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 | |