llvm / llvm/llvm-project

[MLIR][tensor] collapse_shape verifier asserts when static element-count product exceeds int64 range

Open
#221,590 1 comment 0 reactions 0 assignees View on GitHub
mlir
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

Open the contributing 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.