[FIRRTL] Support regions/blocks in LowerLayers

Fix a bug in the `LowerLayers` pass where it would not visit the blocks in
operations with nested regions.  This would cause the pass to produce
invalid IR if the operations within a layerblock contained any operations
which contained blocks which had captures outside the block.  E.g., this
prevented the pass from working correctly with when statements.

This is part of enabling inline layers.  The existing lowering of inline
layers produces `sv.ifdef` operations.  If an inline layer is inside a
bind layer, then the inline layer (lowered to an `sv.ifdef`) needs to be
walked to find any captures.

Signed-off-by: Schuyler Eldridge <schuyler.eldridge@sifive.com>

Author: seldridge

lib/Dialect/FIRRTL/Transforms/LowerLayers.cpp

@@ -468,7 +468,7 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
 
     SmallVector<hw::InnerSymAttr> innerSyms;
     SmallVector<RWProbeOp> rwprobes;
-    for (auto &op : llvm::make_early_inc_range(*body)) {
+    auto layerBlockWalkResult = layerBlock.walk([&](Operation *op) {
       // Record any operations inside the layer block which have inner symbols.
       // Theses may have symbol users which need to be updated.
       if (auto symOp = dyn_cast<hw::InnerSymbolOpInterface>(op))
@@ -488,7 +488,7 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
       if (auto instOp = dyn_cast<InstanceOp>(op)) {
         // Ignore instances which this pass did not create.
         if (!createdInstances.contains(instOp))
-          continue;
+          return WalkResult::advance();
 
         LLVM_DEBUG({
           llvm::dbgs()
@@ -497,7 +497,7 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
               << "        instance: " << instOp.getName() << "\n";
         });
         instOp->moveBefore(layerBlock);
-        continue;
+        return WalkResult::advance();
       }
 
       // Handle subfields, subindexes, and subaccesses which are indexing into
@@ -510,24 +510,24 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
       // is exceedingly rare given that subaccesses are almost always unexepcted
       // when this pass runs.)
       if (isa<SubfieldOp, SubindexOp>(op)) {
-        auto input = op.getOperand(0);
+        auto input = op->getOperand(0);
         if (!firrtl::type_cast<FIRRTLBaseType>(input.getType()).isPassive() &&
             !isAncestorOfValueOwner(layerBlock, input))
-          op.moveBefore(layerBlock);
-        continue;
+          op->moveBefore(layerBlock);
+        return WalkResult::advance();
       }
 
       if (auto subOp = dyn_cast<SubaccessOp>(op)) {
         auto input = subOp.getInput();
         if (firrtl::type_cast<FIRRTLBaseType>(input.getType()).isPassive())
-          continue;
+          return WalkResult::advance();
 
         if (!isAncestorOfValueOwner(layerBlock, input) &&
             !isAncestorOfValueOwner(layerBlock, subOp.getIndex())) {
           subOp->moveBefore(layerBlock);
-          continue;
+          return WalkResult::advance();
         }
-        auto diag = op.emitOpError()
+        auto diag = op->emitOpError()
                     << "has a non-passive operand and captures a value defined "
                        "outside its enclosing bind-convention layerblock.  The "
                        "'LowerLayers' pass cannot lower this as it would "
@@ -539,7 +539,7 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
 
       if (auto rwprobe = dyn_cast<RWProbeOp>(op)) {
         rwprobes.push_back(rwprobe);
-        continue;
+        return WalkResult::advance();
       }
 
       if (auto connect = dyn_cast<FConnectLike>(op)) {
@@ -549,22 +549,22 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
         auto dstInLayerBlock = isAncestorOfValueOwner(layerBlock, dst);
         if (!srcInLayerBlock && !dstInLayerBlock) {
           connect->moveBefore(layerBlock);
-          continue;
+          return WalkResult::advance();
         }
         // Create an input port.
         if (!srcInLayerBlock) {
-          createInputPort(src, op.getLoc());
-          continue;
+          createInputPort(src, op->getLoc());
+          return WalkResult::advance();
         }
         // Create an output port.
         if (!dstInLayerBlock) {
           createOutputPort(dst, src);
           if (!isa<RefType>(dst.getType()))
             connect.erase();
-          continue;
+          return WalkResult::advance();
         }
         // Source and destination in layer block.  Nothing to do.
-        continue;
+        return WalkResult::advance();
       }
 
       // Pre-emptively de-squiggle connections that we are creating.  This will
@@ -587,15 +587,20 @@ LogicalResult LowerLayersPass::runOnModuleBody(FModuleOp moduleOp,
                   *refResolve.getResult().getUsers().begin()))
             if (connect.getDest().getParentBlock() != body) {
               refResolve->moveBefore(layerBlock);
-              continue;
+              return WalkResult::advance();
             }
 
       // For any other ops, create input ports for any captured operands.
-      for (auto operand : op.getOperands()) {
+      for (auto operand : op->getOperands()) {
         if (!isAncestorOfValueOwner(layerBlock, operand))
-          createInputPort(operand, op.getLoc());
+          createInputPort(operand, op->getLoc());
       }
-    }
+
+      return WalkResult::advance();
+    });
+
+    if (layerBlockWalkResult.wasInterrupted())
+      return WalkResult::interrupt();
 
     // Create the new module.  This grabs a lock to modify the circuit.
     FModuleOp newModule = buildNewModule(builder, layerBlock, ports);

test/Dialect/FIRRTL/lower-layers.mlir

@@ -224,6 +224,24 @@ firrtl.circuit "Test" {
     }
   }
 
+  // CHECK:      firrtl.module private @CaptureInWhen_A(
+  // CHECK-SAME:   in %a: !firrtl.uint<1>
+  // CHECK-SAME:   in %cond: !firrtl.uint<1>
+  // CHECK-SAME: )
+
+  // CHECK:      firrtl.module @CaptureInWhen(
+  // CHECK:        %a_a, %a_cond = firrtl.instance a
+  // CHECK-NEXT:   firrtl.matchingconnect %a_cond, %cond :
+  // CHECK-NEXT:   firrtl.matchingconnect %a_a, %a :
+  firrtl.module @CaptureInWhen(in %cond: !firrtl.uint<1>) {
+    %a = firrtl.wire : !firrtl.uint<1>
+    firrtl.layerblock @A {
+      firrtl.when %cond : !firrtl.uint<1> {
+        %b = firrtl.node %a : !firrtl.uint<1>
+      }
+    }
+  }
+
   //===--------------------------------------------------------------------===//
   // Connecting/Defining Refs
   //===--------------------------------------------------------------------===//