[FIRRTL][LowerLayers] Update rwprobe operations if possible.

Fixes first example in llvm#7365 . Add error path to LowerLayers so anything that goes wrong can fail the pass.
dtzSiFive · Aug 6, 2024 · cf464af · cf464af
1 parent bec0dea
commit cf464af
Show file tree

Hide file tree

Showing 4 changed files with 147 additions and 24 deletions.
diff --git a/lib/Dialect/FIRRTL/Transforms/LowerLayers.cpp b/lib/Dialect/FIRRTL/Transforms/LowerLayers.cpp
@@ -162,10 +162,10 @@ class LowerLayersPass
                            SmallVectorImpl<PortInfo> &ports);
 
   /// Strip layer colors from the module's interface.
-  InnerRefMap runOnModuleLike(FModuleLike moduleLike);
+  FailureOr<InnerRefMap> runOnModuleLike(FModuleLike moduleLike);
 
   /// Extract layerblocks and strip probe colors from all ops under the module.
-  void runOnModuleBody(FModuleOp moduleOp, InnerRefMap &innerRefMap);
+  LogicalResult runOnModuleBody(FModuleOp moduleOp, InnerRefMap &innerRefMap);
 
   /// Update the module's port types to remove any explicit layer requirements
   /// from any probe types.
@@ -272,26 +272,33 @@ void LowerLayersPass::removeLayersFromPorts(FModuleLike moduleLike) {
   }
 }
 
