[Mlir-commits] [mlir] 0724358 - [mlir][SparseTensor] Fix crash demapping alloc_tensor with a copy operand (#219319)

[llvmlistbot at llvm.org]

Author: Aman Singh
Date: 2026-08-29T14:49:28-07:00
New Revision: 07243585bc8f62978fa6fed969cd8305c3603b9f

URL: https://github.com/llvm/llvm-project/commit/07243585bc8f62978fa6fed969cd8305c3603b9f
DIFF: https://github.com/llvm/llvm-project/commit/07243585bc8f62978fa6fed969cd8305c3603b9f.diff

LOG: [mlir][SparseTensor] Fix crash demapping alloc_tensor with a copy operand (#219319)

`TensorAllocDemapper` reconstructs demapped level sizes for a
`bufferization.alloc_tensor`/`tensor.empty` by pairing each dynamic
result
dimension with an entry from the op's `dynamic_sizes` operand list,
popping via
`ValueRange::front()`.

When an `alloc_tensor` has a `copy` operand instead of explicit dynamic
sizes,
`dynamic_sizes` is legitimately empty — the op's own verifier requires
that the
sizes are implied by the copy operand and must not be specified — so
`front()`
was called on an empty range and asserted:

```
llvm/include/llvm/ADT/STLExtras.h:1253: Assertion `!empty() && "expected non-empty range"' failed.
```

Fix by special-casing `alloc_tensor`'s copy operand: demap the copy
operand
itself and forward it (with no dynamic-size operands) to a freshly
created,
demapped `alloc_tensor`, since the copy operand's type already fully
determines
the result shape.

## Testing

Added the reduced reproducer from the issue to
`mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir`.

Checked on an assertions build that the reproducer aborts with the
assertion
above before the change and succeeds after it, and that
`mlir/test/Dialect/SparseTensor` is 114/114.

Fixes #216223

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
    mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
index 87c715c425866..03c9e57bd223f 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseReinterpretMap.cpp
@@ -19,6 +19,8 @@
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 
+#include <type_traits>
+
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
@@ -614,6 +616,22 @@ struct TensorAllocDemapper : public OpRewritePattern<AllocOp> {
     if (stt.getEncoding().getDimToLvl().getNumSymbols() != 0)
       return failure();
 
+    if constexpr (std::is_same_v<AllocOp, bufferization::AllocTensorOp>) {
+      // `bufferization.alloc_tensor` does not carry any dynamic size
+      // operands when it has a `copy` operand -- the shape (and the
+      // contents) are inherited from `copy` instead. Simply demap the
+      // `copy` operand and forward it to a newly created (demapped)
+      // `alloc_tensor` op.
+      if (Value copy = op.getCopy()) {
+        Value demappedCopy = genDemap(rewriter, stt.getEncoding(), copy);
+        auto allocOp = AllocOp::create(rewriter, loc, stt.getDemappedType(),
+                                       ValueRange{}, demappedCopy);
+        Value t = genRemap(rewriter, stt.getEncoding(), allocOp.getResult());
+        rewriter.replaceOp(op, t);
+        return success();
+      }
+    }
+
     SmallVector<Value> maxDimCrds;
     maxDimCrds.reserve(stt.getDimRank());
     ValueRange dynSz = op.getDynamicSizes();

diff  --git a/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir b/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
index 89ad74bd7000e..6da99cb457340 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_reinterpret_map.mlir
@@ -109,6 +109,18 @@ func.func @sparse_assemble_reinterpret_map(%val : tensor<?xf64>, %pos:tensor<?xi
   return %0 : tensor<2x4xf64, #BSR>
 }
 
+// CHECK-LABEL:   func.func @sparse_alloc_copy_reinterpret_map(
+// CHECK-SAME:        %[[VAL_0:.*]]: tensor<?x?xf32, #[[$remap]]>) -> tensor<?x?xf32, #[[$remap]]> {
+// CHECK:           %[[VAL_1:.*]] = sparse_tensor.reinterpret_map %[[VAL_0]] : tensor<?x?xf32, #[[$remap]]> to tensor<?x?x2x2xf32, #[[$demap]]>
+// CHECK:           %[[VAL_2:.*]] = bufferization.alloc_tensor() copy(%[[VAL_1]]) : tensor<?x?x2x2xf32, #[[$demap]]>
+// CHECK:           %[[VAL_3:.*]] = sparse_tensor.reinterpret_map %[[VAL_2]] : tensor<?x?x2x2xf32, #[[$demap]]> to tensor<?x?xf32, #[[$remap]]>
+// CHECK:           return %[[VAL_3]] : tensor<?x?xf32, #[[$remap]]>
+// CHECK:         }
+func.func @sparse_alloc_copy_reinterpret_map(%arg0: tensor<?x?xf32, #BSR>) -> tensor<?x?xf32, #BSR> {
+  %0 = bufferization.alloc_tensor() copy(%arg0) : tensor<?x?xf32, #BSR>
+  return %0 : tensor<?x?xf32, #BSR>
+}
+
 // CHECK-LABEL:   func.func @sparse_disassemble_reinterpret_map(
 // CHECK-SAME:         %[[VAL_0:.*]]: tensor<2x4xf64, #[[$remap]]>,
 // CHECK-SAME:         %[[VAL_1:.*]]: tensor<?xf64>,