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
25using namespace mlir::sparse_tensor;
26
27extern "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
49MLIR_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);
61MLIR_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);
69MLIR_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);
77MLIR_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);
85MLIR_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);
94MLIR_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);
103MLIR_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`.
109MLIR_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.
115MLIR_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.
131MLIR_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);
141MLIR_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.
155MLIR_CRUNNERUTILS_EXPORT index_type sparseLvlSize(void *tensor, index_type l);
156
157/// Tensor-storage method to get the size of the given dimension.
158MLIR_CRUNNERUTILS_EXPORT index_type sparseDimSize(void *tensor, index_type d);
159
160/// Tensor-storage method to finalize lexicographic insertions.
161MLIR_CRUNNERUTILS_EXPORT void endLexInsert(void *tensor);
162
163/// Releases the memory for the tensor-storage object.
164MLIR_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.
168MLIR_CRUNNERUTILS_EXPORT char *getTensorFilename(index_type id);
169
170/// Returns the number of stored elements for the sparse tensor being read.
171MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderNSE(void *p);
172
173/// Releases the SparseTensorReader and closes the associated file.
174MLIR_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.
179MLIR_CRUNNERUTILS_EXPORT void *createSparseTensorWriter(char *filename);
180
181/// Finalizes the outputing of a sparse tensor to a file and releases the
182/// SparseTensorWriter.
183MLIR_CRUNNERUTILS_EXPORT void delSparseTensorWriter(void *p);
184
185} // extern "C"
186
187#endif // MLIR_EXECUTIONENGINE_SPARSETENSORRUNTIME_H
188

source code of mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h