[Mlir-commits] [mlir] 860cbeb - [mlir][sparse] add more asserts to sparse support lib

[Aart Bik]

Author: Aart Bik
Date: 2021-09-16T10:13:29-07:00
New Revision: 860cbeb15992a385d9a3fffb3f1a8443cc99b663

URL: https://github.com/llvm/llvm-project/commit/860cbeb15992a385d9a3fffb3f1a8443cc99b663
DIFF: https://github.com/llvm/llvm-project/commit/860cbeb15992a385d9a3fffb3f1a8443cc99b663.diff

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

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

Reviewed By: jsetoain

Differential Revision: https://reviews.llvm.org/D109861

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/SparseUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/SparseUtils.cpp b/mlir/lib/ExecutionEngine/SparseUtils.cpp
index 20b8da5c71b16..3497fe5b49b3f 100644
--- a/mlir/lib/ExecutionEngine/SparseUtils.cpp
+++ b/mlir/lib/ExecutionEngine/SparseUtils.cpp
@@ -87,7 +87,7 @@ 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);
   }
@@ -100,8 +100,6 @@ struct SparseTensorCOO {
   }
   /// 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.
@@ -135,7 +133,6 @@ struct SparseTensorCOO {
   }
   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
@@ -216,15 +213,22 @@ class SparseTensorStorage : public SparseTensorStorageBase {
 
   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; }
@@ -248,6 +252,7 @@ class SparseTensorStorage : public SparseTensorStorageBase {
       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;
   }
 
@@ -271,12 +276,15 @@ class SparseTensorStorage : public SparseTensorStorageBase {
     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;
@@ -312,7 +320,9 @@ class SparseTensorStorage : public SparseTensorStorageBase {
   /// 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.