[mlir][sparse] Using the name "dimSizes" more consistently

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 /// A version of `operator*` on `uint64_t` which checks for overflows.
 static inline uint64_t checkedMul(uint64_t lhs, uint64_t rhs) {
   assert((lhs == 0 || rhs <= std::numeric_limits<uint64_t>::max() / lhs) &&
          "Integer overflow");
   return lhs * rhs;
 }
 
 // TODO: adjust this so it can be used by `openSparseTensorCOO` too.
-// That version doesn't have the permutation, and the `sizes` are
+// That version doesn't have the permutation, and the `dimSizes` are
 // a pointer/C-array rather than `std::vector`.
 //
-/// Asserts that the `sizes` (in target-order) under the `perm` (mapping
+/// 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`.
 static inline void
-assertPermutedSizesMatchShape(const std::vector<uint64_t> &sizes, uint64_t rank,
-                              const uint64_t *perm, const uint64_t *shape) {
+assertPermutedSizesMatchShape(const std::vector<uint64_t> &dimSizes,
+                              uint64_t rank, const uint64_t *perm,
+                              const uint64_t *shape) {
   assert(perm && shape);
-  assert(rank == sizes.size() && "Rank mismatch");
+  assert(rank == dimSizes.size() && "Rank mismatch");
   for (uint64_t r = 0; r < rank; r++)
-    assert((shape[r] == 0 || shape[r] == sizes[perm[r]]) &&
+    assert((shape[r] == 0 || shape[r] == dimSizes[perm[r]]) &&
            "Dimension size mismatch");
 }
 
 /// 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])

 /// A memory-resident sparse tensor in coordinate scheme (collection of
 /// 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 final {
 public:
-  SparseTensorCOO(const std::vector<uint64_t> &szs, uint64_t capacity)
-      : sizes(szs) {
+  SparseTensorCOO(const std::vector<uint64_t> &dimSizes, uint64_t capacity)
+      : dimSizes(dimSizes) {
     if (capacity) {
       elements.reserve(capacity);
       indices.reserve(capacity * getRank());
     }
   }
 
   /// 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();
     uint64_t size = indices.size();
     uint64_t rank = getRank();
-    assert(rank == ind.size());
+    assert(ind.size() == rank && "Element rank mismatch");
     for (uint64_t r = 0; r < rank; r++) {
-      assert(ind[r] < sizes[r]); // within bounds
+      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();
std::sort(elements.begin(), elements.end(),
                   if (e1.indices[r] == e2.indices[r])
                     continue;
                   return e1.indices[r] < e2.indices[r];
                 }
                 return false;
               });
   }
 
-  /// Returns rank.
-  uint64_t getRank() const { return sizes.size(); }
+  /// Get the rank of the tensor.
+  uint64_t getRank() const { return dimSizes.size(); }
 
-  /// Getter for sizes array.
-  const std::vector<uint64_t> &getSizes() const { return sizes; }
+  /// Getter for the dimension-sizes array.
+  const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
 
-  /// Getter for elements array.
+  /// Getter for the elements array.
   const std::vector<Element<V>> &getElements() const { return elements; }
 
   /// Switch into iterator mode.
   void startIterator() {
     iteratorLocked = true;
     iteratorPos = 0;
   }
 
   /// Get the next element.
   const Element<V> *getNext() {
     assert(iteratorLocked && "Attempt to getNext() before startIterator()");
     if (iteratorPos < elements.size())
       return &(elements[iteratorPos++]);
     iteratorLocked = false;
     return nullptr;
   }
 
   /// Factory method. Permutes the original dimensions according to
   /// the given ordering and expects subsequent add() calls to honor
   /// that same ordering for the given indices. The result is a
   /// fully permuted coordinate scheme.
   ///
-  /// Precondition: `sizes` and `perm` must be valid for `rank`.
+  /// Precondition: `dimSizes` and `perm` must be valid for `rank`.
   static SparseTensorCOO<V> *newSparseTensorCOO(uint64_t rank,
-                                                const uint64_t *sizes,
+                                                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++) {
-      assert(sizes[r] > 0 && "Dimension size zero has trivial storage");
-      permsz[perm[r]] = sizes[r];
+      assert(dimSizes[r] > 0 && "Dimension size zero has trivial storage");
+      permsz[perm[r]] = dimSizes[r];
     }
     return new SparseTensorCOO<V>(permsz, capacity);
   }
 
 private:
-  const std::vector<uint64_t> sizes; // per-dimension sizes
-  std::vector<Element<V>> elements;  // all COO elements
-  std::vector<uint64_t> indices;     // shared index pool
+  const std::vector<uint64_t> dimSizes; // per-dimension sizes
+  std::vector<Element<V>> elements;     // all COO elements
+  std::vector<uint64_t> indices;        // shared index pool
   bool iteratorLocked = false;
   unsigned iteratorPos = 0;
 };
 
 // See <https://en.wikipedia.org/wiki/X_Macro>
 //
 // `FOREVERY_SIMPLEX_V` only specifies the non-complex `V` types, because
 // the ABI for complex types has compiler/architecture dependent complexities

 /// 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 {
 public:
   /// Constructs a new storage object.  The `perm` maps the tensor's
   /// semantic-ordering of dimensions to this object's storage-order.
