[mlir][sparse] Factoring out `finalizeSegment` and (generic) `appendIndex`

[wrengr]

Summary

This change introduces two new methods: finalizeSegment and appendIndex; and removes three old methods: endDim, appendCurrentPointer, appendIndex. The two new methods better encapsulate their algorithms, thus allowing to remove repetitious code in several other places.

Depends On D122435

[wrengr]

cleaning up a comment

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

856

this should be at the place of endDim, just to make the delta in this review easier to review.

typically when refactoring
(1) changes, remain in place, only shows delta
(2) move code around, in separate CL, no changes to moved code

910

I liked having appendPointer and appendIndex next to each other, but this pushes it way down. Let's keep it at the original place, which also makes the delta for the review a bit easier than when code moves around like this

[wrengr]

Rebasing for changes to D122435, and addressing comments

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

568

ah, an overflow check, please add that as string at end to make semantics more clear

(and other parts of MLIR will break long before this as I found out earlier ;-)

852

I like this generalization with count as parameter to all appending methods. Just checking, I assume that insert(x.end()) and x.push_back() have similar STL performance?

865

This was of course there before, but for uint64_t I, this is a nop check.
Shall we do the comparison in even higher precision to detect running out of that space
(very unlikely of course, but just observing while here)

956

it would be even shorter if we let appendIndex return the next logical value for full so that this would be

full = appendIndex(d, full, i)

and have all the compressed/dense logic hidden inside the method.

wdyt?

973–980

One of the motivations given for this change is to avoid repeating code. However, in the original code, we only had endDim() pushing zeros into the values array, and I was very used to that invariant. Now both appendIndex and finalizeSegment have this logic. It feels like this part should actuall ycall into appenIndex instead somehow. Do you see a chance to ensure we only add to values in one place?

[wrengr]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

568

Will do. Since I do this check in a couple places over future CLs, I think I'll factor it out into a top-level (inline) function too.

Since the division is expensive (and the error condition rare), does mlir have a cpp flag for "I need extra help debugging, so please turn on even more assertions than usual"?

852

I'm not entirely sure? I've found it notoriously difficult to get much information about stl implementations, since there are multiple vendors/versions of it. The specification https://en.cppreference.com/w/cpp/container/vector/insert requires this usage to be linear in count, but says nothing about the constant factors. I'd imagine in the worst case the overhead when count=1 is just to set up a for-loop that only runs the once (since push_back would also have the same check to prove to itself that the capacity doesn't need resizing). For the cases where count is greater than one, this should be a clear win over calling push_back within a loop: since it need only check the capacity once rather than on every iteration, and since it might possibly do tricks like the following...

Ideally the implementation would use something like memset/memcpy to handle this, since those are usually optimized for architecture-specific issues re alignment and vectorization. But I'm pretty new to C++, so I'm not sure if a random stl vendor could be relied on to do that sort of thing. Once we have some benchmarking tools in place, we can always check to see whether it'd be worth it to roll our own implementation.

For this particular function, the count generalization is just so finalizeSegment doesn't need to repeat the P-validity checking logic; whereas the other call-sites in this and future CLs always use the count=1 default. If vec.insert(vec.end(), 1, val) ends up being notably slower than vec.push_back(val) then I can always go back to the version of the code where this was inlined into finalizeSegment.

865

