[mlir][sparse] add asserts on reading in tensor data

[aartbik]

Summary

Rationale:
Since I made the argument that metadata helps with extra
verification checks, I better actually do that ;-)

[penpornk]

Thank you for the changes!

[aartbik]

In D95072#2511064, @penpornk wrote:

Thank you for the changes!

And thank you for the quick review. Next CL will be more elaborate, it will set up the sparse storage scheme with per-rank dense/sparse annotations for use in MLIR.
Elaborate stuff, but after that anyone can run the code much easier!

mlir/lib/ExecutionEngine/SparseUtils.cpp

 namespace {
 
 /// A sparse tensor element in coordinate scheme (value and indices).
 /// For example, a rank-1 vector element would look like
 ///   ({i}, a[i])
 /// and a rank-5 tensor element like
 ///   ({i,j,k,l,m}, a[i,j,k,l,m])
 struct Element {
-  Element(const std::vector<int64_t> &ind, double val)
+  Element(const std::vector<uint64_t> &ind, double val)
       : indices(ind), value(val){};
-  std::vector<int64_t> indices;
+  std::vector<uint64_t> indices;
   double value;
 };
 
 /// A memory-resident sparse tensor in coordinate scheme (collection of
 /// elements). This data structure is used to read a sparse tensor from
 /// external file 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).
 struct SparseTensor {
 public:
-  SparseTensor(int64_t capacity) : pos(0) { elements.reserve(capacity); }
+  SparseTensor(const std::vector<uint64_t> &szs, uint64_t capacity)
+      : sizes(szs), pos(0) {
+    elements.reserve(capacity);
+  }
   // Add element as indices and value.
-  void add(const std::vector<int64_t> &ind, double val) {
+  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(ind[r] < sizes[r]); // within bounds
     elements.emplace_back(Element(ind, val));
   }
   // Sort elements lexicographically by index.
   void sort() { std::sort(elements.begin(), elements.end(), lexOrder); }
   // Primitive one-time iteration.
   const Element &next() { return elements[pos++]; }
 
 private:
   // Returns true if indices of e1 < indices of e2.
   static bool lexOrder(const Element &e1, const Element &e2) {
     assert(e1.indices.size() == e2.indices.size());
     for (int64_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> elements;
   uint64_t pos;
 };
 
 /// Helper to convert string to lower case.
 static char *toLower(char *token) {
   for (char *c = token; *c; c++)
     *c = tolower(*c);
extern "C" void *openTensorC(char *filename, uint64_t *idata) {
   if (strstr(filename, ".mtx")) {
     readMMEHeader(file, filename, idata);
   } else if (strstr(filename, ".tns")) {
     readExtFROSTTHeader(file, filename, idata);
   } else {
     fprintf(stderr, "Unknown format %s\n", filename);
     exit(1);
   }
-  // Read all nonzero elements.
+  // Prepare sparse tensor object with per-rank dimension sizes
+  // and the number of nonzeros as initial capacity.
   uint64_t rank = idata[0];
   uint64_t nnz = idata[1];
-  SparseTensor *tensor = new SparseTensor(nnz);
-  std::vector<int64_t> indices(rank);
-  double value;
+  std::vector<uint64_t> indices(rank);
+  for (uint64_t r = 0; r < rank; r++)
+    indices[r] = idata[2 + r];
+  SparseTensor *tensor = new SparseTensor(indices, nnz);
+  // Read all nonzero elements.
   for (uint64_t k = 0; k < nnz; k++) {
     for (uint64_t r = 0; r < rank; r++) {
-      if (fscanf(file, "%" PRId64, &indices[r]) != 1) {
+      if (fscanf(file, "%" PRIu64, &indices[r]) != 1) {
         fprintf(stderr, "Cannot find next index in %s\n", filename);
         exit(1);
       }
       indices[r]--; // 0-based index
     }
+    double value;
     if (fscanf(file, "%lg\n", &value) != 1) {
       fprintf(stderr, "Cannot find next value in %s\n", filename);
       exit(1);
     }
     tensor->add(indices, value);
   }
   // Close the file and return sorted tensor.
   fclose(file);