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

[aartbik]

Summary

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.

[aartbik]

cl

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
                       SparseTensorCOO<V> *tensor = nullptr)
       : sizes(szs), rev(getRank()), idx(getRank()), pointers(getRank()),
         indices(getRank()) {
     uint64_t rank = getRank();
     // Store "reverse" permutation.
     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);
         sz = 1;
         allDense = false;
         // Prepare the pointer structure.  We cannot use `appendPointer`
         // here, because `isCompressedDim` won't work until after this
         // preparation has been done.
         pointers[r].push_back(0);
       } else {
         assert(sparsity[r] == DimLevelType::kDense &&
                "singleton not yet supported");
+        sz = checkedMul(sz, sizes[r]);
       }
     }
     // Then assign contents from coordinate scheme tensor if provided.
     if (tensor) {
       // Ensure both preconditions of `fromCOO`.
       assert(tensor->getSizes() == sizes && "Tensor size mismatch");
       tensor->sort();
       // Now actually insert the `elements`.