[MLIR][tensor] collapse_shape verifier asserts when static element-count product exceeds int64 range
- Dominant language
- LLVM
- Stars
- 40.5k
- Forks
- 18.7k
- PR merge metrics
- PR metrics pending
Description
## Describe the bug
The verifier for `tensor.collapse_shape` computes the static element counts of the expanded and collapsed tensor types while checking element preservation. With positive static dimensions whose mathematical product is `2^63`, `mlir-opt --verify-each` aborts in the verifier instead of returning a diagnostic or a failed verification result.
The reproducer is verifier-only: it does not allocate a tensor, run a pass, or launch a backend.
## Expected behavior
The verifier should handle a static element-count product outside the representable positive `int64` range without aborting. It should either reject the input with a regular diagnostic or otherwise return a controlled verification failure.
## Minimal reproducer
Save the following as `collapse-shape-int64-product.mlir`:
```mlir
module {
func.func @identity(%arg0: tensor<4611686018427387904x2xf32>) -> tensor<4611686018427387904x2xf32> {
%0 = tensor.collapse_shape %arg0 [[0], [1]] : tensor<4611686018427387904x2xf32> into tensor<4611686018427387904x2xf32>
return %0 : tensor<4611686018427387904x2xf32>
}
}
```
Run:
```bash
mlir-opt --verify-each collapse-shape-int64-product.mlir
```
Observed result on an assertions-enabled build:
```text
exit status: 134
mlir-opt: ...: Assertion `num.has_value() && "integer overflow in element count computation"' failed.
Key resolved frames from the crash backtrace:
#10 mlir::tensor::CollapseShapeOp::verify() TensorOps.cpp:0:0
#11 mlir::Op::verifyInvariants(mlir::Operation*) TensorDialect.cpp:0:0
#12 (anonymous namespace)::OperationVerifier::verifyOpAndDominance(mlir::Operation&) Verifier.cpp:0:0
```
The optimized build reports `ShapedType::getNumElements` at the assertion site, while the source-level `verifyTensorReshapeOp` callsite is identified in the pinned source links below; neither is emitted as a separate stack frame in this build.
The mathematical element count is `4611686018427387904 * 2 = 9223372036854775808` (`2^63`), which is one above the maximum positive `int64` value.
## Controls
The same identity-shaped operation is accepted for `tensor<2x3xf32>` and for `tensor<2147483648x2xf32>`. A type-only module containing `tensor<4611686018427387904x2xf32>` is also accepted, so the observed abort is tied to the `tensor.collapse_shape` verification path rather than parsing the positive static dimensions alone.
## Relevant implementation
The tensor reshape verifier compares the static element counts of the expanded and collapsed types. The checked helper `ShapedType::tryGetNumElements` can report an unrepresentable product, while the callsite currently uses the asserting `ShapedType::getNumElements` wrapper. In the tested source revision, the callsite is in [`verifyTensorReshapeOp`](https://github.com/llvm/llvm-project/blob/d8145e71418fb1e0a936adfb07dc5317113fc3b6/mlir/lib/Dialect/Tensor/IR/TensorOps.cpp#L2044-L2055), and the checked/asserting helper pair is in [`BuiltinTypeInterfaces.cpp`](https://github.com/llvm/llvm-project/blob/d8145e71418fb1e0a936adfb07dc5317113fc3b6/mlir/lib/IR/BuiltinTypeInterfaces.cpp#L72-L86).
## Environment details
```text
LLVM/MLIR source revision: d8145e71418fb1e0a936adfb07dc5317113fc3b6
mlir-opt: LLVM 23.1.0-rc2, optimized build with assertions
GPU: not required
Execution: verifier-only; no tensor materialization or backend execution
```
The reduced candidate was replayed in an independent invocation and produced the same assertion markers.
## Duplicate search
Related reports include [LLVM issue #178362](https://github.com/llvm/llvm-project/issues/178362), [#204297](https://github.com/llvm/llvm-project/issues/204297), and [#202529](https://github.com/llvm/llvm-project/issues/202529), which cover the same broad `ShapedType` element-count assertion family in `memref.global`, `mem2reg`, or `sroa` paths. [#64638](https://github.com/llvm/llvm-project/issues/64638) is an older large-memref arithmetic case. [#179005](https://github.com/llvm/llvm-project/pull/179005) concerns assertions for unranked tensor reshape operations. [#119866](https://github.com/llvm/llvm-project/issues/119866) is a separate `arith.muli` integer-constant-folding overflow crash during `mlir-opt -canonicalize`, with a different operation and call path.
I did not find a direct duplicate for this `tensor.collapse_shape` verifier callsite in the checked local snapshots or in the live search queries for `tensor.collapse_shape` with element-count overflow terms. Current live search also returned [#179670](https://github.com/llvm/llvm-project/issues/179670), which is a `DenseElementsAttr::reshape` cardinality-mismatch assertion during `--remove-dead-values`, and [#173567](https://github.com/llvm/llvm-project/issues/173567), which is an invalid reassociation-expression assertion; both use different failure mechanisms and callsites.
This report is limited to a verifier crash on a statically typed MLIR input. It does not make a claim about runtime allocation, backend safety, or a confirmed downstream impact.
Contributor guide
Research direction
Start in mlir/lib/Dialect/Tensor/IR/TensorOps.cpp at verifyTensorReshapeOp and compare its element-count handling with ShapedType::tryGetNumElements and getNumElements in mlir/lib/IR/BuiltinTypeInterfaces.cpp. Run mlir-opt --verify-each on the supplied collapse-shape-int64-product.mlir reproducer; done means the input produces a controlled verification failure or diagnostic rather than an assertion abort.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 78/100