[mlir][sparse] Factoring out readAllElements

[wrengr]

Summary

This change enables the SparseTensorCodegen pass to reuse the majority of the readCOO/readSparseTensor code without relying on the SparseTensorCOO/SparseTensorStorage classes. It does however introduce a dynamic function call in the for-loop, so we should perform benchmarking to make sure this refactoring doesn't slow readCOO down too much.

Depends On D138365

[wrengr]

Rebasing to incorporate the bugfix of https://reviews.llvm.org/D138363#3941956. Also doing a bit of redesign of readAllElements since that bugfix means we can perform the dim2lvl mapping in the callback rather than in readAllElements.

(As disscussed offline, the design of readAllElements may not actually work out for the codegen pass, since the type of the callback means we need to be able to form closures in order to close over the COO or other object being inserted into. Nevertheless, the general approach of D138365 can still be used by the codegen pass in order to do the LICM to avoid unnecessary conditionals in the for-loop.)

[wrengr]

git-clang-format

[wrengr]

Abandoning because this is only meaningfully usable if the caller can produce closures to pass to readAllElements (or if they use global variables or similar); hence it'll be difficult for the codegen pass to actually make use of this abstraction. So to solve the code-reuse problem, this should instead be redesigned in a more iterator-style approach.

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

public:
   /// to the `indices` array.
   template <typename V>
   V readCOOElement(uint64_t rank, uint64_t *indices) {
     assert(rank == getRank() && "rank mismatch");
     char *linePtr = readCOOIndices(indices);
     return detail::readCOOValue<V>(&linePtr, isPattern());
   }
 
-  /// Allocates a new COO object for `lvlSizes`, initializes it by reading
-  /// all the elements from the file and applying `dim2lvl` to their indices,
-  /// and then closes the file.
+  /// Reads all `getNNZ()`-many elements from the file, passes each
+  /// element to the callback, and finally closes the file.  N.B.,
+  /// this function passes the dimension-indices to the callback,
+  /// thus the callback must perform the mapping from dimension-indices
+  /// to level-indices.
+  ///
+  /// Preconditions:
+  /// * the file's actual value type can be read as `V`.
+  ///
+  /// Asserts:
+  /// * `isValid()`.
+  //
+  // NOTE: This method is factored out from `readCOO` so that it can be
+  // reused by the codegen-pass.
+  template <typename V>
+  void readAllElements(ElementConsumer<V> yield);
+
+  /// Allocates a new COO object for `lvlSizes`, and initializes it via
+  /// `readAllElements` and applying `dim2lvl` to each element's indices.
   ///
   /// Preconditions:
   /// * `lvlSizes` must be valid for `lvlRank`.
   /// * `dim2lvl` must be valid for `getRank()`.
   /// * `dim2lvl` maps indices valid for `getDimSizes()` to indices
   ///   valid for `lvlSizes`.
   /// * the file's actual value type can be read as `V`.
   ///
   /// Asserts:
   /// * `isValid()`
   /// * `dim2lvl` is a permutation, and therefore also `lvlRank == getRank()`.
   ///   (This requirement will be lifted once we functionalize `dim2lvl`.)
   //
   // NOTE: This method is factored out of `readSparseTensor` primarily to
   // reduce code bloat (since the bulk of the code doesn't care about the
   // `<P,I>` type template parameters).  But we leave it public since it's
   // perfectly reasonable for clients to use.
   template <typename V>
   SparseTensorCOO<V> *readCOO(uint64_t lvlRank, const uint64_t *lvlSizes,
                               const uint64_t *dim2lvl);
 
   /// Allocates a new sparse-tensor storage object with the given encoding,
