diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h b/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h --- a/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h +++ b/mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h @@ -25,6 +25,8 @@ using namespace mlir::sparse_tensor; +using SparseTensorWriter = std::ostream; + extern "C" { //===----------------------------------------------------------------------===// @@ -32,11 +34,7 @@ // Public functions which operate on MLIR buffers (memrefs) to interact // with sparse tensors (which are only visible as opaque pointers externally). // Because these functions deal with memrefs, they should only be used -// by MLIR compiler-generated code (or code similarly guaranteed to remain -// in sync with MLIR; e.g., internal development tools like benchmarks). -// -// Where appropriate, we use macros to generate all variations of these -// functions for each supported primary- and overhead-type. +// by MLIR compiler-generated code (or code that is in sync with MLIR). // //===----------------------------------------------------------------------===// @@ -64,9 +62,6 @@ StridedMemRefType *dim2lvlRef, OverheadType posTp, OverheadType crdTp, PrimaryType valTp, Action action, void *ptr); -// TODO: document what all the arguments are/mean for the functions below, -// especially with regards to "dim"-vs-"lvl" and mappings/permutations. - /// Tensor-storage method to obtain direct access to the values array. #define DECL_SPARSEVALUES(VNAME, V) \ MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparseValues##VNAME( \ @@ -127,6 +122,72 @@ MLIR_SPARSETENSOR_FOREVERY_V(DECL_EXPINSERT) #undef DECL_EXPINSERT +/// Constructs a new SparseTensorReader object, opens the file, reads the +/// header, and validates that the actual contents of the file match +/// the expected `dimShapeRef` and `valTp`. +// +// FIXME: update `SparseTensorCodegenPass` to use +// `_mlir_ciface_createCheckedSparseTensorReader` instead. +MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_createCheckedSparseTensorReader( + char *filename, StridedMemRefType *dimShapeRef, + PrimaryType valTp); + +/// Constructs a new sparse-tensor storage object with the given encoding, +/// initializes it by reading all the elements from the file, and then +/// closes the file. +MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_newSparseTensorFromReader( + void *p, StridedMemRefType *lvlSizesRef, + StridedMemRefType *lvlTypesRef, + StridedMemRefType *lvl2dimRef, + StridedMemRefType *dim2lvlRef, OverheadType posTp, + OverheadType crdTp, PrimaryType valTp); + +/// SparseTensorReader method to copy the dimension-sizes into the +/// provided memref. +// +// FIXME: update `SparseTensorCodegenPass` to use +// `_mlir_ciface_getSparseTensorReaderDimSizes` instead. +MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_copySparseTensorReaderDimSizes( + void *p, StridedMemRefType *dref); + +/// SparseTensorReader method to obtain direct access to the +/// dimension-sizes array. +MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderDimSizes( + StridedMemRefType *out, void *p); + +/// Returns the next element for the sparse tensor being read. +#define DECL_GETNEXT(VNAME, V) \ + MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderNext##VNAME( \ + void *p, StridedMemRefType *dimCoordsRef, \ + StridedMemRefType *vref); +MLIR_SPARSETENSOR_FOREVERY_V(DECL_GETNEXT) +#undef DECL_GETNEXT + +/// Reads the sparse tensor, stores the coordinates and values to the given +/// memrefs. Returns a boolean value to indicate whether the COO elements are +/// sorted. +#define DECL_GETNEXT(VNAME, V, CNAME, C) \ + MLIR_CRUNNERUTILS_EXPORT bool \ + _mlir_ciface_getSparseTensorReaderRead##CNAME##VNAME( \ + void *p, StridedMemRefType *dim2lvlRef, \ + StridedMemRefType *iref, StridedMemRefType *vref) \ + MLIR_SPARSETENSOR_FOREVERY_V_O(DECL_GETNEXT) +#undef DECL_GETNEXT + +/// Outputs the sparse tensor dim-rank, nse, and dim-shape. +MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_outSparseTensorWriterMetaData( + void *p, index_type dimRank, index_type nse, + StridedMemRefType *dimSizesRef); + +/// Outputs an element for the sparse tensor. +#define DECL_OUTNEXT(VNAME, V) \ + MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_outSparseTensorWriterNext##VNAME( \ + void *p, index_type dimRank, \ + StridedMemRefType *dimCoordsRef, \ + StridedMemRefType *vref); +MLIR_SPARSETENSOR_FOREVERY_V(DECL_OUTNEXT) +#undef DECL_OUTNEXT + //===----------------------------------------------------------------------===// // // Public functions which accept only C-style data structures to interact @@ -228,28 +289,8 @@ /// 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. -// -// FIXME: update `SparseTensorCodegenPass` to use -// `_mlir_ciface_createCheckedSparseTensorReader` instead. MLIR_CRUNNERUTILS_EXPORT void *createSparseTensorReader(char *filename); -/// Constructs a new SparseTensorReader object, opens the file, reads the -/// header, and validates that the actual contents of the file match -/// the expected `dimShapeRef` and `valTp`. -MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_createCheckedSparseTensorReader( - char *filename, StridedMemRefType *dimShapeRef, - PrimaryType valTp); - -/// Constructs a new sparse-tensor storage object with the given encoding, -/// initializes it by reading all the elements from the file, and then -/// closes the file. -MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_newSparseTensorFromReader( - void *p, StridedMemRefType *lvlSizesRef, - StridedMemRefType *lvlTypesRef, - StridedMemRefType *lvl2dimRef, - StridedMemRefType *dim2lvlRef, OverheadType posTp, - OverheadType crdTp, PrimaryType valTp); - /// Returns the rank of the sparse tensor being read. MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderRank(void *p); @@ -263,44 +304,10 @@ MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderDimSize(void *p, index_type d); -/// SparseTensorReader method to copy the dimension-sizes into the -/// provided memref. -// -// FIXME: update `SparseTensorCodegenPass` to use -// `_mlir_ciface_getSparseTensorReaderDimSizes` instead. -MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_copySparseTensorReaderDimSizes( - void *p, StridedMemRefType *dref); - -/// SparseTensorReader method to obtain direct access to the -/// dimension-sizes array. -MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderDimSizes( - StridedMemRefType *out, void *p); - /// 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 DECL_GETNEXT(VNAME, V) \ - MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderNext##VNAME( \ - void *p, StridedMemRefType *dimCoordsRef, \ - StridedMemRefType *vref); -MLIR_SPARSETENSOR_FOREVERY_V(DECL_GETNEXT) -#undef DECL_GETNEXT - -/// Reads the sparse tensor, stores the coordinates and values to the given -/// memrefs. Returns a boolean value to indicate whether the COO elements are -/// sorted. -#define DECL_GETNEXT(VNAME, V, CNAME, C) \ - MLIR_CRUNNERUTILS_EXPORT bool \ - _mlir_ciface_getSparseTensorReaderRead##CNAME##VNAME( \ - void *p, StridedMemRefType *dim2lvlRef, \ - StridedMemRefType *iref, StridedMemRefType *vref) \ - MLIR_SPARSETENSOR_FOREVERY_V_O(DECL_GETNEXT) -#undef DECL_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. // @@ -311,20 +318,6 @@ /// SparseTensorWriter. MLIR_CRUNNERUTILS_EXPORT void delSparseTensorWriter(void *p); -/// Outputs the sparse tensor dim-rank, nse, and dim-shape. -MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_outSparseTensorWriterMetaData( - void *p, index_type dimRank, index_type nse, - StridedMemRefType *dimSizesRef); - -/// Outputs an element for the sparse tensor. -#define DECL_OUTNEXT(VNAME, V) \ - MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_outSparseTensorWriterNext##VNAME( \ - void *p, index_type dimRank, \ - StridedMemRefType *dimCoordsRef, \ - StridedMemRefType *vref); -MLIR_SPARSETENSOR_FOREVERY_V(DECL_OUTNEXT) -#undef DECL_OUTNEXT - } // extern "C" #endif // MLIR_EXECUTIONENGINE_SPARSETENSORRUNTIME_H