[CfToHandshake] Structural redesign of the pass

The `CfToHandshake` conversion pass has always (I think) been slightly
weird in its implementation, sequentially triggering conversion patterns
to progressively lower a `func::FuncOp` into a `handshake::FuncOp` in a
way that I am not sure was strictly legal and safe with respect to
MLIR's conversion framework invariants. Sometimes, the pass would also
non-deterministically crash because of a mysterious dialect/operation
registration error, which may have been due to this peculiar design.

This completely redesigns the pass so that the bulk of the conversion
process happens in a single conversion pattern, with every IR change
going through the conversion pattern rewriter. Some piphole conversions
still happen in seperate patterns that can only be triggered after the
bulk conversion has happened. While the outcome of the pass doesn't
change (modulo cosmetic things like the removal of useless merge-like
operations in the entry block), the pass's structure changes
significantly. Gone is the `HandshakeLowering` data-structure along with
its accompanying partial conversion pattern and target; it is largely
replaced with the exported `LowerFuncToHandshake` conversion pattern.
The latter exposes all of its "partial internal conversion steps" as
virtual methods to allow future users to reuse parts of its
implementation when possible while forgoing the rest.

The `handshake::MergeLikeOpInterface` gains an additional method to
query for the data result of its attached operation. This makes querying
this result more reliable from client code, for example this conversion
pass.

include/dynamatic/Conversion/CfToHandshake.h

@@ -16,18 +16,52 @@
 
 #include "dynamatic/Analysis/NameAnalysis.h"
 #include "dynamatic/Dialect/Handshake/HandshakeInterfaces.h"
+#include "dynamatic/Dialect/Handshake/HandshakeOps.h"
 #include "dynamatic/Support/Backedge.h"
 #include "dynamatic/Support/DynamaticPass.h"
 #include "dynamatic/Support/LLVM.h"
+#include "dynamatic/Transforms/FuncMaximizeSSA.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Attributes.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 namespace dynamatic {
 
