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

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-core

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/183948.diff


2 Files Affected:

- (modified) mlir/include/mlir/IR/OpBase.td (+11-2) 
- (modified) mlir/test/mlir-tblgen/types.mlir (+11) 


``````````diff
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>)

``````````

</details>


https://github.com/llvm/llvm-project/pull/183948