[Mlir-commits] [mlir] b7e2044 - [MLIR][ODS] Fix AllElementCountsMatch crash on dynamic shaped types (#183948)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-01T11:59:42+01:00
New Revision: b7e20442d5edc768c986cc1f3363ab76cd9d7281

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

LOG: [MLIR][ODS] Fix AllElementCountsMatch crash on dynamic shaped types (#183948)

The AllElementCountsMatch trait called ShapedType::getNumElements() on
operands or results with dynamic dimensions, which unconditionally
asserts hasStaticShape(). This caused mlir-opt to crash instead of
failing gracefully when the trait was used with dynamically-shaped
types.

Fix this by rewriting AllElementCountsMatch to use an And<> predicate
combining Neg<AnyMatchOperatorPred> (requiring all types to be
statically shaped) with AllMatchSameOperatorPred (requiring equal
element counts). When any type has dynamic dimensions the verification
now fails with a diagnostic instead of crashing.

Update the regression test to expect a verification failure rather than
success when dynamic shapes are present.

Fixes #159740

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpBase.td
    mlir/test/lib/Dialect/Test/TestOps.td
    mlir/test/mlir-tblgen/types.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 7a667d701ab71..7f36e6c74c7f7 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -543,9 +543,19 @@ class AnyMatchOperatorTrait<list<string> names, string operator,
   list<string> values = names;
 }
 
+// Verifies that all named operands/results of a shaped type have the same
+// element count. Fails if any operand or result is not a ShapedType or has
+// dynamic dimensions, since the element count cannot be determined statically.
 class AllElementCountsMatch<list<string> names> :
-    AllMatchSameOperatorTrait<names, ElementCount<"_self">.result,
-                              "element count">;
+    PredOpTrait<"all of {" # !interleave(names, ", ") # "} have same element count",
+        And<[
+            // Fail if any type is not a ShapedType or has dynamic dimensions.
+            Neg<AnyMatchOperatorPred<names,
+                "!::mlir::isa<::mlir::ShapedType>($_self.getType()) || "
+                "!::llvm::cast<::mlir::ShapedType>($_self.getType()).hasStaticShape()">>,
+            // All types are statically shaped; verify element counts match.
+            AllMatchSameOperatorPred<names, ElementCount<"_self">.result>
+        ]>>;
 
 class AllElementTypesMatch<list<string> names> :
     AllMatchSameOperatorTrait<names, ElementType<"_self">.result,

diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 62a374e08ec1c..fe02536a1df5b 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -831,6 +831,13 @@ def OperandZeroAndResultHaveSameElementCount :
   let results = (outs AnyShaped:$res);
 }
 
+def OperandAndResultHaveSameElementCountAnyType :
+    TEST_Op<"operand_and_result_have_same_element_count_any_type",
+            [AllElementCountsMatch<["x", "res"]>]> {
+  let arguments = (ins AnyType:$x);
+  let results = (outs AnyType:$res);
+}
+
 def FourEqualsFive :
     TEST_Op<"four_equals_five", [AllMatch<["5", "4"], "4 equals 5">]>;
 

diff  --git a/mlir/test/mlir-tblgen/types.mlir b/mlir/test/mlir-tblgen/types.mlir
index 7652a87037c92..c2acce0903bf4 100644
--- a/mlir/test/mlir-tblgen/types.mlir
+++ b/mlir/test/mlir-tblgen/types.mlir
@@ -432,6 +432,27 @@ func.func @same_element_count_success(%arg0: tensor<36xi32>, %arg1: tensor<1x2xf
 
 // -----
 
+// Regression test for https://github.com/llvm/llvm-project/issues/159740
+// AllElementCountsMatch should fail (not crash) when operands/results have dynamic shapes.
+func.func @same_element_count_dynamic(%arg0: tensor<2xi32>) {
+  // expected-error at +1 {{all of {x, res} have same element count}}
+  %0 = "test.operand0_and_result_have_same_element_count"(%arg0, %arg0) :
+    (tensor<2xi32>, tensor<2xi32>) -> tensor<?x?xf32>
+  return
+}
+
+// -----
+
+// AllElementCountsMatch should fail (not crash) when operands/results are not
+// ShapedType at all (e.g. a plain integer).
+func.func @same_element_count_non_shaped(%arg0: i32) -> i32 {
+  // expected-error at +1 {{all of {x, res} have same element count}}
+  %0 = "test.operand_and_result_have_same_element_count_any_type"(%arg0) : (i32) -> i32
+  return %0 : i32
+}
+
+// -----
+
 func.func @same_element_count_failure(%arg0: tensor<1xi32>, %arg1: tensor<1x2xf32>) {
   // expected-error at +1 {{all of {x, res} have same element count}}
   "test.operand0_and_result_have_same_element_count"(%arg0, %arg1) : (tensor<1xi32>, tensor<1x2xf32>) -> (tensor<2xi32>)