[mlir][sparse] Add runtime support for reading a COO tensor and writing the data to the given indices and values buffers.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

114–132

It makes me squeamish to have so much code duplication between these macros and the basic ones.

There isn't a great way to handle this due to (intentional) limitations in CPP. But the best option is to rephrase the original x-macros to be variadic: i.e., something like MLIR_SPARSETENSOR_FOREVERY_V(DO,...) DO(F64, double, __VA_ARGS__) etc. Which will at least save you from needing to duplicate the innermost x-macro. However, IIRC Peiming had some issues getting this to work on Windows last time I raised this concern, so you should check with him about what problems he ran into.

If you cannot get the variadic macros to work, then at the very least you should file a bug (assigned to me) and add a FIXME comment with the bug number.

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

397

Please rename to dimRank to match the terminology used elsewhere (and as per D143800 and https://github.com/openxla/stablehlo/pull/1143).

402

Just use !IsSymmetric

406–407

Why use new temp variables rather than just reusing the function parameter variables?

409–413

I like how this helps save code duplication, but I feel it's at the wrong level of abstraction. E.g., if you use the indices and values variables directly rather than using the pIndices and pValues variables, then you could have this function advance the two pointers after writing to them —thus maintaining the invariant that the pointers are always pointing to the next free entry (or one-past-the-end). Doing so keeps both halves of the pointer logic (writing to them and advancing them) together; whereas the current code mixes some of the pointer logic with the isSorted logic.

Having this function advance the pointers after writing to them would also help in clarifying things in the if constexpr (IsSymmetric) conditional, since it would be doing indices[0,1] = indices[-2,-1]; values[0] = values[-1] which (to my mind) helps clarify that it's writing the *next* entry based on the *previous* entry.

438–439

Since IsSymmetric can only be true when dimRank == 2, the pIndices[2] and pIndices[3] memory addresses look like they should be invalid. Yes, I know why the code actually works, but it would be prudent to add a comment explaining why it's okay so that other readers won't get confused.

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

646

please rename to dimRank

[bixia]

Address review comments.

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

114–132

Thanks. Done.

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

38

does it really say, "function-liked"?

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

252

Can we make the NNZ -> NSE or NOE change

(number of stored entries, rather than nonzeros, since we may store a few zeros ;-)

437

I feels like the way we handle symmetry (with the expand option) is starting to convolute our implantation quite a bit. Here we have the whole complication of knowing the number of entries in advance, but still have to compute how many are on the diagonal and thus not counted twice.

perhaps we should abandon the "expand symmetry" option altogether, and simply read in symmetric matrices as lower/upper triangular and start working on adding the proper "property" to our dim/level types

[bixia]

Update after removing expand_symmetry property. Rename nnz to nse.

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

38

see here

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

437

Remove symmetry in https://reviews.llvm.org/D144059

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

114–132

Thanks for doing that :)

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

252

Fwiw I've been using "NSE" in my renaming patch. (I like that that specifically has "stored" in it.)

255

