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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: lijinpei-amd

<details>
<summary>Changes</summary>

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.

---
Full diff: https://github.com/llvm/llvm-project/pull/203289.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp (+1-1) 
- (modified) mlir/test/Dialect/SparseTensor/sparse_pack.mlir (+18) 


``````````diff
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>
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/203289