1 //===- SparseTensorRuntime.h - SparseTensor runtime support lib -*- C++ -*-===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This header file provides the functions which comprise the public API of the 10 // sparse tensor runtime support library for the SparseTensor dialect. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #ifndef MLIR_EXECUTIONENGINE_SPARSETENSORRUNTIME_H 15 #define MLIR_EXECUTIONENGINE_SPARSETENSORRUNTIME_H 16 17 #include "mlir/Dialect/SparseTensor/IR/Enums.h" 18 #include "mlir/ExecutionEngine/CRunnerUtils.h" 19 #include "mlir/ExecutionEngine/Float16bits.h" 20 21 #include <cinttypes> 22 #include <complex> 23 #include <vector> 24 25 using namespace mlir::sparse_tensor; 26 27 extern "C" { 28 29 //===----------------------------------------------------------------------===// 30 // 31 // Public functions which operate on MLIR buffers (memrefs) to interact 32 // with sparse tensors (which are only visible as opaque pointers externally). 33 // Because these functions deal with memrefs, they should only be used 34 // by MLIR compiler-generated code (or code that is in sync with MLIR). 35 // 36 //===----------------------------------------------------------------------===// 37 38 /// This is the "swiss army knife" method for materializing sparse 39 /// tensors into the computation. The types of the `ptr` argument and 40 /// the result depend on the action, as explained in the following table, 41 /// where "STS" means a sparse-tensor-storage object. 42 /// 43 /// Action: `ptr`: Returns: 44 /// --------------------------------------------------------------------------- 45 /// kEmpty - STS, empty 46 /// kFromReader reader STS, input from reader 47 /// kPack buffers STS, from level buffers 48 /// kSortCOOInPlace STS STS, sorted in place 49 MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_newSparseTensor( // NOLINT 50 StridedMemRefType<index_type, 1> *dimSizesRef, 51 StridedMemRefType<index_type, 1> *lvlSizesRef, 52 StridedMemRefType<LevelType, 1> *lvlTypesRef, 53 StridedMemRefType<index_type, 1> *dim2lvlRef, 54 StridedMemRefType<index_type, 1> *lvl2dimRef, OverheadType posTp, 55 OverheadType crdTp, PrimaryType valTp, Action action, void *ptr); 56 57 /// Tensor-storage method to obtain direct access to the values array. 58 #define DECL_SPARSEVALUES(VNAME, V) \ 59 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparseValues##VNAME( \ 60 StridedMemRefType<V, 1> *out, void *tensor); 61 MLIR_SPARSETENSOR_FOREVERY_V(DECL_SPARSEVALUES) 62 #undef DECL_SPARSEVALUES 63 64 /// Tensor-storage method to obtain direct access to the positions array 65 /// for the given level. 66 #define DECL_SPARSEPOSITIONS(PNAME, P) \ 67 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparsePositions##PNAME( \ 68 StridedMemRefType<P, 1> *out, void *tensor, index_type lvl); 69 MLIR_SPARSETENSOR_FOREVERY_O(DECL_SPARSEPOSITIONS) 70 #undef DECL_SPARSEPOSITIONS 71 72 /// Tensor-storage method to obtain direct access to the coordinates array 73 /// for the given level. 74 #define DECL_SPARSECOORDINATES(CNAME, C) \ 75 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparseCoordinates##CNAME( \ 76 StridedMemRefType<C, 1> *out, void *tensor, index_type lvl); 77 MLIR_SPARSETENSOR_FOREVERY_O(DECL_SPARSECOORDINATES) 78 #undef DECL_SPARSECOORDINATES 79 80 /// Tensor-storage method to obtain direct access to the coordinates array 81 /// buffer for the given level (provides an AoS view into the library). 82 #define DECL_SPARSECOORDINATES(CNAME, C) \ 83 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparseCoordinatesBuffer##CNAME( \ 84 StridedMemRefType<C, 1> *out, void *tensor, index_type lvl); 85 MLIR_SPARSETENSOR_FOREVERY_O(DECL_SPARSECOORDINATES) 86 #undef DECL_SPARSECOORDINATES 87 88 /// Tensor-storage method to insert elements in lexicographical 89 /// level-coordinate order. 90 #define DECL_LEXINSERT(VNAME, V) \ 91 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_lexInsert##VNAME( \ 92 void *tensor, StridedMemRefType<index_type, 1> *lvlCoordsRef, \ 93 StridedMemRefType<V, 0> *vref); 94 MLIR_SPARSETENSOR_FOREVERY_V(DECL_LEXINSERT) 95 #undef DECL_LEXINSERT 96 97 /// Tensor-storage method to insert using expansion. 98 #define DECL_EXPINSERT(VNAME, V) \ 99 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_expInsert##VNAME( \ 100 void *tensor, StridedMemRefType<index_type, 1> *lvlCoordsRef, \ 101 StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref, \ 102 StridedMemRefType<index_type, 1> *aref, index_type count); 103 MLIR_SPARSETENSOR_FOREVERY_V(DECL_EXPINSERT) 104 #undef DECL_EXPINSERT 105 106 /// Constructs a new SparseTensorReader object, opens the file, reads the 107 /// header, and validates that the actual contents of the file match 108 /// the expected `dimShapeRef` and `valTp`. 109 MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_createCheckedSparseTensorReader( 110 char *filename, StridedMemRefType<index_type, 1> *dimShapeRef, 111 PrimaryType valTp); 112 113 /// SparseTensorReader method to obtain direct access to the 114 /// dimension-sizes array. 115 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderDimSizes( 116 StridedMemRefType<index_type, 1> *out, void *p); 117 118 /// Reads the sparse tensor, stores the coordinates and values to the given 119 /// memrefs of a COO in AoS format. Returns a boolean to indicate whether 120 /// the COO elements are sorted. 121 #define DECL_READTOBUFFERS(VNAME, V, CNAME, C) \ 122 MLIR_CRUNNERUTILS_EXPORT bool \ 123 _mlir_ciface_getSparseTensorReaderReadToBuffers##CNAME##VNAME( \ 124 void *p, StridedMemRefType<index_type, 1> *dim2lvlRef, \ 125 StridedMemRefType<index_type, 1> *lvl2dimRef, \ 126 StridedMemRefType<C, 1> *cref, StridedMemRefType<V, 1> *vref) \ 127 MLIR_SPARSETENSOR_FOREVERY_V_O(DECL_READTOBUFFERS) 128 #undef DECL_READTOBUFFERS 129 130 /// Outputs the sparse tensor dim-rank, nse, and dim-shape. 131 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_outSparseTensorWriterMetaData( 132 void *p, index_type dimRank, index_type nse, 133 StridedMemRefType<index_type, 1> *dimSizesRef); 134 135 /// Outputs an element for the sparse tensor. 136 #define DECL_OUTNEXT(VNAME, V) \ 137 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_outSparseTensorWriterNext##VNAME( \ 138 void *p, index_type dimRank, \ 139 StridedMemRefType<index_type, 1> *dimCoordsRef, \ 140 StridedMemRefType<V, 0> *vref); 141 MLIR_SPARSETENSOR_FOREVERY_V(DECL_OUTNEXT) 142 #undef DECL_OUTNEXT 143 144 //===----------------------------------------------------------------------===// 145 // 146 // Public functions which accept only C-style data structures to interact 147 // with sparse tensors (which are only visible as opaque pointers externally). 148 // These functions can be used both by MLIR compiler-generated code 149 // as well as by any external runtime that wants to interact with MLIR 150 // compiler-generated code. 151 // 152 //===----------------------------------------------------------------------===// 153 154 /// Tensor-storage method to get the size of the given level. 155 MLIR_CRUNNERUTILS_EXPORT index_type sparseLvlSize(void *tensor, index_type l); 156 157 /// Tensor-storage method to get the size of the given dimension. 158 MLIR_CRUNNERUTILS_EXPORT index_type sparseDimSize(void *tensor, index_type d); 159 160 /// Tensor-storage method to finalize lexicographic insertions. 161 MLIR_CRUNNERUTILS_EXPORT void endLexInsert(void *tensor); 162 163 /// Releases the memory for the tensor-storage object. 164 MLIR_CRUNNERUTILS_EXPORT void delSparseTensor(void *tensor); 165 166 /// Helper function to read a sparse tensor filename from the environment, 167 /// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc. 168 MLIR_CRUNNERUTILS_EXPORT char *getTensorFilename(index_type id); 169 170 /// Returns the number of stored elements for the sparse tensor being read. 171 MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderNSE(void *p); 172 173 /// Releases the SparseTensorReader and closes the associated file. 174 MLIR_CRUNNERUTILS_EXPORT void delSparseTensorReader(void *p); 175 176 /// Creates a SparseTensorWriter for outputting a sparse tensor to a file 177 /// with the given file name. When the file name is empty, std::cout is used. 178 /// Only the extended FROSTT format is supported currently. 179 MLIR_CRUNNERUTILS_EXPORT void *createSparseTensorWriter(char *filename); 180 181 /// Finalizes the outputing of a sparse tensor to a file and releases the 182 /// SparseTensorWriter. 183 MLIR_CRUNNERUTILS_EXPORT void delSparseTensorWriter(void *p); 184 185 } // extern "C" 186 187 #endif // MLIR_EXECUTIONENGINE_SPARSETENSORRUNTIME_H 188