[LLHD] Support conditional drives in mem2reg (#8278)

Extend the LLHD Mem2Reg pass to support conditional drives. This is
pretty trivial since the pass already tracks drive conditions generated
through control flow. All that's needed is to initialize the condition
of a drive's reaching definition that gets propagated across the
lattice.

lib/Dialect/LLHD/Transforms/Mem2Reg.cpp

@@ -40,11 +40,11 @@ static bool isEpsilonDelay(Value value) {
   return false;
 }
 
-/// Check whether an operation is a `llhd.drive` with no enable condition and an
-/// epsilon delay. This corresponds to a blocking assignment in Verilog.
-static bool isUnconditionalBlockingDrive(Operation *op) {
+/// Check whether an operation is a `llhd.drive` with an epsilon delay. This
+/// corresponds to a blocking assignment in Verilog.
+static bool isBlockingDrive(Operation *op) {
   if (auto driveOp = dyn_cast<DrvOp>(op))
-    return !driveOp.getEnable() && isEpsilonDelay(driveOp.getTime());
+    return isEpsilonDelay(driveOp.getTime());
   return false;
 }
 
@@ -278,7 +278,7 @@ struct DriveNode : public OpNode {
             LatticeValue *valueAfter)
       : OpNode(Kind::Drive, op, valueBefore, valueAfter), slot(op.getSignal()),
         value(op.getValue()), def(def) {
-    assert(isUnconditionalBlockingDrive(op));
+    assert(isBlockingDrive(op));
   }
 
   static bool classof(const LatticeNode *n) { return n->kind == Kind::Drive; }
@@ -608,7 +608,7 @@ void Promoter::findPromotableSlots() {
             // Ignore uses outside of the region.
             if (user->getParentRegion() != &region)
               return true;
-            return isa<PrbOp>(user) || isUnconditionalBlockingDrive(user);
+            return isa<PrbOp>(user) || isBlockingDrive(user);
           }))
         continue;
 
@@ -765,12 +765,14 @@ void Promoter::constructLattice() {
 
       // Handle drives.
       if (auto driveOp = dyn_cast<DrvOp>(op)) {
-        if (!isUnconditionalBlockingDrive(&op))
+        if (!isBlockingDrive(&op))
           continue;
         if (!slotOrder.contains(driveOp.getSignal()))
           continue;
-        auto *def =
-            lattice.createDef(driveOp.getValue(), DriveCondition::always());
+        auto condition = DriveCondition::always();
+        if (auto enable = driveOp.getEnable())
+          condition = DriveCondition::conditional(enable);
+        auto *def = lattice.createDef(driveOp.getValue(), condition);
         auto *node = lattice.createNode<DriveNode>(driveOp, def, valueBefore,
                                                    lattice.createValue());
         valueBefore = node->valueAfter;

test/Dialect/LLHD/Transforms/mem2reg.mlir

@@ -440,5 +440,58 @@ hw.module @CaptureAcrossWaits(in %u: i42, in %bool: i1) {
   }
 }
 
+// Conditional drive forwarding.
+// CHECK-LABEL: @ConditionalDrives
+hw.module @ConditionalDrives(in %u: i42, in %v: i42, in %q: i1, in %r: i1) {
+  %0 = llhd.constant_time <0ns, 0d, 1e>
+  %a = llhd.sig %u : i42
+  // CHECK: llhd.process
+  llhd.process {
+    // CHECK-NEXT: llhd.constant_time
+    // CHECK-NEXT: llhd.drv %a, %u after {{%.+}} if %q
+    llhd.drv %a, %u after %0 if %q : !hw.inout<i42>
+    // CHECK-NEXT: llhd.halt
+    llhd.halt
+  }
+  // CHECK: llhd.process
+  llhd.process {
+    // CHECK-NOT: llhd.drv
+    llhd.drv %a, %u after %0 if %q : !hw.inout<i42>
+    // CHECK-NEXT: cf.br ^bb2(%u : i42)
+    cf.br ^bb2
+  ^bb1:
+    // CHECK-NEXT: ^bb1:
+    // CHECK-NOT: llhd.drv
+    llhd.drv %a, %v after %0 if %q : !hw.inout<i42>
+    // CHECK-NEXT: cf.br ^bb2(%v : i42)
+    cf.br ^bb2
+  ^bb2:
+    // CHECK-NEXT: ^bb2([[A:%.+]]: i42):
+    // CHECK-NEXT: llhd.constant_time
+    // CHECK-NEXT: llhd.drv %a, [[A]] after {{%.+}} if %q
+    // CHECK-NEXT: llhd.halt
+    llhd.halt
+  }
+  // CHECK: llhd.process
+  llhd.process {
+    // CHECK-NOT: llhd.drv
+    llhd.drv %a, %u after %0 if %q : !hw.inout<i42>
+    // CHECK-NEXT: cf.br ^bb2(%u, %q : i42, i1)
+    cf.br ^bb2
+  ^bb1:
+    // CHECK-NEXT: ^bb1:
+    // CHECK-NOT: llhd.drv
+    llhd.drv %a, %v after %0 if %r : !hw.inout<i42>
+    // CHECK-NEXT: cf.br ^bb2(%v, %r : i42, i1)
+    cf.br ^bb2
+  ^bb2:
+    // CHECK-NEXT: ^bb2([[A:%.+]]: i42, [[ACOND:%.+]]: i1):
+    // CHECK-NEXT: llhd.constant_time
+    // CHECK-NEXT: llhd.drv %a, [[A]] after {{%.+}} if [[ACOND]]
+    // CHECK-NEXT: llhd.halt
+    llhd.halt
+  }
+}
+
 func.func private @use_i42(%arg0: i42)
 func.func private @use_inout_i42(%arg0: !hw.inout<i42>)