[mlir][sparse] refactoring SparseTensorUtils: (3 of 4) code-cleanup

[wrengr]

Summary

Previously, the SparseTensorUtils.cpp library contained a C++ core implementation, but hid it in an anonymous namespace and only exposed a C-API for accessing it. Now we are factoring out that C++ core into a standalone C++ library so that it can be used directly by downstream clients (per request of one such client). This refactoring has been decomposed into a stack of differentials in order to simplify the code review process, however the full stack of changes should be considered together.

• D133462: Part 1: split one file into several
• D133830: Part 2: Reorder chunks within files
• (this): Part 3: General code cleanup
• D133833: Part 4: Update documentation

This part performs some general code cleanup including:

• making more things const, especially for the targets of pointers
• using preincrement wherever possible (per LLVM style guide)
• adding messages to most assert statments.
• moving argument casting from the core function/method definitions to the CPP wrappers

Depends On D133830

[aartbik]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

78

This idiomatic change really become more personal taste than anything else (and yes, I am aware that for general iterators there is a big advantage, but not for simple integers). But since you seem to feel strongly about it... ;-)

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

780–781

we don't use this kind of style for if-return imlicit else construct

[wrengr]

rebase

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

780–781

I don't follow what you mean

[wrengr]

Rebasing to parent D133462 at Diff 463979.

[wrengr]

git clang-format

[wrengr]

rebase

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

 ///   ({i}, a[i])
 /// and a rank-5 tensor element like
 ///   ({i,j,k,l,m}, a[i,j,k,l,m])
 /// We use pointer to a shared index pool rather than e.g. a direct
 /// vector since that (1) reduces the per-element memory footprint, and
 /// (2) centralizes the memory reservation and (re)allocation to one place.
 template <typename V>
 struct Element final {
-  Element(uint64_t *ind, V val) : indices(ind), value(val){};
-  uint64_t *indices; // pointer into shared index pool
+  Element(const uint64_t *ind, V val) : indices(ind), value(val){};
+  const uint64_t *indices; // pointer into shared index pool
   V value;
 };
 
 /// The type of callback functions which receive an element.  We avoid
 /// packaging the coordinates and value together as an `Element` object
 /// because this helps keep code somewhat cleaner.
 template <typename V>
 using ElementConsumer =