-  /// The `szs` and `sparsity` arrays are already in storage-order.
+  /// The `dimSizes` and `sparsity` arrays are already in storage-order.
   ///
-  /// Precondition: `perm` and `sparsity` must be valid for `szs.size()`.
-  SparseTensorStorageBase(const std::vector<uint64_t> &szs,
+  /// Precondition: `perm` and `sparsity` must be valid for `dimSizes.size()`.
+  SparseTensorStorageBase(const std::vector<uint64_t> &dimSizes,
                           const uint64_t *perm, const DimLevelType *sparsity)
-      : dimSizes(szs), rev(getRank()),
+      : dimSizes(dimSizes), rev(getRank()),
         dimTypes(sparsity, sparsity + getRank()) {
     assert(perm && sparsity);
     const uint64_t rank = getRank();
     // Validate parameters.
     assert(rank > 0 && "Trivial shape is unsupported");
     for (uint64_t r = 0; r < rank; r++) {
       assert(dimSizes[r] > 0 && "Dimension size zero has trivial storage");
       assert((dimTypes[r] == DimLevelType::kDense ||

 /// and annotations to implement all required setup in a general manner.
 template <typename P, typename I, typename V>
 class SparseTensorStorage final : public SparseTensorStorageBase {
   /// Private constructor to share code between the other constructors.
   /// Beware that the object is not necessarily guaranteed to be in a
   /// valid state after this constructor alone; e.g., `isCompressedDim(d)`
   /// doesn't entail `!(pointers[d].empty())`.
   ///
-  /// Precondition: `perm` and `sparsity` must be valid for `szs.size()`.
-  SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
-                      const DimLevelType *sparsity)
-      : SparseTensorStorageBase(szs, perm, sparsity), pointers(getRank()),
+  /// Precondition: `perm` and `sparsity` must be valid for `dimSizes.size()`.
+  SparseTensorStorage(const std::vector<uint64_t> &dimSizes,
+                      const uint64_t *perm, const DimLevelType *sparsity)
+      : SparseTensorStorageBase(dimSizes, perm, sparsity), pointers(getRank()),
         indices(getRank()), idx(getRank()) {}
 
 public:
   /// Constructs a sparse tensor storage scheme with the given dimensions,
   /// permutation, and per-dimension dense/sparse annotations, using
   /// the coordinate scheme tensor for the initial contents if provided.
   ///
-  /// Precondition: `perm` and `sparsity` must be valid for `szs.size()`.
-  SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
-                      const DimLevelType *sparsity, SparseTensorCOO<V> *coo)
-      : SparseTensorStorage(szs, perm, sparsity) {
+  /// Precondition: `perm` and `sparsity` must be valid for `dimSizes.size()`.
+  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++) {
       if (isCompressedDim(r)) {
-        // TODO: Take a parameter between 1 and `sizes[r]`, and multiply
+        // 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 { // Dense dimension.
         sz = checkedMul(sz, getDimSizes()[r]);
       }
     }
     // Then assign contents from coordinate scheme tensor if provided.
     if (coo) {
       // Ensure both preconditions of `fromCOO`.
-      assert(coo->getSizes() == getDimSizes() && "Tensor size mismatch");
+      assert(coo->getDimSizes() == getDimSizes() && "Tensor size mismatch");
       coo->sort();
       // Now actually insert the `elements`.
       const std::vector<Element<V>> &elements = coo->getElements();
       uint64_t nnz = elements.size();
       values.reserve(nnz);
       fromCOO(elements, 0, nnz, 0);
     } else if (allDense) {
       values.resize(sz, 0);
     }
   }
 
   /// Constructs a sparse tensor storage scheme with the given dimensions,
   /// permutation, and per-dimension dense/sparse annotations, using
   /// the given sparse tensor for the initial contents.
   ///
   /// Preconditions:
-  /// * `perm` and `sparsity` must be valid for `szs.size()`.
+  /// * `perm` and `sparsity` must be valid for `dimSizes.size()`.
   /// * The `tensor` must have the same value type `V`.
-  SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
-                      const DimLevelType *sparsity,
+  SparseTensorStorage(const std::vector<uint64_t> &dimSizes,
+                      const uint64_t *perm, const DimLevelType *sparsity,
                       const SparseTensorStorageBase &tensor);
 
   ~SparseTensorStorage() final override = default;
 
   /// Partially specialize these getter methods based on template types.
   void getPointers(std::vector<P> **out, uint64_t d) final override {
     assert(d < getRank());
     *out = &pointers[d];
public:
   /// permutation as is desired for the new sparse tensor storage.
   ///
   /// Precondition: `shape`, `perm`, and `sparsity` must be valid for `rank`.
   static 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->getSizes();
+      const auto &coosz = coo->getDimSizes();
       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];
       }
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 `sizes[d]`, and not previously occurring
+  /// 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)) {
       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(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 `sizes[d]` and not elsewhere occurring in the same segment).
+  /// 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));
     // 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");
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-dimension dense/sparse annotations.
   ///
   /// Preconditions:
   /// (1) the `elements` must be lexicographically sorted.
-  /// (2) the indices of every element are valid for `sizes` (equal rank
+  /// (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();
     assert(d <= rank && hi <= elements.size());
     // Once dimensions are exhausted, insert the numerical values.
     if (d == rank) {
       assert(lo < hi);
public:
   /// * the `tensor` must have the same `V` value type.
   /// * `perm` must be valid for `rank`.
   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 stg-order -> semantic-order
-    const auto &sizes = src.getDimSizes(); // in source storage-order
-    for (uint64_t s = 0; s < rank; s++) {  // `s` source storage-order
-      uint64_t t = perm[rev[s]];           // `t` target-order
+    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
+      uint64_t t = perm[rev[s]];              // `t` target-order
       reord[s] = t;
-      permsz[t] = sizes[s];
+      permsz[t] = dimSizes[s];
     }
   }
 
   virtual ~SparseTensorEnumeratorBase() = default;
 
   // We disallow copying to help avoid leaking the `src` reference.
   // (In addition to avoiding the problem of slicing.)
   SparseTensorEnumeratorBase(const SparseTensorEnumeratorBase &) = delete;

 /// those parameters.
 class SparseTensorNNZ final {
 public:
   /// Allocate the statistics structure for the desired sizes and
   /// sparsity (in the target tensor's storage-order).  This constructor
   /// does not actually populate the statistics, however; for that see
   /// `initialize`.
   ///
-  /// Precondition: `szs` must not contain zeros.
-  SparseTensorNNZ(const std::vector<uint64_t> &szs,
+  /// Precondition: `dimSizes` must not contain zeros.
+  SparseTensorNNZ(const std::vector<uint64_t> &dimSizes,
                   const std::vector<DimLevelType> &sparsity)
-      : dimSizes(szs), dimTypes(sparsity), nnz(getRank()) {
+      : dimSizes(dimSizes), dimTypes(sparsity), nnz(getRank()) {
     assert(dimSizes.size() == dimTypes.size() && "Rank mismatch");
     bool uncompressed = true;
     uint64_t sz = 1; // the product of all `dimSizes` strictly less than `r`.
     for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
       switch (dimTypes[r]) {
       case DimLevelType::kCompressed:
         assert(uncompressed &&
                "Multiple compressed layers not currently supported");
private:
   // All of these are in the target storage-order.
   const std::vector<uint64_t> &dimSizes;
   const std::vector<DimLevelType> &dimTypes;
   std::vector<std::vector<uint64_t>> nnz;
 };
 
 template <typename P, typename I, typename V>
 SparseTensorStorage<P, I, V>::SparseTensorStorage(
-    const std::vector<uint64_t> &szs, const uint64_t *perm,
+    const std::vector<uint64_t> &dimSizes, const uint64_t *perm,
     const DimLevelType *sparsity, const SparseTensorStorageBase &tensor)
-    : SparseTensorStorage(szs, perm, sparsity) {
+    : SparseTensorStorage(dimSizes, perm, sparsity) {
   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`.

 /// Writes the sparse tensor to extended FROSTT format.
 template <typename V>
 static 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 &sizes = coo->getSizes();
+  auto &dimSizes = coo->getDimSizes();
   auto &elements = coo->getElements();
   uint64_t rank = coo->getRank();
   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++)
-    file << sizes[r] << " ";
-  file << sizes[rank - 1] << std::endl;
+    file << dimSizes[r] << " ";
+  file << dimSizes[rank - 1] << std::endl;
   for (uint64_t i = 0; i < nnz; i++) {
     auto &idx = elements[i].indices;
     for (uint64_t r = 0; r < rank; r++)
       file << (idx[r] + 1) << " ";
     file << elements[i].value << std::endl;
   }
   file.flush();
   file.close();
static void fromMLIRSparseTensor(void *tensor, uint64_t *pRank, uint64_t *pNse,
   std::iota(perm.begin(), perm.end(), 0);
   SparseTensorCOO<V> *coo = sparseTensor->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->getSizes()[i];
+    shape[i] = coo->getDimSizes()[i];
 
   V *values = new V[nse];
   uint64_t *indices = new uint64_t[rank * nse];
 
   for (uint64_t i = 0, base = 0; i < nse; i++) {
     values[i] = elements[i].value;
     for (uint64_t j = 0; j < rank; j++)
       indices[base + j] = elements[i].indices[j];