[mlir][sparse] assert fail on mismatch between rank and annotations array

[aartbik]

Summary

Rationale:
Providing the wrong number of sparse/dense annotations was silently
ignored or caused unrelated crashes. This minor change verifies that
the provided number matches the rank.

[bixia]

Thanks for adding the asserts!

mlir/lib/ExecutionEngine/SparseUtils.cpp

 struct SparseTensor {
 public:
   SparseTensor(const std::vector<uint64_t> &szs, uint64_t capacity)
       : sizes(szs), pos(0) {
     elements.reserve(capacity);
   }
   /// Adds element as indices and value.
   void add(const std::vector<uint64_t> &ind, double val) {
-    assert(sizes.size() == ind.size());
-    for (int64_t r = 0, rank = sizes.size(); r < rank; r++)
+    assert(getRank() == ind.size());
+    for (int64_t r = 0, rank = getRank(); r < rank; r++)
       assert(ind[r] < sizes[r]); // within bounds
     elements.emplace_back(Element(ind, val));
   }
   /// Sorts elements lexicographically by index.
   void sort() { std::sort(elements.begin(), elements.end(), lexOrder); }
   /// Primitive one-time iteration.
   const Element &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> &getElements() const { return elements; }
 
 private:
   /// Returns true if indices of e1 < indices of e2.
   static bool lexOrder(const Element &e1, const Element &e2) {

 /// "one-size-fits-all" solution that simply takes an input tensor and
 /// annotations to implement all required setup in a general manner.
 template <typename P, typename I, typename V>
 class SparseTensorStorage : public SparseTensorStorageBase {
 public:
   /// Constructs sparse tensor storage scheme following the given
   /// per-rank dimension dense/sparse annotations.
   SparseTensorStorage(SparseTensor *tensor, bool *sparsity)
-      : sizes(tensor->getSizes()), pointers(sizes.size()),
-        indices(sizes.size()) {
+      : sizes(tensor->getSizes()), pointers(getRank()), indices(getRank()) {
     // Provide hints on capacity.
     // TODO: needs fine-tuning based on sparsity
-    values.reserve(tensor->getElements().size());
-    for (uint64_t d = 0, s = 1, rank = sizes.size(); d < rank; d++) {
-      s *= tensor->getSizes()[d];
+    uint64_t nnz = tensor->getElements().size();
+    values.reserve(nnz);
+    for (uint64_t d = 0, s = 1, rank = getRank(); d < rank; d++) {
+      s *= sizes[d];
       if (sparsity[d]) {
         pointers[d].reserve(s + 1);
         indices[d].reserve(s);
         s = 1;
       }
     }
     // Then setup the tensor.
-    traverse(tensor, sparsity, 0, tensor->getElements().size(), 0);
+    traverse(tensor, sparsity, 0, nnz, 0);
   }
 
   virtual ~SparseTensorStorage() {}
 
+  uint64_t getRank() const { return sizes.size(); }
+
   uint64_t getDimSize(uint64_t d) override { return sizes[d]; }
+
+  // Partially specialize these three methods based on template types.
   void getPointers(std::vector<P> **out, uint64_t d) override {
     *out = &pointers[d];
   }
   void getIndices(std::vector<I> **out, uint64_t d) override {
     *out = &indices[d];
   }
   void getValues(std::vector<V> **out) override { *out = &values; }
 
 private:
   /// Initializes sparse tensor storage scheme from a memory-resident
   /// representation of an external sparse tensor. This method prepares
   /// the pointers and indices arrays under the given per-rank dimension
   /// dense/sparse annotations.
   void traverse(SparseTensor *tensor, bool *sparsity, uint64_t lo, uint64_t hi,
                 uint64_t d) {
     const std::vector<Element> &elements = tensor->getElements();
     // Once dimensions are exhausted, insert the numerical values.
-    if (d == sizes.size()) {
+    if (d == getRank()) {
       values.push_back(lo < hi ? elements[lo].value : 0.0);
       return;
     }
     // Prepare a sparse pointer structure at this dimension.
     if (sparsity[d] && pointers[d].empty())
       pointers[d].push_back(0);
     // Visit all elements in this interval.
     uint64_t full = 0;
private:
   std::vector<uint64_t> sizes; // per-rank dimension sizes
   std::vector<std::vector<P>> pointers;
   std::vector<std::vector<I>> indices;
   std::vector<V> values;
 };
 
 /// Templated reader.
 template <typename P, typename I, typename V>
-void *newSparseTensor(char *filename, bool *sparsity) {
+void *newSparseTensor(char *filename, bool *sparsity, uint64_t size) {
   uint64_t idata[64];
   SparseTensor *t = static_cast<SparseTensor *>(openTensorC(filename, idata));
+  assert(size == t->getRank()); // sparsity array must match rank
   SparseTensorStorageBase *tensor =
       new SparseTensorStorage<P, I, V>(t, sparsity);
   delete t;
   return tensor;
 }
 
 /// Helper to convert string to lower case.
 static char *toLower(char *token) {

 enum TypeEnum : uint64_t { kF64 = 0, kF32 = 1, kU64 = 2, kU32 = 3 };
 
 void *newSparseTensor(char *filename, bool *abase, bool *adata, uint64_t aoff,
                       uint64_t asize, uint64_t astride, uint64_t ptrTp,
                       uint64_t indTp, uint64_t valTp) {
   assert(astride == 1);
   bool *sparsity = abase + aoff;
   if (ptrTp == kU64 && indTp == kU64 && valTp == kF64)
-    return newSparseTensor<uint64_t, uint64_t, double>(filename, sparsity);
+    return newSparseTensor<uint64_t, uint64_t, double>(filename, sparsity,
+                                                       asize);
   if (ptrTp == kU64 && indTp == kU64 && valTp == kF32)
-    return newSparseTensor<uint64_t, uint64_t, float>(filename, sparsity);
+    return newSparseTensor<uint64_t, uint64_t, float>(filename, sparsity,
+                                                      asize);
   if (ptrTp == kU64 && indTp == kU32 && valTp == kF64)
-    return newSparseTensor<uint64_t, uint32_t, double>(filename, sparsity);
+    return newSparseTensor<uint64_t, uint32_t, double>(filename, sparsity,
+                                                       asize);
   if (ptrTp == kU64 && indTp == kU32 && valTp == kF32)
-    return newSparseTensor<uint64_t, uint32_t, float>(filename, sparsity);
+    return newSparseTensor<uint64_t, uint32_t, float>(filename, sparsity,
+                                                      asize);
   if (ptrTp == kU32 && indTp == kU64 && valTp == kF64)
-    return newSparseTensor<uint32_t, uint64_t, double>(filename, sparsity);
+    return newSparseTensor<uint32_t, uint64_t, double>(filename, sparsity,
+                                                       asize);
   if (ptrTp == kU32 && indTp == kU64 && valTp == kF32)
-    return newSparseTensor<uint32_t, uint64_t, float>(filename, sparsity);
+    return newSparseTensor<uint32_t, uint64_t, float>(filename, sparsity,
+                                                      asize);
   if (ptrTp == kU32 && indTp == kU32 && valTp == kF64)
-    return newSparseTensor<uint32_t, uint32_t, double>(filename, sparsity);
+    return newSparseTensor<uint32_t, uint32_t, double>(filename, sparsity,
+                                                       asize);
   if (ptrTp == kU32 && indTp == kU32 && valTp == kF32)
-    return newSparseTensor<uint32_t, uint32_t, float>(filename, sparsity);
+    return newSparseTensor<uint32_t, uint32_t, float>(filename, sparsity,
+                                                      asize);
   fputs("unsupported combination of types\n", stderr);
   exit(1);
 }
 
 uint64_t sparseDimSize(void *tensor, uint64_t d) {
   return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
 }