[mlir][sparse] Extend readCOOIndices to support overhead types beyond index_type.

This is to prepare for implementing the C API for reading a COO tensor to the
given buffers for indices and values.

Reviewed By: wrengr

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

diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
index fa36b7e..ec9a05f 100644 (file)
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
@@ -261,7 +261,19 @@ private:
   /// for the generated library.
   ///
   /// Precondition: `indices` is valid for `getRank()`.
-  char *readCOOIndices(uint64_t *indices);
+  template <typename I>
+  char *readCOOIndices(I *indices) {
+    readLine();
+    // Local variable for tracking the parser's position in the `line` buffer.
+    char *linePtr = line;
+    for (uint64_t dimRank = getRank(), d = 0; d < dimRank; ++d) {
+      // Parse the 1-based index.
+      uint64_t idx = strtoul(linePtr, &linePtr, 10);
+      // Store the 0-based index.
+      indices[d] = static_cast<I>(idx - 1);
+    }
+    return linePtr;
+  }
 
   /// The internal implementation of `readCOO`.  We template over
   /// `IsPattern` and `IsSymmetric` in order to perform LICM without

diff --git a/mlir/lib/ExecutionEngine/SparseTensor/File.cpp b/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
index 19fbaf2..bb227bf 100644 (file)
--- a/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
@@ -53,19 +53,6 @@ void SparseTensorReader::readLine() {
     MLIR_SPARSETENSOR_FATAL("Cannot read next line of %s\n", filename);
 }
 
-char *SparseTensorReader::readCOOIndices(uint64_t *indices) {
-  readLine();
-  // Local variable for tracking the parser's position in the `line` buffer.
-  char *linePtr = line;
-  for (uint64_t rank = getRank(), 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.
 void SparseTensorReader::readHeader() {
   assert(file && "Attempt to readHeader() before openFile()");