[FIRRTL][LowerToHW] Lower contract ops

Lower `firrtl.contract` to `verif.contract` ops. The lowering itself is
pretty straightforward, since the ops are basically equivalent. The only
exception are block arguments in FIRRTL that get replaced with op
results in Verif due to the switch from an SSACFG region with dominance
to a graph region.

The block arguments require a slight change to how nested operations are
handled by the pass. Currently, a post-order walk over the operations is
used. This is problematic since the parent operation does not get an
opportunity to lower its block arguments before its child operations are
lowered, which may need access to those block arguments. Switching to a
pre-order walk does not work, since that wouldn't allow us to modify
the operations during the walk, which we obviously do during the
lowering. This commit therefore adds a worklist that tracks operation
ranges that are yet to be lowered. The worklist allows parent operations
to lower themselves and any of their block arguments, and then add the
nested regions, blocks, or operations onto the worklist for lowering.
It's basically a pre-walk with the ability to mutate the parent
operation before children are lowered.

Block arguments are also no longer automatically assumed to be already
lowered. This was only valid for modules, since we lower the module
ports before we lower the module body. In case of ops like contracts,
the block arguments still need to be lowered. To fix this, modules push
a 1-1 mapping of their block arguments into the map of lowered values,
which allows other ops to provide some other lowering for their block
arguments.

lib/Conversion/FIRRTLToHW/LowerToHW.cpp

@@ -1580,6 +1580,7 @@ struct FIRRTLLowering : public FIRRTLVisitor<FIRRTLLowering, LogicalResult> {
   LogicalResult visitDecl(MemOp op);
   LogicalResult visitDecl(InstanceOp oldInstance);
   LogicalResult visitDecl(VerbatimWireOp op);
+  LogicalResult visitDecl(ContractOp op);
 
   // Unary Ops.
   LogicalResult lowerNoopCast(Operation *op);
@@ -1837,6 +1838,19 @@ struct FIRRTLLowering : public FIRRTLVisitor<FIRRTLLowering, LogicalResult> {
   /// the LTL ops, which were necessary to go from the def-before-use FIRRTL
   /// dialect to the graph-like HW dialect.
   SetVector<Operation *> ltlOpFixupWorklist;
+
+  /// A worklist of operation ranges to be lowered. Parnet operations can push
+  /// their nested operations onto this worklist to be processed after the
+  /// parent operation has handled the region, blocks, and block arguments.
+  SmallVector<std::pair<Block::iterator, Block::iterator>> worklist;
+
+  void addToWorklist(Block &block) {
+    worklist.push_back({block.begin(), block.end()});
+  }
+  void addToWorklist(Region &region) {
+    for (auto &block : llvm::reverse(region))
+      addToWorklist(block);
+  }
 };
 } // end anonymous namespace
 
