MLIR  19.0.0git
Storage.h
Go to the documentation of this file.
1 //===- Storage.h - TACO-flavored sparse tensor representation ---*- 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 file contains definitions for the following classes:
10 //
11 // * `SparseTensorStorageBase`
12 // * `SparseTensorStorage<P, C, V>`
13 //
14 //===----------------------------------------------------------------------===//
15 
16 #ifndef MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H
17 #define MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H
18 
24 
25 namespace mlir {
26 namespace sparse_tensor {
27 
28 //===----------------------------------------------------------------------===//
29 //
30 // SparseTensorStorage Classes
31 //
32 //===----------------------------------------------------------------------===//
33 
34 /// Abstract base class for `SparseTensorStorage<P,C,V>`. This class
35 /// takes responsibility for all the `<P,C,V>`-independent aspects
36 /// of the tensor (e.g., sizes, sparsity, mapping). In addition,
37 /// we use function overloading to implement "partial" method
38 /// specialization, which the C-API relies on to catch type errors
39 /// arising from our use of opaque pointers.
40 ///
41 /// Because this class forms a bridge between the denotational semantics
42 /// of "tensors" and the operational semantics of how we store and
43 /// compute with them, it also distinguishes between two different
44 /// coordinate spaces (and their associated rank, sizes, etc).
45 /// Denotationally, we have the *dimensions* of the tensor represented
46 /// by this object. Operationally, we have the *levels* of the storage
47 /// representation itself.
48 ///
49 /// The *size* of an axis is the cardinality of possible coordinate
50 /// values along that axis (regardless of which coordinates have stored
51 /// element values). As such, each size must be non-zero since if any
52 /// axis has size-zero then the whole tensor would have trivial storage
53 /// (since there are no possible coordinates). Thus we use the plural
54 /// term *sizes* for a collection of non-zero cardinalities, and use
55 /// this term whenever referring to run-time cardinalities. Whereas we
56 /// use the term *shape* for a collection of compile-time cardinalities,
57 /// where zero is used to indicate cardinalities which are dynamic (i.e.,
58 /// unknown/unspecified at compile-time). At run-time, these dynamic
59 /// cardinalities will be inferred from or checked against sizes otherwise
60 /// specified. Thus, dynamic cardinalities always have an "immutable but
61 /// unknown" value; so the term "dynamic" should not be taken to indicate
62 /// run-time mutability.
64 protected:
67 
68 public:
69  /// Constructs a new sparse-tensor storage object with the given encoding.
70  SparseTensorStorageBase(uint64_t dimRank, const uint64_t *dimSizes,
71  uint64_t lvlRank, const uint64_t *lvlSizes,
72  const LevelType *lvlTypes, const uint64_t *dim2lvl,
73  const uint64_t *lvl2dim);
74  virtual ~SparseTensorStorageBase() = default;
75 
76  /// Gets the number of tensor-dimensions.
77  uint64_t getDimRank() const { return dimSizes.size(); }
78 
79  /// Gets the number of storage-levels.
80  uint64_t getLvlRank() const { return lvlSizes.size(); }
81 
82  /// Gets the tensor-dimension sizes array.
83  const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
84 
85  /// Safely looks up the size of the given tensor-dimension.
86  uint64_t getDimSize(uint64_t d) const {
87  assert(d < getDimRank());
88  return dimSizes[d];
89  }
90 
91  /// Gets the storage-level sizes array.
92  const std::vector<uint64_t> &getLvlSizes() const { return lvlSizes; }
93 
94  /// Safely looks up the size of the given storage-level.
95  uint64_t getLvlSize(uint64_t l) const {
96  assert(l < getLvlRank());
97  return lvlSizes[l];
98  }
99 
100  /// Gets the level-types array.
101  const std::vector<LevelType> &getLvlTypes() const { return lvlTypes; }
102 
103  /// Safely looks up the type of the given level.
104  LevelType getLvlType(uint64_t l) const {
105  assert(l < getLvlRank());
106  return lvlTypes[l];
107  }
108 
109  /// Safely checks if the level uses dense storage.
110  bool isDenseLvl(uint64_t l) const { return isDenseLT(getLvlType(l)); }
111 
112  /// Safely checks if the level uses compressed storage.
113  bool isCompressedLvl(uint64_t l) const {
114  return isCompressedLT(getLvlType(l));
115  }
116 
117  /// Safely checks if the level uses loose compressed storage.
118  bool isLooseCompressedLvl(uint64_t l) const {
119  return isLooseCompressedLT(getLvlType(l));
120  }
121 
122  /// Safely checks if the level uses singleton storage.
123  bool isSingletonLvl(uint64_t l) const { return isSingletonLT(getLvlType(l)); }
124 
125  /// Safely checks if the level uses n out of m storage.
126  bool isNOutOfMLvl(uint64_t l) const { return isNOutOfMLT(getLvlType(l)); }
127 
128  /// Safely checks if the level is ordered.
129  bool isOrderedLvl(uint64_t l) const { return isOrderedLT(getLvlType(l)); }
130 
131  /// Safely checks if the level is unique.
132  bool isUniqueLvl(uint64_t l) const { return isUniqueLT(getLvlType(l)); }
133 
134  /// Gets positions-overhead storage for the given level.
135 #define DECL_GETPOSITIONS(PNAME, P) \
136  virtual void getPositions(std::vector<P> **, uint64_t);
138 #undef DECL_GETPOSITIONS
139 
140  /// Gets coordinates-overhead storage for the given level.
141 #define DECL_GETCOORDINATES(INAME, C) \
142  virtual void getCoordinates(std::vector<C> **, uint64_t);
144 #undef DECL_GETCOORDINATES
145 
146  /// Gets primary storage.
147 #define DECL_GETVALUES(VNAME, V) virtual void getValues(std::vector<V> **);
149 #undef DECL_GETVALUES
150 
151  /// Element-wise insertion in lexicographic coordinate order. The first
152  /// argument is the level-coordinates for the value being inserted.
153 #define DECL_LEXINSERT(VNAME, V) virtual void lexInsert(const uint64_t *, V);
155 #undef DECL_LEXINSERT
156 
157  /// Expanded insertion. Note that this method resets the
158  /// values/filled-switch array back to all-zero/false while only
159  /// iterating over the nonzero elements.
160 #define DECL_EXPINSERT(VNAME, V) \
161  virtual void expInsert(uint64_t *, V *, bool *, uint64_t *, uint64_t, \
162  uint64_t);
164 #undef DECL_EXPINSERT
165 
166  /// Finalizes lexicographic insertions.
167  virtual void endLexInsert() = 0;
168 
169 private:
170  const std::vector<uint64_t> dimSizes;
171  const std::vector<uint64_t> lvlSizes;
172  const std::vector<LevelType> lvlTypes;
173  const std::vector<uint64_t> dim2lvlVec;
174  const std::vector<uint64_t> lvl2dimVec;
175 
176 protected:
177  const MapRef map; // non-owning pointers into dim2lvl/lvl2dim vectors
178  const bool allDense;
179 };
180 
181 /// A memory-resident sparse tensor using a storage scheme based on
182 /// per-level sparse/dense annotations. This data structure provides
183 /// a bufferized form of a sparse tensor type. In contrast to generating
184 /// setup methods for each differently annotated sparse tensor, this
185 /// method provides a convenient "one-size-fits-all" solution that simply
186 /// takes an input tensor and annotations to implement all required setup
187 /// in a general manner.
188 template <typename P, typename C, typename V>
190  /// Private constructor to share code between the other constructors.
191  /// Beware that the object is not necessarily guaranteed to be in a
192  /// valid state after this constructor alone; e.g., `isCompressedLvl(l)`
193  /// doesn't entail `!(positions[l].empty())`.
194  SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
195  uint64_t lvlRank, const uint64_t *lvlSizes,
196  const LevelType *lvlTypes, const uint64_t *dim2lvl,
197  const uint64_t *lvl2dim)
198  : SparseTensorStorageBase(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
199  dim2lvl, lvl2dim),
200  positions(lvlRank), coordinates(lvlRank), lvlCursor(lvlRank), coo() {}
201 
202 public:
203  /// Constructs a sparse tensor with the given encoding, and allocates
204  /// overhead storage according to some simple heuristics. When the
205  /// `bool` argument is true and `lvlTypes` are all dense, then this
206  /// ctor will also initialize the values array with zeros. That
207  /// argument should be true when an empty tensor is intended; whereas
208  /// it should usually be false when the ctor will be followed up by
209  /// some other form of initialization.
210  SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
211  uint64_t lvlRank, const uint64_t *lvlSizes,
212  const LevelType *lvlTypes, const uint64_t *dim2lvl,
213  const uint64_t *lvl2dim, SparseTensorCOO<V> *lvlCOO,
214  bool initializeValuesIfAllDense);
215 
216  /// Constructs a sparse tensor with the given encoding, and initializes
217  /// the contents from the COO. This ctor performs the same heuristic
218  /// overhead-storage allocation as the ctor above.
219  SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
220  uint64_t lvlRank, const uint64_t *lvlSizes,
221  const LevelType *lvlTypes, const uint64_t *dim2lvl,
222  const uint64_t *lvl2dim, SparseTensorCOO<V> &lvlCOO);
223 
224  /// Constructs a sparse tensor with the given encoding, and initializes
225  /// the contents from the level buffers. This ctor allocates exactly
226  /// the required amount of overhead storage, not using any heuristics.
227  /// It assumes that the data provided by `lvlBufs` can be directly used to
228  /// interpret the result sparse tensor and performs *NO* integrity test on the
229  /// input data. It also assume that the trailing COO coordinate buffer is
230  /// passed in as a single AoS memory.
231  SparseTensorStorage(uint64_t dimRank, const uint64_t *dimSizes,
232  uint64_t lvlRank, const uint64_t *lvlSizes,
233  const LevelType *lvlTypes, const uint64_t *dim2lvl,
234  const uint64_t *lvl2dim, const intptr_t *lvlBufs);
235 
236  /// Allocates a new empty sparse tensor.
238  newEmpty(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
239  const uint64_t *lvlSizes, const LevelType *lvlTypes,
240  const uint64_t *dim2lvl, const uint64_t *lvl2dim);
241 
242  /// Allocates a new sparse tensor and initializes it from the given COO.
244  newFromCOO(uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
245  const uint64_t *lvlSizes, const LevelType *lvlTypes,
246  const uint64_t *dim2lvl, const uint64_t *lvl2dim,
247  SparseTensorCOO<V> &lvlCOO);
248 
249  /// Allocates a new sparse tensor and initialize it with the data stored level
250  /// buffers directly.
252  packFromLvlBuffers(uint64_t dimRank, const uint64_t *dimSizes,
253  uint64_t lvlRank, const uint64_t *lvlSizes,
254  const LevelType *lvlTypes, const uint64_t *dim2lvl,
255  const uint64_t *lvl2dim, uint64_t srcRank,
256  const intptr_t *buffers);
257 
258  ~SparseTensorStorage() final = default;
259 
260  /// Partially specialize these getter methods based on template types.
261  void getPositions(std::vector<P> **out, uint64_t lvl) final {
262  assert(out && "Received nullptr for out parameter");
263  assert(lvl < getLvlRank());
264  *out = &positions[lvl];
265  }
266  void getCoordinates(std::vector<C> **out, uint64_t lvl) final {
267  assert(out && "Received nullptr for out parameter");
268  assert(lvl < getLvlRank());
269  *out = &coordinates[lvl];
270  }
271  void getValues(std::vector<V> **out) final {
272  assert(out && "Received nullptr for out parameter");
273  *out = &values;
274  }
275 
276  /// Partially specialize lexicographical insertions based on template types.
277  void lexInsert(const uint64_t *lvlCoords, V val) final {
278  assert(lvlCoords);
279  if (allDense) {
280  uint64_t lvlRank = getLvlRank();
281  uint64_t valIdx = 0;
282  // Linearize the address.
283  for (uint64_t l = 0; l < lvlRank; l++)
284  valIdx = valIdx * getLvlSize(l) + lvlCoords[l];
285  values[valIdx] = val;
286  return;
287  }
288  // First, wrap up pending insertion path.
289  uint64_t diffLvl = 0;
290  uint64_t full = 0;
291  if (!values.empty()) {
292  diffLvl = lexDiff(lvlCoords);
293  endPath(diffLvl + 1);
294  full = lvlCursor[diffLvl] + 1;
295  }
296  // Then continue with insertion path.
297  insPath(lvlCoords, diffLvl, full, val);
298  }
299 
300  /// Partially specialize expanded insertions based on template types.
301  void expInsert(uint64_t *lvlCoords, V *values, bool *filled, uint64_t *added,
302  uint64_t count, uint64_t expsz) final {
303  assert((lvlCoords && values && filled && added) && "Received nullptr");
304  if (count == 0)
305  return;
306  // Sort.
307  std::sort(added, added + count);
308  // Restore insertion path for first insert.
309  const uint64_t lastLvl = getLvlRank() - 1;
310  uint64_t c = added[0];
311  assert(c <= expsz);
312  assert(filled[c] && "added coordinate is not filled");
313  lvlCoords[lastLvl] = c;
314  lexInsert(lvlCoords, values[c]);
315  values[c] = 0;
316  filled[c] = false;
317  // Subsequent insertions are quick.
318  for (uint64_t i = 1; i < count; i++) {
319  assert(c < added[i] && "non-lexicographic insertion");
320  c = added[i];
321  assert(c <= expsz);
322  assert(filled[c] && "added coordinate is not filled");
323  lvlCoords[lastLvl] = c;
324  insPath(lvlCoords, lastLvl, added[i - 1] + 1, values[c]);
325  values[c] = 0;
326  filled[c] = false;
327  }
328  }
329 
330  /// Finalizes lexicographic insertions.
331  void endLexInsert() final {
332  if (!allDense) {
333  if (values.empty())
334  finalizeSegment(0);
335  else
336  endPath(0);
337  }
338  }
339 
340  /// Allocates a new COO object and initializes it with the contents.
341  /// Callers must make sure to delete the COO when they're done with it.
343  std::vector<uint64_t> dimCoords(getDimRank());
344  coo = new SparseTensorCOO<V>(getDimSizes(), values.size());
345  toCOO(0, 0, dimCoords);
346  assert(coo->getElements().size() == values.size());
347  return coo;
348  }
349 
350  /// Sort the unordered tensor in place, the method assumes that it is
351  /// an unordered COO tensor.
352  void sortInPlace() {
353  uint64_t nnz = values.size();
354 #ifndef NDEBUG
355  for (uint64_t l = 0; l < getLvlRank(); l++)
356  assert(nnz == coordinates[l].size());
357 #endif
358 
359  // In-place permutation.
360  auto applyPerm = [this](std::vector<uint64_t> &perm) {
361  uint64_t length = perm.size();
362  uint64_t lvlRank = getLvlRank();
363  // Cache for the current level coordinates.
364  std::vector<P> lvlCrds(lvlRank);
365  for (uint64_t i = 0; i < length; i++) {
366  uint64_t current = i;
367  if (i != perm[current]) {
368  for (uint64_t l = 0; l < lvlRank; l++)
369  lvlCrds[l] = coordinates[l][i];
370  V val = values[i];
371  // Deals with a permutation cycle.
372  while (i != perm[current]) {
373  uint64_t next = perm[current];
374  // Swaps the level coordinates and value.
375  for (uint64_t l = 0; l < lvlRank; l++)
376  coordinates[l][current] = coordinates[l][next];
377  values[current] = values[next];
378  perm[current] = current;
379  current = next;
380  }
381  for (uint64_t l = 0; l < lvlRank; l++)
382  coordinates[l][current] = lvlCrds[l];
383  values[current] = val;
384  perm[current] = current;
385  }
386  }
387  };
388 
389  std::vector<uint64_t> sortedIdx(nnz, 0);
390  for (uint64_t i = 0; i < nnz; i++)
391  sortedIdx[i] = i;
392 
393  std::sort(sortedIdx.begin(), sortedIdx.end(),
394  [this](uint64_t lhs, uint64_t rhs) {
395  for (uint64_t l = 0; l < getLvlRank(); l++) {
396  if (coordinates[l][lhs] == coordinates[l][rhs])
397  continue;
398  return coordinates[l][lhs] < coordinates[l][rhs];
399  }
400  assert(lhs == rhs && "duplicate coordinates");
401  return false;
402  });
403 
404  applyPerm(sortedIdx);
405  }
406 
407 private:
408  /// Appends coordinate `crd` to level `lvl`, in the semantically
409  /// general sense. For non-dense levels, that means appending to the
410  /// `coordinates[lvl]` array, checking that `crd` is representable in
411  /// the `C` type; however, we do not verify other semantic requirements
412  /// (e.g., that `crd` is in bounds for `lvlSizes[lvl]`, and not previously
413  /// occurring in the same segment). For dense levels, this method instead
414  /// appends the appropriate number of zeros to the `values` array, where
415  /// `full` is the number of "entries" already written to `values` for this
416  /// segment (aka one after the highest coordinate previously appended).
417  void appendCrd(uint64_t lvl, uint64_t full, uint64_t crd) {
418  if (!isDenseLvl(lvl)) {
419  assert(isCompressedLvl(lvl) || isLooseCompressedLvl(lvl) ||
420  isSingletonLvl(lvl) || isNOutOfMLvl(lvl));
421  coordinates[lvl].push_back(detail::checkOverflowCast<C>(crd));
422  } else { // Dense level.
423  assert(crd >= full && "Coordinate was already filled");
424  if (crd == full)
425  return; // Short-circuit, since it'll be a nop.
426  if (lvl + 1 == getLvlRank())
427  values.insert(values.end(), crd - full, 0);
428  else
429  finalizeSegment(lvl + 1, 0, crd - full);
430  }
431  }
432 
433  /// Computes the assembled-size associated with the `l`-th level,
434  /// given the assembled-size associated with the `(l-1)`-th level.
435  uint64_t assembledSize(uint64_t parentSz, uint64_t l) const {
436  if (isCompressedLvl(l))
437  return positions[l][parentSz];
438  if (isLooseCompressedLvl(l))
439  return positions[l][2 * parentSz - 1];
440  if (isSingletonLvl(l) || isNOutOfMLvl(l))
441  return parentSz; // new size same as the parent
442  assert(isDenseLvl(l));
443  return parentSz * getLvlSize(l);
444  }
445 
446  /// Initializes sparse tensor storage scheme from a memory-resident sparse
447  /// tensor in coordinate scheme. This method prepares the positions and
448  /// coordinates arrays under the given per-level dense/sparse annotations.
449  void fromCOO(const std::vector<Element<V>> &lvlElements, uint64_t lo,
450  uint64_t hi, uint64_t l) {
451  const uint64_t lvlRank = getLvlRank();
452  assert(l <= lvlRank && hi <= lvlElements.size());
453  // Once levels are exhausted, insert the numerical values.
454  if (l == lvlRank) {
455  assert(lo < hi);
456  values.push_back(lvlElements[lo].value);
457  return;
458  }
459  // Visit all elements in this interval.
460  uint64_t full = 0;
461  while (lo < hi) { // If `hi` is unchanged, then `lo < lvlElements.size()`.
462  // Find segment in interval with same coordinate at this level.
463  const uint64_t c = lvlElements[lo].coords[l];
464  uint64_t seg = lo + 1;
465  if (isUniqueLvl(l))
466  while (seg < hi && lvlElements[seg].coords[l] == c)
467  seg++;
468  // Handle segment in interval for sparse or dense level.
469  appendCrd(l, full, c);
470  full = c + 1;
471  fromCOO(lvlElements, lo, seg, l + 1);
472  // And move on to next segment in interval.
473  lo = seg;
474  }
475  // Finalize the sparse position structure at this level.
476  finalizeSegment(l, full);
477  }
478 
479  /// Finalizes the sparse position structure at this level.
480  void finalizeSegment(uint64_t l, uint64_t full = 0, uint64_t count = 1) {
481  if (count == 0)
482  return; // Short-circuit, since it'll be a nop.
483  if (isCompressedLvl(l)) {
484  uint64_t pos = coordinates[l].size();
485  positions[l].insert(positions[l].end(), count,
486  detail::checkOverflowCast<P>(pos));
487  } else if (isLooseCompressedLvl(l)) {
488  // Finish this level, and push pairs for the empty ones, and one
489  // more for next level. Note that this always leaves one extra
490  // unused element at the end.
491  uint64_t pos = coordinates[l].size();
492  positions[l].insert(positions[l].end(), 2 * count,
493  detail::checkOverflowCast<P>(pos));
494  } else if (isSingletonLvl(l) || isNOutOfMLvl(l)) {
495  return; // Nothing to finalize.
496  } else { // Dense dimension.
497  assert(isDenseLvl(l));
498  const uint64_t sz = getLvlSizes()[l];
499  assert(sz >= full && "Segment is overfull");
500  count = detail::checkedMul(count, sz - full);
501  // For dense storage we must enumerate all the remaining coordinates
502  // in this level (i.e., coordinates after the last non-zero
503  // element), and either fill in their zero values or else recurse
504  // to finalize some deeper level.
505  if (l + 1 == getLvlRank())
506  values.insert(values.end(), count, 0);
507  else
508  finalizeSegment(l + 1, 0, count);
509  }
510  }
511 
512  /// Wraps up a single insertion path, inner to outer.
513  void endPath(uint64_t diffLvl) {
514  const uint64_t lvlRank = getLvlRank();
515  const uint64_t lastLvl = lvlRank - 1;
516  assert(diffLvl <= lvlRank);
517  const uint64_t stop = lvlRank - diffLvl;
518  for (uint64_t i = 0; i < stop; i++) {
519  const uint64_t l = lastLvl - i;
520  finalizeSegment(l, lvlCursor[l] + 1);
521  }
522  }
523 
524  /// Continues a single insertion path, outer to inner. The first
525  /// argument is the level-coordinates for the value being inserted.
526  void insPath(const uint64_t *lvlCoords, uint64_t diffLvl, uint64_t full,
527  V val) {
528  const uint64_t lvlRank = getLvlRank();
529  assert(diffLvl <= lvlRank);
530  for (uint64_t l = diffLvl; l < lvlRank; l++) {
531  const uint64_t c = lvlCoords[l];
532  appendCrd(l, full, c);
533  full = 0;
534  lvlCursor[l] = c;
535  }
536  values.push_back(val);
537  }
538 
539  /// Finds the lexicographically first level where the level-coordinates
540  /// in the argument differ from those in the current cursor.
541  uint64_t lexDiff(const uint64_t *lvlCoords) const {
542  const uint64_t lvlRank = getLvlRank();
543  for (uint64_t l = 0; l < lvlRank; l++) {
544  const auto crd = lvlCoords[l];
545  const auto cur = lvlCursor[l];
546  if (crd > cur || (crd == cur && !isUniqueLvl(l)) ||
547  (crd < cur && !isOrderedLvl(l))) {
548  return l;
549  }
550  if (crd < cur) {
551  assert(false && "non-lexicographic insertion");
552  return -1u;
553  }
554  }
555  assert(false && "duplicate insertion");
556  return -1u;
557  }
558 
559  // Performs forall on level entries and inserts into dim COO.
560  void toCOO(uint64_t parentPos, uint64_t l, std::vector<uint64_t> &dimCoords) {
561  if (l == getLvlRank()) {
562  map.pushbackward(lvlCursor.data(), dimCoords.data());
563  assert(coo);
564  assert(parentPos < values.size());
565  coo->add(dimCoords, values[parentPos]);
566  return;
567  }
568  if (isCompressedLvl(l)) {
569  const std::vector<P> &positionsL = positions[l];
570  assert(parentPos + 1 < positionsL.size());
571  const uint64_t pstart = static_cast<uint64_t>(positionsL[parentPos]);
572  const uint64_t pstop = static_cast<uint64_t>(positionsL[parentPos + 1]);
573  const std::vector<C> &coordinatesL = coordinates[l];
574  assert(pstop <= coordinatesL.size());
575  for (uint64_t pos = pstart; pos < pstop; pos++) {
576  lvlCursor[l] = static_cast<uint64_t>(coordinatesL[pos]);
577  toCOO(pos, l + 1, dimCoords);
578  }
579  } else if (isLooseCompressedLvl(l)) {
580  const std::vector<P> &positionsL = positions[l];
581  assert(2 * parentPos + 1 < positionsL.size());
582  const uint64_t pstart = static_cast<uint64_t>(positionsL[2 * parentPos]);
583  const uint64_t pstop =
584  static_cast<uint64_t>(positionsL[2 * parentPos + 1]);
585  const std::vector<C> &coordinatesL = coordinates[l];
586  assert(pstop <= coordinatesL.size());
587  for (uint64_t pos = pstart; pos < pstop; pos++) {
588  lvlCursor[l] = static_cast<uint64_t>(coordinatesL[pos]);
589  toCOO(pos, l + 1, dimCoords);
590  }
591  } else if (isSingletonLvl(l) || isNOutOfMLvl(l)) {
592  assert(parentPos < coordinates[l].size());
593  lvlCursor[l] = static_cast<uint64_t>(coordinates[l][parentPos]);
594  toCOO(parentPos, l + 1, dimCoords);
595  } else { // Dense level.
596  assert(isDenseLvl(l));
597  const uint64_t sz = getLvlSizes()[l];
598  const uint64_t pstart = parentPos * sz;
599  for (uint64_t c = 0; c < sz; c++) {
600  lvlCursor[l] = c;
601  toCOO(pstart + c, l + 1, dimCoords);
602  }
603  }
604  }
605 
606  std::vector<std::vector<P>> positions;
607  std::vector<std::vector<C>> coordinates;
608  std::vector<V> values;
609  std::vector<uint64_t> lvlCursor;
610  SparseTensorCOO<V> *coo;
611 };
612 
613 //===----------------------------------------------------------------------===//
614 //
615 // SparseTensorStorage Factories
616 //
617 //===----------------------------------------------------------------------===//
618 
619 template <typename P, typename C, typename V>
621  uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
622  const uint64_t *lvlSizes, const LevelType *lvlTypes,
623  const uint64_t *dim2lvl, const uint64_t *lvl2dim) {
624  return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes,
625  lvlTypes, dim2lvl, lvl2dim, nullptr,
626  true);
627 }
628 
629 template <typename P, typename C, typename V>
631  uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
632  const uint64_t *lvlSizes, const LevelType *lvlTypes,
633  const uint64_t *dim2lvl, const uint64_t *lvl2dim,
634  SparseTensorCOO<V> &lvlCOO) {
635  return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes,
636  lvlTypes, dim2lvl, lvl2dim, lvlCOO);
637 }
638 
639 template <typename P, typename C, typename V>
641  uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
642  const uint64_t *lvlSizes, const LevelType *lvlTypes,
643  const uint64_t *dim2lvl, const uint64_t *lvl2dim, uint64_t srcRank,
644  const intptr_t *buffers) {
645  return new SparseTensorStorage<P, C, V>(dimRank, dimSizes, lvlRank, lvlSizes,
646  lvlTypes, dim2lvl, lvl2dim, buffers);
647 }
648 
649 //===----------------------------------------------------------------------===//
650 //
651 // SparseTensorStorage Constructors
652 //
653 //===----------------------------------------------------------------------===//
654 
655 template <typename P, typename C, typename V>
657  uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
658  const uint64_t *lvlSizes, const LevelType *lvlTypes,
659  const uint64_t *dim2lvl, const uint64_t *lvl2dim,
660  SparseTensorCOO<V> *lvlCOO, bool initializeValuesIfAllDense)
661  : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
662  dim2lvl, lvl2dim) {
663  assert(!lvlCOO || lvlRank == lvlCOO->getRank());
664  coo = lvlCOO;
665  // Provide hints on capacity of positions and coordinates.
666  // TODO: needs much fine-tuning based on actual sparsity; currently
667  // we reserve position/coordinate space based on all previous dense
668  // levels, which works well up to first sparse level; but we should
669  // really use nnz and dense/sparse distribution.
670  uint64_t sz = 1;
671  for (uint64_t l = 0; l < lvlRank; l++) {
672  if (isCompressedLvl(l)) {
673  positions[l].reserve(sz + 1);
674  positions[l].push_back(0);
675  coordinates[l].reserve(sz);
676  sz = 1;
677  } else if (isLooseCompressedLvl(l)) {
678  positions[l].reserve(2 * sz + 1); // last one unused
679  positions[l].push_back(0);
680  coordinates[l].reserve(sz);
681  sz = 1;
682  } else if (isSingletonLvl(l)) {
683  coordinates[l].reserve(sz);
684  sz = 1;
685  } else if (isNOutOfMLvl(l)) {
686  assert(l == lvlRank - 1 && "unexpected n:m usage");
687  sz = detail::checkedMul(sz, lvlSizes[l]) / 2;
688  coordinates[l].reserve(sz);
689  values.reserve(sz);
690  } else { // Dense level.
691  assert(isDenseLvl(l));
692  sz = detail::checkedMul(sz, lvlSizes[l]);
693  }
694  }
695  if (allDense && initializeValuesIfAllDense)
696  values.resize(sz, 0);
697 }
698 
699 template <typename P, typename C, typename V>
701  uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
702  const uint64_t *lvlSizes, const LevelType *lvlTypes,
703  const uint64_t *dim2lvl, const uint64_t *lvl2dim,
704  SparseTensorCOO<V> &lvlCOO)
705  : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
706  dim2lvl, lvl2dim, nullptr, false) {
707  // Ensure lvlCOO is sorted.
708  assert(lvlRank == lvlCOO.getRank());
709  lvlCOO.sort();
710  // Now actually insert the `elements`.
711  const auto &elements = lvlCOO.getElements();
712  const uint64_t nse = elements.size();
713  assert(values.size() == 0);
714  values.reserve(nse);
715  fromCOO(elements, 0, nse, 0);
716 }
717 
718 template <typename P, typename C, typename V>
720  uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
721  const uint64_t *lvlSizes, const LevelType *lvlTypes,
722  const uint64_t *dim2lvl, const uint64_t *lvl2dim, const intptr_t *lvlBufs)
723  : SparseTensorStorage(dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes,
724  dim2lvl, lvl2dim) {
725  // Note that none of the buffers can be reused because ownership
726  // of the memory passed from clients is not necessarily transferred.
727  // Therefore, all data is copied over into a new SparseTensorStorage.
728  uint64_t trailCOOLen = 0, parentSz = 1, bufIdx = 0;
729  for (uint64_t l = 0; l < lvlRank; l++) {
730  if (!isUniqueLvl(l) && (isCompressedLvl(l) || isLooseCompressedLvl(l))) {
731  // A `(loose)compressed_nu` level marks the start of trailing COO
732  // start level. Since the coordinate buffer used for trailing COO
733  // is passed in as AoS scheme and SparseTensorStorage uses a SoA
734  // scheme, we cannot simply copy the value from the provided buffers.
735  trailCOOLen = lvlRank - l;
736  break;
737  }
738  if (isCompressedLvl(l) || isLooseCompressedLvl(l)) {
739  P *posPtr = reinterpret_cast<P *>(lvlBufs[bufIdx++]);
740  C *crdPtr = reinterpret_cast<C *>(lvlBufs[bufIdx++]);
741  if (isLooseCompressedLvl(l)) {
742  positions[l].assign(posPtr, posPtr + 2 * parentSz);
743  coordinates[l].assign(crdPtr, crdPtr + positions[l][2 * parentSz - 1]);
744  } else {
745  positions[l].assign(posPtr, posPtr + parentSz + 1);
746  coordinates[l].assign(crdPtr, crdPtr + positions[l][parentSz]);
747  }
748  } else if (isSingletonLvl(l)) {
749  assert(0 && "general singleton not supported yet");
750  } else if (isNOutOfMLvl(l)) {
751  assert(0 && "n ouf of m not supported yet");
752  } else {
753  assert(isDenseLvl(l));
754  }
755  parentSz = assembledSize(parentSz, l);
756  }
757 
758  // Handle Aos vs. SoA mismatch for COO.
759  if (trailCOOLen != 0) {
760  uint64_t cooStartLvl = lvlRank - trailCOOLen;
761  assert(!isUniqueLvl(cooStartLvl) &&
762  (isCompressedLvl(cooStartLvl) || isLooseCompressedLvl(cooStartLvl)));
763  P *posPtr = reinterpret_cast<P *>(lvlBufs[bufIdx++]);
764  C *aosCrdPtr = reinterpret_cast<C *>(lvlBufs[bufIdx++]);
765  P crdLen;
766  if (isLooseCompressedLvl(cooStartLvl)) {
767  positions[cooStartLvl].assign(posPtr, posPtr + 2 * parentSz);
768  crdLen = positions[cooStartLvl][2 * parentSz - 1];
769  } else {
770  positions[cooStartLvl].assign(posPtr, posPtr + parentSz + 1);
771  crdLen = positions[cooStartLvl][parentSz];
772  }
773  for (uint64_t l = cooStartLvl; l < lvlRank; l++) {
774  coordinates[l].resize(crdLen);
775  for (uint64_t n = 0; n < crdLen; n++) {
776  coordinates[l][n] = *(aosCrdPtr + (l - cooStartLvl) + n * trailCOOLen);
777  }
778  }
779  parentSz = assembledSize(parentSz, cooStartLvl);
780  }
781 
782  // Copy the values buffer.
783  V *valPtr = reinterpret_cast<V *>(lvlBufs[bufIdx]);
784  values.assign(valPtr, valPtr + parentSz);
785 }
786 
787 } // namespace sparse_tensor
788 } // namespace mlir
789 
790 #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_STORAGE_H
#define MLIR_SPARSETENSOR_FOREVERY_FIXED_O(DO)
Definition: Enums.h:63
#define MLIR_SPARSETENSOR_FOREVERY_V(DO)
Definition: Enums.h:96
#define DECL_LEXINSERT(VNAME, V)
Element-wise insertion in lexicographic coordinate order.
Definition: Storage.h:153
#define DECL_GETPOSITIONS(PNAME, P)
Gets positions-overhead storage for the given level.
Definition: Storage.h:135
#define DECL_EXPINSERT(VNAME, V)
Expanded insertion.
Definition: Storage.h:160
#define DECL_GETCOORDINATES(INAME, C)
Gets coordinates-overhead storage for the given level.
Definition: Storage.h:141
#define DECL_GETVALUES(VNAME, V)
Gets primary storage.
Definition: Storage.h:147
A class for capturing the sparse tensor type map with a compact encoding.
Definition: MapRef.h:33
A memory-resident sparse tensor in coordinate-scheme representation (a collection of Elements).
Definition: COO.h:66
const std::vector< Element< V > > & getElements() const
Gets the elements array.
Definition: COO.h:96
void sort()
Sorts elements lexicographically by coordinates.
Definition: COO.h:133
uint64_t getRank() const
Gets the dimension-rank of the tensor.
Definition: COO.h:90
Abstract base class for SparseTensorStorage<P,C,V>.
Definition: Storage.h:63
uint64_t getLvlRank() const
Gets the number of storage-levels.
Definition: Storage.h:80
const std::vector< LevelType > & getLvlTypes() const
Gets the level-types array.
Definition: Storage.h:101
const std::vector< uint64_t > & getLvlSizes() const
Gets the storage-level sizes array.
Definition: Storage.h:92
bool isSingletonLvl(uint64_t l) const
Safely checks if the level uses singleton storage.
Definition: Storage.h:123
bool isCompressedLvl(uint64_t l) const
Safely checks if the level uses compressed storage.
Definition: Storage.h:113
SparseTensorStorageBase & operator=(const SparseTensorStorageBase &)=delete
SparseTensorStorageBase(const SparseTensorStorageBase &)=default
uint64_t getDimRank() const
Gets the number of tensor-dimensions.
Definition: Storage.h:77
const std::vector< uint64_t > & getDimSizes() const
Gets the tensor-dimension sizes array.
Definition: Storage.h:83
bool isDenseLvl(uint64_t l) const
Safely checks if the level uses dense storage.
Definition: Storage.h:110
uint64_t getLvlSize(uint64_t l) const
Safely looks up the size of the given storage-level.
Definition: Storage.h:95
bool isNOutOfMLvl(uint64_t l) const
Safely checks if the level uses n out of m storage.
Definition: Storage.h:126
uint64_t getDimSize(uint64_t d) const
Safely looks up the size of the given tensor-dimension.
Definition: Storage.h:86
bool isOrderedLvl(uint64_t l) const
Safely checks if the level is ordered.
Definition: Storage.h:129
LevelType getLvlType(uint64_t l) const
Safely looks up the type of the given level.
Definition: Storage.h:104
virtual void endLexInsert()=0
Finalizes lexicographic insertions.
bool isLooseCompressedLvl(uint64_t l) const
Safely checks if the level uses loose compressed storage.
Definition: Storage.h:118
bool isUniqueLvl(uint64_t l) const
Safely checks if the level is unique.
Definition: Storage.h:132
A memory-resident sparse tensor using a storage scheme based on per-level sparse/dense annotations.
Definition: Storage.h:189
void getValues(std::vector< V > **out) final
Definition: Storage.h:271
void getCoordinates(std::vector< C > **out, uint64_t lvl) final
Definition: Storage.h:266
static SparseTensorStorage< P, C, V > * packFromLvlBuffers(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, uint64_t srcRank, const intptr_t *buffers)
Allocates a new sparse tensor and initialize it with the data stored level buffers directly.
Definition: Storage.h:640
void getPositions(std::vector< P > **out, uint64_t lvl) final
Partially specialize these getter methods based on template types.
Definition: Storage.h:261
void expInsert(uint64_t *lvlCoords, V *values, bool *filled, uint64_t *added, uint64_t count, uint64_t expsz) final
Partially specialize expanded insertions based on template types.
Definition: Storage.h:301
SparseTensorCOO< V > * toCOO()
Allocates a new COO object and initializes it with the contents.
Definition: Storage.h:342
void lexInsert(const uint64_t *lvlCoords, V val) final
Partially specialize lexicographical insertions based on template types.
Definition: Storage.h:277
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.
Definition: Storage.h:630
static SparseTensorStorage< P, C, V > * newEmpty(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)
Allocates a new empty sparse tensor.
Definition: Storage.h:620
void sortInPlace()
Sort the unordered tensor in place, the method assumes that it is an unordered COO tensor.
Definition: Storage.h:352
void endLexInsert() final
Finalizes lexicographic insertions.
Definition: Storage.h:331
uint64_t checkedMul(uint64_t lhs, uint64_t rhs)
A version of operator* on uint64_t which guards against overflows (when assertions are enabled).
bool isUniqueLT(LevelType lt)
Definition: Enums.h:391
bool isOrderedLT(LevelType lt)
Definition: Enums.h:388
bool isSingletonLT(LevelType lt)
Definition: Enums.h:384
bool isCompressedLT(LevelType lt)
Definition: Enums.h:378
bool isLooseCompressedLT(LevelType lt)
Definition: Enums.h:381
bool isDenseLT(LevelType lt)
Definition: Enums.h:377
bool isNOutOfMLT(LevelType lt)
Definition: Enums.h:387
Include the generated interface declarations.
This enum defines all the sparse representations supportable by the SparseTensor dialect.
Definition: Enums.h:222