[FIRRTL] Convert LowerMemory to walk

Change LowerMemory to use a walk as opposed to iteration over operations. This is done to make this pass work with operations which have regions and may contain memories, i.e., to make this pass work with layers and whens. Signed-off-by: Schuyler Eldridge <[email protected]>
llvm · Sep 12, 2024 · 9b4512d · 9b4512d
1 parent 008dc9d
commit 9b4512d
Show file tree

Hide file tree

Showing 2 changed files with 48 additions and 13 deletions.
diff --git a/lib/Dialect/FIRRTL/Transforms/LowerMemory.cpp b/lib/Dialect/FIRRTL/Transforms/LowerMemory.cpp
@@ -130,6 +130,9 @@ struct LowerMemoryPass
   /// The set of all memories seen so far.  This is used to "deduplicate"
   /// memories by emitting modules one module for equivalent memories.
   std::map<FirMemory, FMemModuleOp> memories;
+
+  /// A sequence of operations that should be erased later.
+  SmallPtrSet<Operation *, 32> operationsToErase;
 };
 } // end anonymous namespace
 
@@ -322,7 +325,7 @@ void LowerMemoryPass::lowerMemory(MemOp mem, const FirMemory &summary,
     newAnnos.addAnnotations(newMemModAnnos);
     newAnnos.applyToOperation(memInst);
   }
-  mem->erase();
+  operationsToErase.insert(mem);
   ++numLoweredMems;
 }
 
@@ -441,8 +444,8 @@ InstanceOp LowerMemoryPass::emitMemoryInstance(MemOp op, FModuleOp module,
         enConnect->moveAfter(andOp);
         // Erase the old mask connect.
         auto *maskField = maskConnect->getOperand(0).getDefiningOp();
-        maskConnect->erase();
-        maskField->erase();
+        operationsToErase.insert(maskConnect);
+        operationsToErase.insert(maskField);
       };
 
       if (memportKind == MemOp::PortKind::Read) {
@@ -492,27 +495,38 @@ InstanceOp LowerMemoryPass::emitMemoryInstance(MemOp op, FModuleOp module,
   // Update all users of the result of read ports
   for (auto [subfield, result] : returnHolder) {
     subfield->getResult(0).replaceAllUsesWith(inst.getResult(result));
-    subfield->erase();
+    operationsToErase.insert(subfield);
   }
 
   return inst;
 }
 
 LogicalResult LowerMemoryPass::runOnModule(FModuleOp moduleOp,
                                            bool shouldDedup) {
-  for (auto op :
-       llvm::make_early_inc_range(moduleOp.getBodyBlock()->getOps<MemOp>())) {
+  assert(operationsToErase.empty() && "operationsToErase must be empty");
+
+  auto result = moduleOp.walk([&](MemOp op) {
     // Check that the memory has been properly lowered already.
-    if (!type_isa<UIntType>(op.getDataType()))
-      return op->emitError(
-          "memories should be flattened before running LowerMemory");
+    if (!type_isa<UIntType>(op.getDataType())) {
+      op->emitError("memories should be flattened before running LowerMemory");
+      return WalkResult::interrupt();
+    }
 
     auto summary = getSummary(op);
-    if (!summary.isSeqMem())
-      continue;
+    if (summary.isSeqMem())
+      lowerMemory(op, summary, shouldDedup);
+
+    return WalkResult::advance();
+  });
+
+  if (result.wasInterrupted())
+    return failure();
+
+  for (Operation *op : operationsToErase)
+    op->erase();
+
+  operationsToErase.clear();
 
-    lowerMemory(op, summary, shouldDedup);
-  }
   return success();
 }
 

diff --git a/test/Dialect/FIRRTL/lower-memory.mlir b/test/Dialect/FIRRTL/lower-memory.mlir
@@ -275,3 +275,24 @@ firrtl.module @DUT() {
 // CHECK-SAME:  {annotations = [{circt.nonlocal = @[[nla_1]], class = "test1"}]}
 // CHECK-SAME:  @mem0_ext(
 }
+
+// CHECK-LABEL: "MemoryInLayer"
+firrtl.circuit "MemoryInLayer" {
+firrtl.layer @A bind {}
+firrtl.module @MemoryInLayer() {
+
+  // CHECK:      firrtl.layerblock @A
+  // CHECK-NEXT:   firrtl.instance mem
+  firrtl.layerblock @A {
+    %mem_w = firrtl.mem sym @mem Undefined {
+      annotations = [],
+      depth = 2 : i64,
+      name = "mem",
+      portNames = ["w"],
+      readLatency = 1 : i32,
+      writeLatency = 1 : i32
+    } : !firrtl.bundle<addr: uint<1>, en: uint<1>, clk: clock, data: uint<1>, mask: uint<1>>
+  }
+
+}
+}