-  /// initializes it by reading all the elements from the file, and then
-  /// closes the file.  Preconditions/assertions are as per `readCOO`
-  /// and `SparseTensorStorage::newFromCOO`.
+  /// and initializes it via `readAllElements`.  Preconditions/assertions
+  /// are as per `readCOO` and `SparseTensorStorage::newFromCOO`.
   template <typename P, typename I, typename V>
   SparseTensorStorage<P, I, V> *
   readSparseTensor(uint64_t lvlRank, const uint64_t *lvlSizes,
                    const DimLevelType *lvlTypes, const uint64_t *lvl2dim,
                    const uint64_t *dim2lvl) {
     auto *lvlCOO = readCOO<V>(lvlRank, lvlSizes, dim2lvl);
     auto *tensor = SparseTensorStorage<P, I, V>::newFromCOO(
         getRank(), getDimSizes(), lvlRank, lvlTypes, lvl2dim, *lvlCOO);
private:
   /// into the `indices` argument.  Returns the position in the `line`
   /// buffer where the element's value should be parsed from.  This method
   /// has been factored out from `readCOOElement` to minimize code bloat
   /// for the generated library.
   ///
   /// Precondition: `indices` is valid for `getRank()`.
   char *readCOOIndices(uint64_t *indices);
 
-  /// The internal implementation of `readCOO`.  We template over
-  /// `IsPattern` and `IsSymmetric` in order to perform LICM without
-  /// needing to duplicate the source code.
-  //
-  // TODO: We currently take the `dim2lvl` argument as a `PermutationRef`
-  // since that's what `readCOO` creates.  Once we update `readCOO` to
-  // functionalize the mapping, then this helper will just take that
-  // same function.
+  /// The internal implementation of `readAllElements`.  We template over
+  /// `IsPattern` and `IsSymmetric` in order to perform LICM without needing
+  /// to duplicate the source code.
   template <typename V, bool IsPattern, bool IsSymmetric>
-  void readCOOLoop(uint64_t lvlRank, detail::PermutationRef dim2lvl,
-                   SparseTensorCOO<V> *lvlCOO);
+  void readAllElementsImpl(ElementConsumer<V> yield);
 
   /// Reads the MME header of a general sparse matrix of type real.
   void readMMEHeader();
 
   /// Reads the "extended" FROSTT header. Although not part of the
   /// documented format, we assume that the file starts with optional
   /// comments followed by two lines that define the rank, the number of
   /// nonzeros, and the dimensions sizes (one per rank) of the sparse tensor.
SparseTensorCOO<V> *SparseTensorReader::readCOO(uint64_t lvlRank,
   // TODO: This specific implementation does not generalize to arbitrary
   // mappings, but once we functionalize the `dim2lvl` argument we can
   // simply use that function instead.
   const uint64_t dimRank = getRank();
   assert(lvlRank == dimRank && "Rank mismatch");
   detail::PermutationRef d2l(dimRank, dim2lvl);
   // Prepare a COO object with the number of nonzeros as initial capacity.
   auto *lvlCOO = new SparseTensorCOO<V>(lvlRank, lvlSizes, getNNZ());
+  // Parse all the elements and add them to the COO.
+  std::vector<uint64_t> lvlInd(lvlRank);
+  readAllElements<V>(
+      [=, &lvlInd](const std::vector<uint64_t> &dimInd, V value) {
+        d2l.pushforward(dimRank, dimInd.data(), lvlInd.data());
+        lvlCOO->add(lvlInd, value);
+      });
+  return lvlCOO;
+}
+
+template <typename V>
+void SparseTensorReader::readAllElements(ElementConsumer<V> yield) {
+  assert(isValid() && "Attempt to readAllElements() before readHeader()");
   // Do some manual LICM, to avoid assertions in the for-loop.
   const bool IsPattern = isPattern();
   const bool IsSymmetric = (isSymmetric() && getRank() == 2);
   if (IsPattern && IsSymmetric)
-    readCOOLoop<V, true, true>(lvlRank, d2l, lvlCOO);
+    readAllElementsImpl<V, true, true>(yield);
   else if (IsPattern)
-    readCOOLoop<V, true, false>(lvlRank, d2l, lvlCOO);
+    readAllElementsImpl<V, true, false>(yield);
   else if (IsSymmetric)
-    readCOOLoop<V, false, true>(lvlRank, d2l, lvlCOO);
+    readAllElementsImpl<V, false, true>(yield);
   else
-    readCOOLoop<V, false, false>(lvlRank, d2l, lvlCOO);
-  // Close the file and return the COO.
-  closeFile();
-  return lvlCOO;
+    readAllElementsImpl<V, false, false>(yield);
 }
 
 template <typename V, bool IsPattern, bool IsSymmetric>
-void SparseTensorReader::readCOOLoop(uint64_t lvlRank,
-                                     detail::PermutationRef dim2lvl,
-                                     SparseTensorCOO<V> *lvlCOO) {
-  const uint64_t dimRank = getRank();
-  std::vector<uint64_t> dimInd(dimRank);
-  std::vector<uint64_t> lvlInd(lvlRank);
+void SparseTensorReader::readAllElementsImpl(ElementConsumer<V> yield) {
+  std::vector<uint64_t> dimInd(getRank());
   for (uint64_t nnz = getNNZ(), k = 0; k < nnz; ++k) {
     // We inline `readCOOElement` here in order to avoid redundant
     // assertions, since they're guaranteed by the call to `isValid()`
     // and the construction of `dimInd` above.
     char *linePtr = readCOOIndices(dimInd.data());
     const V value = detail::readCOOValue<V, IsPattern>(&linePtr);
-    dim2lvl.pushforward(dimRank, dimInd.data(), lvlInd.data());
     // TODO: <https://github.com/llvm/llvm-project/issues/54179>
-    lvlCOO->add(lvlInd, value);
+    yield(dimInd, value);
     // We currently chose to deal with symmetric matrices by fully
     // constructing them.  In the future, we may want to make symmetry
     // implicit for storage reasons.
     if constexpr (IsSymmetric)
       if (dimInd[0] != dimInd[1]) {
-        // Must recompute `lvlInd`, since arbitrary mappings don't preserve swap.
         std::swap(dimInd[0], dimInd[1]);
-        dim2lvl.pushforward(dimRank, dimInd.data(), lvlInd.data());
-        lvlCOO->add(lvlInd, value);
+        yield(dimInd, value);
       }
   }
+  // Close the file, convert the COO to SparseTensorStorage, and return.
+  closeFile();
 }
 
 /// Writes the sparse tensor to `filename` in extended FROSTT format.
 template <typename V>
 inline void writeExtFROSTT(const SparseTensorCOO<V> &coo,
                            const char *filename) {
   assert(filename && "Got nullptr for filename");
   auto &dimSizes = coo.getDimSizes();