@@ -1847,16 +1861,27 @@ LogicalResult FIRRTLModuleLowering::lowerModuleOperations(
 
 // This is the main entrypoint for the lowering pass.
 LogicalResult FIRRTLLowering::run() {
-  // FIRRTL FModule is a single block because FIRRTL ops are a DAG.  Walk
-  // through each operation, lowering each in turn if we can, introducing
-  // casts if we cannot.
-  auto &body = theModule.getBody();
+  // Mark the module's block arguments are already lowered. This will allow
+  // `getLoweredValue` to return the block arguments as they are.
+  for (auto arg : theModule.getBody().front().getArguments())
+    if (failed(setLowering(arg, arg)))
+      return failure();
 
+  // Add the operations in the body to the worklist and lower all operations
+  // until the worklist is empty. Operations may push their own nested
+  // operations onto the worklist to lower them in turn. The `builder` is
+  // positioned ahead of each operation as it is being lowered.
+  addToWorklist(theModule.getBody());
   SmallVector<Operation *, 16> opsToRemove;
 
-  // Iterate through each operation in the module body, attempting to lower
-  // each of them.  We maintain 'builder' for each invocation.
-  auto result = theModule.walk([&](Operation *op) {
+  while (!worklist.empty()) {
+    auto &[opsIt, opsEnd] = worklist.back();
+    if (opsIt == opsEnd) {
+      worklist.pop_back();
+      continue;
+    }
+    Operation *op = &*opsIt++;
+
     builder.setInsertionPoint(op);
     builder.setLoc(op->getLoc());
     auto done = succeeded(dispatchVisitor(op));
@@ -1872,14 +1897,10 @@ LogicalResult FIRRTLLowering::run() {
         break;
       case LoweringFailure:
         backedgeBuilder.abandon();
-        return WalkResult::interrupt();
+        return failure();
       }
     }
-    return WalkResult::advance();
-  });
-
-  if (result.wasInterrupted())
-    return failure();
+  }
 
   // Replace all backedges with uses of their regular values.  We process them
   // after the module body since the lowering table is too hard to keep up to
@@ -1933,7 +1954,7 @@ LogicalResult FIRRTLLowering::run() {
       if (!isZeroBitFIRRTLType(result.getType()))
         continue;
       if (!zeroI0) {
-        auto builder = OpBuilder::atBlockBegin(&body.front());
+        auto builder = OpBuilder::atBlockBegin(theModule.getBodyBlock());
         zeroI0 = builder.create<hw::ConstantOp>(op->getLoc(),
                                                 builder.getIntegerType(0), 0);
         maybeUnusedValues.insert(zeroI0);
@@ -2082,10 +2103,6 @@ Value FIRRTLLowering::getOrCreateZConstant(Type type) {
 /// unknown width integers.  This returns hw::inout type values if present, it
 /// does not implicitly read from them.
 Value FIRRTLLowering::getPossiblyInoutLoweredValue(Value value) {
-  // Block arguments are considered lowered.
-  if (isa<BlockArgument>(value))
-    return value;
-
   // If we lowered this value, then return the lowered value, otherwise fail.
   if (auto lowering = valueMapping.lookup(value)) {
     assert(!isa<FIRRTLType>(lowering.getType()) &&
@@ -2724,6 +2741,10 @@ void FIRRTLLowering::addIfProceduralBlock(Value cond,
 ///
 FIRRTLLowering::UnloweredOpResult
 FIRRTLLowering::handleUnloweredOp(Operation *op) {
+  // Push nested operations onto the worklist such that they are lowered.
+  for (auto &region : op->getRegions())
+    addToWorklist(region);
+
   // Simply pass through non-FIRRTL operations and consider them already
   // lowered. This allows us to handled partially lowered inputs, and also allow
   // other FIRRTL operations to spawn additional already-lowered operations,
@@ -3397,6 +3418,37 @@ LogicalResult FIRRTLLowering::visitDecl(InstanceOp oldInstance) {
   return success();
 }
 
+LogicalResult FIRRTLLowering::visitDecl(ContractOp oldOp) {
+  SmallVector<Value> inputs;
+  SmallVector<Type> types;
+  for (auto input : oldOp.getInputs()) {
+    auto lowered = getLoweredValue(input);
+    if (!lowered)
+      return failure();
+    inputs.push_back(lowered);
+    types.push_back(lowered.getType());
+  }
+
+  auto newOp = builder.create<verif::ContractOp>(types, inputs);
+  newOp->setDiscardableAttrs(oldOp->getDiscardableAttrDictionary());
+  auto &body = newOp.getBody().emplaceBlock();
+
+  for (auto [newResult, oldResult, oldArg] :
+       llvm::zip(newOp.getResults(), oldOp.getResults(),
+                 oldOp.getBody().getArguments())) {
+    if (failed(setLowering(oldResult, newResult)))
+      return failure();
+    if (failed(setLowering(oldArg, newResult)))
+      return failure();
+  }
+
+  body.getOperations().splice(body.end(),
+                              oldOp.getBody().front().getOperations());
+  addToWorklist(body);
+
+  return success();
+}
+
 //===----------------------------------------------------------------------===//
 // Unary Operations
 //===----------------------------------------------------------------------===//

test/Conversion/FIRRTLToHW/intrinsics.mlir

@@ -124,7 +124,7 @@ firrtl.circuit "Intrinsics" {
     // CHECK-NEXT: verif.assert %a : i1
     firrtl.int.verif.ensure %a : !firrtl.uint<1>
     // CHECK-NEXT: verif.contract
-    verif.contract {
+    firrtl.contract {
       // CHECK-NEXT: verif.require %a : i1
       firrtl.int.verif.require %a : !firrtl.uint<1>
       // CHECK-NEXT: verif.require %a label "hello" : i1

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

@@ -1731,3 +1731,30 @@ firrtl.circuit "Foo" {
   // CHECK-NEXT: hw.instance "Foo" @Foo(a: [[A]]: i42) -> (z: i42)
   firrtl.formal @MyTest2, @Foo {world = "abc"}
 }
+
+// -----
+
+firrtl.circuit "Foo" {
+  // CHECK-LABEL: hw.module @Foo
+  firrtl.module @Foo(in %a: !firrtl.uint<42>, in %b: !firrtl.uint<1337>) {
+    // CHECK: verif.contract {
+    // CHECK-NEXT: }
+    firrtl.contract {
+    }
+
+    // CHECK: [[TMP:%.+]]:2 = verif.contract %a, %b : i42, i1337 {
+    // CHECK-NEXT:   dbg.variable "c2", [[TMP]]#0 : i42
+    // CHECK-NEXT:   dbg.variable "d2", [[TMP]]#1 : i1337
+    // CHECK-NEXT: }
+    %c, %d = firrtl.contract %a, %b : !firrtl.uint<42>, !firrtl.uint<1337> {
+    ^bb0(%c2: !firrtl.uint<42>, %d2: !firrtl.uint<1337>):
+      dbg.variable "c2", %c2 : !firrtl.uint<42>
+      dbg.variable "d2", %d2 : !firrtl.uint<1337>
+    }
+
+    // CHECK: dbg.variable "c", [[TMP]]#0 : i42
+    // CHECK: dbg.variable "d", [[TMP]]#1 : i1337
+    dbg.variable "c", %c : !firrtl.uint<42>
+    dbg.variable "d", %d : !firrtl.uint<1337>
+  }
+}