[LoopScheduleToCalyx] Lower pipeline register to register writes. (#8048

)

include/circt/Dialect/Calyx/CalyxLoweringUtils.h

@@ -29,6 +29,7 @@
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/JSON.h"
 
+#include <optional>
 #include <variant>
 
 namespace circt {
@@ -404,12 +405,13 @@ class ComponentLoweringStateInterface {
   /// Put the name of the callee and the instance of the call into map.
   void addInstance(StringRef calleeName, InstanceOp instanceOp);
 
-  /// Return the group which evaluates the value v. Optionally, caller may
-  /// specify the expected type of the group.
+  /// Returns the evaluating group or None if not found.
   template <typename TGroupOp = calyx::GroupInterface>
-  TGroupOp getEvaluatingGroup(Value v) {
+  std::optional<TGroupOp> findEvaluatingGroup(Value v) {
     auto it = valueGroupAssigns.find(v);
-    assert(it != valueGroupAssigns.end() && "No group evaluating value!");
+    if (it == valueGroupAssigns.end())
+      return std::nullopt;
+
     if constexpr (std::is_same_v<TGroupOp, calyx::GroupInterface>)
       return it->second;
     else {
@@ -419,6 +421,15 @@ class ComponentLoweringStateInterface {
     }
   }
 
+  /// Return the group which evaluates the value v. Optionally, caller may
+  /// specify the expected type of the group.
+  template <typename TGroupOp = calyx::GroupInterface>
+  TGroupOp getEvaluatingGroup(Value v) {
+    std::optional<TGroupOp> group = findEvaluatingGroup<TGroupOp>(v);
+    assert(group.has_value() && "No group evaluating value!");
+    return *group;
+  }
+
   template <typename T, typename = void>
   struct IsFloatingPoint : std::false_type {};
 

lib/Conversion/LoopScheduleToCalyx/LoopScheduleToCalyx.cpp

@@ -1019,11 +1019,11 @@ class BuildPipelineGroups : public calyx::FuncOpPartialLoweringPattern {
 
     // Collect group names for the prologue or epilogue.
     SmallVector<StringAttr> prologueGroups, epilogueGroups;
+    auto &state = getState<ComponentLoweringState>();
 
     auto updatePrologueAndEpilogue = [&](calyx::GroupOp group) {
       // Mark the group for scheduling in the pipeline's block.
-      getState<ComponentLoweringState>().addBlockScheduleable(stage->getBlock(),
-                                                              group);
+      state.addBlockScheduleable(stage->getBlock(), group);
 
       // Add the group to the prologue or epilogue for this stage as
       // necessary. The goal is to fill the pipeline so it will be in steady
@@ -1046,8 +1046,7 @@ class BuildPipelineGroups : public calyx::FuncOpPartialLoweringPattern {
       // Covers the case where there are no values that need to be passed
       // through to the next stage, e.g., some intermediary store.
       for (auto &op : stage.getBodyBlock())
-        if (auto group = getState<ComponentLoweringState>()
-                             .getNonPipelinedGroupFrom<calyx::GroupOp>(&op))
+        if (auto group = state.getNonPipelinedGroupFrom<calyx::GroupOp>(&op))
           updatePrologueAndEpilogue(*group);
     }
 
@@ -1058,26 +1057,51 @@ class BuildPipelineGroups : public calyx::FuncOpPartialLoweringPattern {
       // Get the pipeline register for that result.
       auto pipelineRegister = pipelineRegisters[i];
 
-      // Get the evaluating group for that value.
-      calyx::GroupInterface evaluatingGroup =
-          getState<ComponentLoweringState>().getEvaluatingGroup(value);
-
-      // Remember the final group for this stage result.
       calyx::GroupOp group;
-
-      // Stitch the register in, depending on whether the group was
-      // combinational or sequential.
-      if (auto combGroup =
-              dyn_cast<calyx::CombGroupOp>(evaluatingGroup.getOperation()))
-        group =
-            convertCombToSeqGroup(combGroup, pipelineRegister, value, rewriter);
-      else
-        group =
-            replaceGroupRegister(evaluatingGroup, pipelineRegister, rewriter);
-
-      // Replace the stage result uses with the register out.
-      stage.getResult(i).replaceAllUsesWith(pipelineRegister.getOut());
-
+      // Get the evaluating group for that value.
+      std::optional<calyx::GroupInterface> evaluatingGroup =
+          state.findEvaluatingGroup(value);
+      if (!evaluatingGroup.has_value()) {
+        if (auto opStage =
+                dyn_cast<LoopSchedulePipelineStageOp>(value.getDefiningOp())) {
+          // The pipeline register for this input value needs to be discovered.
+          auto opResult = cast<OpResult>(value);
+          unsigned int opNumber = opResult.getResultNumber();
+          auto &stageRegisters = state.getPipelineRegs(opStage);
+          calyx::RegisterOp opRegister = stageRegisters.find(opNumber)->second;
+          value = opRegister.getOut(); // Pass the `out` wire of this register.
+        }
+        if (value.getDefiningOp<calyx::RegisterOp>() == nullptr) {
+          // We add this for any unhandled cases.
+          llvm::errs() << "unexpected: input value: " << value << ", in stage "
+                       << stage.getStageNumber() << " register " << i
+                       << " is not a register and was not previously "
+                          "evaluated in a Calyx group. Please open an issue.\n";
+          return LogicalResult::failure();
+        }
+        // This is a register's `out` value being written to this pipeline
+        // register. We create a new group to build this assignment.
+        std::string groupName = state.getUniqueName(
+            loweringState().blockName(pipelineRegister->getBlock()));
+        group = calyx::createGroup<calyx::GroupOp>(
+            rewriter, state.getComponentOp(), pipelineRegister->getLoc(),
+            groupName);
+        calyx::buildAssignmentsForRegisterWrite(
+            rewriter, group, state.getComponentOp(), pipelineRegister, value);
+      } else {
+        // This was previously evaluated. Stitch the register in, depending on
+        // whether the group was combinational or sequential.
+        auto combGroup =
+            dyn_cast<calyx::CombGroupOp>(evaluatingGroup->getOperation());
+        group = combGroup == nullptr
+                    ? replaceGroupRegister(*evaluatingGroup, pipelineRegister,
+                                           rewriter)
+                    : convertCombToSeqGroup(combGroup, pipelineRegister, value,
+                                            rewriter);
+
+        // Replace the stage result uses with the register out.
+        stage.getResult(i).replaceAllUsesWith(pipelineRegister.getOut());
+      }
       updatePrologueAndEpilogue(group);
     }
 

test/Conversion/LoopScheduleToCalyx/convert_pipeline.mlir

@@ -267,3 +267,47 @@ module {
   }
 }
 
+// -----
+
+// Pipeline register to pipeline register writes are valid.
+
+//CHECK:      calyx.group @bb0_3 {
+//CHECK-NEXT:   calyx.assign %stage_1_register_0_reg.in = %stage_0_register_0_reg.out : i32
+//CHECK-NEXT:   calyx.assign %stage_1_register_0_reg.write_en = %true : i1
+//CHECK-NEXT:   calyx.group_done %stage_1_register_0_reg.done : i1
+//CHECK-NEXT: }
+//CHECK-NEXT: calyx.group @bb0_4 {
+//CHECK-NEXT:   calyx.assign %stage_2_register_0_reg.in = %stage_1_register_0_reg.out : i32
+//CHECK-NEXT:   calyx.assign %stage_2_register_0_reg.write_en = %true : i1
+//CHECK-NEXT:   calyx.group_done %stage_2_register_0_reg.done : i1
+//CHECK-NEXT: }
+module {
+  func.func @foo(%arg0: i32, %arg1: memref<30x30xi32>) attributes {llvm.linkage = #llvm.linkage<external>} {
+    %alloca = memref.alloca() : memref<30x30xi32>
+    cf.br ^bb4
+  ^bb4:
+    %0 = arith.constant 0 : index
+    %c1_5 = arith.constant 1 : index
+    loopschedule.pipeline II =  2 trip_count =  20 iter_args(%arg5 = %0) : (index) -> () {
+      %6 = arith.cmpi ult, %arg5, %0 : index
+      loopschedule.register %6 : i1
+    } do {
+      %6:2 = loopschedule.pipeline.stage start = 0 {
+        %10 = memref.load %arg1[%0, %0] : memref<30x30xi32>
+        %11 = arith.addi %arg5, %c1_5 : index
+        loopschedule.register %10, %11 : i32, index
+      } : i32, index
+      %7 = loopschedule.pipeline.stage start = 1 {
+        loopschedule.register %6#0 : i32
+      } : i32
+      %8 = loopschedule.pipeline.stage start = 2 {
+        loopschedule.register %7 : i32
+      } : i32
+      loopschedule.terminator iter_args(%6#1), results() : (index) -> ()
+    }
+    cf.br ^bb6
+  ^bb6:
+    return
+  }
+}
+