[mlir][sparse] Add MLIR interface for sparse tensor file input and output.

Reviewed By: aartbik, wrengr

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

3 files changed, 350 insertions, 0 deletions

diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h b/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h
index 94022b64fca7..9da0300d903d 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h
@@ -212,6 +212,64 @@ MLIR_SPARSETENSOR_FOREVERY_V(DECL_CONVERTTOMLIRSPARSETENSOR)
 MLIR_SPARSETENSOR_FOREVERY_V(DECL_CONVERTFROMMLIRSPARSETENSOR)
 #undef DECL_CONVERTFROMMLIRSPARSETENSOR
 
+/// Creates a SparseTensorReader for reading a sparse tensor from a file with
+/// the given file name. This opens the file and read the header meta data based
+/// of the sparse tensor format derived from the suffix of the file name.
+MLIR_CRUNNERUTILS_EXPORT void *createSparseTensorReader(char *filename);
+
+/// Returns the rank of the sparse tensor being read.
+MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderRank(void *p);
+
+/// Returns the is_symmetric bit for the sparse tensor being read.
+MLIR_CRUNNERUTILS_EXPORT bool getSparseTensorReaderIsSymmetric(void *p);
+
+/// Returns the number of non-zero values for the sparse tensor being read.
+MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderNNZ(void *p);
+
+/// Returns the size of a dimension for the sparse tensor being read.
+MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderDimSize(void *p,
+                                                                 index_type d);
+
+/// Returns all dimension sizes for the sparse tensor being read.
+MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderDimSizes(
+    void *p, StridedMemRefType<index_type, 1> *dref);
+
+/// Releases the SparseTensorReader. This also closes the file associated with
+/// the reader.
+MLIR_CRUNNERUTILS_EXPORT void delSparseTensorReader(void *p);
+
+/// Returns the next element for the sparse tensor being read.
+#define IMPL_GETNEXT(VNAME, V)                                                 \
+  MLIR_CRUNNERUTILS_EXPORT V _mlir_ciface_getSparseTensorReaderNext##VNAME(    \
+      void *p, StridedMemRefType<index_type, 1> *iref);
+MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
+#undef IMPL_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);
+
+/// Finalizes the outputing of a sparse tensor to a file and releases the
+/// SparseTensorWriter.
+MLIR_CRUNNERUTILS_EXPORT void delSparseTensorWriter(void *p);
+
+/// Outputs the sparse tensor rank, nnz and shape.
+MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_outSparseTensorWriterMetaData(
+    void *p, index_type rank, index_type nnz,
+    StridedMemRefType<index_type, 1> *dref);
+
+/// Outputs an element for the sparse tensor.
+#define IMPL_OUTNEXT(VNAME, V)                                                 \
+  MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_outSparseTensorWriterNext##VNAME( \
+      void *p, index_type rank, StridedMemRefType<index_type, 1> *iref,        \
+      V value);
+MLIR_SPARSETENSOR_FOREVERY_V(IMPL_OUTNEXT)
+#undef IMPL_OUTNEXT
+
 } // extern "C"
 
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSORRUNTIME_H
diff --git a/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp b/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
index 3ae6d8d011e3..84aa114533d8 100644
--- a/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp
@@ -578,6 +578,107 @@ MLIR_SPARSETENSOR_FOREVERY_V(IMPL_CONVERTTOMLIRSPARSETENSOR)
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_CONVERTFROMMLIRSPARSETENSOR)
 #undef IMPL_CONVERTFROMMLIRSPARSETENSOR
 
