[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

545

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

545

Thanks! Add Reader/Writer to the names.

[bixia]

Rebase.

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

617

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

617

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 ↗

(On Diff #466822)

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 ↗

(On Diff #466822)

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 ↗

(On Diff #466822)

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

244 ↗

(On Diff #466822)

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

581–587

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

592

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 ↗

(On Diff #466822)

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 ↗

(On Diff #466822)

Add
This also closes the file associated with
the reader.

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

581–587

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

592

Removed, that was a mistake.

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
 
+/// Creates a SparseTensorReader for reading a sparse tensor from a file with
+/// the given file name.
+void *createSparseTensorReader(char *filename) {
+  SparseTensorReader *stfile = new SparseTensorReader(filename);
+  stfile->openFile();
+  stfile->readHeader();
+  return static_cast<void *>(stfile);
+}
+
+/// Returns the rank of the sparse tensor being read.
+index_type getSparseTensorRank(void *p) {
+  return static_cast<SparseTensorReader *>(p)->getRank();
+}
+
+/// Returns the is_symmetric bit for the sparse tensor being read.
+bool getSparseTensorIsSymmetric(void *p) {
+  return static_cast<SparseTensorReader *>(p)->isSymmetric();
+}
+
+/// Returns the number of non-zero values for the sparse tensor being read.
+index_type getSparseTensorNNZ(void *p) {
+  return static_cast<SparseTensorReader *>(p)->getNNZ();
+}
+
+/// Returns the size of a dimension for the sparse tensor being read.
+index_type getSparseTensorDimSize(void *p, index_type d) {
+  return static_cast<SparseTensorReader *>(p)->getDimSize(d);
+}
+
+/// Returns all dimension sizes for the sparse tensor being read.
+void _mlir_ciface_getSparseTensorDimSizes(
+    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];
+}
+
+/// Releases the SparseTensorReader.
+void delSparseTensorReader(void *p) {
+  delete static_cast<SparseTensorReader *>(p);
+}
+
+/// Returns the next element for the sparse tensor being read.
+#define IMPL_GETNEXT(VNAME, V)                                                 \
+  V _mlir_ciface_getSparseTensorNext##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 (uint64_t 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;
+
+/// 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.
+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);
+}
+
+/// Finalizes the outputing of a sparse tensor to a file and releases the
+/// SparseTensorWriter.
+void delSparseTensorWriter(void *p) {
+  SparseTensorWriter *file = static_cast<SparseTensorWriter *>(p);
+  file->flush();
+  assert(file->good());
+  if (file != &std::cout)
+    delete file;
+}
+
+/// Outputs the sparse tensor rank, nnz and shape.
+void _mlir_ciface_outSparseTensorMetaData(
+    void *p, index_type rank, index_type nnz,
+    StridedMemRefType<index_type, 1> *dref) {
+  assert(p && dref);
+  assert(dref->strides[0] == 1);
+  index_type *dimSizes = dref->data + dref->offset;
+  SparseTensorWriter &file = *static_cast<SparseTensorWriter *>(p);
+  file << rank << " " << nnz << std::endl;
+  for (uint64_t r = 0; r < rank - 1; ++r)
+    file << dimSizes[r] << " ";
+  file << dimSizes[rank - 1] << std::endl;
+}
+
+/// Outputs an element for the sparse tensor.
+#define IMPL_OUTNEXT(VNAME, V)                                                 \
+  void _mlir_ciface_outSparseTensorNext##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 @getSparseTensorRank(!TensorReader) -> (index)
+  func.func private @getSparseTensorNNZ(!TensorReader) -> (index)
+  func.func private @getSparseTensorIsSymmetric(!TensorReader) -> (i1)
+  func.func private @getSparseTensorDimSizes(!TensorReader, memref<?xindex>) -> () attributes { llvm.emit_c_interface }
+  func.func private @getSparseTensorNextF32(!TensorReader, memref<?xindex>) -> (f32) attributes { llvm.emit_c_interface }
+
+  func.func private @createSparseTensorWriter(!Filename) -> (!TensorWriter)
+  func.func private @delSparseTensorWriter(!TensorWriter)
+  func.func private @outSparseTensorMetaData(!TensorWriter, index, index, memref<?xindex>) -> () attributes { llvm.emit_c_interface }
+  func.func private @outSparseTensorNextF32(!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 @getSparseTensorRank(%tensor) : (!TensorReader) -> index
+    %nnz = call @getSparseTensorNNZ(%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 @getSparseTensorNextF32(%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 @getSparseTensorRank(%tensor) : (!TensorReader) -> index
+    vector.print %rank : index
+    %nnz = call @getSparseTensorNNZ(%tensor) : (!TensorReader) -> index
+    vector.print %nnz : index
+    %symmetric = call @getSparseTensorIsSymmetric(%tensor) : (!TensorReader) -> i1
+    vector.print %symmetric : i1
+    %dimSizes = memref.alloc(%rank) : memref<?xindex>
+    func.call @getSparseTensorDimSizes(%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 %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 @getSparseTensorRank(%tensor0) : (!TensorReader) -> index
+    %nnz = call @getSparseTensorNNZ(%tensor0) : (!TensorReader) -> index
+    %dimSizes = memref.alloc(%rank) : memref<?xindex>
+    func.call @getSparseTensorDimSizes(%tensor0,%dimSizes) : (!TensorReader, memref<?xindex>) -> ()
+    call @outSparseTensorMetaData(%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 @getSparseTensorNextF32(%tensor0, %indices) : (!TensorReader, memref<?xindex>) -> f32
+      func.call @outSparseTensorNextF32(%tensor1, %rank, %indices, %v) : (!TensorWriter, index, memref<?xindex>, f32) -> ()
+    }
+
+    // Release the resources.
+    call @delSparseTensorReader(%tensor0) : (!TensorReader) -> ()
+    call @delSparseTensorWriter(%tensor1) : (!TensorWriter) -> ()
+    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
+  }
+}