impl branch interface in wait op, use simplify region, adapt tests

mvpant · mvpant · commit 3ecbcc1bff64 · 2025-08-25T18:28:24.000+03:00
diff --git a/include/circt/Dialect/LLHD/IR/LLHDStructureOps.td b/include/circt/Dialect/LLHD/IR/LLHDStructureOps.td
@@ -144,6 +144,7 @@ def WaitOp : LLHDOp<"wait", [
   AttrSizedOperandSegments,
   HasParent<"ProcessOp">,
   Terminator,
+  DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>
 ]> {
   let summary = "Suspend execution of a process";
   let description = [{
diff --git a/lib/Dialect/LLHD/IR/LLHDOps.cpp b/lib/Dialect/LLHD/IR/LLHDOps.cpp
@@ -546,6 +546,15 @@ static LogicalResult verifyYieldResults(Operation *op,
   return success();
 }
 
+SuccessorOperands WaitOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return SuccessorOperands(getDestOperandsMutable());
+}
+
+Block *WaitOp::getSuccessorForOperands(ArrayRef<Attribute>) {
+  return getDest();
+}
+
 LogicalResult WaitOp::verify() {
   return verifyYieldResults(*this, getYieldOperands());
 }
diff --git a/lib/Dialect/LLHD/Transforms/LowerProcesses.cpp b/lib/Dialect/LLHD/Transforms/LowerProcesses.cpp
@@ -288,116 +288,39 @@ struct LowerProcessesPass
   void runOnOperation() override;
 };
 
-static std::optional<DenseMap<BlockArgument, Value>>
-getBlockArgsToPrune(Block &block) {
-  DenseMap<BlockArgument, Value> result;
-  SmallPtrSet<Value, 4> branchOperands;
-
-  for (auto &arg : block.getArguments()) {
-    for (auto it = block.pred_begin(); it != block.pred_end(); ++it) {
-      Block *predecessor = *it;
-      if (auto branchOp =
-              dyn_cast<BranchOpInterface>(predecessor->getTerminator())) {
-        SuccessorOperands successorOperands =
-            branchOp.getSuccessorOperands(it.getSuccessorIndex());
-
-        if (Value operand = successorOperands[arg.getArgNumber()]) {
-          branchOperands.insert(operand);
-        } else {
-          return {};
-        }
-      } else if (auto waitOp = dyn_cast<WaitOp>(predecessor->getTerminator())) {
-        auto destOperands = waitOp.getDestOperands();
-        branchOperands.insert(destOperands[arg.getArgNumber()]);
-      } else {
-        return {};
-      }
-    }
-
-    if (branchOperands.size() == 1) {
-      result.insert(std::make_pair(arg, *branchOperands.begin()));
-    }
-
-    branchOperands.clear();
-  }
-
-  return result;
-}
-
-static bool simplifyBlocksArgs(ProcessOp &processOp) {
-  SmallVector<Block *> worklist;
-  for (auto &block : processOp.getBody()) {
-    auto waitOp = dyn_cast<WaitOp>(block.getTerminator());
-    if (!waitOp)
-      continue;
-
-    auto dstOperands = waitOp.getDestOperands();
-    if (dstOperands.empty())
-      continue;
-
-    if (llvm::none_of(dstOperands,
-                      [](auto operand) { return isa<BlockArgument>(operand); }))
-      continue;
-
-    worklist.push_back(&block);
-  }
-
-  while (!worklist.empty()) {
-    auto *block = worklist.pop_back_val();
-    auto blockArgsToPrune = getBlockArgsToPrune(*block);
-    if (!blockArgsToPrune.has_value() || blockArgsToPrune->empty())
-      continue;
-
-    llvm::BitVector argsToErase(block->getNumArguments());
-    for (auto &[blockArg, replaceValue] : *blockArgsToPrune) {
-      LLVM_DEBUG(llvm::dbgs()
-                 << "replace" << blockArg << " with " << replaceValue << "\n");
-
-      argsToErase.set(blockArg.getArgNumber());
-    }
-
-    auto argNums = llvm::to_vector(argsToErase.set_bits());
-    for (auto it = block->pred_begin(); it != block->pred_end(); ++it) {
-      Block *predecessor = *it;
-      if (auto branchOp =
-              dyn_cast<BranchOpInterface>(predecessor->getTerminator())) {
-        SuccessorOperands successorOperands =
-            branchOp.getSuccessorOperands(it.getSuccessorIndex());
-
-        const auto argOff = successorOperands.getProducedOperandCount();
-        for (const auto argNum : llvm::reverse(argNums)) {
-          int start = argOff + argNum;
-          successorOperands.erase(start);
-        }
-      } else if (auto waitOp = dyn_cast<WaitOp>(predecessor->getTerminator())) {
-        auto destOperands = waitOp.getDestOperandsMutable();
-        for (const auto argNum : llvm::reverse(argNums)) {
-          destOperands.erase(argNum);
-        }
-      } else {
-        llvm_unreachable("Unexpected predecessor terminator");
-      }
-    }
-
-    for (auto &[blockArg, replaceValue] : *blockArgsToPrune) {
-      blockArg.replaceAllUsesWith(replaceValue);
-    }
-
-    block->eraseArguments(argsToErase);
-
-    for (auto *successor : block->getSuccessors())
-      worklist.push_back(successor);
-  }
-
-  return true;
+// Perform cleanup on the process operation prior to lowering.
+// This can remove redundant operands passed to the successor of llhd.wait.
+//
+// ^bb0
+//  cf.br ^bb1(%clock : i1)
+// ^bb1(%1: i1): // pred: ^bb0, ^bb2
+//  llhd.wait yield (...), (%clock : i1), ^bb2(%1 : i1)
+// ^bb2(%2: i1): // pred: ^bb1
+//  cf.cond_br %true, ^bb3, ^bb1(%2 : i1)
+// ^bb3:
+//
+// The above IR can be rewritten as:
+//
+// ^bb0
+//  cf.br ^bb1
+// ^bb1: // pred: ^bb0, ^bb2
+//  llhd.wait yield (...), (%clock : i1), ^bb2
+// ^bb2: // pred: ^bb1
+//  cf.cond_br %true, ^bb3, ^bb1
+// ^bb3:
+//
+static void simplifyProcess(ProcessOp &processOp) {
+  OpBuilder builder(processOp);
+  IRRewriter rewriter(builder);
+  (void)simplifyRegions(rewriter, processOp->getRegions());
 }
 
 } // namespace
 
 void LowerProcessesPass::runOnOperation() {
   SmallVector<ProcessOp> processOps(getOperation().getOps<ProcessOp>());
   for (auto processOp : processOps) {
-    simplifyBlocksArgs(processOp);
+    simplifyProcess(processOp);
     Lowering(processOp).lower();
   }
 }
diff --git a/test/Dialect/LLHD/Transforms/lower-processes.mlir b/test/Dialect/LLHD/Transforms/lower-processes.mlir
@@ -158,50 +158,57 @@ hw.module @SkipIfNoWaits() {
 }
 
 // CHECK-LABEL: @SkipIfWaitHasDestinationOperands(
-hw.module @SkipIfWaitHasDestinationOperands(in %a: i42) {
+hw.module @SkipIfWaitHasDestinationOperands(in %a: i42, in %b: i42) {
   // CHECK: llhd.process
-  llhd.process {
-    cf.br ^bb1
-  ^bb1:
-    llhd.wait ^bb2(%a : i42)
-  ^bb2(%0: i42):
-    cf.br ^bb1
+  // CHECK: llhd.wait yield ({{.*}} : i42), ^bb2({{.*}} : i42)
+  %0 = llhd.process -> i42 {
+    cf.br ^bb1(%a : i42)
+  ^bb1(%1: i42):
+    %false = hw.constant false
+    %c100 = hw.constant 100 : i42
+    cf.cond_br %false, ^bb2(%1 : i42), ^bb3(%c100 : i42)
+  ^bb2(%2: i42):
+    llhd.wait yield (%2 : i42), ^bb2(%2 : i42)
+  ^bb3(%3: i42):
+    %4 = comb.add %3, %b : i42
+    cf.br ^bb1(%4 : i42)
   }
 }
 
 // CHECK-LABEL: @SkipIfEntryAndWaitConvergeInWrongSpot(
 hw.module @SkipIfEntryAndWaitConvergeInWrongSpot(in %a: i42) {
   // CHECK: llhd.process
+  %c100 = hw.constant 100 : i42
   llhd.process {
-    cf.br ^bb2  // skip logic after wait
+    cf.br ^bb2(%c100 : i42)
   ^bb1:
     %0 = comb.add %a, %a : i42
-    cf.br ^bb2
-  ^bb2:
-    llhd.wait ^bb1
+    cf.br ^bb2(%0 : i42)
+  ^bb2(%1 : i42):
+    llhd.wait (%1 : i42), ^bb1
   }
 }
 
 // CHECK-LABEL: @SkipIfEntryAndWaitConvergeWithDifferentBlockArgs(
 hw.module @SkipIfEntryAndWaitConvergeWithDifferentBlockArgs(in %a: i42, in %b: i42) {
   // CHECK: llhd.process
-  llhd.process {
+  %0 = llhd.process -> i42 {
     cf.br ^bb2(%a : i42)
   ^bb1:
     cf.br ^bb2(%b : i42)
   ^bb2(%0: i42):
-    llhd.wait ^bb1
+    llhd.wait yield (%0 : i42), ^bb1
   }
 }
 
 // CHECK-LABEL: @SkipIfValueUnobserved(
 hw.module @SkipIfValueUnobserved(in %a: i42) {
   // CHECK: llhd.process
-  llhd.process {
+  %0 = llhd.process -> i42 {
     cf.br ^bb1
   ^bb1:
     %0 = comb.add %a, %a : i42
-    llhd.wait ^bb1
+    llhd.wait yield (%0 : i42), ^bb1
   }
 }
 
@@ -246,19 +253,19 @@ hw.module @PruneWaitOperands(in %clock : i1, in %f1 : i2, in %f2 : i3) {
   %c100_i10 = hw.constant 100 : i10
   %c110_i10 = hw.constant 110 : i10
   %true = hw.constant true
-  // CHECK:      llhd.process -> i1 {
+  // CHECK:      llhd.process -> i1, i8, i2, i10, i3 {
   // CHECK-NEXT:   cf.br ^bb1(%c10_i8, %c100_i10 : i8, i10)
   // CHECK-NEXT: ^bb1(%1: i8, %2: i10):
-  // CHECK-NEXT:   llhd.wait yield (%false : i1), (%clock : i1), ^bb2
+  // CHECK-NEXT:   llhd.wait yield (%clock, %1, %f1, %2, %f2 : i1, i8, i2, i10, i3), (%clock : i1), ^bb2
   // CHECK-NEXT: ^bb2:
   // CHECK:        cf.cond_br {{.*}}, ^bb3, ^bb1(%c20_i8, %c110_i10 : i8, i10)
   // CHECK-NEXT: ^bb3:
   // CHECK:        cf.cond_br {{.*}}, ^bb1(%c10_i8, %c100_i10 : i8, i10), ^bb1(%c20_i8, %c110_i10 : i8, i10)
   // CHECK-NEXT: }
-  %0 = llhd.process -> i1 {
+  %0:5 = llhd.process -> i1, i8, i2, i10, i3 {
     cf.br ^bb1(%clock, %c10_i8, %f1, %c100_i10, %f2 : i1, i8, i2, i10, i3)
   ^bb1(%1: i1, %2: i8, %3: i2, %4: i10, %5: i3):
-    llhd.wait yield (%false : i1), (%clock : i1), ^bb2(%2, %1, %3, %5, %4 : i8, i1, i2, i3, i10)
+    llhd.wait yield (%1, %2, %3, %4, %5 : i1, i8, i2, i10, i3), (%clock : i1), ^bb2(%2, %1, %3, %5, %4 : i8, i1, i2, i3, i10)
   ^bb2(%6: i8, %7: i1, %8: i2, %9: i3, %10: i10):
     %15 = comb.xor bin %7, %true : i1
     %16 = comb.and bin %15, %clock : i1
