[mlir][sparse] Adding assertions for overhead storage types

[wrengr]

Summary

Fixes https://bugs.llvm.org/show_bug.cgi?id=52314 aka https://github.com/llvm/llvm-project/issues/51656

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

250–262

remove this big block, the comment goes in too much detail

if you want to keep the sizes[r] > 0 assert, put it in loop at L269

261–266

remove this comment; too much detail and not relevant for the next lines of code

433–439

just revert to original comment

all details is in the two new methods addPointer/addIndex

if you want to keep some of the notes, I would suggest moving them there, but keeping the comments outside those methods as in the original

457–458

remove comment; that is a good note for future work, but too distracting for reading this block of code

539–540

remove comment

[aartbik]

Please note some build issues, perhaps C++20 issue?

project/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
/var/lib/buildkite-agent/builds/llvm-project/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp:436:22: error: no member named 'numeric_limits' in namespace 'std'

assert(p <= std::numeric_limits<P>::max() &&
            ~~~~~^

/usr/include/assert.h:93:27: note: expanded from macro 'assert'

(static_cast <bool> (expr)                                         \
                     ^~~~

[wrengr]

Adding import for std::numeric_limits. And removing all explanatory documentation

[wrengr]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

250–262

I'm fine removing the comment from the code, but I think it should still be documented *somewhere* (since it's nontrivial to deduce what the worst case is and why). Does mlir have any standard practice for where to put such documentation?

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

29

Ah, that was it ;-)

250–262

In general, the commenting style in this file is concise, so I would adapt to that (as a rule of thumb, mirror the style you find in existing code, if it over communicates, then follow that, if it is terse, keep it that way). For now, I feel the comments go into a lot of detail that are not related to the code where they were placed.

Some of these feel like they belong in a design doc. If you feel very strongly about keeping them, move all information to the addPointer() and addIndex() methods. I still prefer conciseness, but would be okay having some of the stuff you have elsewhere at that central place, like the decision to assert on each insertion rather than doing it up front, for instance.

[aartbik]

Yeah, much better. One last nit on the partial loop merging you did

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

258–259

I would remove this comment, we don't have plans for such a subclass, and the printed string on failure already explains what is wrong

268

merging the loops seems unrelated to this change; also if you feel strongly about having one loop, we should also merge L251, and change the comment structure a bit

[wrengr]

removing more comments

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

268

It's less about having a single loop, more that you wanted the assert moved into the loop and I don't see why the assert would be any more closely related to reserving space than the push_back

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 
 #include <algorithm>
 #include <cassert>
 #include <cctype>
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
+#include <limits>
 #include <numeric>
 #include <vector>
 
 //===----------------------------------------------------------------------===//
 //
 // Internal support for storing and reading sparse tensors.
 //
 // The following memory-resident sparse storage schemes are supported:
public:
   /// permutation, and per-dimension dense/sparse annotations, using
   /// the coordinate scheme tensor for the initial contents if provided.
   SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
                       const DimLevelType *sparsity,
                       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
     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 *= 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 `addPointer`
+        // 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");
       }
     }
-    // Prepare sparse pointer structures for all dimensions.
-    for (uint64_t r = 0; r < rank; r++)
-      if (sparsity[r] == DimLevelType::kCompressed)
-        pointers[r].push_back(0);
     // Then assign contents from coordinate scheme tensor if provided.
     if (tensor) {
-      // Lexicographically sort the tensor, to ensure precondition of `fromCOO`.
+      // Ensure both preconditions of `fromCOO`.
+      assert(tensor->getSizes() == sizes && "Tensor size mismatch");
       tensor->sort();
+      // Now actually insert the `elements`.
       const std::vector<Element<V>> &elements = tensor->getElements();
       uint64_t nnz = elements.size();
       values.reserve(nnz);
       fromCOO(elements, 0, nnz, 0);
     } else if (allDense) {
       values.resize(sz, 0);
     }
   }
if (tensor) {
       for (uint64_t r = 0; r < rank; r++)
         permsz[perm[r]] = sizes[r];
       n = new SparseTensorStorage<P, I, V>(permsz, perm, sparsity);
     }
     return n;
   }
 
 private:
+  /// Appends the next free position of `indices[d]` to `pointers[d]`.
+  /// Thus, when called after inserting the last element of a segment,
+  /// it will append the position where the next segment begins.
+  inline void addPointer(uint64_t d) {
+    assert(isCompressedDim(d)); // Entails `d < getRank()`.
+    uint64_t p = indices[d].size();
+    assert(p <= std::numeric_limits<P>::max() &&
+           "Pointer value is too large for the P-type");
+    pointers[d].push_back(p); // Here is where we convert to `P`.
+  }
+
+  /// Appends the given index to `indices[d]`.
+  inline void addIndex(uint64_t d, uint64_t i) {
+    assert(isCompressedDim(d)); // Entails `d < getRank()`.
+    assert(i <= std::numeric_limits<I>::max() &&
+           "Index value is too large for the I-type");
+    indices[d].push_back(i); // Here is where we convert to `I`.
+  }
+
   /// Initializes sparse tensor storage scheme from a memory-resident sparse
   /// tensor in coordinate scheme. This method prepares the pointers and
   /// indices arrays under the given per-dimension dense/sparse annotations.
