[Mlir-commits] [mlir] 175b9af - [mlir][sparse] avoid reserving dense storage for ptr/idx

[Aart Bik]

Author: Aart Bik
Date: 2022-04-05T17:40:01-07:00
New Revision: 175b9af484f483c3423ab2f78db5de7e25b64c31

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

LOG: [mlir][sparse] avoid reserving dense storage for ptr/idx

This avoids a rather big bug where we were reserving
dense space for the ptx/idx in the first sparse dimension.
For example, using CSR for a 140874 x 140874 matrix with
3977139 nonzero would reserve the full 19845483876 space.
This revision fixes this for now, but we need to revisit
the reservation heuristic to make this better.

Reviewed By: bixia

Differential Revision: https://reviews.llvm.org/D123166

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
index 03c54657d8ebe..fbc858466e064 100644
--- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
@@ -262,12 +262,14 @@ class SparseTensorStorage : public SparseTensorStorageBase {
     for (uint64_t r = 0; r < rank; r++)
       rev[perm[r]] = r;
     // Provide hints on capacity of pointers and indices.
-    // TODO: needs fine-tuning based on sparsity
+    // TODO: needs much fine-tuning based on actual sparsity; currently
+    //       we reserve pointer/index space based on all previous dense
+    //       dimensions, which works well up to first sparse dim; but
+    //       we should really use nnz and dense/sparse distribution.
     bool allDense = true;
     uint64_t sz = 1;
     for (uint64_t r = 0; r < rank; r++) {
       assert(sizes[r] > 0 && "Dimension size zero has trivial storage");
-      sz = checkedMul(sz, sizes[r]);
       if (sparsity[r] == DimLevelType::kCompressed) {
         pointers[r].reserve(sz + 1);
         indices[r].reserve(sz);
@@ -280,6 +282,7 @@ class SparseTensorStorage : public SparseTensorStorageBase {
       } else {
         assert(sparsity[r] == DimLevelType::kDense &&
                "singleton not yet supported");
+        sz = checkedMul(sz, sizes[r]);
       }
     }
     // Then assign contents from coordinate scheme tensor if provided.