[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.

lib/Dialect/FIRRTL/Transforms/LowerLayers.cpp

@@ -465,6 +465,7 @@ LogicalResult 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.
@@ -550,6 +551,11 @@ LogicalResult 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();
@@ -664,6 +670,22 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
                               << splice.second << "\n";);
     }
 
+    // Update RWProbe operations.
+    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 "

test/Dialect/FIRRTL/lower-layers-errors.mlir

@@ -14,3 +14,18 @@ firrtl.circuit "NonPassiveSubaccess" {
     }
   }
 }
+
+// -----
+
+firrtl.circuit "RWProbeCantMove" {
+  firrtl.layer @A bind { }
+  firrtl.module @RWProbeCantMove() 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 <@RWProbeCantMove::@sym> : !firrtl.rwprobe<uint<5>>
+    }
+  }
+}
+

test/Dialect/FIRRTL/lower-layers.mlir

@@ -717,3 +717,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>
+    }
+  }
+}

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)
+