-  /// Precondition: the `elements` must be lexicographically sorted.
+  ///
+  /// Preconditions:
+  /// (1) the `elements` must be lexicographically sorted.
+  /// (2) the indices of every element are valid for `sizes` (equal rank
+  ///     and pointwise less-than).
   void fromCOO(const std::vector<Element<V>> &elements, uint64_t lo,
                uint64_t hi, uint64_t d) {
     // Once dimensions are exhausted, insert the numerical values.
     assert(d <= getRank());
     if (d == getRank()) {
       assert(lo < hi && hi <= elements.size());
       values.push_back(elements[lo].value);
       return;
     }
     // Visit all elements in this interval.
     uint64_t full = 0;
     while (lo < hi) {
       assert(lo < elements.size() && hi <= elements.size());
       // Find segment in interval with same index elements in this dimension.
       uint64_t i = elements[lo].indices[d];
       uint64_t seg = lo + 1;
       while (seg < hi && elements[seg].indices[d] == i)
         seg++;
       // Handle segment in interval for sparse or dense dimension.
       if (isCompressedDim(d)) {
-        indices[d].push_back(i);
+        addIndex(d, i);
       } else {
         // For dense storage we must fill in all the zero values between
         // the previous element (when last we ran this for-loop) and the
         // current element.
         for (; full < i; full++)
           endDim(d + 1);
         full++;
       }
       fromCOO(elements, lo, seg, d + 1);
       // And move on to next segment in interval.
       lo = seg;
     }
     // Finalize the sparse pointer structure at this dimension.
     if (isCompressedDim(d)) {
-      pointers[d].push_back(indices[d].size());
+      addPointer(d);
     } else {
       // For dense storage we must fill in all the zero values after
       // the last element.
       for (uint64_t sz = sizes[d]; full < sz; full++)
         endDim(d + 1);
     }
   }
 
private:
   }
 
   /// Ends a deeper, never seen before dimension.
   void endDim(uint64_t d) {
     assert(d <= getRank());
     if (d == getRank()) {
       values.push_back(0);
     } else if (isCompressedDim(d)) {
-      pointers[d].push_back(indices[d].size());
+      addPointer(d);
     } else {
       for (uint64_t full = 0, sz = sizes[d]; full < sz; full++)
         endDim(d + 1);
     }
   }
 
   /// Wraps up a single insertion path, inner to outer.
   void endPath(uint64_t diff) {
     uint64_t rank = getRank();
     assert(diff <= rank);
     for (uint64_t i = 0; i < rank - diff; i++) {
       uint64_t d = rank - i - 1;
       if (isCompressedDim(d)) {
-        pointers[d].push_back(indices[d].size());
+        addPointer(d);
       } else {
         for (uint64_t full = idx[d] + 1, sz = sizes[d]; full < sz; full++)
           endDim(d + 1);
       }
     }
   }
 
   /// Continues a single insertion path, outer to inner.
   void insPath(const uint64_t *cursor, uint64_t diff, uint64_t top, V val) {
     uint64_t rank = getRank();
     assert(diff < rank);
     for (uint64_t d = diff; d < rank; d++) {
       uint64_t i = cursor[d];
       if (isCompressedDim(d)) {
-        indices[d].push_back(i);
+        addIndex(d, i);
       } else {
         for (uint64_t full = top; full < i; full++)
           endDim(d + 1);
       }
       top = 0;
       idx[d] = i;
     }
     values.push_back(val);
   }
 
   /// Finds the lexicographic differing dimension.
   uint64_t lexDiff(const uint64_t *cursor) {
     for (uint64_t r = 0, rank = getRank(); r < rank; r++)
       if (cursor[r] > idx[r])
         return r;
       else
         assert(cursor[r] == idx[r] && "non-lexicographic insertion");
     assert(0 && "duplication insertion");
     return -1u;
   }
 
   /// Returns true if dimension is compressed.
   inline bool isCompressedDim(uint64_t d) const {
+    assert(d < getRank());
     return (!pointers[d].empty());
   }
 
 private:
   std::vector<uint64_t> sizes; // per-dimension sizes
   std::vector<uint64_t> rev;   // "reverse" permutation
   std::vector<uint64_t> idx;   // index cursor
   std::vector<std::vector<P>> pointers;