[mlir][sparse] Making way for SparseTensorRuntime to support non-permutations

Systematically updates the SparseTensorRuntime to properly distinguish tensor-dimensions from storage-levels (and their associated ranks, shapes, sizes, indices, etc).  With a few exceptions which are noted in the code, this ensures the runtime has all the **semantic** changes necessary to support non-permutations.

(Whereas **operationally**, since we're still using `std::vector<uing64_t>` to represent the mappings, there's no way to pass in any interesting non-permutations.  Changing the representation to `std::function` will be done in a separate differential.)

Depends On D137680

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D137681

1 files changed, 17 insertions, 11 deletions

diff --git a/mlir/lib/ExecutionEngine/SparseTensor/File.cpp b/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
index b5c940b65345..e750bf77b89c 100644
--- a/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
@@ -47,18 +47,24 @@ void SparseTensorReader::closeFile() {
   }
 }
 
-// TODO(wrengr/bixia): figure out how to reorganize the element-parsing
-// loop of `openSparseTensorCOO` into methods of this class, so we can
-// avoid leaking access to the `line` pointer (both for general hygiene
-// and because we can't mark it const due to the second argument of
-// `strtoul`/`strtoud` being `char * *restrict` rather than
-// `char const* *restrict`).
-//
 /// Attempts to read a line from the file.
-char *SparseTensorReader::readLine() {
-  if (fgets(line, kColWidth, file))
-    return line;
-  MLIR_SPARSETENSOR_FATAL("Cannot read next line of %s\n", filename);
+void SparseTensorReader::readLine() {
+  if (!fgets(line, kColWidth, file))
+    MLIR_SPARSETENSOR_FATAL("Cannot read next line of %s\n", filename);
+}
+
+char *SparseTensorReader::readCOOIndices(uint64_t rank, uint64_t *indices) {
+  assert(rank == getRank() && "Rank mismatch");
+  readLine();
+  // Local variable for tracking the parser's position in the `line` buffer.
+  char *linePtr = line;
+  for (uint64_t r = 0; r < rank; ++r) {
+    // Parse the 1-based index.
+    uint64_t idx = strtoul(linePtr, &linePtr, 10);
+    // Store the 0-based index.
+    indices[r] = idx - 1;
+  }
+  return linePtr;
 }
 
 /// Reads and parses the file's header.