+void *createSparseTensorReader(char *filename) {
+  SparseTensorReader *stfile = new SparseTensorReader(filename);
+  stfile->openFile();
+  stfile->readHeader();
+  return static_cast<void *>(stfile);
+}
+
+index_type getSparseTensorReaderRank(void *p) {
+  return static_cast<SparseTensorReader *>(p)->getRank();
+}
+
+bool getSparseTensorReaderIsSymmetric(void *p) {
+  return static_cast<SparseTensorReader *>(p)->isSymmetric();
+}
+
+index_type getSparseTensorReaderNNZ(void *p) {
+  return static_cast<SparseTensorReader *>(p)->getNNZ();
+}
+
+index_type getSparseTensorReaderDimSize(void *p, index_type d) {
+  return static_cast<SparseTensorReader *>(p)->getDimSize(d);
+}
+
+void _mlir_ciface_getSparseTensorReaderDimSizes(
+    void *p, StridedMemRefType<index_type, 1> *dref) {
+  assert(p && dref);
+  assert(dref->strides[0] == 1);
+  index_type *dimSizes = dref->data + dref->offset;
+  SparseTensorReader &file = *static_cast<SparseTensorReader *>(p);
+  const index_type *sizes = file.getDimSizes();
+  index_type rank = file.getRank();
+  for (uint64_t r = 0; r < rank; ++r)
+    dimSizes[r] = sizes[r];
+}
+
+void delSparseTensorReader(void *p) {
+  delete static_cast<SparseTensorReader *>(p);
+}
+
+#define IMPL_GETNEXT(VNAME, V)                                                 \
+  V _mlir_ciface_getSparseTensorReaderNext##VNAME(                             \
+      void *p, StridedMemRefType<index_type, 1> *iref) {                       \
+    assert(p &&iref);                                                          \
+    assert(iref->strides[0] == 1);                                             \
+    index_type *indices = iref->data + iref->offset;                           \
+    SparseTensorReader *stfile = static_cast<SparseTensorReader *>(p);         \
+    index_type rank = stfile->getRank();                                       \
+    char *linePtr = stfile->readLine();                                        \
+    for (index_type r = 0; r < rank; ++r) {                                    \
+      uint64_t idx = strtoul(linePtr, &linePtr, 10);                           \
+      indices[r] = idx - 1;                                                    \
+    }                                                                          \
+    return detail::readCOOValue<V>(&linePtr, stfile->isPattern());             \
+  }
+MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
+#undef IMPL_GETNEXT
+
+void *createSparseTensorWriter(char *filename) {
+  SparseTensorWriter *file =
+      (filename[0] == 0) ? &std::cout : new std::ofstream(filename);
+  *file << "# extended FROSTT format\n";
+  return static_cast<void *>(file);
+}
+
+void delSparseTensorWriter(void *p) {
+  SparseTensorWriter *file = static_cast<SparseTensorWriter *>(p);
+  file->flush();
+  assert(file->good());
+  if (file != &std::cout)
+    delete file;
+}
+
+void _mlir_ciface_outSparseTensorWriterMetaData(
+    void *p, index_type rank, index_type nnz,
+    StridedMemRefType<index_type, 1> *dref) {
+  assert(p && dref);
+  assert(dref->strides[0] == 1);
+  assert(rank != 0);
+  index_type *dimSizes = dref->data + dref->offset;
+  SparseTensorWriter &file = *static_cast<SparseTensorWriter *>(p);
+  file << rank << " " << nnz << std::endl;
+  for (index_type r = 0; r < rank - 1; ++r)
+    file << dimSizes[r] << " ";
+  file << dimSizes[rank - 1] << std::endl;
+}
+
+#define IMPL_OUTNEXT(VNAME, V)                                                 \
+  void _mlir_ciface_outSparseTensorWriterNext##VNAME(                          \
+      void *p, index_type rank, StridedMemRefType<index_type, 1> *iref,        \
+      V value) {                                                               \
+    assert(p &&iref);                                                          \
+    assert(iref->strides[0] == 1);                                             \
+    index_type *indices = iref->data + iref->offset;                           \
+    SparseTensorWriter &file = *static_cast<SparseTensorWriter *>(p);          \
+    for (uint64_t r = 0; r < rank; ++r)                                        \
+      file << (indices[r] + 1) << " ";                                         \
+    file << value << std::endl;                                                \
+  }
+MLIR_SPARSETENSOR_FOREVERY_V(IMPL_OUTNEXT)
+#undef IMPL_OUTNEXT
+
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_io.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_io.mlir
new file mode 100644
index 000000000000..e30072ddc9d6
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_io.mlir
@@ -0,0 +1,191 @@
+// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: TENSOR0="%mlir_src_dir/test/Integration/data/wide.mtx" \
+// RUN: TENSOR1="" \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_lib_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+!Filename = !llvm.ptr<i8>
+!TensorReader = !llvm.ptr<i8>
+!TensorWriter = !llvm.ptr<i8>
+
+module {
+
+  func.func private @getTensorFilename(index) -> (!Filename)
+
+  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 @getSparseTensorReaderDimSizes(!TensorReader, memref<?xindex>)
+    -> () attributes { llvm.emit_c_interface }
+  func.func private @getSparseTensorReaderNextF32(!TensorReader, memref<?xindex>)
+    -> (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>, 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 @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>) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : 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>
+    %vs = memref.alloc(%nnz) : memref<?xf32>
+    %indices = memref.alloc(%rank) : memref<?xindex>
+    scf.for %i = %c0 to %nnz step %c1 {
+      %v = func.call @getSparseTensorReaderNextF32(%tensor, %indices)
+        : (!TensorReader, memref<?xindex>) -> 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>
+  }
+
+  // Reads a COO tensor from the given file name and prints its content.
+  func.func @readTensorFileAndDump(%fileName: !Filename) {
+    %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 @getSparseTensorReaderDimSizes(%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>) -> ()
+    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 %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) {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+
+    %tensor0 = call @createSparseTensorReader(%fileName0)
+      : (!Filename) -> (!TensorReader)
+    %tensor1 = call @createSparseTensorWriter(%fileName1)
+      : (!Filename) -> (!TensorWriter)
+
+    %rank = call @getSparseTensorReaderRank(%tensor0) : (!TensorReader) -> index
+    %nnz = call @getSparseTensorReaderNNZ(%tensor0) : (!TensorReader) -> index
+    %dimSizes = memref.alloc(%rank) : memref<?xindex>
+    func.call @getSparseTensorReaderDimSizes(%tensor0,%dimSizes)
+      : (!TensorReader, memref<?xindex>) -> ()
+    call @outSparseTensorWriterMetaData(%tensor1, %rank, %nnz, %dimSizes)
+      : (!TensorWriter, index, index, memref<?xindex>) -> ()
+
+    //TODO: handle isSymmetric.
+    // Assume rank == 2.
+    %indices = memref.alloc(%rank) : memref<?xindex>
+    scf.for %i = %c0 to %nnz step %c1 {
+      %v = func.call @getSparseTensorReaderNextF32(%tensor0, %indices)
+        : (!TensorReader, memref<?xindex>) -> f32
+      func.call @outSparseTensorWriterNextF32(%tensor1, %rank, %indices, %v)
+        : (!TensorWriter, index, memref<?xindex>, f32) -> ()
+    }
+
+    // Release the resources.
+    call @delSparseTensorReader(%tensor0) : (!TensorReader) -> ()
+    call @delSparseTensorWriter(%tensor1) : (!TensorWriter) -> ()
+    memref.dealloc %dimSizes : memref<?xindex>
+    memref.dealloc %indices : memref<?xindex>
+
+    return
+  }
+
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %fileName0 = call @getTensorFilename(%c0) : (index) -> (!Filename)
+    %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: ( 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
+    // CHECK: 4 256
+    // CHECK: 1 1 -1
+    // CHECK: 1 127 2
+    // CHECK: 1 128 -3
+    // CHECK: 1 255 4
+    // CHECK: 2 2 -5
+    // CHECK: 2 254 6
+    // CHECK: 3 3 -7
+    // CHECK: 4 1 8
+    // CHECK: 4 2 -9
+    // CHECK: 4 4 10
+    // CHECK: 4 99 -11
+    // CHECK: 4 127 12
+    // CHECK: 4 128 -13
+    // CHECK: 4 129 14
+    // CHECK: 4 250 -15
+    // CHECK: 4 254 16
+    // CHECK: 4 256 -17
+    call @createTensorFileFrom(%fileName0, %fileName1)
+      : (!Filename, !Filename) -> ()
+
+    return
+  }
+}