[Mlir-commits] [mlir] [mlir][arith] Fix folds crashing on dynamic-shaped tensors (PR #212072)

[Victor Perez]

https://github.com/victor-eds updated https://github.com/llvm/llvm-project/pull/212072

>From c7c2738cf82566478c10742c1aeb009c394ee6fb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?V=C3=ADctor=20P=C3=A9rez=20Carrasco?=
 <victor.pc.upm at gmail.com>
Date: Sun, 26 Jul 2026 09:11:22 -0700
Subject: [PATCH] [mlir][arith] Fix folds crashing on dynamic-shaped tensors
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Builder::getZeroAttr asserts when asked to build a constant of a
ranked tensor with a dynamic shape (DenseElementsAttr requires a static
shape). Three self-identity folds/patterns reached it without a guard and
aborted on dynamic-shaped operands:

  arith.xori %x, %x : tensor<?xi32>
  arith.subui_extended %x, %x : tensor<?xi32>, tensor<?xi1>
  arith.subi (arith.subi %a, %b), %a : tensor<?xi32>   (SubISubILHSRHSLHS)

Guard each with a static-shape check (as arith.subi(x,x) already does) so
they leave the op unfolded instead of crashing. The TableGen pattern gets
a new HasStaticOrScalarShape constraint.

Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>
---
 .../Dialect/Arith/IR/ArithCanonicalization.td | 12 ++++++-
 mlir/lib/Dialect/Arith/IR/ArithOps.cpp        | 14 ++++++--
 mlir/test/Dialect/Arith/canonicalize.mlir     | 34 +++++++++++++++++++
 3 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
index bfb8e5cc34218..2615bcca23e35 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
+++ b/mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
@@ -15,6 +15,15 @@ include "mlir/Dialect/Arith/IR/ArithOps.td"
 // Create zero attribute of type matching the argument's type.
 def GetZeroAttr : NativeCodeCall<"$_builder.getZeroAttr($0.getType())">;
 
+// A value whose type is a scalar or has a static shape, i.e. one for which a
+// splat constant can be built. Guards patterns that materialize a constant via
+// GetZeroAttr, which asserts on a dynamically-shaped type.
+def IsScalarOrHasStaticShape :
+    Constraint<CPred<"!::llvm::isa<::mlir::ShapedType>($0.getType()) || "
+                     "::llvm::cast<::mlir::ShapedType>($0.getType())"
+                     ".hasStaticShape()">,
+               "type is a scalar or has a static shape">;
+
 // Add two integer attributes and create a new one with the result.
 def AddIntAttrs : NativeCodeCall<"addIntegerAttrs($_builder, $0, $1, $2)">;
 
@@ -193,7 +202,8 @@ def SubILHSSubConstantLHS :
 def SubISubILHSRHSLHS :
     Pat<(Arith_SubIOp:$res (Arith_SubIOp $x, $y, $ovf1), $x, $ovf2),
         (Arith_SubIOp (Arith_ConstantOp (GetZeroAttr $y)), $y,
-            (MergeOverflow $ovf1, $ovf2))>;
+            (MergeOverflow $ovf1, $ovf2)),
+        [(IsScalarOrHasStaticShape $y)]>;
 
 //===----------------------------------------------------------------------===//
 // MulSIExtendedOp
diff --git a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
index 674577e5764d8..7f138b30eb8dd 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -539,6 +539,11 @@ arith::SubUIExtendedOp::fold(FoldAdaptor adaptor,
 
   // subui_extended(x, x) -> 0, false
   if (getLhs() == getRhs()) {
+    // A dynamically-shaped result cannot be a constant; bail before
+    // getZeroAttr, which would assert on a non-static shape.
+    auto shapedType = dyn_cast<ShapedType>(getDiff().getType());
+    if (shapedType && !shapedType.hasStaticShape())
+      return failure();
     Builder builder(getContext());
     auto zeroDiff = builder.getZeroAttr(getDiff().getType());
     auto falseValue = builder.getZeroAttr(borrowTy);
@@ -1153,8 +1158,13 @@ OpFoldResult arith::XOrIOp::fold(FoldAdaptor adaptor) {
   if (matchPattern(adaptor.getRhs(), m_Zero()))
     return getLhs();
   /// xor(x, x) -> 0
-  if (getLhs() == getRhs())
-    return Builder(getContext()).getZeroAttr(getType());
+  if (getLhs() == getRhs()) {
+    // A dynamically-shaped result cannot be a constant; bail before
+    // getZeroAttr, which would assert on a non-static shape.
+    auto shapedType = dyn_cast<ShapedType>(getType());
+    if (!shapedType || shapedType.hasStaticShape())
+      return Builder(getContext()).getZeroAttr(getType());
+  }
   /// xor(xor(x, a), a) -> x
   /// xor(xor(a, x), a) -> x
   if (arith::XOrIOp prev = getLhs().getDefiningOp<arith::XOrIOp>()) {
diff --git a/mlir/test/Dialect/Arith/canonicalize.mlir b/mlir/test/Dialect/Arith/canonicalize.mlir
index d11b00e1c1e24..d48939bd3ab32 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -4263,3 +4263,37 @@ func.func @convertf_fold_f8() -> f8E5M2 {
   return %result : f8E5M2
 }
 
+
+// -----
+
+// Self-identity folds and patterns must not build a constant of a dynamic
+// shape (which would assert); they must leave the op unfolded.
+// CHECK-LABEL: func @xori_self_dynamic
+//       CHECK:   arith.xori
+func.func @xori_self_dynamic(%arg0 : tensor<?xi32>) -> tensor<?xi32> {
+  %0 = arith.xori %arg0, %arg0 : tensor<?xi32>
+  return %0 : tensor<?xi32>
+}
+
+// -----
+
+// CHECK-LABEL: func @subui_extended_self_dynamic
+//       CHECK:   arith.subui_extended
+func.func @subui_extended_self_dynamic(%arg0 : tensor<?xi32>)
+    -> (tensor<?xi32>, tensor<?xi1>) {
+  %low, %bo = arith.subui_extended %arg0, %arg0
+      : tensor<?xi32>, tensor<?xi1>
+  return %low, %bo : tensor<?xi32>, tensor<?xi1>
+}
+
+// -----
+
+// CHECK-LABEL: func @subi_subi_lhs_rhs_lhs_dynamic
+//       CHECK:   arith.subi
+//       CHECK:   arith.subi
+func.func @subi_subi_lhs_rhs_lhs_dynamic(%a : tensor<?xi32>, %b : tensor<?xi32>)
+    -> tensor<?xi32> {
+  %0 = arith.subi %a, %b : tensor<?xi32>
+  %1 = arith.subi %0, %a : tensor<?xi32>
+  return %1 : tensor<?xi32>
+}