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

[llvmlistbot at llvm.org]

Author: Victor Perez
Date: 2026-07-26T19:20:03-04:00
New Revision: 677f13efef9c48acdfba73703344c5234a04ae9e

URL: https://github.com/llvm/llvm-project/commit/677f13efef9c48acdfba73703344c5234a04ae9e
DIFF: https://github.com/llvm/llvm-project/commit/677f13efef9c48acdfba73703344c5234a04ae9e.diff

LOG: [mlir][arith] Fix folds crashing on dynamic-shaped tensors (#212072)

`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:

```mlir
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 `IsScalarOrHasStaticShape` constraint.

Signed-off-by: Víctor Pérez Carrasco <victor.pc.upm at gmail.com>

Added: 
    

Modified: 
    mlir/lib/Dialect/Arith/IR/ArithCanonicalization.td
    mlir/lib/Dialect/Arith/IR/ArithOps.cpp
    mlir/test/Dialect/Arith/canonicalize.mlir

Removed: 
    


################################################################################
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 e260ca02509de..7e4ffeda2bd7e 100644
--- a/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
+++ b/mlir/lib/Dialect/Arith/IR/ArithOps.cpp
@@ -552,6 +552,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);
@@ -1244,8 +1249,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 5352f777109de..baf0d859b287f 100644
--- a/mlir/test/Dialect/Arith/canonicalize.mlir
+++ b/mlir/test/Dialect/Arith/canonicalize.mlir
@@ -4382,3 +4382,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>
+}