[CfToHandshake] Do not replace the same value twice

experimental/lib/Conversion/FtdCfToHandshake.cpp

@@ -113,7 +113,7 @@ struct FtdCfToHandshakePass
 };
 } // namespace
 
-using ArgReplacements = DenseMap<BlockArgument, OpResult>;
+using BlockArgToMergeResult = DenseMap<BlockArgument, OpResult>;
 
 static void channelifyMuxes(handshake::FuncOp &funcOp) {
   // Considering each mux that was added, the inputs and output values must be
@@ -281,7 +281,7 @@ LogicalResult ftd::FtdLowerFuncToHandshake::matchAndRewrite(
 
   // Stores mapping from each value that passes through a merge-like
   // operation to the data result of that merge operation
-  ArgReplacements argReplacements;
+  BlockArgToMergeResult argReplacements;
 
   // Currently, the following 2 functions do nothing but construct the network
   // of CMerges in complete isolation from the rest of the components
@@ -308,7 +308,8 @@ LogicalResult ftd::FtdLowerFuncToHandshake::matchAndRewrite(
   // Create the memory interface according to the algorithm from FPGA'23. This
   // functions introduce new data dependencies that are then passed to FTD for
   // correctly delivering data between them like any real data dependencies
-  if (failed(verifyAndCreateMemInterfaces(funcOp, rewriter, memInfo)))
+  if (failed(verifyAndCreateMemInterfaces(funcOp, rewriter, memInfo,
+                                          argReplacements)))
     return failure();
 
   // Convert the constants and undefined values from the `arith` dialect to

include/dynamatic/Conversion/CfToHandshake.h

@@ -27,6 +27,8 @@
 
 namespace dynamatic {
 
+using BlockArgToMergeResult = DenseMap<BlockArgument, OpResult>;
+
 /// 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.
@@ -129,10 +131,10 @@ class LowerFuncToHandshake : public DynOpConversionPattern<mlir::func::FuncOp> {
   /// - 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.
-  virtual LogicalResult
-  verifyAndCreateMemInterfaces(handshake::FuncOp funcOp,
-                               ConversionPatternRewriter &rewriter,
-                               MemInterfacesInfo &memInfo) const;
+  virtual LogicalResult verifyAndCreateMemInterfaces(
+      handshake::FuncOp funcOp, ConversionPatternRewriter &rewriter,
+      MemInterfacesInfo &memInfo,
+      const BlockArgToMergeResult &argReplacements) const;
 
   /// Sets an integer "bb" attribute on each operation to identify the basic
   /// block from which the operation originates in the std-level IR.

lib/Conversion/CfToHandshake/CfToHandshake.cpp

@@ -132,6 +132,36 @@ mergeFuncResults(handshake::FuncOp funcOp, ConversionPatternRewriter &rewriter,
   return results;
 }
 
+/// This function finds the control merge result of a given Block. We
+/// need to use this function to connect the control network to the memory
+/// interface ops (mc, lsq).
+///
+/// Important remarks:
+/// - This function works only after we called addMergeOps (i.e., after the
+/// merge results become available).
+/// - The initial block of a dataflow circuit does not have a control merge.
+/// Therefore, if the block is the initial block, it returns the "start" channel
+/// of the function.
+/// - Why not just use getBlockControl(Block *block)? That function directly
+/// uses the block argument of the block, which will be eventually removed
+/// during the block inlining and become invalid. This function directly uses
+/// the value that will present in the final IR.
+static Value
+getBlockControlMergeOrFuncControl(const BlockArgToMergeResult &argReplacements,
+                                  Block *block) {
+  Operation *parentOp = block->getParentOp();
+
+  auto funcOp = cast<handshake::FuncOp>(parentOp);
+
+  assert(funcOp && "The parent Op of this block must be a handshake::FuncOp");
+
+  // Check if the block is the first block:
+  if (&(funcOp.getBlocks().front()) == block) {
+    return block->getArguments().back();
+  }
+  return argReplacements.at(block->getArguments().back());
+}
+
 LogicalResult LowerFuncToHandshake::computeLinearDominance(
     DenseMap<Block *, DenseSet<Block *>> &dominations,
     llvm::MapVector<Block *, SmallVector<handshake::MemPortOpInterface>>
@@ -208,8 +238,6 @@ CfToHandshakeTypeConverter::CfToHandshakeTypeConverter() {
 // LowerFuncToHandshake
 //===-----------------------------------------------------------------------==//
 
-using ArgReplacements = DenseMap<BlockArgument, OpResult>;
-
 LogicalResult LowerFuncToHandshake::matchAndRewrite(
     func::FuncOp lowerFuncOp, OpAdaptor /*adaptor*/,
     ConversionPatternRewriter &rewriter) const {
@@ -232,8 +260,11 @@ LogicalResult LowerFuncToHandshake::matchAndRewrite(
     return success();
 
   // Stores mapping from each value that passes through a merge-like operation
-  // to the data result of that merge operation
-  ArgReplacements argReplacements;
+  // to the data result of that merge operation.
+  //
+  // This mapping will be used when inlining all the blocks into one (MILR needs
+  // to know what the block arguments will become when inlining the blocks).
+  BlockArgToMergeResult argReplacements;
   addMergeOps(funcOp, rewriter, argReplacements);
   addBranchOps(funcOp, rewriter);
 
@@ -247,7 +278,8 @@ LogicalResult LowerFuncToHandshake::matchAndRewrite(
   // tagged with the BB they belong to (required by memory interface
   // instantiation logic)
   idBasicBlocks(funcOp, rewriter);
-  if (failed(verifyAndCreateMemInterfaces(funcOp, rewriter, memInfo)))
+  if (failed(verifyAndCreateMemInterfaces(funcOp, rewriter, memInfo,
+                                          argReplacements)))
     return failure();
 
   idBasicBlocks(funcOp, rewriter);
@@ -572,9 +604,9 @@ void LowerFuncToHandshake::insertMerge(BlockArgument blockArg,
   }
 }
 
-void LowerFuncToHandshake::addMergeOps(handshake::FuncOp funcOp,
-                                       ConversionPatternRewriter &rewriter,
-                                       ArgReplacements &argReplacements) const {
+void LowerFuncToHandshake::addMergeOps(
+    handshake::FuncOp funcOp, ConversionPatternRewriter &rewriter,
+    BlockArgToMergeResult &argReplacements) const {
   // Create backedge builder to manage operands of merge operations between
   // insertion and reconnection
   BackedgeBuilder edgeBuilder(rewriter, funcOp.getLoc());
@@ -886,7 +918,8 @@ LogicalResult LowerFuncToHandshake::convertMemoryOps(
 
 LogicalResult LowerFuncToHandshake::verifyAndCreateMemInterfaces(
     handshake::FuncOp funcOp, ConversionPatternRewriter &rewriter,
-    MemInterfacesInfo &memInfo) const {
+    MemInterfacesInfo &memInfo,
+    const BlockArgToMergeResult &argReplacements) const {
 
   if (memInfo.empty())
     return success();
@@ -913,15 +946,19 @@ LogicalResult LowerFuncToHandshake::verifyAndCreateMemInterfaces(
   auto returns = funcOp.getOps<func::ReturnOp>();
   assert(!returns.empty() && "no returns in function");
   if (std::distance(returns.begin(), returns.end()) == 1) {
-    ctrlEnd = getBlockControl((*returns.begin())->getBlock());
+    Value lastBlockCtrlArg = getBlockControlMergeOrFuncControl(
+        argReplacements, (*returns.begin())->getBlock());
+
+    ctrlEnd = lastBlockCtrlArg;
   } else {
     // Merge the control signals of all blocks with a return to create a control
     // representing the final control flow decision
     SmallVector<Value> controls;
     func::ReturnOp lastRetOp;
     for (func::ReturnOp retOp : returns) {
       lastRetOp = retOp;
-      controls.push_back(getBlockControl(retOp->getBlock()));
+      controls.push_back(getBlockControlMergeOrFuncControl(argReplacements,
+                                                           retOp->getBlock()));
     }
     rewriter.setInsertionPointToStart(lastRetOp->getBlock());
     auto mergeOp =
@@ -938,7 +975,8 @@ LogicalResult LowerFuncToHandshake::verifyAndCreateMemInterfaces(
   // format for the memory interface builder
   DenseMap<unsigned, Value> ctrlVals;
   for (auto [blockIdx, block] : llvm::enumerate(funcOp))
-    ctrlVals.insert({blockIdx, getBlockControl(&block)});
+    ctrlVals.insert(
+        {blockIdx, getBlockControlMergeOrFuncControl(argReplacements, &block)});
 
   // Each memory region is independent from the others
   for (auto &[memref, memAccesses] : memInfo) {
@@ -1009,7 +1047,7 @@ void LowerFuncToHandshake::idBasicBlocks(
 
 LogicalResult LowerFuncToHandshake::flattenAndTerminate(
     handshake::FuncOp funcOp, ConversionPatternRewriter &rewriter,
-    const ArgReplacements &argReplacements) const {
+    const BlockArgToMergeResult &argReplacements) const {
   // Erase all cf-level terminators, accumulating operands to func-level returns
   // as we go
   SmallVector<SmallVector<Value>> returnsOperands;
@@ -1034,18 +1072,14 @@ LogicalResult LowerFuncToHandshake::flattenAndTerminate(
 
   // Inline all non-entry blocks into the entry block, erasing them as we go
   Operation *lastOp = &funcOp.front().back();
-  for (Block &block : llvm::make_early_inc_range(funcOp)) {
-    if (block.isEntryBlock())
-      continue;
+  for (Block &block : llvm::make_early_inc_range(llvm::drop_begin(funcOp))) {
 
     // Replace all block arguments with the data result of merge-like
     // operations; this effectively connects all merges to the rest of the
     // circuit
     SmallVector<Value> replacements;
     for (BlockArgument blockArg : block.getArguments()) {
-      Value mergeRes = argReplacements.at(blockArg);
-      replacements.push_back(mergeRes);
-      rewriter.replaceAllUsesWith(blockArg, mergeRes);
+      replacements.push_back(argReplacements.at(blockArg));
     }
     rewriter.inlineBlockBefore(&block, lastOp, replacements);
   }
@@ -1079,7 +1113,8 @@ LogicalResult LowerFuncToHandshake::flattenAndTerminate(
       }
     }
   }
-  endOprds.push_back(getBlockControl(funcOp.getBodyBlock()));
+  endOprds.push_back(getBlockControlMergeOrFuncControl(argReplacements,
+                                                       funcOp.getBodyBlock()));
 
   auto endOp = rewriter.create<handshake::EndOp>(lastOp->getLoc(), endOprds);
   endOp->setAttr(BB_ATTR_NAME, rewriter.getUI32IntegerAttr(exitBlockID));