Hopefully when I=uint64_t the assertion will get optimized away entirely (since the compiler should be able to detect that it's always true). Lifting both operands to a higher precision wouldn't change anything, it'd still always be true, just with a few more zeros at front

973–980

finalizeSegment only ever pushes a bunch of zeros to values and/or calls appendPointer, exactly like endDim used to. Logically finalizeSegment is doing the same thing as endDim did, it's just doing so more generally and more efficiently. Re generality: endDim could only handle the case of completely empty dimensions (initial full is zero), whereas finalizeSegment can handle cases where the initial full isn't zero, as occurs in endPath and the part of appendIndex that handles insPath. Re efficiency: endDim had a bunch of for-loops to repeatedly call the same function (not depending on the induction variable), whereas finalizeSegment accumulates those repeated calls into its count parameter— so we only call finalizeSegment at most rank *times* (rather than at most rank *depth*), and when we hit the basis case all the count-many calls to push_back get combined into a single insert call (thus saving overhead re checking for capacity resizing, improving memory locality, etc).

Whereas appendIndex will write an arbitrary index to the indices array (if compressed), which is a completely different thing. The main change here over the previous version of appendIndex is just that now it can also handle the case where the dimension has dense storage: in which case we insert the appropriate number of zero values. Previously this situation was handled by insPath, but I think it makes more sense as part of appendIndex since the zero insertion is logically being done for the same reasons as writing to the indices array. As for code duplication, appendIndex calls finalizeSegment to perform the zero padding (just as previously insPath called endDim to perform the zero padding).

There is only a slight amount of code duplication, namely the one line for handling the d+1==rank case. This occurs because I adjusted the inductive hypothesis, so both finalizeSegment and appendIndex only take in d such that d<rank (matching the convention of most if not all the other methods); as opposed to the previous situation where d<=rank was allowed. While the duplication could be removed by reverting the IH to d<=rank, I think the d<rank IH is the right way to go. What would it mean for client code to call finalizeSegment(rank) ? There is no rank dimension, so there are no segments of that dimension to be finalized, so I can't think of anything sensible that it could mean.

[wrengr]

Addressing comments

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

956

I like the idea, though I'm not entirely sure of the semantics; i.e., how to document the semantics so as to retain the appropriate encapsulation. I've taken a first pass at it, so let me know if you think it'd be better phrased another way.

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

568

I am not aware of such a flag, but others may know.

But here, rather than passing by reference, how about returning the result so we can do

sz = checkedMulAssign(sz, sizes[r]);

but also more general in a result that is not the same as "lhs".

other = checkedMulAssign(sz, sizes[r]);

852

Thanks for checking. I am not overly worried, and we can indeed always profile and improve later.

865

Ah, yes, never mind. I was hoping to somehow detect "about to wrap-around" from i to i+1, but here we already deal with the new index value, so that will not work.

882

this is technically index set [full..full+i], but with space for the new current element right?
I also would not say cardinality of i+1, since that is only when we start at 0?

972

please put brackets around the subtraction just for my sanity ;-)

[wrengr]

addressing comments

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

568

Will do. Though it's worth noting that I phrased it like operator*= rather than operator* for a couple reasons: (1) I feel like it helps clarify/allow the lhs!=0 precondition, whereas for operator* I think it'd be more appropriate to check whether lhs==0 rather than assuming it; (2) it's really just intended for the case where we're accumulating a bunch of things and thus overflow becomes more likely, rather than being a more general solution for detecting overflows elsewhere.

882

Why would it be [full..full+i]? If full means the number of entries already written, then before calling the method we've already written full many entries (for indices [0..full-1]); within the method we write another i-full many entries (for indices [full..i-1]), thus together we've written a total of [0..full-1]++[full..i-1] == [0..i-1].

In any case, this weird off-by-one relationship —between the i parameter to appendIndex vs the return value which becomes the full for the next iteration of the while-loop in fromCOO— is what I meant when I complained about the semantics being unclear / hard to explain. If full really meant the number of entries/zeros written, then we would return full + (i - full) == i as the new number of entries/zeros written (which matches the [0..i-1] derived above); but we don't, we return i+1 instead. And if full really meant the next unwritten index (since it starts at 0 in the fromCOO while-loop, yet we haven't written the 0th index yet), then the compressed case should also return i+1 (since the ith index has now been written); but it doesn't, it returns full instead. This inconsistency was already there in the original code, so maybe it's a bug rather than intentional; I was just making sure to retain the previous behavior through this refactoring is all

[wrengr]

Re-adjusting the return value of appendIndex.

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

882

Fwiw, I just checked and having the compressed case return i+1 passes all our tests. Which makes sense since appendIndex ignores the full for compressed dimensions anyways. Of course, the uniformity of that means that we could just have full = i + 1; at the callsite in fromCOO, rather than returning anything from this method.

Conversely, having the dense case return i causes the assertion at the end of toCOO to fail for three tests, and six others fail to FileCheck. (All this was expected, since I played around with that before.)

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

568

Yeah, I figured there would be limited use cases but the reference parameter may still surprise the future reader. thanks for the change

882

Ai, I hate it when typos obscure my point ;-). I obviously meant that we wrote out [full, i], as in the original for loop, so I understand that afterwards the full [0..i-1] are written, but one method invocation may write less than that if full was > 0 to start with. But since you reverted back to no return, point is mood.

