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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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

>From d2546ea0a500c1644cd0f444602dcbd6328ec675 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Sat, 28 Feb 2026 10:10:09 -0800
Subject: [PATCH] [MLIR][ODS] Fix AllElementCountsMatch crash on dynamic shaped
 types

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
---
 mlir/include/mlir/IR/OpBase.td   | 13 +++++++++++--
 mlir/test/mlir-tblgen/types.mlir | 11 +++++++++++
 2 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/IR/OpBase.td b/mlir/include/mlir/IR/OpBase.td
index 7a667d701ab71..629d3f530dca1 100644
--- a/mlir/include/mlir/IR/OpBase.td
+++ b/mlir/include/mlir/IR/OpBase.td
@@ -543,9 +543,18 @@ 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 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 has dynamic dimensions.
+            Neg<AnyMatchOperatorPred<names,
+                "!::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/mlir-tblgen/types.mlir b/mlir/test/mlir-tblgen/types.mlir
index 7652a87037c92..ba1bd03bc3076 100644
--- a/mlir/test/mlir-tblgen/types.mlir
+++ b/mlir/test/mlir-tblgen/types.mlir
@@ -432,6 +432,17 @@ 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
+}
+
+// -----
+
 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>)