public:
   /// fully permuted coordinate scheme.
   ///
   /// Precondition: `dimSizes` and `perm` must be valid for `rank`.
   static SparseTensorCOO<V> *newSparseTensorCOO(uint64_t rank,
                                                 const uint64_t *dimSizes,
                                                 const uint64_t *perm,
                                                 uint64_t capacity = 0) {
     std::vector<uint64_t> permsz(rank);
-    for (uint64_t r = 0; r < rank; r++) {
+    for (uint64_t r = 0; r < rank; ++r) {
       assert(dimSizes[r] > 0 && "Dimension size zero has trivial storage");
       permsz[perm[r]] = dimSizes[r];
     }
     return new SparseTensorCOO<V>(permsz, capacity);
   }
 
   /// Get the rank of the tensor.
   uint64_t getRank() const { return dimSizes.size(); }
 
   /// Getter for the dimension-sizes array.
   const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
 
   /// Getter for the elements array.
   const std::vector<Element<V>> &getElements() const { return elements; }
 
   /// Adds element as indices and value.
   void add(const std::vector<uint64_t> &ind, V val) {
     assert(!iteratorLocked && "Attempt to add() after startIterator()");
-    uint64_t *base = indices.data();
+    const uint64_t *base = indices.data();
     uint64_t size = indices.size();
     uint64_t rank = getRank();
     assert(ind.size() == rank && "Element rank mismatch");
-    for (uint64_t r = 0; r < rank; r++) {
+    for (uint64_t r = 0; r < rank; ++r) {
       assert(ind[r] < dimSizes[r] && "Index is too large for the dimension");
       indices.push_back(ind[r]);
     }
     // This base only changes if indices were reallocated. In that case, we
     // need to correct all previous pointers into the vector. Note that this
     // only happens if we did not set the initial capacity right, and then only
     // for every internal vector reallocation (which with the doubling rule
     // should only incur an amortized linear overhead).
-    uint64_t *newBase = indices.data();
+    const uint64_t *newBase = indices.data();
     if (newBase != base) {
-      for (uint64_t i = 0, n = elements.size(); i < n; i++)
+      for (uint64_t i = 0, n = elements.size(); i < n; ++i)
         elements[i].indices = newBase + (elements[i].indices - base);
       base = newBase;
     }
     // Add element as (pointer into shared index pool, value) pair.
     elements.emplace_back(base + size, val);
   }
 
   /// Sorts elements lexicographically by index.
   void sort() {
     assert(!iteratorLocked && "Attempt to sort() after startIterator()");
     // TODO: we may want to cache an `isSorted` bit, to avoid
     // unnecessary/redundant sorting.
     uint64_t rank = getRank();
     std::sort(elements.begin(), elements.end(),
               [rank](const Element<V> &e1, const Element<V> &e2) {
-                for (uint64_t r = 0; r < rank; r++) {
+                for (uint64_t r = 0; r < rank; ++r) {
                   if (e1.indices[r] == e2.indices[r])
                     continue;
                   return e1.indices[r] < e2.indices[r];
                 }
                 return false;
               });
   }
 

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

 
 #include "mlir/ExecutionEngine/SparseTensor/Storage.h"
 
 #include <fstream>
 
 namespace mlir {
 namespace sparse_tensor {
 
+// TODO: benchmark whether to keep various methods inline vs moving them
+// off to the cpp file.
+
+// TODO: consider distinguishing separate classes for before vs
+// after reading the header; so as to statically avoid the need
+// to `assert(isValid())`.
+
 /// This class abstracts over the information stored in file headers,
 /// as well as providing the buffers and methods for parsing those headers.
 class SparseTensorFile final {
 public:
   enum class ValueKind {
     kInvalid = 0,
     kPattern = 1,
     kReal = 2,
     kInteger = 3,
     kComplex = 4,
     kUndefined = 5
   };
 
-  explicit SparseTensorFile(char *filename) : filename(filename) {
+  explicit SparseTensorFile(const char *filename) : filename(filename) {
     assert(filename && "Received nullptr for filename");
   }
 
   // Disallows copying, to avoid duplicating the `file` pointer.
   SparseTensorFile(const SparseTensorFile &) = delete;
   SparseTensorFile &operator=(const SparseTensorFile &) = delete;
 
   // This dtor tries to avoid leaking the `file`.  (Though it's better
public:
   /// Gets the dimension-sizes array.  The pointer itself is always
   /// valid; however, the values stored therein are only valid after
   /// parsing the header.
   const uint64_t *getDimSizes() const { return idata + 2; }
 
   /// Safely gets the size of the given dimension.  Is only valid
   /// after parsing the header.
   uint64_t getDimSize(uint64_t d) const {
-    assert(d < getRank());
+    assert(d < getRank() && "Dimension out of bounds");
     return idata[2 + d];
   }
 
   /// Asserts the shape subsumes the actual dimension sizes.  Is only
   /// valid after parsing the header.
   void assertMatchesShape(uint64_t rank, const uint64_t *shape) const;
 
 private:
   void readMMEHeader();
   void readExtFROSTTHeader();
 
   static constexpr int kColWidth = 1025;
-  const char *filename;
+  const char *const filename;
   FILE *file = nullptr;
   ValueKind valueKind_ = ValueKind::kInvalid;
   bool isSymmetric_ = false;
   uint64_t idata[512];
   char line[kColWidth];
 };
 
 namespace detail {

 }
 
 } // namespace detail
 
 /// Reads a sparse tensor with the given filename into a memory-resident
 /// sparse tensor in coordinate scheme.
 template <typename V>
 inline SparseTensorCOO<V> *
-openSparseTensorCOO(char *filename, uint64_t rank, const uint64_t *shape,
+openSparseTensorCOO(const char *filename, uint64_t rank, const uint64_t *shape,
                     const uint64_t *perm, PrimaryType valTp) {
   SparseTensorFile stfile(filename);
   stfile.openFile();
   stfile.readHeader();
   // Check tensor element type against the value type in the input file.
   SparseTensorFile::ValueKind valueKind = stfile.getValueKind();
   bool tensorIsInteger =
       (valTp >= PrimaryType::kI64 && valTp <= PrimaryType::kI8);
   bool tensorIsReal = (valTp >= PrimaryType::kF64 && valTp <= PrimaryType::kI8);
   if ((valueKind == SparseTensorFile::ValueKind::kReal && tensorIsInteger) ||
       (valueKind == SparseTensorFile::ValueKind::kComplex && tensorIsReal)) {
     MLIR_SPARSETENSOR_FATAL(
         "Tensor element type %d not compatible with values in file %s\n",
         static_cast<int>(valTp), filename);
   }
   stfile.assertMatchesShape(rank, shape);
   // Prepare sparse tensor object with per-dimension sizes
   // and the number of nonzeros as initial capacity.
   uint64_t nnz = stfile.getNNZ();
   auto *coo = SparseTensorCOO<V>::newSparseTensorCOO(rank, stfile.getDimSizes(),
                                                      perm, nnz);
   // Read all nonzero elements.
   std::vector<uint64_t> indices(rank);
-  for (uint64_t k = 0; k < nnz; k++) {
+  for (uint64_t k = 0; k < nnz; ++k) {
     char *linePtr = stfile.readLine();
-    for (uint64_t r = 0; r < rank; r++) {
+    for (uint64_t r = 0; r < rank; ++r) {
+      // Parse the 1-based index.
       uint64_t idx = strtoul(linePtr, &linePtr, 10);
-      // Add 0-based index.
+      // Add the 0-based index.
       indices[perm[r]] = idx - 1;
     }
     detail::readCOOValue(coo, indices, &linePtr, stfile.isPattern(),
                          stfile.isSymmetric() && indices[0] != indices[1]);
   }
   // Close the file and return tensor.
   stfile.closeFile();
   return coo;
 }
 
-/// Writes the sparse tensor to `dest` in extended FROSTT format.
+/// Writes the sparse tensor to `filename` in extended FROSTT format.
 template <typename V>
-inline void outSparseTensor(void *tensor, void *dest, bool sort) {
-  assert(tensor && dest);
-  auto coo = static_cast<SparseTensorCOO<V> *>(tensor);
-  if (sort)
-    coo->sort();
-  char *filename = static_cast<char *>(dest);
-  auto &dimSizes = coo->getDimSizes();
-  auto &elements = coo->getElements();
-  uint64_t rank = coo->getRank();
-  uint64_t nnz = elements.size();
+inline void writeExtFROSTT(const SparseTensorCOO<V> &coo,
+                           const char *filename) {
+  assert(filename && "Got nullptr for filename");
+  auto &dimSizes = coo.getDimSizes();
+  auto &elements = coo.getElements();
+  const uint64_t rank = coo.getRank();
+  const uint64_t nnz = elements.size();
   std::fstream file;
   file.open(filename, std::ios_base::out | std::ios_base::trunc);
   assert(file.is_open());
   file << "; extended FROSTT format\n" << rank << " " << nnz << std::endl;
-  for (uint64_t r = 0; r < rank - 1; r++)
+  for (uint64_t r = 0; r < rank - 1; ++r)
     file << dimSizes[r] << " ";
   file << dimSizes[rank - 1] << std::endl;
-  for (uint64_t i = 0; i < nnz; i++) {
+  for (uint64_t i = 0; i < nnz; ++i) {
     auto &idx = elements[i].indices;
-    for (uint64_t r = 0; r < rank; r++)
+    for (uint64_t r = 0; r < rank; ++r)
       file << (idx[r] + 1) << " ";
     file << elements[i].value << std::endl;
   }
   file.flush();
   file.close();
   assert(file.good());
 }
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_FILE_H

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

 #include "mlir/ExecutionEngine/SparseTensor/COO.h"
 #include "mlir/ExecutionEngine/SparseTensor/CheckedMul.h"
 #include "mlir/ExecutionEngine/SparseTensor/Enums.h"
 #include "mlir/ExecutionEngine/SparseTensor/ErrorHandling.h"
 
 namespace mlir {
 namespace sparse_tensor {
 
+//===----------------------------------------------------------------------===//
 // This forward decl is sufficient to split `SparseTensorStorageBase` into
 // its own header, but isn't sufficient for `SparseTensorStorage` to join it.
 template <typename V>
 class SparseTensorEnumeratorBase;
 
+// These macros ensure consistent error messages, without risk of incuring
+// an additional method call to do so.
+#define ASSERT_VALID_DIM(d)                                                    \
+  assert(d < getRank() && "Dimension index is out of bounds");
+#define ASSERT_COMPRESSED_DIM(d)                                               \
+  assert(isCompressedDim(d) && "Dimension is not compressed");
+#define ASSERT_DENSE_DIM(d) assert(isDenseDim(d) && "Dimension is not dense");
+
 /// Abstract base class for `SparseTensorStorage<P,I,V>`.  This class
 /// takes responsibility for all the `<P,I,V>`-independent aspects
 /// of the tensor (e.g., shape, sparsity, permutation).  In addition,
 /// we use function overloading to implement "partial" method
 /// specialization, which the C-API relies on to catch type errors
 /// arising from our use of opaque pointers.
 class SparseTensorStorageBase {
 protected:
public:
   /// Get the rank of the tensor.
   uint64_t getRank() const { return dimSizes.size(); }
 
   /// Getter for the dimension-sizes array, in storage-order.
   const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
 
   /// Safely lookup the size of the given (storage-order) dimension.
   uint64_t getDimSize(uint64_t d) const {
-    assert(d < getRank());
+    ASSERT_VALID_DIM(d);
     return dimSizes[d];
   }
 
   /// Getter for the "reverse" permutation, which maps this object's
   /// storage-order to the tensor's semantic-order.
   const std::vector<uint64_t> &getRev() const { return rev; }
 
   /// Getter for the dimension-types array, in storage-order.
   const std::vector<DimLevelType> &getDimTypes() const { return dimTypes; }
 
   /// Safely check if the (storage-order) dimension uses dense storage.
   bool isDenseDim(uint64_t d) const {
-    assert(d < getRank());
+    ASSERT_VALID_DIM(d);
     return dimTypes[d] == DimLevelType::kDense;
   }
 
   /// Safely check if the (storage-order) dimension uses compressed storage.
   bool isCompressedDim(uint64_t d) const {
-    assert(d < getRank());
+    ASSERT_VALID_DIM(d);
     switch (dimTypes[d]) {
     case DimLevelType::kCompressed:
     case DimLevelType::kCompressedNu:
     case DimLevelType::kCompressedNo:
     case DimLevelType::kCompressedNuNo:
       return true;
     default:
       return false;
     }
   }
 
   /// Safely check if the (storage-order) dimension uses singleton storage.
   bool isSingletonDim(uint64_t d) const {
-    assert(d < getRank());
+    ASSERT_VALID_DIM(d);
     switch (dimTypes[d]) {
     case DimLevelType::kSingleton:
     case DimLevelType::kSingletonNu:
     case DimLevelType::kSingletonNo:
     case DimLevelType::kSingletonNuNo:
       return true;
     default:
       return false;
     }
   }
 
   /// Safely check if the (storage-order) dimension is ordered.
   bool isOrderedDim(uint64_t d) const {
-    assert(d < getRank());
+    ASSERT_VALID_DIM(d);
     switch (dimTypes[d]) {
     case DimLevelType::kCompressedNo:
     case DimLevelType::kCompressedNuNo:
     case DimLevelType::kSingletonNo:
     case DimLevelType::kSingletonNuNo:
       return false;
     default:
       return true;
     }
   }
 
   /// Safely check if the (storage-order) dimension is unique.
   bool isUniqueDim(uint64_t d) const {
-    assert(d < getRank());
+    ASSERT_VALID_DIM(d);
     switch (dimTypes[d]) {
     case DimLevelType::kCompressedNu:
     case DimLevelType::kCompressedNuNo:
     case DimLevelType::kSingletonNu:
     case DimLevelType::kSingletonNuNo:
       return false;
     default:
       return true;
#define DECL_EXPINSERT(VNAME, V)                                               \
   FOREVERY_V(DECL_EXPINSERT)
 #undef DECL_EXPINSERT
 
   /// Finishes insertion.
   virtual void endInsert() = 0;
 
 private:
   const std::vector<uint64_t> dimSizes;
-  std::vector<uint64_t> rev;
+  std::vector<uint64_t> rev; // conceptually `const`
   const std::vector<DimLevelType> dimTypes;
 };
 
+//===----------------------------------------------------------------------===//
 // This forward decl is necessary for defining `SparseTensorStorage`,
 // but isn't sufficient for splitting it off.
 template <typename P, typename I, typename V>
 class SparseTensorEnumerator;
 
 /// A memory-resident sparse tensor using a storage scheme based on
 /// per-dimension sparse/dense annotations. This data structure provides a
 /// bufferized form of a sparse tensor type. In contrast to generating setup
public:
   newSparseTensor(uint64_t rank, const uint64_t *shape, const uint64_t *perm,
                   const DimLevelType *sparsity,
                   const SparseTensorStorageBase *source);
 
   ~SparseTensorStorage() final = default;
 
   /// Partially specialize these getter methods based on template types.
   void getPointers(std::vector<P> **out, uint64_t d) final {
-    assert(d < getRank());
+    ASSERT_VALID_DIM(d);
     *out = &pointers[d];
   }
   void getIndices(std::vector<I> **out, uint64_t d) final {
-    assert(d < getRank());
+    ASSERT_VALID_DIM(d);
     *out = &indices[d];
   }
   void getValues(std::vector<V> **out) final { *out = &values; }
 
   /// Partially specialize lexicographical insertions based on template types.
   void lexInsert(const uint64_t *cursor, V val) final {
     // First, wrap up pending insertion path.
     uint64_t diff = 0;
void expInsert(uint64_t *cursor, V *values, bool *filled, uint64_t *added,
                  uint64_t count) final {
     if (count == 0)
       return;
     // Sort.
     std::sort(added, added + count);
     // Restore insertion path for first insert.
     const uint64_t lastDim = getRank() - 1;
     uint64_t index = added[0];
+    assert(filled[index] && "added index is not filled");
     cursor[lastDim] = index;
     lexInsert(cursor, values[index]);
-    assert(filled[index]);
     values[index] = 0;
     filled[index] = false;
     // Subsequent insertions are quick.
-    for (uint64_t i = 1; i < count; i++) {
+    for (uint64_t i = 1; i < count; ++i) {
       assert(index < added[i] && "non-lexicographic insertion");
       index = added[i];
+      assert(filled[index] && "added index is not filled");
       cursor[lastDim] = index;
       insPath(cursor, lastDim, added[i - 1] + 1, values[index]);
-      assert(filled[index]);
       values[index] = 0;
       filled[index] = false;
     }
   }
 
   /// Finalizes lexicographic insertions.
   void endInsert() final {
     if (values.empty())
public:
   }
 
 private:
   /// Appends an arbitrary new position to `pointers[d]`.  This method
   /// checks that `pos` is representable in the `P` type; however, it
   /// does not check that `pos` is semantically valid (i.e., larger than
   /// the previous position and smaller than `indices[d].capacity()`).
   void appendPointer(uint64_t d, uint64_t pos, uint64_t count = 1) {
-    assert(isCompressedDim(d));
+    ASSERT_COMPRESSED_DIM(d);
     assert(pos <= std::numeric_limits<P>::max() &&
            "Pointer value is too large for the P-type");
     pointers[d].insert(pointers[d].end(), count, static_cast<P>(pos));
   }
 
   /// Appends index `i` to dimension `d`, in the semantically general
   /// sense.  For non-dense dimensions, that means appending to the
   /// `indices[d]` array, checking that `i` is representable in the `I`
   /// type; however, we do not verify other semantic requirements (e.g.,
   /// that `i` is in bounds for `dimSizes[d]`, and not previously occurring
   /// in the same segment).  For dense dimensions, this method instead
   /// appends the appropriate number of zeros to the `values` array,
   /// where `full` is the number of "entries" already written to `values`
   /// for this segment (aka one after the highest index previously appended).
   void appendIndex(uint64_t d, uint64_t full, uint64_t i) {
     if (isCompressedDim(d) || isSingletonDim(d)) {
       assert(i <= std::numeric_limits<I>::max() &&
              "Index value is too large for the I-type");
       indices[d].push_back(static_cast<I>(i));
     } else { // Dense dimension.
-      assert(isDenseDim(d));
+      ASSERT_DENSE_DIM(d);
       assert(i >= full && "Index was already filled");
       if (i == full)
         return; // Short-circuit, since it'll be a nop.
       if (d + 1 == getRank())
         values.insert(values.end(), i - full, 0);
       else
         finalizeSegment(d + 1, 0, i - full);
     }
   }
 
   /// Writes the given coordinate to `indices[d][pos]`.  This method
   /// checks that `i` is representable in the `I` type; however, it
   /// does not check that `i` is semantically valid (i.e., in bounds
   /// for `dimSizes[d]` and not elsewhere occurring in the same segment).
   void writeIndex(uint64_t d, uint64_t pos, uint64_t i) {
-    assert(isCompressedDim(d));
+    ASSERT_COMPRESSED_DIM(d);
     // Subscript assignment to `std::vector` requires that the `pos`-th
     // entry has been initialized; thus we must be sure to check `size()`
     // here, instead of `capacity()` as would be ideal.
     assert(pos < indices[d].size() && "Index position is out of bounds");
     assert(i <= std::numeric_limits<I>::max() &&
            "Index value is too large for the I-type");
     indices[d][pos] = static_cast<I>(i);
   }
private:
   /// indices arrays under the given per-dimension dense/sparse annotations.
   ///
   /// Preconditions:
   /// (1) the `elements` must be lexicographically sorted.
   /// (2) the indices of every element are valid for `dimSizes` (equal rank
   ///     and pointwise less-than).
   void fromCOO(const std::vector<Element<V>> &elements, uint64_t lo,
                uint64_t hi, uint64_t d) {
-    uint64_t rank = getRank();
+    const uint64_t rank = getRank();
     assert(d <= rank && hi <= elements.size());
     // Once dimensions are exhausted, insert the numerical values.
     if (d == rank) {
       assert(lo < hi);
       values.push_back(elements[lo].value);
       return;
     }
     // Visit all elements in this interval.
     uint64_t full = 0;
     while (lo < hi) { // If `hi` is unchanged, then `lo < elements.size()`.
       // Find segment in interval with same index elements in this dimension.
-      uint64_t i = elements[lo].indices[d];
+      const uint64_t i = elements[lo].indices[d];
       uint64_t seg = lo + 1;
-      bool merge = isUniqueDim(d);
-      while (merge && seg < hi && elements[seg].indices[d] == i)
-        seg++;
+      if (isUniqueDim(d))
+        while (seg < hi && elements[seg].indices[d] == i)
+          ++seg;
       // Handle segment in interval for sparse or dense dimension.
       appendIndex(d, full, i);
       full = i + 1;
       fromCOO(elements, lo, seg, d + 1);
       // And move on to next segment in interval.
       lo = seg;
     }
     // Finalize the sparse pointer structure at this dimension.
     finalizeSegment(d, full);
   }
 
   /// Finalize the sparse pointer structure at this dimension.
   void finalizeSegment(uint64_t d, uint64_t full = 0, uint64_t count = 1) {
     if (count == 0)
       return; // Short-circuit, since it'll be a nop.
     if (isCompressedDim(d)) {
       appendPointer(d, indices[d].size(), count);
     } else if (isSingletonDim(d)) {
       return;
     } else { // Dense dimension.
-      assert(isDenseDim(d));
+      ASSERT_DENSE_DIM(d);
       const uint64_t sz = getDimSizes()[d];
       assert(sz >= full && "Segment is overfull");
       count = detail::checkedMul(count, sz - full);
       // For dense storage we must enumerate all the remaining coordinates
       // in this dimension (i.e., coordinates after the last non-zero
       // element), and either fill in their zero values or else recurse
       // to finalize some deeper dimension.
       if (d + 1 == getRank())
         values.insert(values.end(), count, 0);
       else
         finalizeSegment(d + 1, 0, count);
     }
   }
 
   /// 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++) {
+    const uint64_t rank = getRank();
+    assert(diff <= rank && "Dimension-diff is out of bounds");
+    for (uint64_t i = 0; i < rank - diff; ++i) {
       const uint64_t d = rank - i - 1;
       finalizeSegment(d, idx[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];
+    ASSERT_VALID_DIM(diff);
+    const uint64_t rank = getRank();
+    for (uint64_t d = diff; d < rank; ++d) {
+      const uint64_t i = cursor[d];
       appendIndex(d, top, i);
       top = 0;
       idx[d] = i;
     }
     values.push_back(val);
   }
 
   /// Finds the lexicographic differing dimension.
   uint64_t lexDiff(const uint64_t *cursor) const {
-    for (uint64_t r = 0, rank = getRank(); r < rank; r++)
+    const uint64_t rank = getRank();
+    for (uint64_t r = 0; 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;
   }
 
   // Allow `SparseTensorEnumerator` to access the data-members (to avoid
   // the cost of virtual-function dispatch in inner loops), without
   // making them public to other client code.
   friend class SparseTensorEnumerator<P, I, V>;
 
   std::vector<std::vector<P>> pointers;
   std::vector<std::vector<I>> indices;
   std::vector<V> values;
   std::vector<uint64_t> idx; // index cursor for lexicographic insertion.
 };
 
+#undef ASSERT_COMPRESSED_DIM
+#undef ASSERT_VALID_DIM
+
+//===----------------------------------------------------------------------===//
 /// A (higher-order) function object for enumerating the elements of some
 /// `SparseTensorStorage` under a permutation.  That is, the `forallElements`
 /// method encapsulates the loop-nest for enumerating the elements of
 /// the source tensor (in whatever order is best for the source tensor),
 /// and applies a permutation to the coordinates/indices before handing
 /// each element to the callback.  A single enumerator object can be
 /// freely reused for several calls to `forallElements`, just so long
 /// as each call is sequential with respect to one another.
public:
   SparseTensorEnumeratorBase(const SparseTensorStorageBase &tensor,
                              uint64_t rank, const uint64_t *perm)
       : src(tensor), permsz(src.getRev().size()), reord(getRank()),
         cursor(getRank()) {
     assert(perm && "Received nullptr for permutation");
     assert(rank == getRank() && "Permutation rank mismatch");
     const auto &rev = src.getRev();           // source-order -> semantic-order
     const auto &dimSizes = src.getDimSizes(); // in source storage-order
-    for (uint64_t s = 0; s < rank; s++) {     // `s` source storage-order
+    for (uint64_t s = 0; s < rank; ++s) {     // `s` source storage-order
       uint64_t t = perm[rev[s]];              // `t` target-order
       reord[s] = t;
       permsz[t] = dimSizes[s];
     }
   }
 
   virtual ~SparseTensorEnumeratorBase() = default;
 

 
 protected:
   const SparseTensorStorageBase &src;
   std::vector<uint64_t> permsz; // in target order.
   std::vector<uint64_t> reord;  // source storage-order -> target order.
   std::vector<uint64_t> cursor; // in target order.
 };
 
+//===----------------------------------------------------------------------===//
 template <typename P, typename I, typename V>
 class SparseTensorEnumerator final : public SparseTensorEnumeratorBase<V> {
   using Base = SparseTensorEnumeratorBase<V>;
+  using StorageImpl = SparseTensorStorage<P, I, V>;
 
 public:
   /// Constructs an enumerator with the given permutation for mapping
   /// the semantic-ordering of dimensions to the desired target-ordering.
   ///
   /// Precondition: `perm` must be valid for `rank`.
-  SparseTensorEnumerator(const SparseTensorStorage<P, I, V> &tensor,
-                         uint64_t rank, const uint64_t *perm)
+  SparseTensorEnumerator(const StorageImpl &tensor, uint64_t rank,
+                         const uint64_t *perm)
       : Base(tensor, rank, perm) {}
 
   ~SparseTensorEnumerator() final = default;
 
   void forallElements(ElementConsumer<V> yield) final {
     forallElements(yield, 0, 0);
   }
 
 private:
   /// The recursive component of the public `forallElements`.
   void forallElements(ElementConsumer<V> yield, uint64_t parentPos,
                       uint64_t d) {
     // Recover the `<P,I,V>` type parameters of `src`.
-    const auto &src =
-        static_cast<const SparseTensorStorage<P, I, V> &>(this->src);
+    const auto &src = static_cast<const StorageImpl &>(this->src);
     if (d == Base::getRank()) {
       assert(parentPos < src.values.size() &&
              "Value position is out of bounds");
       // TODO: <https://github.com/llvm/llvm-project/issues/54179>
       yield(this->cursor, src.values[parentPos]);
     } else if (src.isCompressedDim(d)) {
       // Look up the bounds of the `d`-level segment determined by the
       // `d-1`-level position `parentPos`.
       const std::vector<P> &pointersD = src.pointers[d];
       assert(parentPos + 1 < pointersD.size() &&
              "Parent pointer position is out of bounds");
       const uint64_t pstart = static_cast<uint64_t>(pointersD[parentPos]);
       const uint64_t pstop = static_cast<uint64_t>(pointersD[parentPos + 1]);
       // Loop-invariant code for looking up the `d`-level coordinates/indices.
       const std::vector<I> &indicesD = src.indices[d];
       assert(pstop <= indicesD.size() && "Index position is out of bounds");
       uint64_t &cursorReordD = this->cursor[this->reord[d]];
-      for (uint64_t pos = pstart; pos < pstop; pos++) {
+      for (uint64_t pos = pstart; pos < pstop; ++pos) {
         cursorReordD = static_cast<uint64_t>(indicesD[pos]);
         forallElements(yield, pos, d + 1);
       }
     } else if (src.isSingletonDim(d)) {
       MLIR_SPARSETENSOR_FATAL("unsupported dimension level type");
     } else { // Dense dimension.
-      assert(src.isDenseDim(d));
+      assert(src.isDenseDim(d)); // TODO: reuse the ASSERT_DENSE_DIM message
       const uint64_t sz = src.getDimSizes()[d];
       const uint64_t pstart = parentPos * sz;
       uint64_t &cursorReordD = this->cursor[this->reord[d]];
-      for (uint64_t i = 0; i < sz; i++) {
+      for (uint64_t i = 0; i < sz; ++i) {
         cursorReordD = i;
         forallElements(yield, pstart + i, d + 1);
       }
     }
   }
 };
 
+//===----------------------------------------------------------------------===//
 /// Statistics regarding the number of nonzero subtensors in
 /// a source tensor, for direct sparse=>sparse conversion a la
 /// <https://arxiv.org/abs/2001.02609>.
 ///
 /// N.B., this class stores references to the parameters passed to
 /// the constructor; thus, objects of this class must not outlive
 /// those parameters.
 class SparseTensorNNZ final {
private:
   const std::vector<DimLevelType> &dimTypes;
   std::vector<std::vector<uint64_t>> nnz;
 };
 
 //===----------------------------------------------------------------------===//
 // Definitions of the ctors and factories of `SparseTensorStorage<P,I,V>`.
 
 namespace detail {
-
-// TODO: try to unify this with `SparseTensorFile::assertMatchesShape`
-// which is used by `openSparseTensorCOO`.  It's easy enough to resolve
-// the `std::vector` vs pointer mismatch for `dimSizes`; but it's trickier
-// to resolve the presence/absence of `perm` (without introducing extra
-// overhead), so perhaps the code duplication is unavoidable.
-//
 /// Asserts that the `dimSizes` (in target-order) under the `perm` (mapping
 /// semantic-order to target-order) are a refinement of the desired `shape`
 /// (in semantic-order).
 ///
 /// Precondition: `perm` and `shape` must be valid for `rank`.
-inline void assertPermutedSizesMatchShape(const std::vector<uint64_t> &dimSizes,
-                                          uint64_t rank, const uint64_t *perm,
-                                          const uint64_t *shape) {
-  assert(perm && shape);
-  assert(rank == dimSizes.size() && "Rank mismatch");
-  for (uint64_t r = 0; r < rank; r++)
-    assert((shape[r] == 0 || shape[r] == dimSizes[perm[r]]) &&
-           "Dimension size mismatch");
-}
-
+void assertPermutedSizesMatchShape(const std::vector<uint64_t> &dimSizes,
+                                   const uint64_t *shape);
 } // namespace detail
 
 template <typename P, typename I, typename V>
 SparseTensorStorage<P, I, V> *SparseTensorStorage<P, I, V>::newSparseTensor(
     uint64_t rank, const uint64_t *shape, const uint64_t *perm,
     const DimLevelType *sparsity, SparseTensorCOO<V> *coo) {
-  SparseTensorStorage<P, I, V> *n = nullptr;
   if (coo) {
     const auto &coosz = coo->getDimSizes();
+#ifndef NDEBUG
     detail::assertPermutedSizesMatchShape(coosz, rank, perm, shape);
-    n = new SparseTensorStorage<P, I, V>(coosz, perm, sparsity, coo);
-  } else {
-    std::vector<uint64_t> permsz(rank);
-    for (uint64_t r = 0; r < rank; r++) {
-      assert(shape[r] > 0 && "Dimension size zero has trivial storage");
-      permsz[perm[r]] = shape[r];
-    }
-    // We pass the null `coo` to ensure we select the intended constructor.
-    n = new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, coo);
-  }
-  return n;
+#endif
+    return new SparseTensorStorage<P, I, V>(coosz, perm, sparsity, coo);
+  }
+  // else
+  for (uint64_t r = 0; r < rank; ++r) {
+  return new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, coo);
 }
 
 template <typename P, typename I, typename V>
 SparseTensorStorage<P, I, V> *SparseTensorStorage<P, I, V>::newSparseTensor(
     uint64_t rank, const uint64_t *shape, const uint64_t *perm,
     const DimLevelType *sparsity, const SparseTensorStorageBase *source) {
   assert(source && "Got nullptr for source");
   SparseTensorEnumeratorBase<V> *enumerator;
   source->newEnumerator(&enumerator, rank, perm);
   const auto &permsz = enumerator->permutedSizes();
+#ifndef NDEBUG
   detail::assertPermutedSizesMatchShape(permsz, rank, perm, shape);
+#endif
   auto *tensor =
       new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, *source);
   delete enumerator;
   return tensor;
 }
 
 template <typename P, typename I, typename V>
 SparseTensorStorage<P, I, V>::SparseTensorStorage(
     const std::vector<uint64_t> &dimSizes, const uint64_t *perm,
     const DimLevelType *sparsity, SparseTensorCOO<V> *coo)
     : SparseTensorStorage(dimSizes, perm, sparsity) {
   // Provide hints on capacity of pointers and indices.
   // TODO: needs much fine-tuning based on actual sparsity; currently
   //       we reserve pointer/index space based on all previous dense
   //       dimensions, which works well up to first sparse dim; but
   //       we should really use nnz and dense/sparse distribution.
   bool allDense = true;
   uint64_t sz = 1;
-  for (uint64_t r = 0, rank = getRank(); r < rank; r++) {
+  for (uint64_t r = 0, rank = getRank(); r < rank; ++r) {
     if (isCompressedDim(r)) {
       // TODO: Take a parameter between 1 and `dimSizes[r]`, and multiply
       // `sz` by that before reserving. (For now we just use 1.)
       pointers[r].reserve(sz + 1);
       pointers[r].push_back(0);
       indices[r].reserve(sz);
       sz = 1;
       allDense = false;
     } else if (isSingletonDim(r)) {
       indices[r].reserve(sz);
       sz = 1;
       allDense = false;
     } else { // Dense dimension.
-      assert(isDenseDim(r));
+      ASSERT_DENSE_DIM(r);
       sz = detail::checkedMul(sz, getDimSizes()[r]);
     }
   }
   // Then assign contents from coordinate scheme tensor if provided.
   if (coo) {
     // Ensure both preconditions of `fromCOO`.
     assert(coo->getDimSizes() == getDimSizes() && "Tensor size mismatch");
     coo->sort();
SparseTensorStorage<P, I, V>::SparseTensorStorage(
   SparseTensorEnumeratorBase<V> *enumerator;
   tensor.newEnumerator(&enumerator, getRank(), perm);
   {
     // Initialize the statistics structure.
     SparseTensorNNZ nnz(getDimSizes(), getDimTypes());
     nnz.initialize(*enumerator);
     // Initialize "pointers" overhead (and allocate "indices", "values").
     uint64_t parentSz = 1; // assembled-size (not dimension-size) of `r-1`.
-    for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
+    for (uint64_t rank = getRank(), r = 0; r < rank; ++r) {
       if (isCompressedDim(r)) {
         pointers[r].reserve(parentSz + 1);
         pointers[r].push_back(0);
         uint64_t currentPos = 0;
         nnz.forallIndices(r, [this, &currentPos, r](uint64_t n) {
           currentPos += n;
           appendPointer(r, currentPos);
         });
for (uint64_t rank = getRank(), r = 0; r < rank; ++r) {
       if (isCompressedDim(r))
         indices[r].resize(parentSz, 0);
     }
     values.resize(parentSz, 0); // Both allocate and zero-initialize.
   }
   // The yieldPos loop
   enumerator->forallElements([this](const std::vector<uint64_t> &ind, V val) {
     uint64_t parentSz = 1, parentPos = 0;
-    for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
+    for (uint64_t rank = getRank(), r = 0; r < rank; ++r) {
       if (isCompressedDim(r)) {
         // If `parentPos == parentSz` then it's valid as an array-lookup;
         // however, it's semantically invalid here since that entry
         // does not represent a segment of `indices[r]`.  Moreover, that
         // entry must be immutable for `assembledSize` to remain valid.
         assert(parentPos < parentSz && "Pointers position is out of bounds");
         const uint64_t currentPos = pointers[r][parentPos];
         // This increment won't overflow the `P` type, since it can't
         // exceed the original value of `pointers[r][parentPos+1]`
         // which was already verified to be within bounds for `P`
         // when it was written to the array.
         pointers[r][parentPos]++;
         writeIndex(r, currentPos, ind[r]);
         parentPos = currentPos;
       } else if (isSingletonDim(r)) {
         // the new parentPos equals the old parentPos.
       } else { // Dense dimension.
-        assert(isDenseDim(r));
+        ASSERT_DENSE_DIM(r);
         parentPos = parentPos * getDimSizes()[r] + ind[r];
       }
       parentSz = assembledSize(parentSz, r);
     }
     assert(parentPos < values.size() && "Value position is out of bounds");
     values[parentPos] = val;
   });
   // No longer need the enumerator, so we'll delete it ASAP.
   delete enumerator;
   // The finalizeYieldPos loop
-  for (uint64_t parentSz = 1, rank = getRank(), r = 0; r < rank; r++) {
+  for (uint64_t parentSz = 1, rank = getRank(), r = 0; r < rank; ++r) {
     if (isCompressedDim(r)) {
       assert(parentSz == pointers[r].size() - 1 &&
              "Actual pointers size doesn't match the expected size");
       // Can't check all of them, but at least we can check the last one.
       assert(pointers[r][parentSz - 1] == pointers[r][parentSz] &&
              "Pointers got corrupted");
       // TODO: optimize this by using `memmove` or similar.
-      for (uint64_t n = 0; n < parentSz; n++) {
+      for (uint64_t n = 0; n < parentSz; ++n) {
         const uint64_t parentPos = parentSz - n;
         pointers[r][parentPos] = pointers[r][parentPos - 1];
       }
       pointers[r][0] = 0;
     }
     parentSz = assembledSize(parentSz, r);
   }
 }
 
 } // namespace sparse_tensor
 } // namespace mlir
 
+#undef ASSERT_DENSE_DIM
+
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H

mlir/lib/ExecutionEngine/SparseTensor/File.cpp

 // to simplify benchmarking, testing, and debugging MLIR code operating on
 // sparse tensors.  However, the provided functionality is **not** part of
 // core MLIR itself.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/SparseTensor/File.h"
 
-#ifdef MLIR_SPARSETENSOR_DEFINE_FUNCTIONS // We are building this library
-
 #include <cctype>
 #include <cstring>
 
 using namespace mlir::sparse_tensor;
 
 /// Opens the file for reading.
 void SparseTensorFile::openFile() {
   if (file)
void SparseTensorFile::readHeader() {
   assert(isValid() && "Failed to read the header");
 }
 
 /// Asserts the shape subsumes the actual dimension sizes.  Is only
 /// valid after parsing the header.
 void SparseTensorFile::assertMatchesShape(uint64_t rank,
                                           const uint64_t *shape) const {
   assert(rank == getRank() && "Rank mismatch");
-  for (uint64_t r = 0; r < rank; r++)
+  for (uint64_t r = 0; r < rank; ++r)
     assert((shape[r] == 0 || shape[r] == idata[2 + r]) &&
            "Dimension size mismatch");
 }
 
 /// Helper to convert string to lower case.
 static inline char *toLower(char *token) {
   for (char *c = token; *c; ++c)
     *c = tolower(*c);
while (true) {
     readLine();
     if (line[0] != '#')
       break;
   }
   // Next line contains RANK and NNZ.
   if (sscanf(line, "%" PRIu64 "%" PRIu64 "\n", idata, idata + 1) != 2)
     MLIR_SPARSETENSOR_FATAL("Cannot find metadata in %s\n", filename);
   // Followed by a line with the dimension sizes (one per rank).
-  for (uint64_t r = 0; r < idata[0]; r++)
+  for (uint64_t r = 0; r < idata[0]; ++r)
     if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1)
       MLIR_SPARSETENSOR_FATAL("Cannot find dimension size %s\n", filename);
   readLine(); // end of line
   // The FROSTT format does not define the data type of the nonzero elements.
   valueKind_ = ValueKind::kUndefined;
 }
-
-#endif // MLIR_SPARSETENSOR_DEFINE_FUNCTIONS

mlir/lib/ExecutionEngine/SparseTensor/NNZ.cpp

 // to simplify benchmarking, testing, and debugging MLIR code operating on
 // sparse tensors.  However, the provided functionality is **not** part of
 // core MLIR itself.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/SparseTensor/Storage.h"
 
-#ifdef MLIR_SPARSETENSOR_DEFINE_FUNCTIONS // We are building this library
-
 using namespace mlir::sparse_tensor;
 
 SparseTensorNNZ::SparseTensorNNZ(const std::vector<uint64_t> &dimSizes,
                                  const std::vector<DimLevelType> &sparsity)
     : dimSizes(dimSizes), dimTypes(sparsity), nnz(getRank()) {
   assert(dimSizes.size() == dimTypes.size() && "Rank mismatch");
   bool uncompressed = true;
   (void)uncompressed;
default:
       MLIR_SPARSETENSOR_FATAL("unsupported dimension level type");
     }
     sz = detail::checkedMul(sz, dimSizes[r]);
   }
 }
 
 void SparseTensorNNZ::forallIndices(uint64_t stopDim,
                                     SparseTensorNNZ::NNZConsumer yield) const {
-  assert(stopDim < getRank() && "Stopping-dimension is out of bounds");
+  assert(stopDim < getRank() && "Dimension out of bounds");
   assert(dimTypes[stopDim] == DimLevelType::kCompressed &&
          "Cannot look up non-compressed dimensions");
   forallIndices(yield, stopDim, 0, 0);
 }
 
 void SparseTensorNNZ::add(const std::vector<uint64_t> &ind) {
   uint64_t parentPos = 0;
-  for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
+  for (uint64_t rank = getRank(), r = 0; r < rank; ++r) {
     if (dimTypes[r] == DimLevelType::kCompressed)
       nnz[r][parentPos]++;
     parentPos = parentPos * dimSizes[r] + ind[r];
   }
 }
 
 void SparseTensorNNZ::forallIndices(SparseTensorNNZ::NNZConsumer yield,
                                     uint64_t stopDim, uint64_t parentPos,
                                     uint64_t d) const {
   assert(d <= stopDim);
   if (d == stopDim) {
     assert(parentPos < nnz[d].size() && "Cursor is out of range");
     yield(nnz[d][parentPos]);
   } else {
     const uint64_t sz = dimSizes[d];
     const uint64_t pstart = parentPos * sz;
     for (uint64_t i = 0; i < sz; i++)
       forallIndices(yield, stopDim, pstart + i, d + 1);
   }
 }
-
-#endif // MLIR_SPARSETENSOR_DEFINE_FUNCTIONS

mlir/lib/ExecutionEngine/SparseTensor/Storage.cpp

 //===- StorageBase.cpp - TACO-flavored sparse tensor representation -------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains method definitions for `SparseTensorStorageBase`.
 // In particular we want to ensure that the default implementations of
 // the "partial method specialization" trick aren't inline (since there's
 // no benefit).  Though this also helps ensure that we avoid weak-vtables:
 // <https://llvm.org/docs/CodingStandards.html#provide-a-virtual-method-anchor-for-classes-in-headers>
 //
+// (This file also contains the definition of `assertPermutedSizesMatchShape`
+// which is used by `SparseTensorStorage` factories.)
+//
 // This file is part of the lightweight runtime support library for sparse
 // tensor manipulations.  The functionality of the support library is meant
 // to simplify benchmarking, testing, and debugging MLIR code operating on
 // sparse tensors.  However, the provided functionality is **not** part of
 // core MLIR itself.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/SparseTensor/Storage.h"
 
-#ifdef MLIR_SPARSETENSOR_DEFINE_FUNCTIONS // We are building this library
-
 using namespace mlir::sparse_tensor;
 
 SparseTensorStorageBase::SparseTensorStorageBase(
     const std::vector<uint64_t> &dimSizes, const uint64_t *perm,
     const DimLevelType *sparsity)
     : dimSizes(dimSizes), rev(getRank()),
       dimTypes(sparsity, sparsity + getRank()) {
-  assert(perm && sparsity);
+  assert(perm && "Got nullptr for permutation");
+  assert(sparsity && "Got nullptr for sparsity");
   const uint64_t rank = getRank();
   // Validate parameters.
   assert(rank > 0 && "Trivial shape is unsupported");
-  for (uint64_t r = 0; r < rank; r++) {
+  for (uint64_t r = 0; r < rank; ++r) {
     assert(dimSizes[r] > 0 && "Dimension size zero has trivial storage");
     assert((isDenseDim(r) || isCompressedDim(r) || isSingletonDim(r)) &&
            "Unsupported DimLevelType");
   }
   // Construct the "reverse" (i.e., inverse) permutation.
-  for (uint64_t r = 0; r < rank; r++)
+  // TODO: should move this computation off to the codegen
+  for (uint64_t r = 0; r < rank; ++r)
     rev[perm[r]] = r;
 }
 
 // Helper macro for generating error messages when some
 // `SparseTensorStorage<P,I,V>` is cast to `SparseTensorStorageBase`
 // and then the wrong "partial method specialization" is called.
 #define FATAL_PIV(NAME)                                                        \
   MLIR_SPARSETENSOR_FATAL("<P,I,V> type mismatch for: " #NAME);
void SparseTensorStorageBase::expInsert(uint64_t *, V *, bool *, uint64_t *, \
                                           uint64_t) {                          \
     FATAL_PIV("expInsert" #VNAME);                                             \
   }
 FOREVERY_V(IMPL_EXPINSERT)
 #undef IMPL_EXPINSERT
 
 #undef FATAL_PIV
 
-#endif // MLIR_SPARSETENSOR_DEFINE_FUNCTIONS
+// TODO: try to unify this with `SparseTensorFile::assertMatchesShape`
+// (which is used by `openSparseTensorCOO`).  It's easy enough to resolve
+// the `std::vector` vs pointer mismatch for `dimSizes`; but it's trickier
+// to resolve the presence/absence of `perm` (without introducing extra
+// overhead), so perhaps the code duplication is unavoidable?
+void mlir::sparse_tensor::detail::assertPermutedSizesMatchShape(
+    const std::vector<uint64_t> &dimSizes, uint64_t rank, const uint64_t *perm,
+    const uint64_t *shape) {
+  assert(perm && "Got nullptr for permutation");
+  assert(shape && "Got nullptr for shape");
+  assert(rank == dimSizes.size() && "Rank mismatch");
+  for (uint64_t d = 0; d < rank; ++d)
+    assert((shape[d] == 0 || shape[d] == dimSizes[perm[d]]) &&
+           "Dimension size mismatch");
+}

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 //      to interact with MLIR compiler-generated code.
 //
 // In both cases (I) and (II), the SparseTensorStorage format is externally
 // only visible as an opaque pointer.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/SparseTensorUtils.h"
+
+#ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
+
 #include "mlir/ExecutionEngine/SparseTensor/COO.h"
 #include "mlir/ExecutionEngine/SparseTensor/ErrorHandling.h"
 #include "mlir/ExecutionEngine/SparseTensor/File.h"
 #include "mlir/ExecutionEngine/SparseTensor/Storage.h"
 
-#ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
-
-#include <algorithm>
-#include <cassert>
-#include <cctype>
-#include <cstdio>
-#include <cstdlib>
-#include <cstring>
-#include <fstream>
-#include <functional>
-#include <iostream>
-#include <limits>
 #include <numeric>
 
 using namespace mlir::sparse_tensor;
 
 namespace {
 
 /// Initializes sparse tensor from an external COO-flavored format.
 template <typename V>
 static SparseTensorStorage<uint64_t, uint64_t, V> *
-toMLIRSparseTensor(uint64_t rank, uint64_t nse, uint64_t *shape, V *values,
-                   uint64_t *indices, uint64_t *perm, uint8_t *sparse) {
-  const DimLevelType *sparsity = (DimLevelType *)(sparse);
+toMLIRSparseTensor(uint64_t rank, uint64_t nse, const uint64_t *shape,
+                   const V *values, const uint64_t *indices,
+                   const uint64_t *perm, const DimLevelType *sparsity) {
 #ifndef NDEBUG
   // Verify that perm is a permutation of 0..(rank-1).
   std::vector<uint64_t> order(perm, perm + rank);
   std::sort(order.begin(), order.end());
   for (uint64_t i = 0; i < rank; ++i)
     if (i != order[i])
       MLIR_SPARSETENSOR_FATAL("Not a permutation of 0..%" PRIu64 "\n", rank);
 
#endif
   auto *tensor = SparseTensorStorage<uint64_t, uint64_t, V>::newSparseTensor(
       rank, shape, perm, sparsity, coo);
   delete coo;
   return tensor;
 }
 
 /// Converts a sparse tensor to an external COO-flavored format.
 template <typename V>
-static void fromMLIRSparseTensor(void *tensor, uint64_t *pRank, uint64_t *pNse,
-                                 uint64_t **pShape, V **pValues,
-                                 uint64_t **pIndices) {
-  assert(tensor);
-  auto sparseTensor =
-      static_cast<SparseTensorStorage<uint64_t, uint64_t, V> *>(tensor);
-  uint64_t rank = sparseTensor->getRank();
+static void
+fromMLIRSparseTensor(const SparseTensorStorage<uint64_t, uint64_t, V> *tensor,
+                     uint64_t *pRank, uint64_t *pNse, uint64_t **pShape,
+                     V **pValues, uint64_t **pIndices) {
+  assert(tensor && "Received nullptr for tensor");
+  uint64_t rank = tensor->getRank();
   std::vector<uint64_t> perm(rank);
   std::iota(perm.begin(), perm.end(), 0);
-  SparseTensorCOO<V> *coo = sparseTensor->toCOO(perm.data());
+  SparseTensorCOO<V> *coo = tensor->toCOO(perm.data());
 
   const std::vector<Element<V>> &elements = coo->getElements();
   uint64_t nse = elements.size();
 
   uint64_t *shape = new uint64_t[rank];
   for (uint64_t i = 0; i < rank; i++)
     shape[i] = coo->getDimSizes()[i];
 

 }
 
 void endInsert(void *tensor) {
   return static_cast<SparseTensorStorageBase *>(tensor)->endInsert();
 }
 
 #define IMPL_OUTSPARSETENSOR(VNAME, V)                                         \
   void outSparseTensor##VNAME(void *coo, void *dest, bool sort) {              \
-    return outSparseTensor<V>(coo, dest, sort);                                \
+    assert(coo && "Got nullptr for COO object");                               \
+    auto &coo_ = *static_cast<SparseTensorCOO<V> *>(coo);                      \
+    if (sort)                                                                  \
+      coo_.sort();                                                             \
+    return writeExtFROSTT(coo_, static_cast<char *>(dest));                    \
   }
 FOREVERY_V(IMPL_OUTSPARSETENSOR)
 #undef IMPL_OUTSPARSETENSOR
 
 void delSparseTensor(void *tensor) {
   delete static_cast<SparseTensorStorageBase *>(tensor);
 }
 
void readSparseTensorShape(char *filename, std::vector<uint64_t> *out) {
   stfile.readHeader();
   stfile.closeFile();
   const uint64_t rank = stfile.getRank();
   const uint64_t *dimSizes = stfile.getDimSizes();
   out->reserve(rank);
   out->assign(dimSizes, dimSizes + rank);
 }
 
+// We can't use `static_cast` here because `DimLevelType` is an enum-class.
 // TODO: generalize beyond 64-bit indices.
 #define IMPL_CONVERTTOMLIRSPARSETENSOR(VNAME, V)                               \
   void *convertToMLIRSparseTensor##VNAME(                                      \
       uint64_t rank, uint64_t nse, uint64_t *shape, V *values,                 \
       uint64_t *indices, uint64_t *perm, uint8_t *sparse) {                    \
     return toMLIRSparseTensor<V>(rank, nse, shape, values, indices, perm,      \
-                                 sparse);                                      \
+                                 reinterpret_cast<DimLevelType *>(sparse));    \
   }
 FOREVERY_V(IMPL_CONVERTTOMLIRSPARSETENSOR)
 #undef IMPL_CONVERTTOMLIRSPARSETENSOR
 
 // TODO: Currently, values are copied from SparseTensorStorage to
 // SparseTensorCOO, then to the output.  We may want to reduce the number
 // of copies.
 //
 // TODO: generalize beyond 64-bit indices, no dim ordering, all dimensions
 // compressed
 #define IMPL_CONVERTFROMMLIRSPARSETENSOR(VNAME, V)                             \
   void convertFromMLIRSparseTensor##VNAME(void *tensor, uint64_t *pRank,       \
                                           uint64_t *pNse, uint64_t **pShape,   \
                                           V **pValues, uint64_t **pIndices) {  \
-    fromMLIRSparseTensor<V>(tensor, pRank, pNse, pShape, pValues, pIndices);   \
+    fromMLIRSparseTensor<V>(                                                   \
+        static_cast<SparseTensorStorage<uint64_t, uint64_t, V> *>(tensor),     \
+        pRank, pNse, pShape, pValues, pIndices);                               \
   }
 FOREVERY_V(IMPL_CONVERTFROMMLIRSPARSETENSOR)
 #undef IMPL_CONVERTFROMMLIRSPARSETENSOR
 
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS