[Mlir-commits] [mlir] [mlir][tensor] Fix element-count overflow assertion in reshape verifier (PR #222117)

[Stanislav Bardyuk]

https://github.com/kodlan created https://github.com/llvm/llvm-project/pull/222117

verifyTensorReshapeOp (shared by tensor.collapse_shape and
tensor.expand_shape) compares the static element counts of the two types
with ShapedType::getNumElements(), which asserts when the product does not
fit in int64_t since 8756292abb00. So a type like
tensor<4611686018427387904x2xf32> (2^63 elements) makes the verifier abort
in an assertions build instead of reporting anything:

```mlir
%0 = tensor.collapse_shape %arg0 [[0], [1]] : tensor<4611686018427387904x2xf32> into tensor<4611686018427387904x2xf32>
```

Use the checked ShapedType::tryGetNumElements() instead and reject the op
with a regular diagnostic when either count overflows. Rejecting rather
than skipping the check keeps the empty-reassociation case
(`collapse_shape [] : tensor<4611686018427387904x2xf32> into tensor<f32>`)
invalid, since that count comparison is the only check that catches it,
and it avoids relying on wrapped multiplication in release builds. It also
catches grouped cases whose per-group product wraps to a plausible value
(e.g. `[[0, 1]] : tensor<4294967296x4294967296xf32> into tensor<0xf32>`).

The per-group products in reshapeLikeShapesAreCompatible and
inferCollapsedType, and the count check in tensor.reshape, still use plain
multiplication (no assertion, but wrap on overflow); left for a follow-up.

Fixes #221590

>From 866dd48c31b1a72c50469a9a8aee1ee464cad8b5 Mon Sep 17 00:00:00 2001
From: Stanislav Bardyuk <sbardyuk at google.com>
Date: Tue, 8 Sep 2026 19:43:34 +0000
Subject: [PATCH] [mlir][tensor] Fix element-count overflow assertion in
 reshape verifier

verifyTensorReshapeOp compared the static element counts of the expanded
and collapsed types with ShapedType::getNumElements(), which asserts on
int64 overflow. Use tryGetNumElements() and reject the op with a
diagnostic when either count is not representable.

Fixes #221590
---
 mlir/lib/Dialect/Tensor/IR/TensorOps.cpp | 13 ++++++++-----
 mlir/test/Dialect/Tensor/invalid.mlir    | 24 ++++++++++++++++++++++++
 2 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
index fb14c1953fd22..3084a769458cd 100644
--- a/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
+++ b/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp
@@ -2071,12 +2071,15 @@ static LogicalResult verifyTensorReshapeOp(TensorReshapeOp op,
 
   // Reshape must preserve the number of elements when statically known.
   if (expandedType.hasStaticShape() && collapsedType.hasStaticShape()) {
-    int64_t expandedNumElements = expandedType.getNumElements();
-    int64_t collapsedNumElements = collapsedType.getNumElements();
-    if (expandedNumElements != collapsedNumElements) {
+    std::optional<int64_t> expandedNumElements =
+        ShapedType::tryGetNumElements(expandedType.getShape());
+    std::optional<int64_t> collapsedNumElements =
+        ShapedType::tryGetNumElements(collapsedType.getShape());
+    if (!expandedNumElements || !collapsedNumElements)
+      return op.emitOpError("number of elements exceeds the int64 range");
+    if (*expandedNumElements != *collapsedNumElements)
       return op.emitOpError("number of elements must be preserved: ")
-             << expandedNumElements << " != " << collapsedNumElements;
-    }
+             << *expandedNumElements << " != " << *collapsedNumElements;
   }
 
   auto maps = op.getReassociationMaps();
diff --git a/mlir/test/Dialect/Tensor/invalid.mlir b/mlir/test/Dialect/Tensor/invalid.mlir
index a526d7ed61722..8fa233f993640 100644
--- a/mlir/test/Dialect/Tensor/invalid.mlir
+++ b/mlir/test/Dialect/Tensor/invalid.mlir
@@ -723,3 +723,27 @@ func.func @no_fold_invalid_collapse() -> tensor<i64> {
     return %0 : tensor<i64>
 }
 
+// -----
+
+func.func @collapse_shape_num_elements_overflow(%arg0: tensor<4611686018427387904x2xf32>) {
+  // expected-error at +1 {{'tensor.collapse_shape' op number of elements exceeds the int64 range}}
+  %0 = tensor.collapse_shape %arg0 [[0], [1]] : tensor<4611686018427387904x2xf32> into tensor<4611686018427387904x2xf32>
+  return
+}
+
+// -----
+
+func.func @collapse_shape_num_elements_overflow_to_scalar(%arg0: tensor<4611686018427387904x2xf32>) {
+  // expected-error at +1 {{'tensor.collapse_shape' op number of elements exceeds the int64 range}}
+  %0 = tensor.collapse_shape %arg0 [] : tensor<4611686018427387904x2xf32> into tensor<f32>
+  return
+}
+
+// -----
+
+func.func @expand_shape_num_elements_overflow(%arg0: tensor<4611686018427387904x2xf32>) {
+  // expected-error at +1 {{'tensor.expand_shape' op number of elements exceeds the int64 range}}
+  %0 = tensor.expand_shape %arg0 [[0], [1]] output_shape [4611686018427387904, 2] : tensor<4611686018427387904x2xf32> into tensor<4611686018427387904x2xf32>
+  return
+}
+