[FIRRTL] Support layers in MergeConnections

lib/Dialect/FIRRTL/Transforms/MergeConnections.cpp

@@ -25,6 +25,7 @@
 #include "circt/Dialect/FIRRTL/FIRRTLUtils.h"
 #include "circt/Dialect/FIRRTL/Passes.h"
 #include "circt/Support/Debug.h"
+#include "mlir/IR/Iterators.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
@@ -253,30 +254,67 @@ bool MergeConnection::peelConnect(MatchingConnectOp connect) {
 bool MergeConnection::run() {
   ImplicitLocOpBuilder theBuilder(moduleOp.getLoc(), moduleOp.getContext());
   builder = &theBuilder;
+
+  // Block worklist that tracks the current position within a block.
+  SmallVector<std::pair<Block::iterator, Block::iterator>> worklist;
+
+  // Walk the IR in order, top-to-bottom, stepping into blocks as they are
+  // found.  This is basically the same as `moduleOp.walk`, however, it allows
+  // for visiting operations that are inserted _after_ the current operation.
+  // Using the existing `walk` does not do this.
   auto *body = moduleOp.getBodyBlock();
-  // Merge connections by forward iterations.
-  for (auto it = body->begin(), e = body->end(); it != e;) {
-    auto connectOp = dyn_cast<MatchingConnectOp>(*it);
-    if (!connectOp) {
-      it++;
-      continue;
+  worklist.push_back({body->begin(), body->end()});
+  while (!worklist.empty()) {
+    auto &[it, e] = worklist.back();
+
+    // Merge connections by forward iterations.
+    bool opWithBlocks = false;
+    while (it != e) {
+      // Add blocks to the stack such that they will be pulled off in-order.
+      for (auto &region : llvm::reverse(it->getRegions()))
+        for (auto &block : llvm::reverse(region.getBlocks())) {
+          worklist.push_back({block.begin(), block.end()});
+          opWithBlocks = true;
+        }
+
+      // We found one or more blocks.  Stop and go process these blocks.
+      if (opWithBlocks) {
+        ++it;
+        break;
+      }
+
+      // This operation does not have blocks.  Process it normally.
+      auto connectOp = dyn_cast<MatchingConnectOp>(*it);
+      if (!connectOp) {
+        ++it;
+        continue;
+      }
+      builder->setInsertionPointAfter(connectOp);
+      builder->setLoc(connectOp.getLoc());
+      bool removeOp = peelConnect(connectOp);
+      ++it;
+      if (removeOp)
+        connectOp.erase();
     }
-    builder->setInsertionPointAfter(connectOp);
-    builder->setLoc(connectOp.getLoc());
-    bool removeOp = peelConnect(connectOp);
-    ++it;
-    if (removeOp)
-      connectOp.erase();
+
+    // We found a block and added to the worklist.
+    if (opWithBlocks)
+      continue;
+
+    // We finished processing a block.
+    worklist.pop_back();
   }
 
   // Clean up dead operations introduced by this pass.
-  for (auto &op : llvm::make_early_inc_range(llvm::reverse(*body)))
-    if (isa<SubfieldOp, SubindexOp, InvalidValueOp, ConstantOp, BitCastOp,
-            CatPrimOp>(op))
-      if (op.use_empty()) {
-        changed = true;
-        op.erase();
-      }
+  moduleOp.walk<mlir::WalkOrder::PostOrder, mlir::ReverseIterator>(
+      [&](Operation *op) {
+        if (isa<SubfieldOp, SubindexOp, InvalidValueOp, ConstantOp, BitCastOp,
+                CatPrimOp>(op))
+          if (op->use_empty()) {
+            changed = true;
+            op->erase();
+          }
+      });
 
   return changed;
 }

test/Dialect/FIRRTL/merge-connections.mlir

@@ -2,6 +2,7 @@
 // RUN: circt-opt --pass-pipeline='builtin.module(firrtl.circuit(firrtl.module(merge-connections{aggressive-merging=true})))' %s | FileCheck %s --check-prefixes=AGGRESSIVE,COMMON
 
 firrtl.circuit "Test"   {
+  firrtl.layer @A bind {}
   // circuit Test :
   //   module Test :
   //     input a : {c: {clock: Clock, valid:UInt<1>}[2]}
@@ -29,6 +30,36 @@ firrtl.circuit "Test"   {
      firrtl.matchingconnect %11, %10 : !firrtl.uint<1>
   }
 
+  // This is the same as @Test except it puts the destination into a layer.
+  //
+  // COMMON-LABEL: firrtl.module @Layers(
+  // COMMON-NEXT:    firrtl.layerblock @A {
+  // COMMON-NEXT:      %b = firrtl.wire
+  // COMMON-NEXT:      firrtl.matchingconnect %b, %a
+  // COMMON-NEXT:    }
+  // COMMON-NEXT:  }
+  firrtl.module @Layers(in %a: !firrtl.vector<bundle<clock: clock, valid: uint<1>>, 2>) {
+    firrtl.layerblock @A {
+      %b = firrtl.wire : !firrtl.vector<bundle<clock: clock, valid: uint<1>>, 2>
+      %0 = firrtl.subindex %a[0] : !firrtl.vector<bundle<clock: clock, valid: uint<1>>, 2>
+      %1 = firrtl.subindex %b[0] : !firrtl.vector<bundle<clock: clock, valid: uint<1>>, 2>
+      %2 = firrtl.subfield %0[clock] : !firrtl.bundle<clock: clock, valid: uint<1>>
+      %3 = firrtl.subfield %1[clock] : !firrtl.bundle<clock: clock, valid: uint<1>>
+      firrtl.matchingconnect %3, %2 : !firrtl.clock
+      %4 = firrtl.subfield %0[valid] : !firrtl.bundle<clock: clock, valid: uint<1>>
+      %5 = firrtl.subfield %1[valid] : !firrtl.bundle<clock: clock, valid: uint<1>>
+      firrtl.matchingconnect %5, %4 : !firrtl.uint<1>
+      %6 = firrtl.subindex %a[1] : !firrtl.vector<bundle<clock: clock, valid: uint<1>>, 2>
+      %7 = firrtl.subindex %b[1] : !firrtl.vector<bundle<clock: clock, valid: uint<1>>, 2>
+      %8 = firrtl.subfield %6[clock] : !firrtl.bundle<clock: clock, valid: uint<1>>
+      %9 = firrtl.subfield %7[clock] : !firrtl.bundle<clock: clock, valid: uint<1>>
+      firrtl.matchingconnect %9, %8 : !firrtl.clock
+      %10 = firrtl.subfield %6[valid] : !firrtl.bundle<clock: clock, valid: uint<1>>
+      %11 = firrtl.subfield %7[valid] : !firrtl.bundle<clock: clock, valid: uint<1>>
+      firrtl.matchingconnect %11, %10 : !firrtl.uint<1>
+    }
+  }
+
   // circuit Bar :
   //   module Bar :
   //     output a : {b: UInt<1>, c:UInt<1>}