-/// This class is strongly inspired by CIRCT's own `HandshakeLowering` class. It
-/// provides all the conversion steps necessary to concert a func-level function
-/// into a matching handshake-level function.
-class HandshakeLowering {
+/// Converts cf-level types to those expected by handshake-level IR. Right now
+/// these are the sames but this will change as soon as the new type system is
+/// integrated.
+class CfToHandshakeTypeConverter : public TypeConverter {
 public:
+  CfToHandshakeTypeConverter();
+};
+
+/// Converts a func-level function into a handshake-level function. The function
+/// signature gets an extra control-only argument to represent the starting
+/// point of the control network. If the function did not return any result, a
+/// control-only result is added to signal function completion. All of the
+/// pattern's intermediate conversion steps are virtual, allowing other passes
+/// to reuse part of the conversion while defining custom behavior.
+class LowerFuncToHandshake : public OpConversionPattern<mlir::func::FuncOp> {
+public:
+  /// Constructs from a type converter and a reference to the englobing pass's
+  /// name analysis (to name new operations as they are inserted into the IR).
+  LowerFuncToHandshake(TypeConverter &typeConverter, NameAnalysis &namer,
+                       MLIRContext *ctx);
+
+  LogicalResult
+  matchAndRewrite(mlir::func::FuncOp funcOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override;
+
+  /// Strategy to use when putting the matched func-level function into maximal
+  /// SSA form.
+  class FuncSSAStrategy : public dynamatic::SSAMaximizationStrategy {
+    /// Filters out block arguments of type MemRefType
+    bool maximizeArgument(BlockArgument arg) override;
+
+    /// Filters out allocation operations
+    bool maximizeOp(Operation &op) override;
+  };
+
   /// Groups memory operations by interface and group for a given memory region.
   struct MemAccesses {
     /// Memory operations for a simple memory controller, grouped by
@@ -37,57 +71,42 @@ class HandshakeLowering {
     llvm::MapVector<unsigned, SmallVector<Operation *>> lsqPorts;
   };
 
-  /// Groups information to "rewire the IR" around a particular merge-like
-  /// operation.
-  struct MergeOpInfo {
-    /// The merge-like operation under consideration.
-    handshake::MergeLikeOpInterface mergeLikeOp;
-    /// The original block argument that the merge-like operation "replaces".
-    BlockArgument blockArg;
-    /// All data operands to the merge-like operation that need to be resolved
-    /// during branch insertion.
-    SmallVector<Backedge> dataEdges;
-    /// An optional index operand that needs to be resolved for mux-like
-    /// operations.
-    std::optional<Backedge> indexEdge{};
-  };
-
-  /// Groups information to rewire the IR around merge-like operations by owning
-  /// basic block (which must still exist).
-  using BlockOps = DenseMap<Block *, std::vector<MergeOpInfo>>;
-
   /// Stores a mapping between memory regions (identified by the function
   /// argument they correspond to) and the set of memory operations referencing
   /// them.
   using MemInterfacesInfo = llvm::MapVector<Value, MemAccesses>;
 
-  /// Constructor simply takes the region being lowered and a reference to the
-  /// top-level name analysis.
-  explicit HandshakeLowering(Region &region, NameAnalysis &nameAnalysis)
-      : region(region), nameAnalysis(nameAnalysis) {}
+  /// Creates a Handshake-level equivalent to the matched func-level function,
+  /// returning it on success. A `nullptr` return value indicates a failure.
+  virtual FailureOr<handshake::FuncOp>
+  lowerSignature(mlir::func::FuncOp funcOp,
+                 ConversionPatternRewriter &rewriter) const;
 
-  /// Creates the control-only network by adding a control-only argument to the
-  /// region's entry block and forwarding it through all basic blocks.
-  LogicalResult createControlNetwork(ConversionPatternRewriter &rewriter);
+  /// Produces the list of named attributes that the Handshake function's
+  /// builder will be passed during signature lowering.
+  virtual SmallVector<mlir::NamedAttribute>
+  deriveNewAttributes(mlir::func::FuncOp funcOp) const;
 
   /// Adds merge-like operations after all block arguments within the region,
-  /// then removes all block arguments and corresponding branch operands. This
-  /// always succeeds.
-  LogicalResult addMergeOps(ConversionPatternRewriter &rewriter);
+  /// then removes all block arguments and corresponding branch operands.
+  virtual void
+  addMergeOps(handshake::FuncOp funcOp, ConversionPatternRewriter &rewriter,
+              DenseMap<BlockArgument, OpResult> &blockArgReplacements) const;
 
   /// Adds handshake-level branch-like operations before all cf-level
   /// branch-like terminators within the region. This needs to happen after
   /// merge-insertion because it also replaces data operands of merge-like
   /// operations with the result value(s) of inserted branch-like operations.
-  /// This always succeeds.
-  LogicalResult addBranchOps(ConversionPatternRewriter &rewriter);
+  virtual void addBranchOps(handshake::FuncOp funcOp,
+                            ConversionPatternRewriter &rewriter) const;
 
   /// Identifies all memory interfaces and their associated operations in the
-  /// function, replaces all load/store-like operations by their handshake
+  /// function, converts all load/store-like operations by their handshake
   /// counterparts, and fills `memInfo` with information about which operations
   /// use which interface.
-  LogicalResult replaceMemoryOps(ConversionPatternRewriter &rewriter,
-                                 MemInterfacesInfo &memInfo);
+  virtual LogicalResult convertMemoryOps(handshake::FuncOp funcOp,
+                                         ConversionPatternRewriter &rewriter,
+                                         MemInterfacesInfo &memInfo) const;
 
   /// Verifies that LSQ groups derived from input IR annotations make sense
   /// (check for linear dominance property within each group and cross-group
@@ -101,99 +120,61 @@ class HandshakeLowering {
   /// - Both a `handhsake::MemoryControllerOp` and `handhsake::LSQOp` will be
   /// instantiated if some but not all of its accesses indicate that they should
   /// connect to an LSQ.
-  LogicalResult
-  verifyAndCreateMemInterfaces(ConversionPatternRewriter &rewriter,
-                               MemInterfacesInfo &memInfo);
-
-  /// Converts each `func::CallOp` operation to an equivalent
-  /// `handshake::InstanceOp` operation. This always succeeds.
-  LogicalResult convertCalls(ConversionPatternRewriter &rewriter);
-
-  /// Connect constants to the rest of the circuit. Constants are triggered by a
-  /// source if their successor is not a branch/return or memory operation.
-  /// Otherwise they are triggered by the control-only network.
-  LogicalResult connectConstants(ConversionPatternRewriter &rewriter);
-
-  /// Replaces undefined operations (mlir::LLVM::UndefOp) with a default "0"
-  /// constant triggered by the enclosing block's control merge.
-  LogicalResult replaceUndefinedValues(ConversionPatternRewriter &rewriter);
+  virtual LogicalResult
+  verifyAndCreateMemInterfaces(handshake::FuncOp funcOp,
+                               ConversionPatternRewriter &rewriter,
+                               MemInterfacesInfo &memInfo) const;
 
   /// Sets an integer "bb" attribute on each operation to identify the basic
   /// block from which the operation originates in the std-level IR.
-  LogicalResult idBasicBlocks(ConversionPatternRewriter &rewriter);
+  virtual void idBasicBlocks(handshake::FuncOp funcOp,
+                             ConversionPatternRewriter &rewriter) const;
 
   /// Creates the region's return network by sequentially moving all blocks'
   /// operations to the entry block, replacing func::ReturnOp's with
   /// handshake::ReturnOp's, deleting all block terminators and non-entry
   /// blocks, merging the results of all return statements, and creating the
   /// region's end operation.
-  LogicalResult createReturnNetwork(ConversionPatternRewriter &rewriter);
-
-  /// Returns the entry control value for operations contained within this
-  /// block.
-  Value getBlockEntryControl(Block *block) const {
-    auto it = blockControls.find(block);
-    assert(it != blockControls.end() &&
-           "No block entry control value registerred for this block!");
-    return it->second;
-  }
-
-  /// Set the control value of a basic block.
-  void setBlockEntryControl(Block *block, Value v) {
-    blockControls[block] = v;
-  };
-
-  /// Returns a reference to the region being lowered.
-  Region &getRegion() { return region; }
+  virtual void flattenAndTerminate(
+      handshake::FuncOp funcOp, ConversionPatternRewriter &rewriter,
+      const DenseMap<BlockArgument, OpResult> &blockArgReplacements) const;
 
-protected:
-  /// The region being lowered.
-  Region &region;
-  /// Start point of the control-only network
-  BlockArgument startCtrl;
-
-  /// Inserts a merge-like operation in the IR for the block argument and
-  /// returns information necessary to rewire the IR around the new operation
-  /// once all merges have been inserted. A control-merge is inserted for
-  /// control-only (data-less) arguments. For other types of arguments, a
-  /// non-deterministic merge is inserted for blocks with 0 or a single
-  /// predecessor while a mux is inserted for blocks with multiple predecessors.
-  MergeOpInfo insertMerge(BlockArgument blockArg, BackedgeBuilder &edgeBuilder,
-                          ConversionPatternRewriter &rewriter);
+  /// Returns the value representing the block's control signal.
+  virtual Value getBlockControl(Block *block) const;
 
 private:
-  /// Associates basic blocks of the region being lowered to their respective
-  /// control value.
-  DenseMap<Block *, Value> blockControls;
   /// Name analysis to name new memory operations as they are created and keep
   /// reference accesses in memory dependencies consistent.
-  NameAnalysis &nameAnalysis;
-};
+  NameAnalysis &namer;
 
-/// Pointer to function lowering a region using a conversion pattern rewriter.
-using RegionLoweringFunc =
-    llvm::function_ref<LogicalResult(Region &, ConversionPatternRewriter &)>;
-
-/// Partially lowers a region using a provided lowering function.
-LogicalResult partiallyLowerRegion(const RegionLoweringFunc &loweringFunc,
-                                   Region &region);
-
-/// Runs a partial lowering method on an instance of the class the method
-/// belongs to. We need two variadic template parameters because arguments
-/// provided to this function may be slightly different but convertible to the
-/// arguments expected by the partial lowering method. Success status is
-/// forwarded from the partial lowering method.
-template <typename T, typename... TArgs1, typename... TArgs2>
-static LogicalResult runPartialLowering(
-    T &instance,
-    LogicalResult (T::*memberFunc)(ConversionPatternRewriter &, TArgs2...),
-    TArgs1 &...args) {
-  return partiallyLowerRegion(
-      [&](Region &, ConversionPatternRewriter &rewriter) -> LogicalResult {
-        return (instance.*memberFunc)(rewriter, args...);
-      },
-      instance.getRegion());
-}
+  /// Groups information to "rewire the IR" around a particular merge-like
+  /// operation.
+  struct MergeOpInfo {
+    /// The original block argument that the merge-like operation "replaces".
+    BlockArgument blockArg;
+    /// The merge-like operation under consideration.
+    handshake::MergeLikeOpInterface op = nullptr;
+    /// Each vector entry represent a data operand to the merge as
+    /// 1. the backedge that was inserted to temporarily represent it,
+    /// 2. the predecessor block from which the data should come,
+    /// 3. a boolean indicating whether this is the first operand to come from
+    ///    the associated block
+    SmallVector<std::tuple<Backedge, Block *, bool>, 2> operands;
+    /// An optional index operand that needs to be resolved for mux-like
+    /// operations.
+    std::optional<Backedge> indexEdge{};
+
+    /// Constructs from the block argument the future merge-like operation will
+    /// replace.
+    MergeOpInfo(BlockArgument blockArg) : blockArg(blockArg) {}
+  };
+
+  /// Inserts a merge-like operation in the IR for the provided block argument.
+  /// Stores information about the merge-like operation in the last argument.
+  void insertMerge(BlockArgument blockArg, ConversionPatternRewriter &rewriter,
+                   BackedgeBuilder &edgeBuilder,
+                   LowerFuncToHandshake::MergeOpInfo &iMerge) const;
+};
 
 #define GEN_PASS_DECL_CFTOHANDSHAKE
 #define GEN_PASS_DEF_CFTOHANDSHAKE

include/dynamatic/Dialect/Handshake/HandshakeInterfaces.td

@@ -124,6 +124,11 @@ def MergeLikeOpInterface : OpInterface<"MergeLikeOpInterface"> {
       }],
       "mlir::OperandRange", "getDataOperands", (ins)
     >,
+    InterfaceMethod<[{
+        Returns the operation's result representing the selected data operand.
+      }],
+      "mlir::OpResult", "getDataResult", (ins)
+    >,
   ];
 
   let verify = [{

include/dynamatic/Dialect/Handshake/HandshakeOps.td

@@ -354,7 +354,8 @@ def LazyForkOp : Handshake_Op<"lazy_fork", [
 }
 
 def MergeOp : Handshake_Op<"merge", [
-  Pure, MergeLikeOpInterface, SOSTInterface, SameOperandsAndResultType
+  Pure, SOSTInterface, SameOperandsAndResultType,
+  DeclareOpInterfaceMethods<MergeLikeOpInterface, ["getDataResult"]>
 ]> {
   let summary = "merge operation";
   let description = [{
@@ -375,7 +376,8 @@ def MergeOp : Handshake_Op<"merge", [
 }
 
 def MuxOp : Handshake_Op<"mux", [
-  Pure, MergeLikeOpInterface,
+  Pure,
+  DeclareOpInterfaceMethods<MergeLikeOpInterface, ["getDataResult"]>,
   DeclareOpInterfaceMethods<InferTypeOpInterface, ["inferReturnTypes"]>,
   DeclareOpInterfaceMethods<ControlInterface, ["isControl"]>,
   DeclareOpInterfaceMethods<NamedIOInterface, ["getOperandName"]>
@@ -405,7 +407,8 @@ def MuxOp : Handshake_Op<"mux", [
 }
 
 def ControlMergeOp : Handshake_Op<"control_merge", [
-  Pure, MergeLikeOpInterface, HasClock, SOSTInterface,
+  Pure, HasClock, SOSTInterface,
+  DeclareOpInterfaceMethods<MergeLikeOpInterface, ["getDataResult"]>,
   DeclareOpInterfaceMethods<NamedIOInterface, ["getResultName"]>
 ]> {
   let summary = "control merge operation";