[Mlir-commits] [mlir] ae1c507 - [mlir][sparse_tensor] Fix out-of-bounds read in SparseAssembleOpConverter (#203289)

[llvmlistbot at llvm.org]

Author: lijinpei-amd
Date: 2026-06-13T14:20:02-07:00
New Revision: ae1c5074a0c6ec128fe824d456314e547b57165a

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

LOG: [mlir][sparse_tensor] Fix out-of-bounds read in SparseAssembleOpConverter (#203289)

The assemble codegen loop iterates over the level rank but asserted on
`getDimShape()[lvl]`, which is sized by the dimension rank. Index the
level shape instead, matching the loop bound and the next line.

Fixes #203225.

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
    mlir/test/Dialect/SparseTensor/sparse_pack.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 0ca63b46e25a8..2a57fcd0b36a8 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -1360,7 +1360,7 @@ struct SparseAssembleOpConverter : public OpConversionPattern<AssembleOp> {
     Level trailCOORank = stt.getLvlRank() - trailCOOStart;
     // Sets up SparseTensorSpecifier.
     for (Level lvl = 0, lvlRank = stt.getLvlRank(); lvl < lvlRank; lvl++) {
-      assert(ShapedType::isStatic(stt.getDimShape()[lvl]));
+      assert(ShapedType::isStatic(stt.getLvlShape()[lvl]));
 
       // Sets up the level size.
       auto lvlSize = constantIndex(rewriter, loc, stt.getLvlShape()[lvl]);

diff  --git a/mlir/test/Dialect/SparseTensor/sparse_pack.mlir b/mlir/test/Dialect/SparseTensor/sparse_pack.mlir
index ebbcc5fc7c7cf..d3c40573f439c 100644
--- a/mlir/test/Dialect/SparseTensor/sparse_pack.mlir
+++ b/mlir/test/Dialect/SparseTensor/sparse_pack.mlir
@@ -76,3 +76,21 @@ func.func @sparse_unpack(%sp : tensor<100x100xf64, #COO>,
                                  -> (tensor<2xindex>, tensor<6x2xi32>), tensor<6xf64>, (index, index), index
   return %rd, %rp, %ri : tensor<6xf64>, tensor<2xindex>, tensor<6x2xi32>
 }
+
+// Tests that the codegen does not crash when the sparse encoding has more
+// levels than dimensions (e.g. a blocked map using floordiv/mod).
+// See https://github.com/llvm/llvm-project/issues/203225
+
+#BSR = #sparse_tensor.encoding<{
+  map = (d0, d1) -> (d0 floordiv 2 : dense, d1 floordiv 2 : compressed, d0 mod 2 : dense, d1 mod 2 : dense)
+}>
+
+// Just check that codegen succeeds (no out-of-bounds crash); the exact
+// lowering is not the point of this regression test.
+// CHECK-LABEL:   func.func @sparse_pack_blocked
+func.func @sparse_pack_blocked(%values: tensor<?xf64>, %pos: tensor<?xindex>, %coordinates: tensor<?xindex>)
+                    -> tensor<2x4xf64, #BSR> {
+  %0 = sparse_tensor.assemble (%pos, %coordinates), %values
+     : (tensor<?xindex>, tensor<?xindex>), tensor<?xf64> to tensor<2x4xf64, #BSR>
+  return %0 : tensor<2x4xf64, #BSR>
+}