[mlir][sparse] add more asserts to sparse support lib

[aartbik]

Summary

We are having issues running the integration test of the sparse compiler
on AArch64 (crashing in the lib). This revision adds more assertions.

[mehdi_amini]

A UBSAN/ASAN build may also help for debugging this issue?

[aartbik]

In D109861#3003010, @mehdi_amini wrote:

A UBSAN/ASAN build may also help for debugging this issue?

Yeah, I have been running asan, tsan and msan runs (the latter less useful with JIT unfortunately) but on x86.
Javier is very kindly helping me on an Aarch64 machine and can reproduce the crash (IR is identical to x86 code, which does not crash).

This revision was just a very first start of the debug process, I don't think it will find anything.

But with Javier's help, we will get to the bottom of this!

[jsetoain]

Healthy checks to have in place, these should catch corruption in the mtx files cleanly. I'm working on a SAN build myself, hopefully it will help pinpoint the source of offending accesses. Thanks!

mlir/lib/ExecutionEngine/SparseUtils.cpp

 /// elements). This data structure is used to read a sparse tensor from
 /// any external format into memory and sort the elements lexicographically
 /// by indices before passing it back to the client (most packed storage
 /// formats require the elements to appear in lexicographic index order).
 template <typename V>
 struct SparseTensorCOO {
 public:
   SparseTensorCOO(const std::vector<uint64_t> &szs, uint64_t capacity)
-      : sizes(szs), pos(0) {
+      : sizes(szs) {
     if (capacity)
       elements.reserve(capacity);
   }
   /// Adds element as indices and value.
   void add(const std::vector<uint64_t> &ind, V val) {
     assert(getRank() == ind.size());
     for (uint64_t r = 0, rank = getRank(); r < rank; r++)
       assert(ind[r] < sizes[r]); // within bounds
     elements.emplace_back(ind, val);
   }
   /// Sorts elements lexicographically by index.
   void sort() { std::sort(elements.begin(), elements.end(), lexOrder); }
-  /// Primitive one-time iteration.
-  const Element<V> &next() { return elements[pos++]; }
   /// Returns rank.
   uint64_t getRank() const { return sizes.size(); }
   /// Getter for sizes array.
   const std::vector<uint64_t> &getSizes() const { return sizes; }
   /// Getter for elements array.
   const std::vector<Element<V>> &getElements() const { return elements; }
 
   /// Factory method. Permutes the original dimensions according to
for (uint64_t r = 0, rank = e1.indices.size(); r < rank; r++) {
       if (e1.indices[r] == e2.indices[r])
         continue;
       return e1.indices[r] < e2.indices[r];
     }
     return false;
   }
   std::vector<uint64_t> sizes; // per-rank dimension sizes
   std::vector<Element<V>> elements;
-  uint64_t pos;
 };
 
 /// Abstract base class of sparse tensor storage. Note that we use
 /// function overloading to implement "partial" method specialization.
 class SparseTensorStorageBase {
 public:
   enum DimLevelType : uint8_t { kDense = 0, kCompressed = 1, kSingleton = 2 };
 
for (uint64_t d = 0, rank = getRank(); d < rank; d++)
       if (sparsity[d] == kCompressed)
         pointers[d].push_back(0);
     // Then setup the tensor.
     fromCOO(tensor, sparsity, 0, nnz, 0);
   }
 
   virtual ~SparseTensorStorage() {}
 
+  /// Get the rank of the tensor.
   uint64_t getRank() const { return sizes.size(); }
 
-  uint64_t getDimSize(uint64_t d) override { return sizes[d]; }
+  /// Get the size in the given dimension of the tensor.
+  uint64_t getDimSize(uint64_t d) override {
+    assert(d < getRank());
+    return sizes[d];
+  }
 
   // Partially specialize these three methods based on template types.
   void getPointers(std::vector<P> **out, uint64_t d) override {
+    assert(d < getRank());
     *out = &pointers[d];
   }
   void getIndices(std::vector<I> **out, uint64_t d) override {
+    assert(d < getRank());
     *out = &indices[d];
   }
   void getValues(std::vector<V> **out) override { *out = &values; }
 
   /// Returns this sparse tensor storage scheme as a new memory-resident
   /// sparse tensor in coordinate scheme with the given dimension order.
   SparseTensorCOO<V> *toCOO(uint64_t *perm) {
     // Restore original order of the dimension sizes and allocate coordinate
     // scheme with desired new ordering specified in perm.
     uint64_t size = getRank();
     std::vector<uint64_t> orgsz(size);
     for (uint64_t r = 0; r < size; r++)
       orgsz[rev[r]] = sizes[r];
     SparseTensorCOO<V> *tensor = SparseTensorCOO<V>::newSparseTensorCOO(
         size, orgsz.data(), perm, values.size());
     // Populate coordinate scheme restored from old ordering and changed with
     // new ordering. Rather than applying both reorderings during the recursion,
     // we compute the combine permutation in advance.
     std::vector<uint64_t> reord(size);
     for (uint64_t r = 0; r < size; r++)
       reord[r] = perm[rev[r]];
     std::vector<uint64_t> idx(size);
     toCOO(tensor, reord, idx, 0, 0);
+    assert(tensor->getElements().size() == values.size());
     return tensor;
   }
 
   /// Factory method. Expects a coordinate scheme that respects the same
   /// permutation as is desired for the new sparse storage scheme.
   static SparseTensorStorage<P, I, V> *
   newSparseTensor(SparseTensorCOO<V> *t, uint8_t *sparsity, uint64_t *perm) {
     t->sort(); // sort lexicographically
     SparseTensorStorage<P, I, V> *n =
         new SparseTensorStorage<P, I, V>(t, sparsity, perm);
     delete t;
     return n;
   }
 
 private:
   /// 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-rank dimension dense/sparse annotations.
   void fromCOO(SparseTensorCOO<V> *tensor, uint8_t *sparsity, uint64_t lo,
                uint64_t hi, uint64_t d) {
     const std::vector<Element<V>> &elements = tensor->getElements();
     // Once dimensions are exhausted, insert the numerical values.
     if (d == getRank()) {
+      assert(lo >= hi || lo < elements.size());
       values.push_back(lo < hi ? elements[lo].value : 0);
       return;
     }
+    assert(d < getRank());
     // 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.
       unsigned idx = elements[lo].indices[d];
       unsigned seg = lo + 1;
       while (seg < hi && elements[seg].indices[d] == idx)
         seg++;
       // Handle segment in interval for sparse or dense dimension.
       if (sparsity[d] == kCompressed) {
         indices[d].push_back(idx);
if (sparsity[d] == kCompressed) {
         fromCOO(tensor, sparsity, 0, 0, d + 1); // pass empty
     }
   }
 
   /// Stores the sparse tensor storage scheme into a memory-resident sparse
   /// tensor in coordinate scheme.
   void toCOO(SparseTensorCOO<V> *tensor, std::vector<uint64_t> &reord,
              std::vector<uint64_t> &idx, uint64_t pos, uint64_t d) {
+    assert(d <= getRank());
     if (d == getRank()) {
+      assert(pos < values.size());
       tensor->add(idx, values[pos]);
     } else if (pointers[d].empty()) {
       // Dense dimension.
       for (uint64_t i = 0, sz = sizes[d], off = pos * sz; i < sz; i++) {
         idx[reord[d]] = i;
         toCOO(tensor, reord, idx, off + i, d + 1);
       }
     } else {