[LowerToHW] Fix memories moved to testbench bug

Fixes a bug where memories instantiated under the DUT would be moved to
the testbench directory.

This bug was introduced in b40db34 where the set of DUT modules was
precomputed by recording all FIRRTL modules under the DUT.  The `isInDUT`
function would then query this precomputed set.  However, this meant that
queries needed to be FIRRTL modules and not HW modules.  Memories, because
they were visited once FIRRTL module bodies were moved to HW module
bodies, would then incorrectly always return false from the `isInDUT`
method because the _new_ module was being used instead of the _old_
module.

This is fixed by making the `isInDUT` function do the right thing in both
cases.  A new mapping of HW module to FIRRTL module is added.  The
`isInDUT` will use this to query the DUT module set if it is given a new
HW module.

Add one test of this behavior.

Fixes #6775.

Signed-off-by: Schuyler Eldridge <[email protected]>

lib/Conversion/FIRRTLToHW/LowerToHW.cpp

@@ -268,6 +268,11 @@ struct CircuitLoweringState {
     return it != oldToNewModuleMap.end() ? it->second : nullptr;
   }
 
+  Operation *getOldModule(Operation *newModule) {
+    auto it = newToOldModuleMap.find(newModule);
+    return it != newToOldModuleMap.end() ? it->second : nullptr;
+  }
+
   // Process remaining annotations and emit warnings on unprocessed annotations
   // still remaining in the annoSet.
   void processRemainingAnnotations(Operation *op, const AnnotationSet &annoSet);
@@ -299,8 +304,11 @@ struct CircuitLoweringState {
 
   // Return true if this module is the DUT or is instantiated by the DUT.
   // Returns false if the module is not instantiated by the DUT or is
-  // instantiated under a bind.
+  // instantiated under a bind.  This will accept either an old FIRRTL module or
+  // a new HW module.
   bool isInDUT(igraph::ModuleOpInterface child) {
+    if (auto hwModule = dyn_cast<hw::HWModuleOp>(child.getOperation()))
+      child = cast<igraph::ModuleOpInterface>(getOldModule(hwModule));
     return dutModules.contains(child);
   }
 
@@ -330,14 +338,18 @@ struct CircuitLoweringState {
   CircuitLoweringState(const CircuitLoweringState &) = delete;
   void operator=(const CircuitLoweringState &) = delete;
 
+  /// Mapping of FModuleOp to HWModuleOp
   DenseMap<Operation *, Operation *> oldToNewModuleMap;
 
+  /// Mapping of HWModuleOp to FModuleOp
+  DenseMap<Operation *, Operation *> newToOldModuleMap;
+
   /// Cache of module symbols.  We need to test hirarchy-based properties to
   /// lower annotaitons.
   InstanceGraph &instanceGraph;
 
-  /// The set of modules that are instantiated under the DUT.  This is
-  /// precomputed as a module being under the DUT may rely on knowledge of
+  /// The set of old FIRRTL modules that are instantiated under the DUT.  This
+  /// is precomputed as a module being under the DUT may rely on knowledge of
   /// properties of the instance and is not suitable for querying in the
   /// parallel execution region of this pass when the backing instances may
   /// already be erased.
@@ -599,6 +611,7 @@ void FIRRTLModuleLowering::runOnOperation() {
                 return failure();
 
               state.oldToNewModuleMap[&op] = loweredMod;
+              state.newToOldModuleMap[loweredMod] = &op;
               modulesToProcess.push_back(loweredMod);
               // Lower all the alias types.
               module.walk([&](Operation *op) {
@@ -616,6 +629,7 @@ void FIRRTLModuleLowering::runOnOperation() {
               if (!loweredMod)
                 return failure();
               state.oldToNewModuleMap[&op] = loweredMod;
+              state.newToOldModuleMap[loweredMod] = &op;
               return success();
             })
             .Case<FMemModuleOp>([&](auto memModule) {
@@ -624,6 +638,7 @@ void FIRRTLModuleLowering::runOnOperation() {
               if (!loweredMod)
                 return failure();
               state.oldToNewModuleMap[&op] = loweredMod;
+              state.newToOldModuleMap[loweredMod] = &op;
               return success();
             })
             .Default([&](Operation *op) {

test/Conversion/FIRRTLToHW/lower-to-hw.mlir

@@ -1593,6 +1593,7 @@ firrtl.circuit "Directories" attributes {
   // CHECK:       hw.module private @BoundUnderDUT
   // CHECK-SAME:    output_file = #hw.output_file<"testbench{{/|\\\\}}"
   firrtl.module private @BoundUnderDUT() {}
+  // CHECK:       hw.module private @DUT
   firrtl.module private @DUT() attributes {
     annotations = [
       {
@@ -1601,7 +1602,28 @@ firrtl.circuit "Directories" attributes {
     ]
   } {
     firrtl.instance boundUnderDUT {lowerToBind} @BoundUnderDUT()
+    // Memories in the DUT shouldn't be moved into the testbench.
+    // See: https://github.com/llvm/circt/issues/6775
+    // CHECK:     seq.firmem
+    // CHECK-NOT:   output_file
+    %mem_r = firrtl.mem Undefined {
+      depth = 2 : i64,
+      name = "mem",
+      portNames = ["r"],
+      prefix = "",
+      readLatency = 1 : i32,
+      writeLatency = 1 : i32
+    } : !firrtl.bundle<addr: uint<1>, en: uint<1>, clk: clock, data flip: uint<32>>
+    %c0_clock = firrtl.specialconstant 0 : !firrtl.clock
+    %0 = firrtl.subfield %mem_r[clk] : !firrtl.bundle<addr: uint<1>, en: uint<1>, clk: clock, data flip: uint<32>>
+    firrtl.strictconnect %0, %c0_clock : !firrtl.clock
+    %c0_ui1 = firrtl.constant 0 : !firrtl.uint<1>
+    %1 = firrtl.subfield %mem_r[en] : !firrtl.bundle<addr: uint<1>, en: uint<1>, clk: clock, data flip: uint<32>>
+    firrtl.strictconnect %1, %c0_ui1 : !firrtl.uint<1>
+    %2 = firrtl.subfield %mem_r[addr] : !firrtl.bundle<addr: uint<1>, en: uint<1>, clk: clock, data flip: uint<32>>
+    firrtl.strictconnect %2, %c0_ui1 : !firrtl.uint<1>
   }
+  // CHECK:       hw.module @Directories
   firrtl.module @Directories() {
     firrtl.instance dut @DUT()
     firrtl.instance dut_A @Directories_A()