[wrengr]

rebase

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 /// and annotations to implement all required setup in a general manner.
 template <typename P, typename I, typename V>
 class SparseTensorStorage : protected EnumerableSparseTensorStorage<V> {
 private:
   /// Partial constructor for sharing code between the public
   /// constructors.  This only does the bare minimum setup for the
   /// other constructors, thus the resulting object is not guaranteed
   /// to be in a valid state.  In particular `isCompressedDim` (and
-  /// hence `appendCurrentPointer` etc) will not work.
+  /// hence `appendPointer` etc) will not work.
   ///
   /// Precondition: `perm` and `sparsity` must be valid for `szs.size()`.
   SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
                       const DimLevelType *sparsity)
       : sizes(szs), rev(getRank()), idx(getRank()), pointers(getRank()),
         indices(getRank()) {
     // Validate parameters.
-    for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
+    uint64_t rank = getRank();
+    assert(rank > 0 && "Rank zero is unsupported");
+    for (uint64_t r = 0; r < rank; r++) {
       assert(sizes[r] > 0 && "Dimension size zero has trivial storage");
       assert((sparsity[r] == DimLevelType::kDense ||
               sparsity[r] == DimLevelType::kCompressed) &&
              "Unsupported DimLevelType");
     }
     // Store "reverse" permutation.
     for (uint64_t rank = getRank(), r = 0; r < rank; r++)
       rev[perm[r]] = r;
   }
 
 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> *tensor)
       : SparseTensorStorage(szs, perm, sparsity) {
     // Provide hints on capacity of pointers and indices.
     // TODO: needs fine-tuning based on sparsity
     bool allDense = true;
     uint64_t sz = 1;
     for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
       // TODO(wrengr): this division is expensive, but replacing it
       // with ffs/clz or intrinsics is non-portable; does mlir have
       // any cpp macro to enable this only when we want "paranoia mode"
       // for extra debugging info?  Or should we just rely on Clang's
       // -fsanitize=integer ?
       assert(sizes[r] <= std::numeric_limits<uint64_t>::max() / sz);
       sz *= sizes[r];
       if (sparsity[r] == DimLevelType::kCompressed) {
         pointers[r].reserve(sz + 1);
for (uint64_t i = 1; i < count; i++) {
       values[index] = 0.0;
       filled[index] = false;
     }
   }
 
   /// Finalizes lexicographic insertions.
   void endInsert() override {
     if (values.empty())
-      endDim(0);
+      finalizeSegment(0);
     else
       endPath(0);
   }
 
   /// Returns this sparse tensor storage scheme as a new memory-resident
   /// sparse tensor in coordinate scheme with the given dimension order.
   ///
   /// Precondition: `perm` must be valid for `getRank()`.
newSparseTensor(uint64_t rank, const uint64_t *shape, const uint64_t *perm,
     assert(source && "Got nullptr for source");
     SparseTensorEnumerator<V> enumerator(*source, rank, perm);
     const auto &permsz = enumerator.permutedSizes();
     assertPermutedSizesMatchShape(permsz, rank, perm, shape);
     return new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, *source);
   }
 
 private:
-  /// Appends the next free position of `indices[d]` to `pointers[d]`.
-  /// Thus, when called after inserting the last element of a segment,
-  /// it will append the position where the next segment begins.
-  inline void appendCurrentPointer(uint64_t d) {
-    assert(isCompressedDim(d)); // Entails `d < getRank()`.
-    appendPointer(d, indices[d].size());
-  }
-
   /// 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()`).
   inline void appendPointer(uint64_t d, uint64_t pos) {
     // This first assertion is trivially satisfied by both use sites;
     // but we leave it in just to be safe.
     assert(isCompressedDim(d)); // Entails `d < getRank()`.
     assert(pos <= std::numeric_limits<P>::max() &&
            "Pointer value is too large for the P-type");
     pointers[d].push_back(pos); // Here is where we convert to `P`.
   }
 