-InnerRefMap LowerLayersPass::runOnModuleLike(FModuleLike moduleLike) {
+FailureOr<InnerRefMap>
+LowerLayersPass::runOnModuleLike(FModuleLike moduleLike) {
   LLVM_DEBUG({
     llvm::dbgs() << "Module: " << moduleLike.getModuleName() << "\n";
     llvm::dbgs() << "  Examining Layer Blocks:\n";
   });
 
   // Strip away layers from the interface of the module-like op.
   InnerRefMap innerRefMap;
-  TypeSwitch<Operation *, void>(moduleLike.getOperation())
-      .Case<FModuleOp>([&](auto op) {
-        op.setLayers({});
-        removeLayersFromPorts(op);
-        runOnModuleBody(op, innerRefMap);
-      })
-      .Case<FExtModuleOp, FIntModuleOp, FMemModuleOp>([&](auto op) {
-        op.setLayers({});
-        removeLayersFromPorts(op);
-      })
-      .Case<ClassOp, ExtClassOp>([](auto) {})
-      .Default([](auto) { assert(0 && "unknown module-like op"); });
+  auto result =
+      TypeSwitch<Operation *, LogicalResult>(moduleLike.getOperation())
+          .Case<FModuleOp>([&](auto op) {
+            op.setLayers({});
+            removeLayersFromPorts(op);
+            return runOnModuleBody(op, innerRefMap);
+          })
+          .Case<FExtModuleOp, FIntModuleOp, FMemModuleOp>([&](auto op) {
+            op.setLayers({});
+            removeLayersFromPorts(op);
+            return success();
+          })
+          .Case<ClassOp, ExtClassOp>([](auto) { return success(); })
+          .Default(
+              [](auto *op) { return op->emitError("unknown module-like op"); });
+
+  if (failed(result))
+    return failure();
 
   return innerRefMap;
 }
@@ -305,8 +312,8 @@ void LowerLayersPass::lowerInlineLayerBlock(LayerOp layer,
   layerBlock.erase();
 }
 
-void LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
-                                      InnerRefMap &innerRefMap) {
+LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
+                                               InnerRefMap &innerRefMap) {
   CircuitOp circuitOp = moduleOp->getParentOfType<CircuitOp>();
   StringRef circuitName = circuitOp.getName();
   hw::InnerSymbolNamespace ns(moduleOp);
@@ -327,7 +334,7 @@ void LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
   //    this layer block to the new module.
   // 4. Instantiate the new module outside the layer block and hook it up.
   // 5. Erase the layer block.
-  moduleOp.walk<mlir::WalkOrder::PostOrder>([&](Operation *op) {
+  auto result = moduleOp.walk<mlir::WalkOrder::PostOrder>([&](Operation *op) {
     // Strip layer requirements from any op that might represent a probe.
     if (auto wire = dyn_cast<WireOp>(op)) {
       removeLayersFromValue(wire.getResult());
@@ -465,6 +472,7 @@ void LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
     };
 
     SmallVector<hw::InnerSymAttr> innerSyms;
+    SmallVector<RWProbeOp> rwprobes;
     for (auto &op : llvm::make_early_inc_range(*body)) {
       // Record any operations inside the layer block which have inner symbols.
       // Theses may have symbol users which need to be updated.
@@ -510,6 +518,11 @@ void LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
         continue;
       }
 
+      if (auto rwprobe = dyn_cast<RWProbeOp>(op)) {
+        rwprobes.push_back(rwprobe);
+        continue;
+      }
+
       if (auto connect = dyn_cast<FConnectLike>(op)) {
         auto src = connect.getSrc();
         auto dst = connect.getDest();
@@ -624,6 +637,24 @@ void LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
                               << splice.second << "\n";);
     }
 
+    // Update RWProbe operations.
+    if (!rwprobes.empty()) {
+      for (auto rwprobe : rwprobes) {
+        auto target = rwprobe.getTarget();
+        auto mapped = innerRefMap.find(target);
+        if (mapped == innerRefMap.end())
+          return rwprobe.emitError("rwprobe target not moved"),
+                 WalkResult::interrupt();
+
+        if (mapped->second.second != newModule.getModuleNameAttr())
+          return rwprobe.emitError("rwprobe target refers to different module"),
+                 WalkResult::interrupt();
+
+        rwprobe.setTargetAttr(
+            hw::InnerRefAttr::get(mapped->second.second, target.getName()));
+      }
+    }
+
     // Connect instance ports to values.
     assert(ports.size() == connectValues.size() &&
            "the number of instance ports and values to connect to them must be "
@@ -656,6 +687,7 @@ void LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
 
     return WalkResult::advance();
   });
+  return success(!result.wasInterrupted());
 }
 
 void LowerLayersPass::createMacroDecls(CircuitNamespace &ns, OpBuilder &b,
@@ -718,18 +750,25 @@ void LowerLayersPass::runOnOperation() {
     }
   }
 
-  auto mergeMaps = [](auto &&a, auto &&b) {
-    for (auto bb : b)
-      a.insert(bb);
+  auto mergeMaps = [](auto &&a, auto &&b) -> FailureOr<InnerRefMap> {
+    if (failed(a) || failed(b))
+      return failure();
+    for (auto bb : *b)
+      a->insert(bb);
     return std::forward<decltype(a)>(a);
   };
 
   // Lower the layer blocks of each module.
   SmallVector<FModuleLike> modules(
       circuitOp.getBodyBlock()->getOps<FModuleLike>());
-  auto innerRefMap =
-      transformReduce(circuitOp.getContext(), modules, InnerRefMap{}, mergeMaps,
-                      [this](FModuleLike mod) { return runOnModuleLike(mod); });
+  auto failureOrInnerRefMap = transformReduce(
+      circuitOp.getContext(), modules, FailureOr<InnerRefMap>(InnerRefMap{}),
+      mergeMaps, [this](FModuleLike mod) -> FailureOr<InnerRefMap> {
+        return runOnModuleLike(mod);
+      });
+  if (failed(failureOrInnerRefMap))
+    return signalPassFailure();
+  auto &innerRefMap = *failureOrInnerRefMap;
 
   // Rewrite any hw::HierPathOps which have namepaths that contain rewritting
   // inner refs.

diff --git a/test/Dialect/FIRRTL/lower-layers-errors.mlir b/test/Dialect/FIRRTL/lower-layers-errors.mlir
@@ -0,0 +1,14 @@
+// RUN: circt-opt -firrtl-lower-layers -split-input-file -verify-diagnostics %s
+
+firrtl.circuit "RWProbeCrossLayer" {
+  firrtl.layer @A bind { }
+  firrtl.module @RWProbeCrossLayer() attributes {layers = [@A]} {
+    %z = firrtl.constant 0 : !firrtl.uint<5>
+    %w = firrtl.node sym @sym %z : !firrtl.uint<5>
+    firrtl.layerblock @A {
+      // expected-error @below {{rwprobe target not moved}}
+      %rw = firrtl.ref.rwprobe <@RWProbeCrossLayer::@sym> : !firrtl.rwprobe<uint<5>>
+    }
+  }
+}
+
diff --git a/test/Dialect/FIRRTL/lower-layers.mlir b/test/Dialect/FIRRTL/lower-layers.mlir
@@ -694,3 +694,34 @@ firrtl.circuit "Foo" attributes {
 //
 // CHECK:       sv.verbatim
 // CHECK-SAME:    #hw.output_file<"testbench{{/|\\\\}}layers_Foo_A.sv", excludeFromFileList>
+
+// -----
+
+// Check rwprobe ops are updated.
+// CHECK-LABEL: circuit "RWTH"
+firrtl.circuit "RWTH" {
+  firrtl.layer @T  bind { }
+  firrtl.module @RWTH() attributes {convention = #firrtl<convention scalarized>, layers = [@T]} {
+    %d_p = firrtl.instance d @DUT(out p: !firrtl.rwprobe<uint<1>, @T>)
+    %one = firrtl.constant 1 : !firrtl.uint<1>
+    firrtl.ref.force_initial %one, %d_p, %one: !firrtl.uint<1>, !firrtl.rwprobe<uint<1>, @T>, !firrtl.uint<1>
+  }
+//      CHECK:    firrtl.module private @DUT_T(out %p: !firrtl.rwprobe<uint<1>>) {
+// CHECK-NEXT:      %w = firrtl.wire sym @[[SYM:.+]] : !firrtl.uint<1>
+// CHECK-NEXT:      %0 = firrtl.ref.rwprobe <@DUT_T::@[[SYM]]> : !firrtl.rwprobe<uint<1>>
+// CHECK-NEXT:      firrtl.ref.define %p, %0 : !firrtl.rwprobe<uint<1>>
+// CHECK-NEXT:    }
+// CHECK-NEXT:    firrtl.module @DUT(out %p: !firrtl.rwprobe<uint<1>>) attributes {convention = #firrtl<convention scalarized>} {
+// CHECK-NEXT:      %t_p = firrtl.instance t sym @t {lowerToBind, output_file = #hw.output_file<"layers_RWTH_T.sv", excludeFromFileList>} @DUT_T(out p: !firrtl.rwprobe<uint<1>>)
+// CHECK-NEXT:      firrtl.ref.define %p, %t_p : !firrtl.rwprobe<uint<1>>
+// CHECK-NEXT:    }
+
+  firrtl.module @DUT(out %p: !firrtl.rwprobe<uint<1>, @T>) attributes {convention = #firrtl<convention scalarized>} {
+    firrtl.layerblock @T {
+      %w = firrtl.wire sym @sym : !firrtl.uint<1>
+      %0 = firrtl.ref.rwprobe <@DUT::@sym> : !firrtl.rwprobe<uint<1>>
+      %1 = firrtl.ref.cast %0 : (!firrtl.rwprobe<uint<1>>) -> !firrtl.rwprobe<uint<1>, @T>
+      firrtl.ref.define %p, %1 : !firrtl.rwprobe<uint<1>, @T>
+    }
+  }
+}
diff --git a/test/firtool/layers-rwprobe.fir b/test/firtool/layers-rwprobe.fir
@@ -0,0 +1,39 @@
+; RUN: firtool %s | FileCheck %s
+
+FIRRTL version 4.0.0
+
+; Check use of colored rwprobe in enablelayers.
+; Check rwprobe of hardware in a layer.
+
+; Order isn't critical, but for simplicity just check lines.
+
+;      CHECK: module TH();
+; CHECK-NEXT:   `ifndef SYNTHESIS
+; CHECK-NEXT:     initial
+; CHECK-NEXT:       force TH.d.t.w = 1'h1;
+; CHECK-NEXT:   `endif // not def SYNTHESIS
+; CHECK-NEXT:   DUT d ();
+; CHECK-NEXT: endmodule
+
+;      CHECK: module DUT_T();
+; CHECK-NEXT:   wire w = 1'h0;
+; CHECK-NEXT: endmodule
+;
+;      CHECK: module DUT();
+; CHECK-NEXT: endmodule
+
+circuit TH:
+  layer T, bind:
+
+  public module TH enablelayer T:
+    inst d of DUT
+    force_initial(d.p, UInt<1>(1))
+
+  public module DUT:
+    output p : RWProbe<UInt<1>, T>
+
+    layerblock T:
+      wire w : UInt<1>
+      connect w, UInt<1>(0)
+      define p = rwprobe(w)
+