[SCFToCalyx] Add combinational group inlining [11/12] (#1862)

This commit introduces combinational group inlining. The conversion strategy of SCFToCalyx is to create a distinct CombGroupOp for each combinational operator in the source program. This implies that the created CombGroupOp is analogous to the SSA value which the original operation represented.

In Calyx, a source operand for a `calyx.assign` must be fully defined within the group that performs the assignment. This means that the source operand is either the output of a sequential component, or a fully defined combinational assignment chain. This commit achieves the latter by recursively inlining use-def chains op combinational logic, by backtracking through the CombGroupOps which generated the combinational assignments of the original program.

lib/Conversion/SCFToCalyx/SCFToCalyx.cpp

@@ -1234,6 +1234,85 @@ class BuildControl : public FuncOpPartialLoweringPattern {
   }
 };
 
+/// This pass recursively inlines use-def chains of combinational logic (from
+/// non-stateful groups) into groups referenced in the control schedule.
+class InlineCombGroups
+    : public PartialLoweringPattern<calyx::GroupInterface,
+                                    OpInterfaceRewritePattern> {
+public:
+  InlineCombGroups(MLIRContext *context, LogicalResult &resRef,
+                   ProgramLoweringState &pls)
+      : PartialLoweringPattern(context, resRef), pls(pls) {}
+
+  LogicalResult partiallyLower(calyx::GroupInterface originGroup,
+                               PatternRewriter &rewriter) const override {
+    auto &state = pls.compLoweringState(
+        originGroup->getParentOfType<calyx::ComponentOp>());
+
+    /// Filter groups which are not part of the control schedule.
+    if (SymbolTable::symbolKnownUseEmpty(originGroup.symName(),
+                                         state.getComponentOp().getControlOp()))
+      return success();
+
+    /// Maintain a set of the groups which we've inlined so far. The group
+    /// itself is implicitly inlined.
+    llvm::SmallSetVector<Operation *, 8> inlinedGroups;
+    inlinedGroups.insert(originGroup);
+
+    /// Starting from the matched originGroup, we traverse use-def chains of
+    /// combinational logic, and inline assignments from the defining
+    /// combinational groups.
+    recurseInlineCombGroups(
+        rewriter, state, inlinedGroups, originGroup, originGroup,
+        /*disable inlining of the originGroup itself*/ false);
+    return success();
+  }
+
+private:
+  void
+  recurseInlineCombGroups(PatternRewriter &rewriter,
+                          ComponentLoweringState &state,
+                          llvm::SmallSetVector<Operation *, 8> &inlinedGroups,
+                          calyx::GroupInterface originGroup,
+                          calyx::GroupInterface recGroup, bool doInline) const {
+    inlinedGroups.insert(recGroup);
+    for (auto assignOp : recGroup.getBody()->getOps<calyx::AssignOp>()) {
+      if (doInline) {
+        /// Inline the assignment into the originGroup.
+        auto clonedAssignOp = rewriter.clone(*assignOp.getOperation());
+        clonedAssignOp->moveBefore(originGroup.getBody(),
+                                   originGroup.getBody()->end());
+      }
+      auto src = assignOp.src();
+      auto srcDefOp = src.getDefiningOp();
+
+      /// Things which stop recursive inlining (or in other words, what
+      /// breaks combinational paths).
+      /// - Component inputs
+      /// - Register and memory reads
+      /// - Constant ops (constant ops are not evaluated by any group)
+      /// - 'While' return values (these are registers, however, 'while'
+      ///   return values have at the current point of conversion not yet
+      ///   been rewritten to their register outputs, see comment in
+      ///   LateSSAReplacement)
+      if (src.isa<BlockArgument>() ||
+          isa<calyx::RegisterOp, calyx::MemoryOp, hw::ConstantOp, ConstantOp,
+              scf::WhileOp>(srcDefOp))
+        continue;
+
+      auto srcCombGroup = state.getEvaluatingGroup<calyx::CombGroupOp>(src);
+      assert(srcCombGroup && "expected combinational group");
+      if (inlinedGroups.count(srcCombGroup))
+        continue;
+
+      recurseInlineCombGroups(rewriter, state, inlinedGroups, originGroup,
+                              srcCombGroup, true);
+    }
+  }
+
+  ProgramLoweringState &pls;
+};
+
 /// LateSSAReplacement contains various functions for replacing SSA values that
 /// were not replaced during op construction.
 class LateSSAReplacement : public FuncOpPartialLoweringPattern {
@@ -1665,6 +1744,10 @@ void SCFToCalyxPass::runOnOperation() {
   /// basic block in the source function.
   addOncePattern<BuildControl>(loweringPatterns, funcMap, *loweringState);
 
+  /// This pass recursively inlines use-def chains of combinational logic (from
+  /// non-stateful groups) into groups referenced in the control schedule.
+  addOncePattern<InlineCombGroups>(loweringPatterns, *loweringState);
+
   /// This pattern performs various SSA replacements that must be done
   /// after control generation.
   addOncePattern<LateSSAReplacement>(loweringPatterns, funcMap, *loweringState);