[LLHD] Simplify single-block llhd.process terminated by llhd.halt

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

@@ -205,6 +205,7 @@ def HaltOp : LLHDOp<"halt", [
     attr-dict
   }];
   let hasVerifier = 1;
+  let hasCanonicalizeMethod = 1;
 }
 
 //===----------------------------------------------------------------------===//

lib/Dialect/LLHD/IR/LLHDOps.cpp

@@ -487,11 +487,72 @@ DrvOp::ensureOnlySafeAccesses(const MemorySlot &slot,
 //===----------------------------------------------------------------------===//
 
 LogicalResult ProcessOp::canonicalize(ProcessOp op, PatternRewriter &rewriter) {
-  if (op.getBody().hasOneBlock() && op.getNumResults() == 0) {
-    auto &block = op.getBody().front();
-    if (block.getOperations().size() == 1 && isa<HaltOp>(block.getTerminator()))
-      rewriter.eraseOp(op);
+  if (!op.getBody().hasOneBlock())
+    return failure();
+
+  auto &block = op.getBody().front();
+  if (op.getNumResults() == 0) {
+    rewriter.eraseOp(op);
+    return success();
+  }
+
+  auto haltOp = dyn_cast<HaltOp>(block.getTerminator());
+  if (!haltOp)
+    return failure();
+
+  // Only constants and halt terminator are expected in a single block.
+  if (!llvm::all_of(block.without_terminator(), [](auto &bodyOp) {
+        return bodyOp.template hasTrait<OpTrait::ConstantLike>();
+      }))
+    return failure();
+
+  llvm::BitVector resultsToErase(op.getNumResults());
+  for (auto [i, operand] : llvm::enumerate(haltOp.getYieldOperands())) {
+    auto *defOp = operand.getDefiningOp();
+    if (defOp && defOp->hasTrait<OpTrait::ConstantLike>()) {
+      // If the constant is available outside the process, use it directly;
+      // otherwise move it outside.
+      if (!defOp->getParentRegion()->isProperAncestor(&op.getBody())) {
+        defOp->moveBefore(op);
+      }
+      op.getResult(i).replaceAllUsesWith(operand);
+      resultsToErase.set(i);
+    }
+  }
+
+  const auto countResultsToErase = resultsToErase.count();
+  if (countResultsToErase == 0)
+    return failure();
+
+  // Just remove whole process operation if we can replace all results.
+  if (countResultsToErase == op.getNumResults()) {
+    rewriter.eraseOp(op);
+    return success();
+  }
+
+  // Otherwise, prune only those that are safe to remove.
+  haltOp->eraseOperands(resultsToErase);
+
+  SmallVector<Type> resultTypes;
+  resultsToErase.flip();
+  for (auto i : resultsToErase.set_bits()) {
+    resultTypes.push_back(op.getResult(i).getType());
+  }
+
+  auto newProcessOp = ProcessOp::create(rewriter, op.getLoc(), resultTypes,
+                                        op->getOperands(), op->getAttrs());
+  newProcessOp.getBody().takeBody(op.getBody());
+
+  // Update old results with new values, accounting for pruned halt operands.
+  unsigned off = 0;
+  for (auto oldResult : op.getResults()) {
+    if (resultsToErase.test(oldResult.getResultNumber())) {
+      oldResult.replaceAllUsesWith(newProcessOp.getResult(off));
+      ++off;
+    }
   }
+
+  rewriter.eraseOp(op);
   return success();
 }
 
@@ -554,6 +615,55 @@ LogicalResult WaitOp::verify() {
 // HaltOp
 //===----------------------------------------------------------------------===//
 
+static LogicalResult deduplicateYieldOperands(PatternRewriter &rewriter,
+                                              HaltOp &op) {
+  auto processOp = dyn_cast<ProcessOp>(op->getParentOp());
+  if (!processOp || !processOp.getBody().hasOneBlock())
+    return failure();
+
+  auto yieldOperands = op.getYieldOperands();
+  llvm::SmallDenseMap<Value, unsigned> uniqueOperands;
+  llvm::SmallDenseMap<unsigned, unsigned> origToNewPos;
+  llvm::BitVector operandsToErase(yieldOperands.size());
+
+  for (auto [operandNo, operand] : llvm::enumerate(yieldOperands)) {
+    if (!uniqueOperands.contains(operand)) {
+      const auto newPos = uniqueOperands.size();
+      uniqueOperands.insert(std::make_pair(operand, newPos));
+      origToNewPos.insert(std::make_pair(operandNo, newPos));
+    } else {
+      auto firstOperandPos = uniqueOperands.lookup(operand);
+      origToNewPos.insert(std::make_pair(operandNo, firstOperandPos));
+      operandsToErase.set(operandNo);
+    }
+  }
+
+  if (uniqueOperands.size() == yieldOperands.size())
+    return failure();
+
+  op->eraseOperands(operandsToErase);
+  SmallVector<Type> resultTypes = llvm::to_vector(op->getOperandTypes());
+
+  rewriter.setInsertionPoint(processOp);
+  auto newProcessOp =
+      ProcessOp::create(rewriter, processOp->getLoc(), resultTypes,
+                        processOp->getOperands(), processOp->getAttrs());
+  newProcessOp.getBody().takeBody(processOp.getBody());
+
+  for (auto oldResult : processOp.getResults()) {
+    auto newResultPos = origToNewPos.lookup(oldResult.getResultNumber());
+    auto newResult = newProcessOp.getResult(newResultPos);
+    rewriter.replaceAllUsesWith(oldResult, newResult);
+  }
+
+  rewriter.eraseOp(processOp);
+  return success();
+}
+
+LogicalResult HaltOp::canonicalize(HaltOp op, PatternRewriter &rewriter) {
+  return deduplicateYieldOperands(rewriter, op);
+}
+
 LogicalResult HaltOp::verify() {
   return verifyYieldResults(*this, getYieldOperands());
 }

test/Dialect/LLHD/Canonicalization/processes.mlir

@@ -35,3 +35,60 @@ hw.module @InlineCombinational(in %a: i42, in %b: i42, in %c: i8917, out u: i42,
   }
   hw.output %0, %1 : i42, i9001
 }
+
+// CHECK-LABEL: hw.module @IgnoreMultiBlockHalt
+hw.module @IgnoreMultiBlockHalt(in %a : i1, in %b : i1, out v1 : i1, out v2 : i1) {
+  // CHECK: llhd.halt %a, %a
+  %0:2 = llhd.process -> i1, i1 {
+  ^bb0:
+    cf.br ^bb1
+  ^bb1:
+    cf.cond_br %b, ^bb1, ^bb2
+  ^bb2:
+    %true = hw.constant true
+    llhd.halt %a, %a : i1, i1
+  }
+  hw.output %0#0, %0#1 : i1, i1
+}
+
+// CHECK-LABEL: hw.module @DeduplicateHaltOperands0
+hw.module @DeduplicateHaltOperands0(in %a : i1, in %b : i1,
+                                    out v1 : i1, out v2 : i1, out v3 : i1, out v4 : i1) {
+  // CHECK:      %0:2 = llhd.process -> i1, i1 {
+  // CHECK-NEXT:   llhd.halt %a, %b : i1, i1
+  // CHECK-NEXT: }
+  // CHECK-NEXT: hw.output %0#0, %0#1, %0#0, %0#1
+  %false = hw.constant false
+  %0:6 = llhd.process -> i1, i1, i1, i1, i1, i1 {
+    %true = hw.constant true
+    llhd.halt %false, %a, %b, %a, %true, %b : i1, i1, i1, i1, i1, i1
+  }
+  hw.output %0#1, %0#2, %0#3, %0#5 : i1, i1, i1, i1
+}
+
+// CHECK-LABEL: hw.module @CanonProcessHalt0
+hw.module @CanonProcessHalt0(out v1 : i1, out v2 : i1) {
+  // CHECK-NOT: llhd.halt
+  // CHECK: hw.output %false, %true
+  %false = hw.constant false
+  %0:2 = llhd.process -> i1, i1 {
+    %true = hw.constant true
+    llhd.halt %false, %true : i1, i1
+  }
+  hw.output %0#0, %0#1 : i1, i1
+}
+
+// CHECK-LABEL: hw.module @CanonProcessHalt1
+hw.module @CanonProcessHalt1(in %a : i1, in %b : i1,
+                             out v1 : i1, out v2 : i1, out v3 : i1, out v4 : i1) {
+  // CHECK:      %0:2 = llhd.process -> i1, i1 {
+  // CHECK-NEXT:   llhd.halt %a, %b : i1, i1
+  // CHECK-NEXT: }
+  // CHECK-NEXT: hw.output %0#1, %false, %0#0, %true
+  %0:4 = llhd.process -> i1, i1, i1, i1 {
+    %false = hw.constant false
+    %true = hw.constant true
+    llhd.halt %false, %a, %true, %b : i1, i1, i1, i1
+  }
+  hw.output %0#3, %0#0, %0#1, %0#2 : i1, i1, i1, i1
+}