[Mlir-commits] [mlir] e800c66 - [mlir][sparse] Extend readCOOIndices to support overhead types beyond index_type.

[llvmlistbot at llvm.org]

Author: bixia1
Date: 2023-02-13T14:51:59-08:00
New Revision: e800c664f0e925f6acb30dfa054906e3059ce125

URL: https://github.com/llvm/llvm-project/commit/e800c664f0e925f6acb30dfa054906e3059ce125
DIFF: https://github.com/llvm/llvm-project/commit/e800c664f0e925f6acb30dfa054906e3059ce125.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
    mlir/lib/ExecutionEngine/SparseTensor/File.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
index fa36b7e393eeb..ec9a05f8b1246 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
@@ -261,7 +261,19 @@ class SparseTensorReader final {
   /// 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 19fbaf23687b0..bb227bf2c2dd4 100644
--- 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()");