From 9b26fd4b19fd75ae94f63a62c8dcf7f36f6784e9 Mon Sep 17 00:00:00 2001
From: Fabian Schuiki <fabian@schuiki.ch>
Date: Wed, 26 Feb 2025 12:14:30 -0800
Subject: [PATCH] [LLHD] Support conditional drives in mem2reg

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   | 20 +++++----
 test/Dialect/LLHD/Transforms/mem2reg.mlir | 53 +++++++++++++++++++++++
 2 files changed, 64 insertions(+), 9 deletions(-)

diff --git a/lib/Dialect/LLHD/Transforms/Mem2Reg.cpp b/lib/Dialect/LLHD/Transforms/Mem2Reg.cpp
index 98a0bebe679f..668432b14b07 100644
--- a/lib/Dialect/LLHD/Transforms/Mem2Reg.cpp
+++ b/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;
diff --git a/test/Dialect/LLHD/Transforms/mem2reg.mlir b/test/Dialect/LLHD/Transforms/mem2reg.mlir
index c6b24b636b45..615539c8edd3 100644
--- a/test/Dialect/LLHD/Transforms/mem2reg.mlir
+++ b/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>)