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

[bixia]

Format the test.

[wrengr]

Be sure to add these functions to the SparseTensorUtils.h header as well, and be sure to annotate them with the MLIR_CRUNNERUTILS_EXPORT macro, otherwise the generated code won't actually be allowed to call them.

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

546

I think these various wrappers for calling SparseTensorReader methods should be renamed to getSparseTensorReader{Rank,IsSymmetric,NNZ,...} —both to make it clear that they work on SparseTensorReader objects, and to avoid future/potential name clashes with similarly named method wrappers for other classes.

Or if you think those names are too long/ugly, then we should try to come up with some alternative naming scheme and apply it to all the functions exposed by this file. Maybe something like the naming schema st_{class}_{method}, where the st_ very succinctly indicates these are part of the sparse_tensor dialect, and splitting the class vs method names helps avoid some other irregularities in our current function names. Thus, we'd have st_reader_create, st_reader_getRank, etc for the stuff in this differential, but also st_storage_getValues, st_storage_getPointers, etc for stuff that's already there from before. (Or as a slight variation we could use the schema ST{Class}_{method}, i.e.: STReader_create, STStorage_getValues, etc.)

[wrengr]

Oh, also be sure to actually rename the SparseTensorFile class to SparseTensorReader

Update: Ah nevermind. I see that's being done in D135477

[bixia]

Add Reader/Writer to routine names for make it clear that the routines are for SparseTensorReader/Writer.
Add the routines to SparseTensorUtils.h and annotate the routines with MLIR_CRUNNERUTILS_EXPORT.

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

546

Thanks! Add Reader/Writer to the names.

[bixia]

Rebase.

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

618

minor nit, use index_type for the r var in the for loop as well, just to make sure we don't get issues on platforms with surprising index types

618

can we assert(rank != 0) before doing the rank-1 test?
or, iterate over all ranks and add spaces safe for the last one

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

166

Yeah, this is a neat feature to allow this kind of easy testing!

[bixia]

Use index_type to look through rank. Add assert(rank!=0) before doing rank-1.

[wrengr]

Beware that you'll need to rebase this for D135613

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

208–209

It's worth mentioning that this not only creates the object itself, but also opens the file and parses the header. (Both so that folks don't go looking for functions to do those steps, but also to make it clear that this can exit due to I/O or parsing errors.)

229

Might be worth mentioning that this also closes the file (just so folks don't go looking for a function to do that step).

239

Please switch this to the using syntax instead, to match the style used in the rest of include/mlir/ExecutionEngine/SparseTensor/*

244

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

582–588

I'd feel better if this code and its replica in openSparseTensorCOO were consolidated into a single method, in order to avoid the code duplication— especially since there's that 0-indexing vs 1-indexing conversion, so we really want to avoid the potential maintenance issues.

I can write up a quick CL for making that change in File.h if you don't want to do it in this CL

593

You shouldn't need this typedef here, since it's already given in the header file.

[wrengr]

Forgot to hit the accept button :) Though do be sure to rebase after addressing the documentation nits

[bixia]

Replace typedef with using. Modify function documentation. Remove an unnecessary typedef.

[bixia]

Rebase.

[wrengr]

In D135613#3852960, @NathanielMcVicar wrote:

This appears to have broken the C interface. I suppose the corresponding rename will be required there. https://lab.llvm.org/buildbot/#/builders/13/builds/26975/steps/6/logs/stdio

C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/ExecutionEngine/SparseTensorRuntime.h(245): error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/ExecutionEngine/SparseTensorRuntime.h(245): warning C4190: '_mlir_ciface_getSparseTensorReaderNextF16' has C-linkage specified, but returns UDT 'f16' which is incompatible with C
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\include\mlir/ExecutionEngine/Float16bits.h(35): note: see declaration of 'f16'

[bixia]

Sorry about this. Trying to fix it with https://reviews.llvm.org/D135811

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

208–209

Add
This opens the file and read the header meta data based
of the sparse tensor format derived from the suffix of the file name.

229

Add
This also closes the file associated with
the reader.

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

582–588

Good point. I will do this in a follow up CL.

593

Removed, that was a mistake.

mlir/include/mlir/ExecutionEngine/SparseTensorUtils.h

 /// returned from `convertToMLIRSparseTensor`.
 #define DECL_CONVERTFROMMLIRSPARSETENSOR(VNAME, V)                             \
   MLIR_CRUNNERUTILS_EXPORT void convertFromMLIRSparseTensor##VNAME(            \
       void *tensor, uint64_t *pRank, uint64_t *pNse, uint64_t **pShape,        \
       V **pValues, uint64_t **pIndices);
 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.
+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.
+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
+
+typedef std::ostream SparseTensorWriter;
+
+/// 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_SPARSETENSORUTILS_H

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

void convertFromMLIRSparseTensor##VNAME(void *tensor, uint64_t *pRank,       \
                                           V **pValues, uint64_t **pIndices) {  \
     fromMLIRSparseTensor<V>(                                                   \
         static_cast<SparseTensorStorage<uint64_t, uint64_t, V> *>(tensor),     \
         pRank, pNse, pShape, pValues, pIndices);                               \
   }
 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
+
+typedef std::ostream SparseTensorWriter;
+
+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

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

+// 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
+  }
+}