[Mlir-commits] [mlir] [mlir][emitc] Fix crash in form-expressions when identity cast is folded (PR #183894)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/183894

When the --form-expressions pass runs applyPatternsGreedily, the greedy rewriter calls CastOpInterface::foldTrait on ops inside ExpressionOp bodies. For an identity cast (e.g. `emitc.cast i32 to i32`), this fold succeeds and replaces the op's result with its operand (a block argument), leaving the ExpressionOp body in the form `{ yield %arg }`.

Subsequently, if the cast expression's result had its use count reduced to one (e.g. because a dead non-EmitC use was eliminated), FoldExpressionOp would select this expression as a candidate to fold into an outer expression. It then calls usedExpression.getRootOp(), which returns nullptr because the body yields a block argument rather than an op result. The subsequent mapper.lookup(nullptr) crashes with an assertion.

Fix by adding a FoldTrivialExpressionOp canonicalization pattern that handles ExpressionOps whose body yields a block argument directly: such an expression is a passthrough and can be replaced by the corresponding operand value. This canonicalization fires during greedy rewriting before FoldExpressionOp encounters the invalid state.

Fixes #179844

>From 6056c149d44627662ccdbcdbd04a77a859e047b1 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Fri, 27 Feb 2026 13:04:54 -0800
Subject: [PATCH] [mlir][emitc] Fix crash in form-expressions when identity
 cast is folded

When the --form-expressions pass runs applyPatternsGreedily, the greedy
rewriter calls CastOpInterface::foldTrait on ops inside ExpressionOp
bodies. For an identity cast (e.g. `emitc.cast i32 to i32`), this fold
succeeds and replaces the op's result with its operand (a block
argument), leaving the ExpressionOp body in the form `{ yield %arg }`.

Subsequently, if the cast expression's result had its use count reduced
to one (e.g. because a dead non-EmitC use was eliminated), FoldExpressionOp
would select this expression as a candidate to fold into an outer
expression. It then calls usedExpression.getRootOp(), which returns
nullptr because the body yields a block argument rather than an op
result. The subsequent mapper.lookup(nullptr) crashes with an assertion.

Fix by adding a FoldTrivialExpressionOp canonicalization pattern that
handles ExpressionOps whose body yields a block argument directly: such
an expression is a passthrough and can be replaced by the corresponding
operand value. This canonicalization fires during greedy rewriting before
FoldExpressionOp encounters the invalid state.

Fixes #179844
---
 mlir/lib/Dialect/EmitC/IR/EmitC.cpp           | 22 ++++++++++++++-
 mlir/test/Dialect/EmitC/form-expressions.mlir | 28 +++++++++++++++++++
 2 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Dialect/EmitC/IR/EmitC.cpp b/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
index fa8a159b50d98..6979f34c1e047 100644
--- a/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
+++ b/mlir/lib/Dialect/EmitC/IR/EmitC.cpp
@@ -461,11 +461,31 @@ struct RemoveRecurringExpressionOperands
   }
 };
 
+/// If an ExpressionOp body yields a block argument directly (no root op),
+/// this means a contained op was folded away (e.g., an identity cast whose
+/// in/out types match). Canonicalize by replacing the expression with the
+/// corresponding operand value.
+struct FoldTrivialExpressionOp : public OpRewritePattern<ExpressionOp> {
+  using OpRewritePattern<ExpressionOp>::OpRewritePattern;
+  LogicalResult matchAndRewrite(ExpressionOp expressionOp,
+                                PatternRewriter &rewriter) const override {
+    auto yieldOp = cast<YieldOp>(expressionOp.getBody()->getTerminator());
+    Value yieldedValue = yieldOp.getResult();
+    auto blockArg = dyn_cast_if_present<BlockArgument>(yieldedValue);
+    if (!blockArg)
+      return failure();
+    rewriter.replaceOp(expressionOp,
+                       expressionOp.getOperand(blockArg.getArgNumber()));
+    return success();
+  }
+};
+
 } // namespace
 
 void ExpressionOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                MLIRContext *context) {
-  results.add<RemoveRecurringExpressionOperands>(context);
+  results.add<RemoveRecurringExpressionOperands, FoldTrivialExpressionOp>(
+      context);
 }
 
 ParseResult ExpressionOp::parse(OpAsmParser &parser, OperationState &result) {
diff --git a/mlir/test/Dialect/EmitC/form-expressions.mlir b/mlir/test/Dialect/EmitC/form-expressions.mlir
index 58eac4381ccb7..f5c002bba84a1 100644
--- a/mlir/test/Dialect/EmitC/form-expressions.mlir
+++ b/mlir/test/Dialect/EmitC/form-expressions.mlir
@@ -227,3 +227,31 @@ func.func @expression_with_constant(%arg0: i32) -> i32 {
   %a = emitc.mul %arg0, %c42 : (i32, i32) -> i32
   return %a : i32
 }
+
+// Regression test for https://github.com/llvm/llvm-project/issues/179844:
+// An identity cast (same in/out type) inside an expression is folded away
+// by CastOpInterface::foldTrait, leaving the expression body yielding a block
+// argument. FoldTrivialExpressionOp must canonicalize such expressions before
+// FoldExpressionOp tries to fold them.
+//
+// CHECK-LABEL: func.func @identity_cast_folded(
+// CHECK-SAME:                                  %[[ARG0:.*]]: i32) -> i32 {
+// CHECK:         %[[EXPR:.*]] = emitc.expression %[[ARG0]] : (i32) -> i32 {
+// CHECK:           %[[RES:.*]] = bitwise_and %[[ARG0]], %[[ARG0]] : (i32, i32) -> i32
+// CHECK:           yield %[[RES]] : i32
+// CHECK:         }
+// CHECK:         return %[[EXPR]] : i32
+// CHECK:       }
+
+func.func @identity_cast_folded(%arg0: i32) -> i32 {
+  // emitc.cast i32->i32 is an identity cast; CastOpInterface folds it away.
+  // After folding, FoldTrivialExpressionOp must canonicalize the expression
+  // that wraps this cast (which just yields its block arg) before
+  // FoldExpressionOp tries to fold it into the bitwise_and expression.
+  %0 = emitc.cast %arg0 : i32 to i32
+  // The bitwise_and uses both the cast result and the original arg, giving
+  // two uses of the cast expression's result (one from this op, one from the
+  // identity comparison below that keeps the cast live during canonicalization).
+  %1 = emitc.bitwise_and %0, %arg0 : (i32, i32) -> i32
+  return %1 : i32
+}