1 //===- SparseTensorUtils.h - Enums shared with the runtime ------*- 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 enums and functions which comprise the 10 // public API of `ExecutionEngine/SparseUtils.cpp`. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #ifndef MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_ 15 #define MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_ 16 17 #include "mlir/ExecutionEngine/CRunnerUtils.h" 18 #include "mlir/ExecutionEngine/Float16bits.h" 19 20 #include <cinttypes> 21 #include <complex> 22 #include <vector> 23 24 extern "C" { 25 26 //===----------------------------------------------------------------------===// 27 // 28 // Typedefs and enums. These are required to be public so that they 29 // can be shared with `Transforms/SparseTensorConversion.cpp`, since 30 // they define the arguments to the public functions declared later on. 31 // 32 // This section also defines x-macros <https://en.wikipedia.org/wiki/X_Macro> 33 // so that we can generate variations of the public functions for each 34 // supported primary- and/or overhead-type. 35 // 36 //===----------------------------------------------------------------------===// 37 38 /// This type is used in the public API at all places where MLIR expects 39 /// values with the built-in type "index". For now, we simply assume that 40 /// type is 64-bit, but targets with different "index" bit widths should 41 /// link with an alternatively built runtime support library. 42 // TODO: support such targets? 43 using index_type = uint64_t; 44 45 /// Encoding of overhead types (both pointer overhead and indices 46 /// overhead), for "overloading" @newSparseTensor. 47 enum class OverheadType : uint32_t { 48 kIndex = 0, 49 kU64 = 1, 50 kU32 = 2, 51 kU16 = 3, 52 kU8 = 4 53 }; 54 55 // This x-macro calls its argument on every overhead type which has 56 // fixed-width. It excludes `index_type` because that type is often 57 // handled specially (e.g., by translating it into the architecture-dependent 58 // equivalent fixed-width overhead type). 59 #define FOREVERY_FIXED_O(DO) \ 60 DO(64, uint64_t) \ 61 DO(32, uint32_t) \ 62 DO(16, uint16_t) \ 63 DO(8, uint8_t) 64 65 // This x-macro calls its argument on every overhead type, including 66 // `index_type`. 67 #define FOREVERY_O(DO) \ 68 FOREVERY_FIXED_O(DO) \ 69 DO(0, index_type) 70 71 // These are not just shorthands but indicate the particular 72 // implementation used (e.g., as opposed to C99's `complex double`, 73 // or MLIR's `ComplexType`). 74 using complex64 = std::complex<double>; 75 using complex32 = std::complex<float>; 76 77 /// Encoding of the elemental type, for "overloading" @newSparseTensor. 78 enum class PrimaryType : uint32_t { 79 kF64 = 1, 80 kF32 = 2, 81 kF16 = 3, 82 kBF16 = 4, 83 kI64 = 5, 84 kI32 = 6, 85 kI16 = 7, 86 kI8 = 8, 87 kC64 = 9, 88 kC32 = 10 89 }; 90 91 // This x-macro includes all `V` types. 92 #define FOREVERY_V(DO) \ 93 DO(F64, double) \ 94 DO(F32, float) \ 95 DO(F16, f16) \ 96 DO(BF16, bf16) \ 97 DO(I64, int64_t) \ 98 DO(I32, int32_t) \ 99 DO(I16, int16_t) \ 100 DO(I8, int8_t) \ 101 DO(C64, complex64) \ 102 DO(C32, complex32) 103 104 /// The actions performed by @newSparseTensor. 105 enum class Action : uint32_t { 106 kEmpty = 0, 107 kFromFile = 1, 108 kFromCOO = 2, 109 kSparseToSparse = 3, 110 kEmptyCOO = 4, 111 kToCOO = 5, 112 kToIterator = 6 113 }; 114 115 /// This enum mimics `SparseTensorEncodingAttr::DimLevelType` for 116 /// breaking dependency cycles. `SparseTensorEncodingAttr::DimLevelType` 117 /// is the source of truth and this enum should be kept consistent with it. 118 enum class DimLevelType : uint8_t { 119 kDense = 0, 120 kCompressed = 1, 121 kSingleton = 2 122 }; 123 124 //===----------------------------------------------------------------------===// 125 // 126 // Public functions which operate on MLIR buffers (memrefs) to interact 127 // with sparse tensors (which are only visible as opaque pointers externally). 128 // Because these functions deal with memrefs, they should only be used 129 // by MLIR compiler-generated code (or code similarly guaranteed to remain 130 // in sync with MLIR; e.g., internal development tools like benchmarks). 131 // 132 // Where appropriate, we use macros to generate all variations of these 133 // functions for each supported primary- and overhead-type. 134 // 135 //===----------------------------------------------------------------------===// 136 137 /// The @newSparseTensor function for constructing a new sparse tensor. 138 /// This is the "swiss army knife" method for materializing sparse 139 /// tensors into the computation. The types of the `ptr` argument and 140 /// the result depend on the action, as explained in the following table 141 /// (where "STS" means a sparse-tensor-storage object, and "COO" means 142 /// a coordinate-scheme object). 143 /// 144 /// Action: `ptr`: Returns: 145 /// kEmpty unused STS, empty 146 /// kEmptyCOO unused COO, empty 147 /// kFromFile char* filename STS, read from the file 148 /// kFromCOO COO STS, copied from the COO source 149 /// kToCOO STS COO, copied from the STS source 150 /// kSparseToSparse STS STS, copied from the STS source 151 /// kToIterator STS COO-Iterator, call @getNext to use 152 MLIR_CRUNNERUTILS_EXPORT void * 153 _mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT 154 StridedMemRefType<index_type, 1> *sref, 155 StridedMemRefType<index_type, 1> *pref, 156 OverheadType ptrTp, OverheadType indTp, 157 PrimaryType valTp, Action action, void *ptr); 158 159 /// Tensor-storage method to obtain direct access to the values array. 160 #define DECL_SPARSEVALUES(VNAME, V) \ 161 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparseValues##VNAME( \ 162 StridedMemRefType<V, 1> *out, void *tensor); 163 FOREVERY_V(DECL_SPARSEVALUES) 164 #undef DECL_SPARSEVALUES 165 166 /// Tensor-storage method to obtain direct access to the pointers array 167 /// for the given dimension. 168 #define DECL_SPARSEPOINTERS(PNAME, P) \ 169 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparsePointers##PNAME( \ 170 StridedMemRefType<P, 1> *out, void *tensor, index_type d); 171 FOREVERY_O(DECL_SPARSEPOINTERS) 172 #undef DECL_SPARSEPOINTERS 173 174 /// Tensor-storage method to obtain direct access to the indices array 175 /// for the given dimension. 176 #define DECL_SPARSEINDICES(INAME, I) \ 177 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_sparseIndices##INAME( \ 178 StridedMemRefType<I, 1> *out, void *tensor, index_type d); 179 FOREVERY_O(DECL_SPARSEINDICES) 180 #undef DECL_SPARSEINDICES 181 182 /// Coordinate-scheme method for adding a new element. 183 #define DECL_ADDELT(VNAME, V) \ 184 MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_addElt##VNAME( \ 185 void *coo, StridedMemRefType<V, 0> *vref, \ 186 StridedMemRefType<index_type, 1> *iref, \ 187 StridedMemRefType<index_type, 1> *pref); 188 FOREVERY_V(DECL_ADDELT) 189 #undef DECL_ADDELT 190 191 /// Coordinate-scheme method for getting the next element while iterating. 192 #define DECL_GETNEXT(VNAME, V) \ 193 MLIR_CRUNNERUTILS_EXPORT bool _mlir_ciface_getNext##VNAME( \ 194 void *coo, StridedMemRefType<index_type, 1> *iref, \ 195 StridedMemRefType<V, 0> *vref); 196 FOREVERY_V(DECL_GETNEXT) 197 #undef DECL_GETNEXT 198 199 /// Tensor-storage method to insert elements in lexicographical index order. 200 #define DECL_LEXINSERT(VNAME, V) \ 201 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_lexInsert##VNAME( \ 202 void *tensor, StridedMemRefType<index_type, 1> *cref, \ 203 StridedMemRefType<V, 0> *vref); 204 FOREVERY_V(DECL_LEXINSERT) 205 #undef DECL_LEXINSERT 206 207 /// Tensor-storage method to insert using expansion. 208 #define DECL_EXPINSERT(VNAME, V) \ 209 MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_expInsert##VNAME( \ 210 void *tensor, StridedMemRefType<index_type, 1> *cref, \ 211 StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref, \ 212 StridedMemRefType<index_type, 1> *aref, index_type count); 213 FOREVERY_V(DECL_EXPINSERT) 214 #undef DECL_EXPINSERT 215 216 //===----------------------------------------------------------------------===// 217 // 218 // Public functions which accept only C-style data structures to interact 219 // with sparse tensors (which are only visible as opaque pointers externally). 220 // These functions can be used both by MLIR compiler-generated code 221 // as well as by any external runtime that wants to interact with MLIR 222 // compiler-generated code. 223 // 224 //===----------------------------------------------------------------------===// 225 226 /// Tensor-storage method to get the size of the given dimension. 227 MLIR_CRUNNERUTILS_EXPORT index_type sparseDimSize(void *tensor, index_type d); 228 229 /// Tensor-storage method to finalize lexicographic insertions. 230 MLIR_CRUNNERUTILS_EXPORT void endInsert(void *tensor); 231 232 /// Coordinate-scheme method to write to file in extended FROSTT format. 233 #define DECL_OUTSPARSETENSOR(VNAME, V) \ 234 MLIR_CRUNNERUTILS_EXPORT void outSparseTensor##VNAME(void *coo, void *dest, \ 235 bool sort); 236 FOREVERY_V(DECL_OUTSPARSETENSOR) 237 #undef DECL_OUTSPARSETENSOR 238 239 /// Releases the memory for the tensor-storage object. 240 void delSparseTensor(void *tensor); 241 242 /// Releases the memory for the coordinate-scheme object. 243 #define DECL_DELCOO(VNAME, V) \ 244 MLIR_CRUNNERUTILS_EXPORT void delSparseTensorCOO##VNAME(void *coo); 245 FOREVERY_V(DECL_DELCOO) 246 #undef DECL_DELCOO 247 248 /// Helper function to read a sparse tensor filename from the environment, 249 /// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc. 250 MLIR_CRUNNERUTILS_EXPORT char *getTensorFilename(index_type id); 251 252 /// Helper function to read the header of a file and return the 253 /// shape/sizes, without parsing the elements of the file. 254 MLIR_CRUNNERUTILS_EXPORT void readSparseTensorShape(char *filename, 255 std::vector<uint64_t> *out); 256 257 /// Initializes sparse tensor from a COO-flavored format expressed using 258 /// C-style data structures. The expected parameters are: 259 /// 260 /// rank: rank of tensor 261 /// nse: number of specified elements (usually the nonzeros) 262 /// shape: array with dimension size for each rank 263 /// values: a "nse" array with values for all specified elements 264 /// indices: a flat "nse * rank" array with indices for all specified elements 265 /// perm: the permutation of the dimensions in the storage 266 /// sparse: the sparsity for the dimensions 267 /// 268 /// For example, the sparse matrix 269 /// | 1.0 0.0 0.0 | 270 /// | 0.0 5.0 3.0 | 271 /// can be passed as 272 /// rank = 2 273 /// nse = 3 274 /// shape = [2, 3] 275 /// values = [1.0, 5.0, 3.0] 276 /// indices = [ 0, 0, 1, 1, 1, 2] 277 #define DECL_CONVERTTOMLIRSPARSETENSOR(VNAME, V) \ 278 MLIR_CRUNNERUTILS_EXPORT void *convertToMLIRSparseTensor##VNAME( \ 279 uint64_t rank, uint64_t nse, uint64_t *shape, V *values, \ 280 uint64_t *indices, uint64_t *perm, uint8_t *sparse); 281 FOREVERY_V(DECL_CONVERTTOMLIRSPARSETENSOR) 282 #undef DECL_CONVERTTOMLIRSPARSETENSOR 283 284 /// Converts a sparse tensor to COO-flavored format expressed using 285 /// C-style data structures. The expected output parameters are pointers 286 /// for these values: 287 /// 288 /// rank: rank of tensor 289 /// nse: number of specified elements (usually the nonzeros) 290 /// shape: array with dimension size for each rank 291 /// values: a "nse" array with values for all specified elements 292 /// indices: a flat "nse * rank" array with indices for all specified elements 293 /// 294 /// The input is a pointer to `SparseTensorStorage<P, I, V>`, typically 295 /// returned from `convertToMLIRSparseTensor`. 296 #define DECL_CONVERTFROMMLIRSPARSETENSOR(VNAME, V) \ 297 MLIR_CRUNNERUTILS_EXPORT void convertFromMLIRSparseTensor##VNAME( \ 298 void *tensor, uint64_t *pRank, uint64_t *pNse, uint64_t **pShape, \ 299 V **pValues, uint64_t **pIndices); 300 FOREVERY_V(DECL_CONVERTFROMMLIRSPARSETENSOR) 301 #undef DECL_CONVERTFROMMLIRSPARSETENSOR 302 303 } // extern "C" 304 305 #endif // MLIR_EXECUTIONENGINE_SPARSETENSORUTILS_H_ 306