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

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-02T11:57:40+01:00
New Revision: 00af181a4eddad6ad2b0ea2e435ecde5f4033314

URL: https://github.com/llvm/llvm-project/commit/00af181a4eddad6ad2b0ea2e435ecde5f4033314
DIFF: https://github.com/llvm/llvm-project/commit/00af181a4eddad6ad2b0ea2e435ecde5f4033314.diff

LOG: [mlir][emitc] Fix crash in form-expressions when identity cast is folded (#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

Added: 
    

Modified: 
    mlir/lib/Dialect/EmitC/IR/EmitC.cpp
    mlir/test/Dialect/EmitC/form-expressions.mlir

Removed: 
    


################################################################################
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
+}