20 #ifndef MLIR_EXECUTIONENGINE_SPARSETENSOR_FILE_H
21 #define MLIR_EXECUTIONENGINE_SPARSETENSOR_FILE_H
29 namespace sparse_tensor {
37 struct is_complex<std::complex<T>> final :
public std::true_type {};
42 template <
typename V,
bool IsPattern>
43 inline std::enable_if_t<!is_complex<V>::value, V>
readValue(
char **linePtr) {
47 if constexpr (IsPattern)
49 return strtod(*linePtr, linePtr);
55 template <
typename V,
bool IsPattern>
56 inline std::enable_if_t<is_complex<V>::value, V>
readValue(
char **linePtr) {
61 if constexpr (IsPattern)
63 double re = strtod(*linePtr, linePtr);
64 double im = strtod(*linePtr, linePtr);
73 inline V
readValue(
char **linePtr,
bool isPattern) {
74 return isPattern ? readValue<V, true>(linePtr) : readValue<V, false>(linePtr);
102 assert(filename &&
"Received nullptr for filename");
113 const uint64_t *dimShape,
120 "Tensor element type %d not compatible with values in file %s\n",
121 static_cast<int>(valTp), filename);
153 assert(
isValid() &&
"Attempt to isPattern() before readHeader()");
160 assert(
isValid() &&
"Attempt to isSymmetric() before readHeader()");
167 assert(
isValid() &&
"Attempt to getRank() before readHeader()");
174 assert(
isValid() &&
"Attempt to getNSE() before readHeader()");
186 assert(d <
getRank() &&
"Dimension out of bounds");
197 template <
typename P,
typename I,
typename V>
200 const LevelType *lvlTypes,
const uint64_t *dim2lvl,
201 const uint64_t *lvl2dim) {
202 const uint64_t dimRank =
getRank();
203 MapRef map(dimRank, lvlRank, dim2lvl, lvl2dim);
204 auto *lvlCOO = readCOO<V>(map, lvlSizes);
206 dimRank,
getDimSizes(), lvlRank, lvlSizes, lvlTypes, dim2lvl, lvl2dim,
215 template <
typename C,
typename V>
216 bool readToBuffers(uint64_t lvlRank,
const uint64_t *dim2lvl,
217 const uint64_t *lvl2dim, C *lvlCoordinates, V *values);
226 template <
typename C>
227 char *readCoords(C *dimCoords) {
230 char *linePtr = line;
231 for (uint64_t dimRank =
getRank(), d = 0; d < dimRank; ++d) {
233 uint64_t c = strtoul(linePtr, &linePtr, 10);
235 dimCoords[d] =
static_cast<C
>(c - 1);
241 template <
typename V>
242 SparseTensorCOO<V> *readCOO(
const MapRef &map,
const uint64_t *lvlSizes);
246 template <
typename V,
bool IsPattern>
247 void readCOOLoop(
const MapRef &map, SparseTensorCOO<V> *coo);
252 template <
typename C,
typename V,
bool IsPattern>
253 bool readToBuffersLoop(
const MapRef &map, C *lvlCoordinates, V *values);
256 void readMMEHeader();
262 void readExtFROSTTHeader();
264 static constexpr
int kColWidth = 1025;
265 const char *
const filename;
266 FILE *file =
nullptr;
268 bool isSymmetric_ =
false;
270 char line[kColWidth];
279 template <
typename V>
280 SparseTensorCOO<V> *SparseTensorReader::readCOO(
const MapRef &map,
281 const uint64_t *lvlSizes) {
282 assert(
isValid() &&
"Attempt to readCOO() before readHeader()");
284 auto *coo =
new SparseTensorCOO<V>(map.getLvlRank(), lvlSizes,
getNSE());
287 readCOOLoop<V, true>(map, coo);
289 readCOOLoop<V, false>(map, coo);
295 template <
typename V,
bool IsPattern>
296 void SparseTensorReader::readCOOLoop(
const MapRef &map,
297 SparseTensorCOO<V> *coo) {
298 const uint64_t dimRank = map.getDimRank();
299 const uint64_t lvlRank = map.getLvlRank();
301 std::vector<uint64_t> dimCoords(dimRank);
302 std::vector<uint64_t> lvlCoords(lvlRank);
303 for (uint64_t k = 0, nse =
getNSE(); k < nse; k++) {
304 char *linePtr = readCoords(dimCoords.data());
305 const V value = detail::readValue<V, IsPattern>(&linePtr);
306 map.pushforward(dimCoords.data(), lvlCoords.data());
307 coo->add(lvlCoords, value);
311 template <
typename C,
typename V>
313 const uint64_t *dim2lvl,
314 const uint64_t *lvl2dim,
315 C *lvlCoordinates, V *values) {
316 assert(
isValid() &&
"Attempt to readCOO() before readHeader()");
319 isPattern() ? readToBuffersLoop<C, V, true>(map, lvlCoordinates, values)
320 : readToBuffersLoop<C, V, false>(map, lvlCoordinates, values);
325 template <
typename C,
typename V,
bool IsPattern>
326 bool SparseTensorReader::readToBuffersLoop(
const MapRef &map, C *lvlCoordinates,
330 const uint64_t nse =
getNSE();
332 std::vector<C> dimCoords(dimRank);
333 bool isSorted =
false;
335 const auto readNextElement = [&]() {
336 linePtr = readCoords<C>(dimCoords.data());
338 *values = detail::readValue<V, IsPattern>(&linePtr);
342 C *prevLvlCoords = lvlCoordinates - lvlRank;
343 for (uint64_t l = 0; l < lvlRank; ++l) {
344 if (prevLvlCoords[l] != lvlCoordinates[l]) {
345 if (prevLvlCoords[l] > lvlCoordinates[l])
351 lvlCoordinates += lvlRank;
356 for (uint64_t n = 1; n < nse; ++n)
#define MLIR_SPARSETENSOR_FATAL(...)
This macro helps minimize repetition of the printf-and-exit idiom, as well as ensuring that we print ...
A class for capturing the sparse tensor type map with a compact encoding.
void pushforward(const T *in, T *out) const
uint64_t getLvlRank() const
uint64_t getDimRank() const
This class abstracts over the information stored in file headers, as well as providing the buffers an...
void assertMatchesShape(uint64_t rank, const uint64_t *shape) const
Asserts the shape subsumes the actual dimension sizes.
bool isPattern() const
Gets the MME "pattern" property setting.
void closeFile()
Closes the file.
SparseTensorReader & operator=(const SparseTensorReader &)=delete
uint64_t getDimSize(uint64_t d) const
Safely gets the size of the given dimension.
SparseTensorReader(const SparseTensorReader &)=delete
void readHeader()
Reads and parses the file's header.
const uint64_t * getDimSizes() const
Gets the dimension-sizes array.
bool canReadAs(PrimaryType valTy) const
Checks if the file's ValueKind can be converted into the given tensor PrimaryType.
uint64_t getNSE() const
Gets the number of stored elements.
bool isValid() const
Checks if a header has been successfully read.
ValueKind getValueKind() const
Returns the stored value kind.
SparseTensorStorage< P, I, V > * readSparseTensor(uint64_t lvlRank, const uint64_t *lvlSizes, const LevelType *lvlTypes, const uint64_t *dim2lvl, const uint64_t *lvl2dim)
Allocates a new sparse-tensor storage object with the given encoding, initializes it by reading all t...
bool readToBuffers(uint64_t lvlRank, const uint64_t *dim2lvl, const uint64_t *lvl2dim, C *lvlCoordinates, V *values)
Reads the COO tensor from the file, stores the coordinates and values to the given buffers,...
bool isSymmetric() const
Gets the MME "symmetric" property setting.
static SparseTensorReader * create(const char *filename, uint64_t dimRank, const uint64_t *dimShape, PrimaryType valTp)
Factory method to allocate a new reader, open the file, read the header, and validate that the actual...
uint64_t getRank() const
Gets the dimension-rank of the tensor.
SparseTensorReader(const char *filename)
void openFile()
Opens the file for reading.
A memory-resident sparse tensor using a storage scheme based on per-level sparse/dense annotations.
static SparseTensorStorage< P, C, V > * newFromCOO(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank, const uint64_t *lvlSizes, const LevelType *lvlTypes, const uint64_t *dim2lvl, const uint64_t *lvl2dim, SparseTensorCOO< V > &lvlCOO)
Allocates a new sparse tensor and initializes it from the given COO.
std::enable_if_t<!is_complex< V >::value, V > readValue(char **linePtr)
Returns an element-value of non-complex type.
PrimaryType
Encoding of the elemental type, for "overloading" @newSparseTensor.
LevelType
This enum defines all the sparse representations supportable by the SparseTensor dialect.
Include the generated interface declarations.