I find it confusing that this method has the same name as readCOO<I> but very different arguments. You should rename this method to readIntoBuffers (or readInto if you're feeling succinct).

Also, you should have this method take the lvlRank argument in addition to dim2lvl so that it matches the API of readCOO<I>. Also, I follow the style of (a) providing buffer lengths before the buffers themselves, and (b) placing out-parameters last. Thus, the arguments should be in the order (lvlRank, dim2lvl, maxNse, indices, values)

302

Ditto my comments at readCOO<I,V>, this method should:

• be named readIntoBuffersLoop.
• take the lvlRank for API symmetry (and for the future-looking case, since we may/will eventually want to read files in via a non-permutation).
• use the argument order (lvlRank, dim2lvl, maxNse, indices, values).

In addition, you should change the type of dim2lvl to be detail::PermutationRef, so that it matches the API of readCOOLoop<I>. When you move the construction of the PermutationRef to readIntoBuffers<I,V> be sure to also add the assert(lvlRank == dimRank && "Rank mismatch"); as well.

Ideally both readCOO<I> and readIntoBuffers<I,V> should take PermutationRef rather than a pointer: because we want to push construction of PermutationRef as far up the callstack as possible, since doing so helps make more of the code typesafe and improves the C++ API. (Honestly I'm not entirely sure why I didn't do that in the first place; but I can make that change in a separate CL.)

390

You should add an assertion here that nnz <= maxNse, to ensure we catch errors as early as possible.

You should also rename the local variable to nse so that it matches the function parameter. Don't worry about the mismatch between the local variable and the method name; I'll be fixing that in my big renaming CL

398

Why split this out as a separate function? It would be cleaner to combine it with getOneElement so that it maintains the invariant that indices/values are always pointing to free memory.

I'm guessing you're worried about the sortedness check. If so, then the cleanest thing would be to move that sortedness check into getOneElement so that the function maintains that invariant as well! To make that work all you have to do is set isSorted=false before getting the first element (so that the conditional fails and all we do is read and advance pointers), and then set isSorted=true before entering the loop.

Rewriting getOneElement to include both advanceOneElement and the sortedness check will greatly help readability and reasoning. It'll help not just because it's maintaining invariants (which always helps), but because it's using the proper inductive hypothesis. It's precisely because that's the correct IH that (a) the function maintains a bunch of invariants, and (b) the rest of the code is just a simple loop calling the function plus a basis case outside of the loop.

411

Once you add the assert(nse <= maxNse) above, then you don't need this assertion anymore

[bixia]

Reorder parameters. Rename local nnz to nse. Combines two lambdas.

[bixia]

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

255

Rename to readToBuffers.
Remove maxNse, and move the assert to the client. Add doc to specify precondition.
Didn't add lvlRank and dim2lvl to this routine through. Unlike readCOO, which allocate the COO and needs those two values, the routine here doesn't need to allocate the COO and doesn't need those parameters. Let me know if you don't agree on this.

302

Rename to readToBuffersLoop.
Use PermutationRef instead.

390

Rename nnz to nse.
Assertion is added to caller.

[bixia]

Trying to fix the build.

[aartbik]

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

299

we don't template over isSymmetric?

405

this is only safe when we already have one element.

You do have the unrolled for loop below, but I don't see how that helps?

408

it feels like

 if (prevIndices[d] != indices[d]) { 
    if (prevIndices[d] >= indices[d])
       isSorted = false;
    break;
}

avoids writing true to the isSorted field for every element.
In fact, the == case denotes a duplicate

417

is this one reading outside the buffer?

[bixia]

Avoid writing isSorted with the same value. Add comments to the code.

[bixia]

Split a macro into variadic version and non-variadic version to avoid warnings.

[aartbik]

I think you missed some of my comments?

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

256

Technically this should be a *Level, since it maps dims (as index now) to levels in the array.

@wrengr do we want a typedef for dim2lvl and lvl2dim too just to be complete?

299

did you see my comment. The comment on L297 still mentions IsSymmetric, which is gone now

[bixia]

Remove IsSymmetric in comment.

[wrengr]

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

252

please use "coordinates" to match the new naming convention

255

Since readToBuffers does still construct a PermutationRef, it should have the same assert(lvlRank == dimRank && "Rank mismatch") assertion that readCOO does. While the PermutationRef ctor checks that dim2lvl is in fact a permutation, having this assertion before calling the ctor does two things: First, it ensures consistent error messages between the runtime/codegen passes, which is important for UX. Second, it helps catch errors —dim2lvl can be a permutation of dimRank even if dimRank != lvlRank, so we need to be sure to catch the rank mismatch in order to ensure correctness.

Also, as I mention in other comments below, the readToBuffersLoop does in fact want the lvlRank argument

255

Please rename the I template parameter to C to match the new naming convention.

256

We can, though that'd best be done as a separate CL, just to keep things clean.

Fwiw, I'm not too worried about it at the moment. Once I wrap up my current CL, I'll be getting back to updating the compiler to support non-permutations; and once the compiler is in a position to generate the dim2lvl/lvl2dim functions, that's when I was going to go through and update everything in the RT to take function-pointers in lieu of the current vectors/memrefs/pointers. Thus, so long as the new methods follow the same coding conventions as the pre-existing methods, then it should be easy enough to change at that point.

256

Please rename to lvlCoordinates to match the new naming convention.

(N.B., we do not use the shorter lvlCoords here because this is a shared pool of coords, rather than just the coords for a single element)

301

Please rename the I template parameter to C to match the new naming convention.

302

Please rename to lvlCoordinates to match the new naming convention.

(N.B., we do not use the shorter lvlCoords here because this is a shared pool of coords, rather than just the coords for a single element)

369

Please rename the I template parameter to C to match the new naming convention.

370

Please rename to lvlCoordinates to match the new naming convention.

(N.B., we do not use the shorter lvlCoords here because this is a shared pool of coords, rather than just the coords for a single element)

388

Please rename the I template parameter to C to match the new naming convention.

390

Please rename to lvlCoordinates to match the new naming convention.

(N.B., we do not use the shorter lvlCoords here because this is a shared pool of coords, rather than just the coords for a single element)

394

Please rename to dimCoords to match the new naming convention.

405–406

This comment should be moved up to where you first declare/define the isSorted variable, rather than here

407

This should be lvlRank.

Once we change dim2lvl to be an arbitrary mapping rather than a permutation, then each element will have lvlRank-many coordinates, therefore we need to make sure to advance the lvlCoordinates pointer by lvlRank for each element.

407

Please rename to prevLvlCoords to match the new naming convention

408

This should be lvlRank.

408

This should be l

416

This should be lvlRank.

423

These braces aren't needed for so simple a loop body

mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h

286

please say "coordinates" to match the new naming convention

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

634

use CNAME and C to match the new naming convention

[aartbik]

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

256

We can, though that'd best be done as a separate CL, just to keep things clean.

Yeah, I was not suggesting to add that here, but this was really more a note to you on perhaps adding two more types for the two mappings.
Using something like

void foo(Dim2Lvl dim2lvl)

would be nice ;-=)

[bixia]

Rename indices to lvlCoordinates, add assert(lvlRank == dimRank && "Rank mismatch").

[bixia]

Rename I to C, indices to coordinates.

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

256

Will use the new typedef here when we have it.

256

Ok. Leaving this as it is.

299

We remove the expansion for isSymmetric in https://reviews.llvm.org/D144059 and now only record the property.

299

Sorry for missing this.

405

I add a comment to the code to explain this.

408

Agree. changed to
if (prevIndices[d] > indices[d])

isSorted = false;

411

we don't have maxNse anymore.

417

This reads the first element with isSorted = false, to prepare for the reading of the rest, with isSorted initialized to true. Add a comment to the code.

[wrengr]

Just a few more nits re naming. Once these and the couple remaining ones from before (dimInd->dimCoords; d->l) are finished, then this should be good to go! :)

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

252

Shouldn't that say "reads from the file" or something similar? (Since this isn't using SparseTensorCOO)

397

You already asserted this in readToBuffers, so don't need to do it again here (since readToBuffersLoop is private).

414

This should be just prevLvlCoords since it's intended to refer to the coordinates of a single element (despite the fact that those "coords" live inside the larger "coordinates" buffer). Yeah, I know it looks idiosyncratic here, since the function parameter is being used both for the entire collection of "coordinates" (when viewed from outside the function) as well as being reused to mean just the "coords" of the most-recent element (when viewed from inside the loop).

If you want to avoid the idiosyncrasy, then feel free to (re)introduce a local variable lvlCoords for the pointer that we advance by lvlRank each time this lambda is called. (Yeah, you used to have such a local variable in earlier versions of this CL; but when I suggested removing it, that was before I started on the Big Renaming CL to try and clean all this stuff up ;)

[bixia]

Rename d->l, dimInd->dimCoords, prevlvlCoordinates->prevLvlCoords.

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

252

add "from the file".

[wrengr]

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

379–381

personally I'd rather this be phrased as const bool isSorted = IsPattern ? ... : ...; (or const bool isSorted = isPattern() ? ... : ...;) since that's a lot more succinct and keeps things const

405

should be dimCoords here too

410–421

To avoid the code repetition, it'd be nice to replace this whole conditional and loopnest with isSorted &&= elemLT(Element<_>(prevLvlCoords,_), Element<_>(lvlCoordinates,_)) where const ElementLT<_> elemLT(lvlRank); is defined above the lambda. Though I'm hesitant to suggest that wholeheartedly, because there's the issue of filling in those underscores. The natural thing to do would be to use V and *values/*(values-1), but that incurs the cost of copying V and hence would slow things down. (The unnatural thing to do would be to define struct Empty {}; or similar and then use that instead of V; but that cure is worse than the disease imo.)

So, rather than making that change per se, could you add a TODO comment saying that we should define a new CoordsLT which is like ElementLT but doesn't have the V parameter? That way I can remember to fix that in the future since it'd be trivial enough to define the new comparator. Alternatively, could you file a ticket with that same comment and assign it to me?

412

should be prevLvlCoords here too

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

637

rename to cref (or more ideally lvlCoordinatesRef) to match the new naming scheme

638

You don't need to assert these, since that's already handled by the ASSERT_NO_STRIDE macro.

643–644

These variables should use the full "size" rather than just "s"

646

Does this really need to be ==? The more liberal thing would be to assert(lvlCoordinatesSize >= reader.getNNZ() * lvlRank); which also matches the liberalness of assert(valuesSize >= reader.getNNZ()). If you really do want equality, then you should use the ASSERT_USIZE_EQ macro instead. (And if you really want >= then you should define a new ASSERT_USIZE_GE/ASSERT_USIZE_LE macro that mirrors ASSERT_USIZE_EQ except for the comparison operator used.)

646

This check should be using lvlRank instead, since that's what the buffer actually needs to store. Of course, there's no way to get the lvlRank from the reader itself, so you'll need to pass it in as an additional argument to this function.

647

"enough"

648–649

Since you don't need the ps variable anywhere other than this assertion, you should use the ASSERT_USIZE_EQ macro instead.

654

should be lvlCoordinates to match the name in the readToBuffers method

[bixia]

Address review comments: use ASSERT_USIZE_EQ, rename dimRank to lvlRank, fix comments.

[bixia]

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

410–421

Add a TODO.

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

637

use cref.

646

Relax the check to use <=, and use replace dimRank with lvlRank.

[aartbik]

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

408

The comment "We set isSorted to false to read the first element" feels like you are actively setting it here instead of above (where you also comment on that). I think here it should read

"// Note that isSorted was set to false while reading the first element to...."

[bixia]

Fix comment.

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

 
 // NOTE: Client code will need to include "mlir/ExecutionEngine/Float16bits.h"
 // if they want to use the `MLIR_SPARSETENSOR_FOREVERY_V` macro.
 
 #include <cinttypes>
 #include <complex>
 #include <optional>
 
 namespace mlir {
 namespace sparse_tensor {
 
 /// This type is used in the public API at all places where MLIR expects
 /// values with the built-in type "index".  For now, we simply assume that
 /// type is 64-bit, but targets with different "index" bit widths should
 /// link with an alternatively built runtime support library.
 // TODO: support such targets?
 using index_type = uint64_t;
enum class PrimaryType : uint32_t {
   kI32 = 6,
   kI16 = 7,
   kI8 = 8,
   kC64 = 9,
   kC32 = 10
 };
 
 // This x-macro includes all `V` types.
+// TODO: We currently split out the non-variadic version from the variadic
+// version. Using ##__VA_ARGS__ to avoid the split gives
+//   warning: token pasting of ',' and __VA_ARGS__ is a GNU extension
+//   [-Wgnu-zero-variadic-macro-arguments]
+// and __VA_OPT__(, ) __VA_ARGS__ requires c++20.
 #define MLIR_SPARSETENSOR_FOREVERY_V(DO)                                       \
   DO(F64, double)                                                              \
   DO(F32, float)                                                               \
   DO(F16, f16)                                                                 \
   DO(BF16, bf16)                                                               \
   DO(I64, int64_t)                                                             \
   DO(I32, int32_t)                                                             \
   DO(I16, int16_t)                                                             \
   DO(I8, int8_t)                                                               \
   DO(C64, complex64)                                                           \
   DO(C32, complex32)
 
+// This x-macro includes all `V` types and supports variadic arguments.
+#define MLIR_SPARSETENSOR_FOREVERY_V_VAR(DO, ...)                              \
+  DO(F64, double, __VA_ARGS__)                                                 \
+  DO(F32, float, __VA_ARGS__)                                                  \
+  DO(F16, f16, __VA_ARGS__)                                                    \
+  DO(BF16, bf16, __VA_ARGS__)                                                  \
+  DO(I64, int64_t, __VA_ARGS__)                                                \
+  DO(I32, int32_t, __VA_ARGS__)                                                \
+  DO(I16, int16_t, __VA_ARGS__)                                                \
+  DO(I8, int8_t, __VA_ARGS__)                                                  \
+  DO(C64, complex64, __VA_ARGS__)                                              \
+  DO(C32, complex32, __VA_ARGS__)
+
+// This x-macro calls its argument on every pair of overhead and `V` types.
+#define MLIR_SPARSETENSOR_FOREVERY_V_O(DO)                                     \
+  MLIR_SPARSETENSOR_FOREVERY_V_VAR(DO, 64, uint64_t)                           \
+  MLIR_SPARSETENSOR_FOREVERY_V_VAR(DO, 32, uint32_t)                           \
+  MLIR_SPARSETENSOR_FOREVERY_V_VAR(DO, 16, uint16_t)                           \
+  MLIR_SPARSETENSOR_FOREVERY_V_VAR(DO, 8, uint8_t)                             \
+  MLIR_SPARSETENSOR_FOREVERY_V_VAR(DO, 0, index_type)
+
 constexpr bool isFloatingPrimaryType(PrimaryType valTy) {
   return PrimaryType::kF64 <= valTy && valTy <= PrimaryType::kBF16;
 }
 
 constexpr bool isIntegralPrimaryType(PrimaryType valTy) {
   return PrimaryType::kI64 <= valTy && valTy <= PrimaryType::kI8;
 }
 

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

readSparseTensor(uint64_t lvlRank, const uint64_t *lvlSizes,
                    const uint64_t *dim2lvl) {
     auto *lvlCOO = readCOO<V>(lvlRank, lvlSizes, dim2lvl);
     auto *tensor = SparseTensorStorage<P, I, V>::newFromCOO(
         getRank(), getDimSizes(), lvlRank, lvlTypes, lvl2dim, *lvlCOO);
     delete lvlCOO;
     return tensor;
   }
 
+  /// Reads the COO tensor from the file, stores the coordinates and values to
+  /// the given buffers, returns a boolean value to indicate whether the COO
+  /// elements are sorted.
+  /// Precondition: the buffers should have enough space to hold the elements.
+  template <typename C, typename V>
+  bool readToBuffers(uint64_t lvlRank, const uint64_t *dim2lvl,
+                     C *lvlCoordinates, V *values);
+
 private:
   /// Attempts to read a line from the file.  Is private because there's
   /// no reason for client code to call it.
   void readLine();
 
   /// Reads the next line of the input file and parses the coordinates
   /// into the `indices` argument.  Returns the position in the `line`
   /// buffer where the element's value should be parsed from.  This method
private:
   // 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.
   template <typename V, bool IsPattern>
   void readCOOLoop(uint64_t lvlRank, detail::PermutationRef dim2lvl,
                    SparseTensorCOO<V> *lvlCOO);
 
+  /// The internal implementation of `readToBuffers`.  We template over
+  /// `IsPattern` in order to perform LICM without needing to duplicate the
+  /// source code.
+  template <typename C, typename V, bool IsPattern>
+  bool readToBuffersLoop(uint64_t lvlRank, detail::PermutationRef dim2lvl,
+                         C *lvlCoordinates, V *values);
+
   /// 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.
   void readExtFROSTTHeader();
for (uint64_t nnz = getNNZ(), k = 0; k < nnz; ++k) {
     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);
   }
 }
 
+template <typename C, typename V>
+bool SparseTensorReader::readToBuffers(uint64_t lvlRank,
+                                       const uint64_t *dim2lvl,
+                                       C *lvlCoordinates, V *values) {
+  assert(isValid() && "Attempt to readCOO() before readHeader()");
+  const uint64_t dimRank = getRank();
+  assert(lvlRank == dimRank && "Rank mismatch");
+  detail::PermutationRef d2l(dimRank, dim2lvl);
+  // Do some manual LICM, to avoid assertions in the for-loop.
+  bool isSorted =
+      isPattern()
+          ? readToBuffersLoop<C, V, true>(lvlRank, d2l, lvlCoordinates, values)
+          : readToBuffersLoop<C, V, false>(lvlRank, d2l, lvlCoordinates,
+                                           values);
+
+  // Close the file and return isSorted.
+  closeFile();
+  return isSorted;
+}
+
+template <typename C, typename V, bool IsPattern>
+bool SparseTensorReader::readToBuffersLoop(uint64_t lvlRank,
+                                           detail::PermutationRef dim2lvl,
+                                           C *lvlCoordinates, V *values) {
+  const uint64_t dimRank = getRank();
+  const uint64_t nse = getNNZ();
+  std::vector<C> dimCoords(dimRank);
+  // Read the first element with isSorted=false as a way to avoid accessing its
+  // previous element.
+  bool isSorted = false;
+  char *linePtr;
+  // We inline `readCOOElement` here in order to avoid redundant assertions,
+  // since they're guaranteed by the call to `isValid()` and the construction
+  // of `dimCoords` above.
+  auto readElement = [&]() {
+    linePtr = readCOOIndices<C>(dimCoords.data());
+    dim2lvl.pushforward(dimRank, dimCoords.data(), lvlCoordinates);
+    *values = detail::readCOOValue<V, IsPattern>(&linePtr);
+    if (isSorted) {
+      // Note that isSorted was set to false while reading the first element,
+      // to guarantee the safeness of using prevLvlCoords.
+      C *prevLvlCoords = lvlCoordinates - lvlRank;
+      // TODO: define a new CoordsLT which is like ElementLT but doesn't have
+      // the V parameter, and use it here.
+      for (uint64_t l = 0; l < lvlRank; ++l) {
+        if (prevLvlCoords[l] != lvlCoordinates[l]) {
+          if (prevLvlCoords[l] > lvlCoordinates[l])
+            isSorted = false;
+          break;
+        }
+      }
+    }
+    lvlCoordinates += lvlRank;
+    ++values;
+  };
+  readElement();
+  isSorted = true;
+  for (uint64_t n = 1; n < nse; ++n)
+    readElement();
+
+  return isSorted;
+}
+
 /// 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();
   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)
     file << dimSizes[r] << " ";
   file << dimSizes[rank - 1] << std::endl;

mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h

 /// Returns the next element for the sparse tensor being read.
 #define DECL_GETNEXT(VNAME, V)                                                 \
   MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderNext##VNAME( \
       void *p, StridedMemRefType<index_type, 1> *iref,                         \
       StridedMemRefType<V, 0> *vref);
 MLIR_SPARSETENSOR_FOREVERY_V(DECL_GETNEXT)
 #undef DECL_GETNEXT
 
+/// Reads the sparse tensor, stores the coordinates and values to the given
+/// memrefs. Returns a boolean value to indicate whether the COO elements are
+/// sorted.
+#define DECL_GETNEXT(VNAME, V, CNAME, C)                                       \
+  MLIR_CRUNNERUTILS_EXPORT bool                                                \
+      _mlir_ciface_getSparseTensorReaderRead##CNAME##VNAME(                    \
+          void *p, StridedMemRefType<index_type, 1> *dim2lvlRef,               \
+          StridedMemRefType<C, 1> *iref, StridedMemRefType<V, 1> *vref)        \
+          MLIR_SPARSETENSOR_FOREVERY_V_O(DECL_GETNEXT)
+#undef DECL_GETNEXT
+
 using SparseTensorWriter = std::ostream;
 
 /// Creates a SparseTensorWriter for outputing a sparse tensor to a file with
 /// the given file name. When the file name is empty, std::cout is used.
 //
 // Only the extended FROSTT format is supported currently.
 MLIR_CRUNNERUTILS_EXPORT void *createSparseTensorWriter(char *filename);
 

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

void _mlir_ciface_getSparseTensorReaderNext##VNAME(                          \
     const uint64_t rank = MEMREF_GET_USIZE(iref);                              \
     index_type *indices = MEMREF_GET_PAYLOAD(iref);                            \
     V *value = MEMREF_GET_PAYLOAD(vref);                                       \
     *value = reader.readCOOElement<V>(rank, indices);                          \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
 #undef IMPL_GETNEXT
 
+#define IMPL_GETNEXT(VNAME, V, CNAME, C)                                       \
+  bool _mlir_ciface_getSparseTensorReaderRead##CNAME##VNAME(                   \
+      void *p, StridedMemRefType<index_type, 1> *dim2lvlRef,                   \
+      StridedMemRefType<C, 1> *cref, StridedMemRefType<V, 1> *vref) {          \
+    assert(p);                                                                 \
+    auto &reader = *static_cast<SparseTensorReader *>(p);                      \
+    ASSERT_NO_STRIDE(cref);                                                    \
+    ASSERT_NO_STRIDE(vref);                                                    \
+    ASSERT_NO_STRIDE(dim2lvlRef);                                              \
+    const uint64_t cSize = MEMREF_GET_USIZE(cref);                             \
+    const uint64_t vSize = MEMREF_GET_USIZE(vref);                             \
+    const uint64_t lvlRank = reader.getRank();                                 \
+    assert(vSize *lvlRank <= cSize);                                           \
+    assert(vSize >= reader.getNNZ() && "Not enough space in buffers");         \
+    ASSERT_USIZE_EQ(dim2lvlRef, lvlRank);                                      \
+    (void)cSize;                                                               \
+    (void)vSize;                                                               \
+    (void)lvlRank;                                                             \
+    C *lvlCoordinates = MEMREF_GET_PAYLOAD(cref);                              \
+    V *values = MEMREF_GET_PAYLOAD(vref);                                      \
+    index_type *dim2lvl = MEMREF_GET_PAYLOAD(dim2lvlRef);                      \
+    return reader.readToBuffers<C, V>(lvlRank, dim2lvl, lvlCoordinates,        \
+                                      values);                                 \
+  }
+MLIR_SPARSETENSOR_FOREVERY_V_O(IMPL_GETNEXT)
+#undef IMPL_GETNEXT
+
 void *_mlir_ciface_newSparseTensorFromReader(
     void *p, StridedMemRefType<index_type, 1> *lvlSizesRef,
     StridedMemRefType<DimLevelType, 1> *lvlTypesRef,
     StridedMemRefType<index_type, 1> *lvl2dimRef,
     StridedMemRefType<index_type, 1> *dim2lvlRef, OverheadType ptrTp,
     OverheadType indTp, PrimaryType valTp) {
   assert(p);
   SparseTensorReader &reader = *static_cast<SparseTensorReader *>(p);

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_io.mlir

module {
 
   func.func private @createSparseTensorReader(!Filename) -> (!TensorReader)
   func.func private @delSparseTensorReader(!TensorReader) -> ()
   func.func private @getSparseTensorReaderRank(!TensorReader) -> (index)
   func.func private @getSparseTensorReaderNNZ(!TensorReader) -> (index)
   func.func private @getSparseTensorReaderIsSymmetric(!TensorReader) -> (i1)
   func.func private @copySparseTensorReaderDimSizes(!TensorReader,
     memref<?xindex>) -> () attributes { llvm.emit_c_interface }
+  func.func private @getSparseTensorReaderRead0F32(!TensorReader,
+    memref<?xindex>, memref<?xindex>, memref<?xf32>)
+    -> (i1) attributes { llvm.emit_c_interface }
   func.func private @getSparseTensorReaderNextF32(!TensorReader,
     memref<?xindex>, memref<f32>) -> () attributes { llvm.emit_c_interface }
 
   func.func private @createSparseTensorWriter(!Filename) -> (!TensorWriter)
   func.func private @delSparseTensorWriter(!TensorWriter)
   func.func private @outSparseTensorWriterMetaData(!TensorWriter, index, index,
     memref<?xindex>) -> () attributes { llvm.emit_c_interface }
   func.func private @outSparseTensorWriterNextF32(!TensorWriter, index,
     memref<?xindex>, memref<f32>) -> () attributes { llvm.emit_c_interface }
 
   func.func @dumpi(%arg0: memref<?xindex>) {
     %c0 = arith.constant 0 : index
     %v = vector.transfer_read %arg0[%c0], %c0: memref<?xindex>, vector<17xindex>
     vector.print %v : vector<17xindex>
     return
   }
 
+  func.func @dumpi2(%arg0: memref<?xindex, strided<[2], offset: ?>>) {
+    %c0 = arith.constant 0 : index
+    %v = vector.transfer_read %arg0[%c0], %c0 :
+      memref<?xindex, strided<[2], offset: ?>>, vector<17xindex>
+    vector.print %v : vector<17xindex>
+    return
+  }
+
   func.func @dumpf(%arg0: memref<?xf32>) {
     %c0 = arith.constant 0 : index
     %d0 = arith.constant 0.0 : f32
     %v = vector.transfer_read %arg0[%c0], %d0: memref<?xf32>, vector<17xf32>
     vector.print %v : vector<17xf32>
     return
   }
 
   // Returns the indices and values of the tensor.
   func.func @readTensorFile(%tensor: !TensorReader)
-    -> (memref<?xindex>, memref<?xindex>, memref<?xf32>) {
+    -> (memref<?xindex>, memref<?xf32>, i1) {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
+    %c2 = arith.constant 2 : index
 
     %rank = call @getSparseTensorReaderRank(%tensor) : (!TensorReader) -> index
     %nnz = call @getSparseTensorReaderNNZ(%tensor) : (!TensorReader) -> index
 
     // Assume rank == 2.
-    %x0s = memref.alloc(%nnz) : memref<?xindex>
-    %x1s = memref.alloc(%nnz) : memref<?xindex>
+    %isize = arith.muli %c2, %nnz : index
+    %xs = memref.alloc(%isize) : memref<?xindex>
     %vs = memref.alloc(%nnz) : memref<?xf32>
-    %indices = memref.alloc(%rank) : memref<?xindex>
-    %value = memref.alloca() : memref<f32>
-    scf.for %i = %c0 to %nnz step %c1 {
-      func.call @getSparseTensorReaderNextF32(%tensor, %indices, %value)
-        : (!TensorReader, memref<?xindex>, memref<f32>) -> ()
-      // TODO: can we use memref.subview to avoid the need for the %value
-      //       buffer?
-      %v = memref.load %value[] : memref<f32>
-      memref.store %v, %vs[%i] : memref<?xf32>
-      %i0 = memref.load %indices[%c0] : memref<?xindex>
-      memref.store %i0, %x0s[%i] : memref<?xindex>
-      %i1 = memref.load %indices[%c1] : memref<?xindex>
-      memref.store %i1, %x1s[%i] : memref<?xindex>
-    }
-
-    // Release the resource for the indices.
-    memref.dealloc %indices : memref<?xindex>
-    return %x0s, %x1s, %vs : memref<?xindex>, memref<?xindex>, memref<?xf32>
+    %dim2lvl = memref.alloca(%c2) : memref<?xindex>
+    memref.store %c0, %dim2lvl[%c0] : memref<?xindex>
+    memref.store %c1, %dim2lvl[%c1] : memref<?xindex>
+    %isSorted =func.call @getSparseTensorReaderRead0F32(%tensor, %dim2lvl, %xs, %vs)
+        : (!TensorReader, memref<?xindex>, memref<?xindex>, memref<?xf32>) -> (i1)
+    return %xs, %vs, %isSorted : memref<?xindex>, memref<?xf32>, i1
   }
 
   // Reads a COO tensor from the given file name and prints its content.
   func.func @readTensorFileAndDump(%fileName: !Filename) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c2 = arith.constant 2 : index
     %tensor = call @createSparseTensorReader(%fileName)
       : (!Filename) -> (!TensorReader)
     %rank = call @getSparseTensorReaderRank(%tensor) : (!TensorReader) -> index
     vector.print %rank : index
     %nnz = call @getSparseTensorReaderNNZ(%tensor) : (!TensorReader) -> index
     vector.print %nnz : index
     %symmetric = call @getSparseTensorReaderIsSymmetric(%tensor)
       : (!TensorReader) -> i1
     vector.print %symmetric : i1
     %dimSizes = memref.alloc(%rank) : memref<?xindex>
     func.call @copySparseTensorReaderDimSizes(%tensor, %dimSizes)
       : (!TensorReader, memref<?xindex>) -> ()
     call @dumpi(%dimSizes) : (memref<?xindex>) -> ()
-    %x0s, %x1s, %vs = call @readTensorFile(%tensor)
-      : (!TensorReader) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>)
 
-    call @dumpi(%x0s) : (memref<?xindex>) -> ()
-    call @dumpi(%x1s) : (memref<?xindex>) -> ()
+    %xs, %vs, %isSorted = call @readTensorFile(%tensor)
+      : (!TensorReader) -> (memref<?xindex>, memref<?xf32>, i1)
+    %x0s = memref.subview %xs[%c0][%nnz][%c2]
+      : memref<?xindex> to memref<?xindex, strided<[2], offset: ?>>
+    %x1s = memref.subview %xs[%c1][%nnz][%c2]
+      : memref<?xindex> to memref<?xindex, strided<[2], offset: ?>>
+    vector.print %isSorted : i1
+    call @dumpi2(%x0s) : (memref<?xindex, strided<[2], offset: ?>>) -> ()
+    call @dumpi2(%x1s) : (memref<?xindex, strided<[2], offset: ?>>) -> ()
     call @dumpf(%vs) : (memref<?xf32>) -> ()
 
     // Release the resources.
     call @delSparseTensorReader(%tensor) : (!TensorReader) -> ()
     memref.dealloc %dimSizes : memref<?xindex>
-    memref.dealloc %x0s : memref<?xindex>
-    memref.dealloc %x1s : memref<?xindex>
+    memref.dealloc %xs : memref<?xindex>
     memref.dealloc %vs : memref<?xf32>
 
     return
   }
 
   // Reads a COO tensor from a file with fileName0 and writes its content to
   // another file with fileName1.
   func.func @createTensorFileFrom(%fileName0: !Filename, %fileName1: !Filename) {
func.func @entry() {
     %fileName1 = call @getTensorFilename(%c1) : (index) -> (!Filename)
 
     // Write the sparse tensor data from file through the SparseTensorReader and
     // print the data.
     // CHECK: 2
     // CHECK: 17
     // CHECK: 0
     // CHECK: ( 4, 256, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 )
+    // CHECK: 1
     // CHECK: ( 0, 0, 0, 0, 1, 1, 2, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3 )
     // CHECK: ( 0, 126, 127, 254, 1, 253, 2, 0, 1, 3, 98, 126, 127, 128, 249, 253, 255 )
     // CHECK: ( -1, 2, -3, 4, -5, 6, -7, 8, -9, 10, -11, 12, -13, 14, -15, 16, -17 )
     call @readTensorFileAndDump(%fileName0) : (!Filename) -> ()
 
     // Write the sparse tensor data to std::cout through the SparseTensorWriter.
     // CHECK: # extended FROSTT format
     // CHECK: 2 17
func.func @entry() {
     // CHECK: 4 250 -15
     // CHECK: 4 254 16
     // CHECK: 4 256 -17
     call @createTensorFileFrom(%fileName0, %fileName1)
       : (!Filename, !Filename) -> ()
 
     return
   }
-}
+}
 No newline at end of file