-  /// Appends the given index to `indices[d]`.  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 previously occurring in the same segment).
-  inline void appendIndex(uint64_t d, uint64_t i) {
-    assert(isCompressedDim(d)); // Entails `d < getRank()`.
-    assert(i <= std::numeric_limits<I>::max() &&
-           "Index value is too large for the I-type");
-    indices[d].push_back(i); // Here is where we convert to `I`.
+  /// 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.
+    assert(d < getRank());
+    if (isCompressedDim(d)) {
+      const uint64_t pos = indices[d].size();
+      // Inlined `appendPointer(d, pos)` called `count` times.
+      assert(pos <= std::numeric_limits<P>::max() &&
+             "Pointer value is too large for the P-type");
+      const P p = pos; // Convert the `uint64_t` to `P`.
+      pointers[d].insert(pointers[d].end(), count, p);
+    } else {
+      const uint64_t sz = sizes[d];
+      assert(sz >= full && "Segment is overfull");
+      // Assuming we check for overflows in the constructor, then this
+      // multiply will never overflow.
+      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())
+        // TODO: can we replace this with `memset`, or is this good enough?
+        values.insert(values.end(), count, 0);
+      else
+        finalizeSegment(d + 1, 0, count);
+    }
   }
 
   /// 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).
   inline void writeIndex(uint64_t d, uint64_t pos, uint64_t i) {
     // This first assertion trivially entailed at the one call-site;
     // but just to be safe.
     assert(isCompressedDim(d)); // Entails `d < getRank()`.
     // 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] = i; // Here is where we convert to `I`.
   }
 
+  /// 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
+  /// in the same segment).  For dense dimensions, this method instead
+  /// appends the appropriate number of zeros to the `values` array, where
+  /// `full` indicates the highest index already written for this segment.
+  void appendIndex(uint64_t d, uint64_t full, uint64_t i) {
+    assert(d < getRank());
+    if (isCompressedDim(d)) {
+      assert(i <= std::numeric_limits<I>::max() &&
+             "Index value is too large for the I-type");
+      indices[d].push_back(i); // Here is where we convert `uint64_t` to `I`.
+    } else {
+      assert(i >= full && "Index was already filled");
+      if (i == full)
+        return; // short-circuit nop.
+      if (d + 1 == getRank())
+        // TODO: can we replace this with `memset`, or is this good enough?
+        values.insert(values.end(), i - full, 0);
+      else
+        finalizeSegment(d + 1, 0, i - full);
+    }
+  }
+
   /// 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
   ///     and pointwise less-than).
void fromCOO(const std::vector<Element<V>> &elements, uint64_t lo,
     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];
       uint64_t seg = lo + 1;
       while (seg < hi && elements[seg].indices[d] == i)
         seg++;
       // Handle segment in interval for sparse or dense dimension.
-      if (isCompressedDim(d)) {
-        appendIndex(d, i);
-      } else {
-        // For dense storage we must fill in all the zero values between
-        // the previous element (when last we ran this for-loop) and the
-        // current element.
-        for (; full < i; full++)
-          endDim(d + 1);
-        full++;
-      }
+      appendIndex(d, full, i);
+      if (!isCompressedDim(d)) // i.e., is `DimLevelType::kDense`
+        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.
-    if (isCompressedDim(d)) {
-      appendCurrentPointer(d);
-    } else {
-      // For dense storage we must fill in all the zero values after
-      // the last element.
-      for (uint64_t sz = sizes[d]; full < sz; full++)
-        endDim(d + 1);
-    }
-  }
-
-  /// Ends a deeper, never seen before dimension.
-  void endDim(uint64_t d) {
-    assert(d <= getRank());
-    if (d == getRank()) {
-      values.push_back(0);
-    } else if (isCompressedDim(d)) {
-      appendCurrentPointer(d);
-    } else {
-      for (uint64_t full = 0, sz = sizes[d]; full < sz; full++)
-        endDim(d + 1);
-    }
+    finalizeSegment(d, full);
   }
 
   /// 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++) {
       uint64_t d = rank - i - 1;
-      if (isCompressedDim(d)) {
-        appendCurrentPointer(d);
-      } else {
-        for (uint64_t full = idx[d] + 1, sz = sizes[d]; full < sz; full++)
-          endDim(d + 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];
-      if (isCompressedDim(d)) {
-        appendIndex(d, i);
-      } else {
-        for (uint64_t full = top; full < i; full++)
-          endDim(d + 1);
-      }
+      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 {