[LLHD] Remove redundant destination operands from wait operation

include/circt/Dialect/LLHD/IR/LLHDStructureOps.td

@@ -143,7 +143,7 @@ def CombinationalOp : LLHDOp<"combinational", [
 def WaitOp : LLHDOp<"wait", [
   AttrSizedOperandSegments,
   HasParent<"ProcessOp">,
-  Terminator,
+  Terminator
 ]> {
   let summary = "Suspend execution of a process";
   let description = [{

lib/Dialect/LLHD/Transforms/Deseq.cpp

@@ -505,6 +505,11 @@ TruthTable Deseq::computeBoolean(OpResult value) {
     return result;
   }
 
+  if (auto instanceOp = dyn_cast<hw::InstanceOp>(op)) {
+    // TODO: Should we delve into hw.module itself?
+    return getUnknownBoolean();
+  }
+
   // Otherwise check if the operation depends on any of the triggers. If it
   // does, create a poison value since we don't really know how the trigger
   // affects this boolean. If it doesn't, create an unknown value.

lib/Dialect/LLHD/Transforms/LowerProcesses.cpp

@@ -287,10 +287,121 @@ struct LowerProcessesPass
     : public llhd::impl::LowerProcessesPassBase<LowerProcessesPass> {
   void runOnOperation() override;
 };
+
+static LogicalResult dropRedundantArguments(RewriterBase &rewriter,
+                                            Block &block) {
+  SmallVector<size_t> argsToErase;
+
+  // Go through the arguments of the block.
+  for (auto [argIdx, blockOperand] : llvm::enumerate(block.getArguments())) {
+    bool sameArg = true;
+    Value commonValue;
+
+    // Go through the block predecessor and flag if they pass to the block
+    // different values for the same argument.
+    for (Block::pred_iterator predIt = block.pred_begin(),
+                              predE = block.pred_end();
+         predIt != predE; ++predIt) {
+      Operation *terminator = (*predIt)->getTerminator();
+      auto branchOperand =
+          llvm::TypeSwitch<Operation *, std::optional<Value>>(terminator)
+              .Case([&, argIdx = argIdx](BranchOpInterface branchOp) {
+                auto succIndex = predIt.getSuccessorIndex();
+                SuccessorOperands succOperands =
+                    branchOp.getSuccessorOperands(succIndex);
+                auto branchOperands = succOperands.getForwardedOperands();
+                return branchOperands[argIdx];
+              })
+              .Case([&, argIdx = argIdx](WaitOp waitOp) {
+                auto destOperands = waitOp.getDestOperands();
+                return destOperands[argIdx];
+              })
+              .Default([](Operation *) { return std::nullopt; });
+
+      if (!branchOperand.has_value()) {
+        sameArg = false;
+        break;
+      }
+
+      if (!commonValue) {
+        commonValue = *branchOperand;
+        continue;
+      }
+
+      if (*branchOperand != commonValue) {
+        sameArg = false;
+        break;
+      }
+    }
+
+    // If they are passing the same value, drop the argument.
+    if (commonValue && sameArg) {
+      argsToErase.push_back(argIdx);
+
+      // Remove the argument from the block.
+      rewriter.replaceAllUsesWith(blockOperand, commonValue);
+    }
+  }
+
+  // Remove the arguments.
+  for (size_t argIdx : llvm::reverse(argsToErase)) {
+    block.eraseArgument(argIdx);
+
+    // Remove the argument from the branch ops.
+    for (auto predIt = block.pred_begin(), predE = block.pred_end();
+         predIt != predE; ++predIt) {
+      llvm::TypeSwitch<Operation *, void>((*predIt)->getTerminator())
+          .Case([&](BranchOpInterface branchOp) {
+            auto branch = cast<BranchOpInterface>((*predIt)->getTerminator());
+            unsigned succIndex = predIt.getSuccessorIndex();
+            SuccessorOperands succOperands =
+                branch.getSuccessorOperands(succIndex);
+            succOperands.erase(argIdx);
+          })
+          .Case([&](WaitOp waitOp) {
+            auto destOperands = waitOp.getDestOperandsMutable();
+            destOperands.erase(argIdx);
+          })
+          .Default([](Operation *) {
+            llvm_unreachable("Unexpected predecessor terminator");
+          });
+    }
+  }
+
+  return success(!argsToErase.empty());
+}
+
+static void simplifyProcess(ProcessOp &processOp) {
+  OpBuilder builder(processOp);
+  IRRewriter rewriter(builder);
+  (void)simplifyRegions(rewriter, processOp->getRegions());
+
+  // simplifyRegions does not prune the destination operands of the `llhd.wait`
+  // operation because it does not implement BranchOpInterface, due to its
+  // side-effect semantics. Implementing BranchOpInterface could allow other
+  // optimizations to forward branch operands from the `llhd.wait` operation to
+  // the destination block where execution resumes. However, this would be
+  // invalid, as the SSA value might change across the wait operation.
+  // Therefore, we manually prune the destination block arguments ourselves.
+  for (auto &block : processOp.getBody()) {
+    auto waitOp = dyn_cast<WaitOp>(block.getTerminator());
+    if (!waitOp)
+      continue;
+
+    auto dstOperands = waitOp.getDestOperands();
+    if (dstOperands.empty())
+      continue;
+
+    (void)dropRedundantArguments(rewriter, *waitOp.getDest());
+  }
+}
+
 } // namespace
 
 void LowerProcessesPass::runOnOperation() {
   SmallVector<ProcessOp> processOps(getOperation().getOps<ProcessOp>());
-  for (auto processOp : processOps)
+  for (auto processOp : processOps) {
+    // simplifyProcess(processOp);
     Lowering(processOp).lower();
+  }
 }