18a91bc7bSHarrietAkot //===- SparseTensorUtils.cpp - Sparse Tensor Utils for MLIR execution -----===//
28a91bc7bSHarrietAkot //
38a91bc7bSHarrietAkot // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
48a91bc7bSHarrietAkot // See https://llvm.org/LICENSE.txt for license information.
58a91bc7bSHarrietAkot // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
68a91bc7bSHarrietAkot //
78a91bc7bSHarrietAkot //===----------------------------------------------------------------------===//
88a91bc7bSHarrietAkot //
98a91bc7bSHarrietAkot // This file implements a light-weight runtime support library that is useful
108a91bc7bSHarrietAkot // for sparse tensor manipulations. The functionality provided in this library
118a91bc7bSHarrietAkot // is meant to simplify benchmarking, testing, and debugging MLIR code that
128a91bc7bSHarrietAkot // operates on sparse tensors. The provided functionality is **not** part
138a91bc7bSHarrietAkot // of core MLIR, however.
148a91bc7bSHarrietAkot //
158a91bc7bSHarrietAkot //===----------------------------------------------------------------------===//
168a91bc7bSHarrietAkot 
17845561ecSwren romano #include "mlir/ExecutionEngine/SparseTensorUtils.h"
188a91bc7bSHarrietAkot 
198a91bc7bSHarrietAkot #ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
208a91bc7bSHarrietAkot 
218a91bc7bSHarrietAkot #include <algorithm>
228a91bc7bSHarrietAkot #include <cassert>
238a91bc7bSHarrietAkot #include <cctype>
248a91bc7bSHarrietAkot #include <cstdio>
258a91bc7bSHarrietAkot #include <cstdlib>
268a91bc7bSHarrietAkot #include <cstring>
27efa15f41SAart Bik #include <fstream>
28753fe330Swren romano #include <functional>
29efa15f41SAart Bik #include <iostream>
304d0a18d0Swren romano #include <limits>
318a91bc7bSHarrietAkot #include <numeric>
32736c1b66SAart Bik 
338a91bc7bSHarrietAkot //===----------------------------------------------------------------------===//
348a91bc7bSHarrietAkot //
358a91bc7bSHarrietAkot // Internal support for storing and reading sparse tensors.
368a91bc7bSHarrietAkot //
378a91bc7bSHarrietAkot // The following memory-resident sparse storage schemes are supported:
388a91bc7bSHarrietAkot //
398a91bc7bSHarrietAkot // (a) A coordinate scheme for temporarily storing and lexicographically
408a91bc7bSHarrietAkot //     sorting a sparse tensor by index (SparseTensorCOO).
418a91bc7bSHarrietAkot //
428a91bc7bSHarrietAkot // (b) A "one-size-fits-all" sparse tensor storage scheme defined by
438a91bc7bSHarrietAkot //     per-dimension sparse/dense annnotations together with a dimension
448a91bc7bSHarrietAkot //     ordering used by MLIR compiler-generated code (SparseTensorStorage).
458a91bc7bSHarrietAkot //
468a91bc7bSHarrietAkot // The following external formats are supported:
478a91bc7bSHarrietAkot //
488a91bc7bSHarrietAkot // (1) Matrix Market Exchange (MME): *.mtx
498a91bc7bSHarrietAkot //     https://math.nist.gov/MatrixMarket/formats.html
508a91bc7bSHarrietAkot //
518a91bc7bSHarrietAkot // (2) Formidable Repository of Open Sparse Tensors and Tools (FROSTT): *.tns
528a91bc7bSHarrietAkot //     http://frostt.io/tensors/file-formats.html
538a91bc7bSHarrietAkot //
548a91bc7bSHarrietAkot // Two public APIs are supported:
558a91bc7bSHarrietAkot //
568a91bc7bSHarrietAkot // (I) Methods operating on MLIR buffers (memrefs) to interact with sparse
578a91bc7bSHarrietAkot //     tensors. These methods should be used exclusively by MLIR
588a91bc7bSHarrietAkot //     compiler-generated code.
598a91bc7bSHarrietAkot //
608a91bc7bSHarrietAkot // (II) Methods that accept C-style data structures to interact with sparse
618a91bc7bSHarrietAkot //      tensors. These methods can be used by any external runtime that wants
628a91bc7bSHarrietAkot //      to interact with MLIR compiler-generated code.
638a91bc7bSHarrietAkot //
648a91bc7bSHarrietAkot // In both cases (I) and (II), the SparseTensorStorage format is externally
658a91bc7bSHarrietAkot // only visible as an opaque pointer.
668a91bc7bSHarrietAkot //
678a91bc7bSHarrietAkot //===----------------------------------------------------------------------===//
688a91bc7bSHarrietAkot 
698a91bc7bSHarrietAkot namespace {
708a91bc7bSHarrietAkot 
7103fe15ceSAart Bik static constexpr int kColWidth = 1025;
7203fe15ceSAart Bik 
7372ec2f76Swren romano /// A version of `operator*` on `uint64_t` which checks for overflows.
7472ec2f76Swren romano static inline uint64_t checkedMul(uint64_t lhs, uint64_t rhs) {
7572ec2f76Swren romano   assert((lhs == 0 || rhs <= std::numeric_limits<uint64_t>::max() / lhs) &&
7672ec2f76Swren romano          "Integer overflow");
7772ec2f76Swren romano   return lhs * rhs;
7872ec2f76Swren romano }
7972ec2f76Swren romano 
80774674ceSwren romano // This macro helps minimize repetition of this idiom, as well as ensuring
81774674ceSwren romano // we have some additional output indicating where the error is coming from.
82774674ceSwren romano // (Since `fprintf` doesn't provide a stacktrace, this helps make it easier
83774674ceSwren romano // to track down whether an error is coming from our code vs somewhere else
84774674ceSwren romano // in MLIR.)
85774674ceSwren romano #define FATAL(...)                                                             \
86774674ceSwren romano   {                                                                            \
87774674ceSwren romano     fprintf(stderr, "SparseTensorUtils: " __VA_ARGS__);                        \
88774674ceSwren romano     exit(1);                                                                   \
89774674ceSwren romano   }
90774674ceSwren romano 
918cb33240Swren romano // TODO: adjust this so it can be used by `openSparseTensorCOO` too.
92fa6aed2aSwren romano // That version doesn't have the permutation, and the `dimSizes` are
938cb33240Swren romano // a pointer/C-array rather than `std::vector`.
948cb33240Swren romano //
95fa6aed2aSwren romano /// Asserts that the `dimSizes` (in target-order) under the `perm` (mapping
968cb33240Swren romano /// semantic-order to target-order) are a refinement of the desired `shape`
978cb33240Swren romano /// (in semantic-order).
988cb33240Swren romano ///
998cb33240Swren romano /// Precondition: `perm` and `shape` must be valid for `rank`.
1008cb33240Swren romano static inline void
101fa6aed2aSwren romano assertPermutedSizesMatchShape(const std::vector<uint64_t> &dimSizes,
102fa6aed2aSwren romano                               uint64_t rank, const uint64_t *perm,
103fa6aed2aSwren romano                               const uint64_t *shape) {
1048cb33240Swren romano   assert(perm && shape);
105fa6aed2aSwren romano   assert(rank == dimSizes.size() && "Rank mismatch");
1068cb33240Swren romano   for (uint64_t r = 0; r < rank; r++)
107fa6aed2aSwren romano     assert((shape[r] == 0 || shape[r] == dimSizes[perm[r]]) &&
1088cb33240Swren romano            "Dimension size mismatch");
1098cb33240Swren romano }
1108cb33240Swren romano 
1118a91bc7bSHarrietAkot /// A sparse tensor element in coordinate scheme (value and indices).
1128a91bc7bSHarrietAkot /// For example, a rank-1 vector element would look like
1138a91bc7bSHarrietAkot ///   ({i}, a[i])
1148a91bc7bSHarrietAkot /// and a rank-5 tensor element like
1158a91bc7bSHarrietAkot ///   ({i,j,k,l,m}, a[i,j,k,l,m])
116ccd047cbSAart Bik /// We use pointer to a shared index pool rather than e.g. a direct
117ccd047cbSAart Bik /// vector since that (1) reduces the per-element memory footprint, and
118ccd047cbSAart Bik /// (2) centralizes the memory reservation and (re)allocation to one place.
1198a91bc7bSHarrietAkot template <typename V>
12076944420Swren romano struct Element final {
121ccd047cbSAart Bik   Element(uint64_t *ind, V val) : indices(ind), value(val){};
122ccd047cbSAart Bik   uint64_t *indices; // pointer into shared index pool
1238a91bc7bSHarrietAkot   V value;
1248a91bc7bSHarrietAkot };
1258a91bc7bSHarrietAkot 
126753fe330Swren romano /// The type of callback functions which receive an element.  We avoid
127753fe330Swren romano /// packaging the coordinates and value together as an `Element` object
128753fe330Swren romano /// because this helps keep code somewhat cleaner.
129753fe330Swren romano template <typename V>
130753fe330Swren romano using ElementConsumer =
131753fe330Swren romano     const std::function<void(const std::vector<uint64_t> &, V)> &;
132753fe330Swren romano 
1338a91bc7bSHarrietAkot /// A memory-resident sparse tensor in coordinate scheme (collection of
1348a91bc7bSHarrietAkot /// elements). This data structure is used to read a sparse tensor from
1358a91bc7bSHarrietAkot /// any external format into memory and sort the elements lexicographically
1368a91bc7bSHarrietAkot /// by indices before passing it back to the client (most packed storage
1378a91bc7bSHarrietAkot /// formats require the elements to appear in lexicographic index order).
1388a91bc7bSHarrietAkot template <typename V>
13976944420Swren romano struct SparseTensorCOO final {
1408a91bc7bSHarrietAkot public:
141fa6aed2aSwren romano   SparseTensorCOO(const std::vector<uint64_t> &dimSizes, uint64_t capacity)
142fa6aed2aSwren romano       : dimSizes(dimSizes) {
143ccd047cbSAart Bik     if (capacity) {
1448a91bc7bSHarrietAkot       elements.reserve(capacity);
145ccd047cbSAart Bik       indices.reserve(capacity * getRank());
1468a91bc7bSHarrietAkot     }
147ccd047cbSAart Bik   }
148ccd047cbSAart Bik 
1498a91bc7bSHarrietAkot   /// Adds element as indices and value.
1508a91bc7bSHarrietAkot   void add(const std::vector<uint64_t> &ind, V val) {
1518a91bc7bSHarrietAkot     assert(!iteratorLocked && "Attempt to add() after startIterator()");
152ccd047cbSAart Bik     uint64_t *base = indices.data();
153ccd047cbSAart Bik     uint64_t size = indices.size();
1548a91bc7bSHarrietAkot     uint64_t rank = getRank();
155fa6aed2aSwren romano     assert(ind.size() == rank && "Element rank mismatch");
156ccd047cbSAart Bik     for (uint64_t r = 0; r < rank; r++) {
157fa6aed2aSwren romano       assert(ind[r] < dimSizes[r] && "Index is too large for the dimension");
158ccd047cbSAart Bik       indices.push_back(ind[r]);
1598a91bc7bSHarrietAkot     }
160ccd047cbSAart Bik     // This base only changes if indices were reallocated. In that case, we
161ccd047cbSAart Bik     // need to correct all previous pointers into the vector. Note that this
162ccd047cbSAart Bik     // only happens if we did not set the initial capacity right, and then only
163ccd047cbSAart Bik     // for every internal vector reallocation (which with the doubling rule
164ccd047cbSAart Bik     // should only incur an amortized linear overhead).
165298d2fa1SMehdi Amini     uint64_t *newBase = indices.data();
166298d2fa1SMehdi Amini     if (newBase != base) {
167ccd047cbSAart Bik       for (uint64_t i = 0, n = elements.size(); i < n; i++)
168298d2fa1SMehdi Amini         elements[i].indices = newBase + (elements[i].indices - base);
169298d2fa1SMehdi Amini       base = newBase;
170ccd047cbSAart Bik     }
171ccd047cbSAart Bik     // Add element as (pointer into shared index pool, value) pair.
172ccd047cbSAart Bik     elements.emplace_back(base + size, val);
173ccd047cbSAart Bik   }
174ccd047cbSAart Bik 
1758a91bc7bSHarrietAkot   /// Sorts elements lexicographically by index.
1768a91bc7bSHarrietAkot   void sort() {
1778a91bc7bSHarrietAkot     assert(!iteratorLocked && "Attempt to sort() after startIterator()");
178cf358253Swren romano     // TODO: we may want to cache an `isSorted` bit, to avoid
179cf358253Swren romano     // unnecessary/redundant sorting.
180ccd047cbSAart Bik     uint64_t rank = getRank();
181aff9c89fSwren romano     std::sort(elements.begin(), elements.end(),
182aff9c89fSwren romano               [rank](const Element<V> &e1, const Element<V> &e2) {
183ccd047cbSAart Bik                 for (uint64_t r = 0; r < rank; r++) {
184ccd047cbSAart Bik                   if (e1.indices[r] == e2.indices[r])
185ccd047cbSAart Bik                     continue;
186ccd047cbSAart Bik                   return e1.indices[r] < e2.indices[r];
1878a91bc7bSHarrietAkot                 }
188ccd047cbSAart Bik                 return false;
189ccd047cbSAart Bik               });
190ccd047cbSAart Bik   }
191ccd047cbSAart Bik 
192fa6aed2aSwren romano   /// Get the rank of the tensor.
193fa6aed2aSwren romano   uint64_t getRank() const { return dimSizes.size(); }
194ccd047cbSAart Bik 
195fa6aed2aSwren romano   /// Getter for the dimension-sizes array.
196fa6aed2aSwren romano   const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
197ccd047cbSAart Bik 
198fa6aed2aSwren romano   /// Getter for the elements array.
1998a91bc7bSHarrietAkot   const std::vector<Element<V>> &getElements() const { return elements; }
2008a91bc7bSHarrietAkot 
2018a91bc7bSHarrietAkot   /// Switch into iterator mode.
2028a91bc7bSHarrietAkot   void startIterator() {
2038a91bc7bSHarrietAkot     iteratorLocked = true;
2048a91bc7bSHarrietAkot     iteratorPos = 0;
2058a91bc7bSHarrietAkot   }
206ccd047cbSAart Bik 
2078a91bc7bSHarrietAkot   /// Get the next element.
2088a91bc7bSHarrietAkot   const Element<V> *getNext() {
2098a91bc7bSHarrietAkot     assert(iteratorLocked && "Attempt to getNext() before startIterator()");
2108a91bc7bSHarrietAkot     if (iteratorPos < elements.size())
2118a91bc7bSHarrietAkot       return &(elements[iteratorPos++]);
2128a91bc7bSHarrietAkot     iteratorLocked = false;
2138a91bc7bSHarrietAkot     return nullptr;
2148a91bc7bSHarrietAkot   }
2158a91bc7bSHarrietAkot 
2168a91bc7bSHarrietAkot   /// Factory method. Permutes the original dimensions according to
2178a91bc7bSHarrietAkot   /// the given ordering and expects subsequent add() calls to honor
2188a91bc7bSHarrietAkot   /// that same ordering for the given indices. The result is a
2198a91bc7bSHarrietAkot   /// fully permuted coordinate scheme.
2208d8b566fSwren romano   ///
221fa6aed2aSwren romano   /// Precondition: `dimSizes` and `perm` must be valid for `rank`.
2228a91bc7bSHarrietAkot   static SparseTensorCOO<V> *newSparseTensorCOO(uint64_t rank,
223fa6aed2aSwren romano                                                 const uint64_t *dimSizes,
2248a91bc7bSHarrietAkot                                                 const uint64_t *perm,
2258a91bc7bSHarrietAkot                                                 uint64_t capacity = 0) {
2268a91bc7bSHarrietAkot     std::vector<uint64_t> permsz(rank);
227d83a7068Swren romano     for (uint64_t r = 0; r < rank; r++) {
228fa6aed2aSwren romano       assert(dimSizes[r] > 0 && "Dimension size zero has trivial storage");
229fa6aed2aSwren romano       permsz[perm[r]] = dimSizes[r];
230d83a7068Swren romano     }
2318a91bc7bSHarrietAkot     return new SparseTensorCOO<V>(permsz, capacity);
2328a91bc7bSHarrietAkot   }
2338a91bc7bSHarrietAkot 
2348a91bc7bSHarrietAkot private:
235fa6aed2aSwren romano   const std::vector<uint64_t> dimSizes; // per-dimension sizes
236ccd047cbSAart Bik   std::vector<Element<V>> elements;     // all COO elements
237ccd047cbSAart Bik   std::vector<uint64_t> indices;        // shared index pool
238db6796dfSMehdi Amini   bool iteratorLocked = false;
239db6796dfSMehdi Amini   unsigned iteratorPos = 0;
2408a91bc7bSHarrietAkot };
2418a91bc7bSHarrietAkot 
2428cb33240Swren romano // Forward.
2438cb33240Swren romano template <typename V>
2448cb33240Swren romano class SparseTensorEnumeratorBase;
2458cb33240Swren romano 
246774674ceSwren romano // Helper macro for generating error messages when some
247774674ceSwren romano // `SparseTensorStorage<P,I,V>` is cast to `SparseTensorStorageBase`
248774674ceSwren romano // and then the wrong "partial method specialization" is called.
249774674ceSwren romano #define FATAL_PIV(NAME) FATAL("<P,I,V> type mismatch for: " #NAME);
250774674ceSwren romano 
2518d8b566fSwren romano /// Abstract base class for `SparseTensorStorage<P,I,V>`.  This class
2528d8b566fSwren romano /// takes responsibility for all the `<P,I,V>`-independent aspects
2538d8b566fSwren romano /// of the tensor (e.g., shape, sparsity, permutation).  In addition,
2548d8b566fSwren romano /// we use function overloading to implement "partial" method
2558d8b566fSwren romano /// specialization, which the C-API relies on to catch type errors
2568d8b566fSwren romano /// arising from our use of opaque pointers.
2578a91bc7bSHarrietAkot class SparseTensorStorageBase {
2588a91bc7bSHarrietAkot public:
2598d8b566fSwren romano   /// Constructs a new storage object.  The `perm` maps the tensor's
2608d8b566fSwren romano   /// semantic-ordering of dimensions to this object's storage-order.
261fa6aed2aSwren romano   /// The `dimSizes` and `sparsity` arrays are already in storage-order.
2628d8b566fSwren romano   ///
263fa6aed2aSwren romano   /// Precondition: `perm` and `sparsity` must be valid for `dimSizes.size()`.
264fa6aed2aSwren romano   SparseTensorStorageBase(const std::vector<uint64_t> &dimSizes,
2658d8b566fSwren romano                           const uint64_t *perm, const DimLevelType *sparsity)
266fa6aed2aSwren romano       : dimSizes(dimSizes), rev(getRank()),
2678d8b566fSwren romano         dimTypes(sparsity, sparsity + getRank()) {
268753fe330Swren romano     assert(perm && sparsity);
2698d8b566fSwren romano     const uint64_t rank = getRank();
2708d8b566fSwren romano     // Validate parameters.
2718d8b566fSwren romano     assert(rank > 0 && "Trivial shape is unsupported");
2728d8b566fSwren romano     for (uint64_t r = 0; r < rank; r++) {
2738d8b566fSwren romano       assert(dimSizes[r] > 0 && "Dimension size zero has trivial storage");
2748d8b566fSwren romano       assert((dimTypes[r] == DimLevelType::kDense ||
2758d8b566fSwren romano               dimTypes[r] == DimLevelType::kCompressed) &&
2768d8b566fSwren romano              "Unsupported DimLevelType");
2778d8b566fSwren romano     }
2788d8b566fSwren romano     // Construct the "reverse" (i.e., inverse) permutation.
2798d8b566fSwren romano     for (uint64_t r = 0; r < rank; r++)
2808d8b566fSwren romano       rev[perm[r]] = r;
2818d8b566fSwren romano   }
2828d8b566fSwren romano 
2838d8b566fSwren romano   virtual ~SparseTensorStorageBase() = default;
2848d8b566fSwren romano 
2858d8b566fSwren romano   /// Get the rank of the tensor.
2868d8b566fSwren romano   uint64_t getRank() const { return dimSizes.size(); }
2878d8b566fSwren romano 
2888d8b566fSwren romano   /// Getter for the dimension-sizes array, in storage-order.
2898d8b566fSwren romano   const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
2908d8b566fSwren romano 
2918d8b566fSwren romano   /// Safely lookup the size of the given (storage-order) dimension.
2928d8b566fSwren romano   uint64_t getDimSize(uint64_t d) const {
2938d8b566fSwren romano     assert(d < getRank());
2948d8b566fSwren romano     return dimSizes[d];
2958d8b566fSwren romano   }
2968d8b566fSwren romano 
2978d8b566fSwren romano   /// Getter for the "reverse" permutation, which maps this object's
2988d8b566fSwren romano   /// storage-order to the tensor's semantic-order.
2998d8b566fSwren romano   const std::vector<uint64_t> &getRev() const { return rev; }
3008d8b566fSwren romano 
3018d8b566fSwren romano   /// Getter for the dimension-types array, in storage-order.
3028d8b566fSwren romano   const std::vector<DimLevelType> &getDimTypes() const { return dimTypes; }
3038d8b566fSwren romano 
3048d8b566fSwren romano   /// Safely check if the (storage-order) dimension uses compressed storage.
3058d8b566fSwren romano   bool isCompressedDim(uint64_t d) const {
3068d8b566fSwren romano     assert(d < getRank());
3078d8b566fSwren romano     return (dimTypes[d] == DimLevelType::kCompressed);
3088d8b566fSwren romano   }
3098a91bc7bSHarrietAkot 
3108cb33240Swren romano   /// Allocate a new enumerator.
3111313f5d3Swren romano #define DECL_NEWENUMERATOR(VNAME, V)                                           \
3121313f5d3Swren romano   virtual void newEnumerator(SparseTensorEnumeratorBase<V> **, uint64_t,       \
3131313f5d3Swren romano                              const uint64_t *) const {                         \
314774674ceSwren romano     FATAL_PIV("newEnumerator" #VNAME);                                         \
3158cb33240Swren romano   }
3161313f5d3Swren romano   FOREVERY_V(DECL_NEWENUMERATOR)
3171313f5d3Swren romano #undef DECL_NEWENUMERATOR
3188cb33240Swren romano 
3194f2ec7f9SAart Bik   /// Overhead storage.
320a9a19f59Swren romano #define DECL_GETPOINTERS(PNAME, P)                                             \
321a9a19f59Swren romano   virtual void getPointers(std::vector<P> **, uint64_t) {                      \
322a9a19f59Swren romano     FATAL_PIV("getPointers" #PNAME);                                           \
323774674ceSwren romano   }
324a9a19f59Swren romano   FOREVERY_FIXED_O(DECL_GETPOINTERS)
325a9a19f59Swren romano #undef DECL_GETPOINTERS
326a9a19f59Swren romano #define DECL_GETINDICES(INAME, I)                                              \
327a9a19f59Swren romano   virtual void getIndices(std::vector<I> **, uint64_t) {                       \
328a9a19f59Swren romano     FATAL_PIV("getIndices" #INAME);                                            \
329774674ceSwren romano   }
330a9a19f59Swren romano   FOREVERY_FIXED_O(DECL_GETINDICES)
331a9a19f59Swren romano #undef DECL_GETINDICES
3328a91bc7bSHarrietAkot 
3334f2ec7f9SAart Bik   /// Primary storage.
3341313f5d3Swren romano #define DECL_GETVALUES(VNAME, V)                                               \
335774674ceSwren romano   virtual void getValues(std::vector<V> **) { FATAL_PIV("getValues" #VNAME); }
3361313f5d3Swren romano   FOREVERY_V(DECL_GETVALUES)
3371313f5d3Swren romano #undef DECL_GETVALUES
3388a91bc7bSHarrietAkot 
3394f2ec7f9SAart Bik   /// Element-wise insertion in lexicographic index order.
3401313f5d3Swren romano #define DECL_LEXINSERT(VNAME, V)                                               \
341774674ceSwren romano   virtual void lexInsert(const uint64_t *, V) { FATAL_PIV("lexInsert" #VNAME); }
3421313f5d3Swren romano   FOREVERY_V(DECL_LEXINSERT)
3431313f5d3Swren romano #undef DECL_LEXINSERT
3444f2ec7f9SAart Bik 
3454f2ec7f9SAart Bik   /// Expanded insertion.
3461313f5d3Swren romano #define DECL_EXPINSERT(VNAME, V)                                               \
3471313f5d3Swren romano   virtual void expInsert(uint64_t *, V *, bool *, uint64_t *, uint64_t) {      \
348774674ceSwren romano     FATAL_PIV("expInsert" #VNAME);                                             \
3494f2ec7f9SAart Bik   }
3501313f5d3Swren romano   FOREVERY_V(DECL_EXPINSERT)
3511313f5d3Swren romano #undef DECL_EXPINSERT
3524f2ec7f9SAart Bik 
3534f2ec7f9SAart Bik   /// Finishes insertion.
354f66e5769SAart Bik   virtual void endInsert() = 0;
355f66e5769SAart Bik 
356753fe330Swren romano protected:
357753fe330Swren romano   // Since this class is virtual, we must disallow public copying in
358753fe330Swren romano   // order to avoid "slicing".  Since this class has data members,
359753fe330Swren romano   // that means making copying protected.
360753fe330Swren romano   // <https://github.com/isocpp/CppCoreGuidelines/blob/master/CppCoreGuidelines.md#Rc-copy-virtual>
361753fe330Swren romano   SparseTensorStorageBase(const SparseTensorStorageBase &) = default;
362753fe330Swren romano   // Copy-assignment would be implicitly deleted (because `dimSizes`
363753fe330Swren romano   // is const), so we explicitly delete it for clarity.
364753fe330Swren romano   SparseTensorStorageBase &operator=(const SparseTensorStorageBase &) = delete;
365753fe330Swren romano 
3668a91bc7bSHarrietAkot private:
3678d8b566fSwren romano   const std::vector<uint64_t> dimSizes;
3688d8b566fSwren romano   std::vector<uint64_t> rev;
3698d8b566fSwren romano   const std::vector<DimLevelType> dimTypes;
3708a91bc7bSHarrietAkot };
3718a91bc7bSHarrietAkot 
372774674ceSwren romano #undef FATAL_PIV
373774674ceSwren romano 
374753fe330Swren romano // Forward.
375753fe330Swren romano template <typename P, typename I, typename V>
376753fe330Swren romano class SparseTensorEnumerator;
377753fe330Swren romano 
3788a91bc7bSHarrietAkot /// A memory-resident sparse tensor using a storage scheme based on
3798a91bc7bSHarrietAkot /// per-dimension sparse/dense annotations. This data structure provides a
3808a91bc7bSHarrietAkot /// bufferized form of a sparse tensor type. In contrast to generating setup
3818a91bc7bSHarrietAkot /// methods for each differently annotated sparse tensor, this method provides
3828a91bc7bSHarrietAkot /// a convenient "one-size-fits-all" solution that simply takes an input tensor
3838a91bc7bSHarrietAkot /// and annotations to implement all required setup in a general manner.
3848a91bc7bSHarrietAkot template <typename P, typename I, typename V>
38576944420Swren romano class SparseTensorStorage final : public SparseTensorStorageBase {
3868cb33240Swren romano   /// Private constructor to share code between the other constructors.
3878cb33240Swren romano   /// Beware that the object is not necessarily guaranteed to be in a
3888cb33240Swren romano   /// valid state after this constructor alone; e.g., `isCompressedDim(d)`
3898cb33240Swren romano   /// doesn't entail `!(pointers[d].empty())`.
3908cb33240Swren romano   ///
391fa6aed2aSwren romano   /// Precondition: `perm` and `sparsity` must be valid for `dimSizes.size()`.
392fa6aed2aSwren romano   SparseTensorStorage(const std::vector<uint64_t> &dimSizes,
393fa6aed2aSwren romano                       const uint64_t *perm, const DimLevelType *sparsity)
394fa6aed2aSwren romano       : SparseTensorStorageBase(dimSizes, perm, sparsity), pointers(getRank()),
3958cb33240Swren romano         indices(getRank()), idx(getRank()) {}
3968cb33240Swren romano 
3978a91bc7bSHarrietAkot public:
3988a91bc7bSHarrietAkot   /// Constructs a sparse tensor storage scheme with the given dimensions,
3998a91bc7bSHarrietAkot   /// permutation, and per-dimension dense/sparse annotations, using
4008a91bc7bSHarrietAkot   /// the coordinate scheme tensor for the initial contents if provided.
4018d8b566fSwren romano   ///
402fa6aed2aSwren romano   /// Precondition: `perm` and `sparsity` must be valid for `dimSizes.size()`.
403fa6aed2aSwren romano   SparseTensorStorage(const std::vector<uint64_t> &dimSizes,
404fa6aed2aSwren romano                       const uint64_t *perm, const DimLevelType *sparsity,
405fa6aed2aSwren romano                       SparseTensorCOO<V> *coo)
406fa6aed2aSwren romano       : SparseTensorStorage(dimSizes, perm, sparsity) {
4078a91bc7bSHarrietAkot     // Provide hints on capacity of pointers and indices.
408175b9af4SAart Bik     // TODO: needs much fine-tuning based on actual sparsity; currently
409175b9af4SAart Bik     //       we reserve pointer/index space based on all previous dense
410175b9af4SAart Bik     //       dimensions, which works well up to first sparse dim; but
411175b9af4SAart Bik     //       we should really use nnz and dense/sparse distribution.
412f66e5769SAart Bik     bool allDense = true;
413f66e5769SAart Bik     uint64_t sz = 1;
4148d8b566fSwren romano     for (uint64_t r = 0, rank = getRank(); r < rank; r++) {
4158d8b566fSwren romano       if (isCompressedDim(r)) {
416fa6aed2aSwren romano         // TODO: Take a parameter between 1 and `dimSizes[r]`, and multiply
4178d8b566fSwren romano         // `sz` by that before reserving. (For now we just use 1.)
418f66e5769SAart Bik         pointers[r].reserve(sz + 1);
4198d8b566fSwren romano         pointers[r].push_back(0);
420f66e5769SAart Bik         indices[r].reserve(sz);
421f66e5769SAart Bik         sz = 1;
422f66e5769SAart Bik         allDense = false;
4238d8b566fSwren romano       } else { // Dense dimension.
4248d8b566fSwren romano         sz = checkedMul(sz, getDimSizes()[r]);
4258a91bc7bSHarrietAkot       }
4268a91bc7bSHarrietAkot     }
4278a91bc7bSHarrietAkot     // Then assign contents from coordinate scheme tensor if provided.
4288d8b566fSwren romano     if (coo) {
4294d0a18d0Swren romano       // Ensure both preconditions of `fromCOO`.
430fa6aed2aSwren romano       assert(coo->getDimSizes() == getDimSizes() && "Tensor size mismatch");
4318d8b566fSwren romano       coo->sort();
4324d0a18d0Swren romano       // Now actually insert the `elements`.
4338d8b566fSwren romano       const std::vector<Element<V>> &elements = coo->getElements();
434ceda1ae9Swren romano       uint64_t nnz = elements.size();
4358a91bc7bSHarrietAkot       values.reserve(nnz);
436ceda1ae9Swren romano       fromCOO(elements, 0, nnz, 0);
4371ce77b56SAart Bik     } else if (allDense) {
438f66e5769SAart Bik       values.resize(sz, 0);
4398a91bc7bSHarrietAkot     }
4408a91bc7bSHarrietAkot   }
4418a91bc7bSHarrietAkot 
4428cb33240Swren romano   /// Constructs a sparse tensor storage scheme with the given dimensions,
4438cb33240Swren romano   /// permutation, and per-dimension dense/sparse annotations, using
4448cb33240Swren romano   /// the given sparse tensor for the initial contents.
4458cb33240Swren romano   ///
4468cb33240Swren romano   /// Preconditions:
447fa6aed2aSwren romano   /// * `perm` and `sparsity` must be valid for `dimSizes.size()`.
4488cb33240Swren romano   /// * The `tensor` must have the same value type `V`.
449fa6aed2aSwren romano   SparseTensorStorage(const std::vector<uint64_t> &dimSizes,
450fa6aed2aSwren romano                       const uint64_t *perm, const DimLevelType *sparsity,
4518cb33240Swren romano                       const SparseTensorStorageBase &tensor);
4528cb33240Swren romano 
45376944420Swren romano   ~SparseTensorStorage() final override = default;
4548a91bc7bSHarrietAkot 
455f66e5769SAart Bik   /// Partially specialize these getter methods based on template types.
45676944420Swren romano   void getPointers(std::vector<P> **out, uint64_t d) final override {
4578a91bc7bSHarrietAkot     assert(d < getRank());
4588a91bc7bSHarrietAkot     *out = &pointers[d];
4598a91bc7bSHarrietAkot   }
46076944420Swren romano   void getIndices(std::vector<I> **out, uint64_t d) final override {
4618a91bc7bSHarrietAkot     assert(d < getRank());
4628a91bc7bSHarrietAkot     *out = &indices[d];
4638a91bc7bSHarrietAkot   }
46476944420Swren romano   void getValues(std::vector<V> **out) final override { *out = &values; }
4658a91bc7bSHarrietAkot 
46603fe15ceSAart Bik   /// Partially specialize lexicographical insertions based on template types.
46776944420Swren romano   void lexInsert(const uint64_t *cursor, V val) final override {
4681ce77b56SAart Bik     // First, wrap up pending insertion path.
4691ce77b56SAart Bik     uint64_t diff = 0;
4701ce77b56SAart Bik     uint64_t top = 0;
4711ce77b56SAart Bik     if (!values.empty()) {
4721ce77b56SAart Bik       diff = lexDiff(cursor);
4731ce77b56SAart Bik       endPath(diff + 1);
4741ce77b56SAart Bik       top = idx[diff] + 1;
4751ce77b56SAart Bik     }
4761ce77b56SAart Bik     // Then continue with insertion path.
4771ce77b56SAart Bik     insPath(cursor, diff, top, val);
478f66e5769SAart Bik   }
479f66e5769SAart Bik 
4804f2ec7f9SAart Bik   /// Partially specialize expanded insertions based on template types.
4814f2ec7f9SAart Bik   /// Note that this method resets the values/filled-switch array back
4824f2ec7f9SAart Bik   /// to all-zero/false while only iterating over the nonzero elements.
4834f2ec7f9SAart Bik   void expInsert(uint64_t *cursor, V *values, bool *filled, uint64_t *added,
48476944420Swren romano                  uint64_t count) final override {
4854f2ec7f9SAart Bik     if (count == 0)
4864f2ec7f9SAart Bik       return;
4874f2ec7f9SAart Bik     // Sort.
4884f2ec7f9SAart Bik     std::sort(added, added + count);
4894f2ec7f9SAart Bik     // Restore insertion path for first insert.
4903bf2ba3bSwren romano     const uint64_t lastDim = getRank() - 1;
4914f2ec7f9SAart Bik     uint64_t index = added[0];
4923bf2ba3bSwren romano     cursor[lastDim] = index;
4934f2ec7f9SAart Bik     lexInsert(cursor, values[index]);
4944f2ec7f9SAart Bik     assert(filled[index]);
4954f2ec7f9SAart Bik     values[index] = 0;
4964f2ec7f9SAart Bik     filled[index] = false;
4974f2ec7f9SAart Bik     // Subsequent insertions are quick.
4984f2ec7f9SAart Bik     for (uint64_t i = 1; i < count; i++) {
4994f2ec7f9SAart Bik       assert(index < added[i] && "non-lexicographic insertion");
5004f2ec7f9SAart Bik       index = added[i];
5013bf2ba3bSwren romano       cursor[lastDim] = index;
5023bf2ba3bSwren romano       insPath(cursor, lastDim, added[i - 1] + 1, values[index]);
5034f2ec7f9SAart Bik       assert(filled[index]);
5043bf2ba3bSwren romano       values[index] = 0;
5054f2ec7f9SAart Bik       filled[index] = false;
5064f2ec7f9SAart Bik     }
5074f2ec7f9SAart Bik   }
5084f2ec7f9SAart Bik 
509f66e5769SAart Bik   /// Finalizes lexicographic insertions.
51076944420Swren romano   void endInsert() final override {
5111ce77b56SAart Bik     if (values.empty())
51272ec2f76Swren romano       finalizeSegment(0);
5131ce77b56SAart Bik     else
5141ce77b56SAart Bik       endPath(0);
5151ce77b56SAart Bik   }
516f66e5769SAart Bik 
5178cb33240Swren romano   void newEnumerator(SparseTensorEnumeratorBase<V> **out, uint64_t rank,
51876944420Swren romano                      const uint64_t *perm) const final override {
5198cb33240Swren romano     *out = new SparseTensorEnumerator<P, I, V>(*this, rank, perm);
5208cb33240Swren romano   }
5218cb33240Swren romano 
5228a91bc7bSHarrietAkot   /// Returns this sparse tensor storage scheme as a new memory-resident
5238a91bc7bSHarrietAkot   /// sparse tensor in coordinate scheme with the given dimension order.
5248d8b566fSwren romano   ///
5258d8b566fSwren romano   /// Precondition: `perm` must be valid for `getRank()`.
526753fe330Swren romano   SparseTensorCOO<V> *toCOO(const uint64_t *perm) const {
5278cb33240Swren romano     SparseTensorEnumeratorBase<V> *enumerator;
5288cb33240Swren romano     newEnumerator(&enumerator, getRank(), perm);
529753fe330Swren romano     SparseTensorCOO<V> *coo =
5308cb33240Swren romano         new SparseTensorCOO<V>(enumerator->permutedSizes(), values.size());
5318cb33240Swren romano     enumerator->forallElements([&coo](const std::vector<uint64_t> &ind, V val) {
532753fe330Swren romano       coo->add(ind, val);
533753fe330Swren romano     });
5348d8b566fSwren romano     // TODO: This assertion assumes there are no stored zeros,
5358d8b566fSwren romano     // or if there are then that we don't filter them out.
5368d8b566fSwren romano     // Cf., <https://github.com/llvm/llvm-project/issues/54179>
5378d8b566fSwren romano     assert(coo->getElements().size() == values.size());
5388cb33240Swren romano     delete enumerator;
5398d8b566fSwren romano     return coo;
5408a91bc7bSHarrietAkot   }
5418a91bc7bSHarrietAkot 
5428a91bc7bSHarrietAkot   /// Factory method. Constructs a sparse tensor storage scheme with the given
5438a91bc7bSHarrietAkot   /// dimensions, permutation, and per-dimension dense/sparse annotations,
5448a91bc7bSHarrietAkot   /// using the coordinate scheme tensor for the initial contents if provided.
5458a91bc7bSHarrietAkot   /// In the latter case, the coordinate scheme must respect the same
5468a91bc7bSHarrietAkot   /// permutation as is desired for the new sparse tensor storage.
5478d8b566fSwren romano   ///
5488d8b566fSwren romano   /// Precondition: `shape`, `perm`, and `sparsity` must be valid for `rank`.
5498a91bc7bSHarrietAkot   static SparseTensorStorage<P, I, V> *
550d83a7068Swren romano   newSparseTensor(uint64_t rank, const uint64_t *shape, const uint64_t *perm,
5518d8b566fSwren romano                   const DimLevelType *sparsity, SparseTensorCOO<V> *coo) {
5528a91bc7bSHarrietAkot     SparseTensorStorage<P, I, V> *n = nullptr;
5538d8b566fSwren romano     if (coo) {
554fa6aed2aSwren romano       const auto &coosz = coo->getDimSizes();
5558cb33240Swren romano       assertPermutedSizesMatchShape(coosz, rank, perm, shape);
5568d8b566fSwren romano       n = new SparseTensorStorage<P, I, V>(coosz, perm, sparsity, coo);
5578a91bc7bSHarrietAkot     } else {
5588a91bc7bSHarrietAkot       std::vector<uint64_t> permsz(rank);
559d83a7068Swren romano       for (uint64_t r = 0; r < rank; r++) {
560d83a7068Swren romano         assert(shape[r] > 0 && "Dimension size zero has trivial storage");
561d83a7068Swren romano         permsz[perm[r]] = shape[r];
562d83a7068Swren romano       }
5638cb33240Swren romano       // We pass the null `coo` to ensure we select the intended constructor.
5648cb33240Swren romano       n = new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, coo);
5658a91bc7bSHarrietAkot     }
5668a91bc7bSHarrietAkot     return n;
5678a91bc7bSHarrietAkot   }
5688a91bc7bSHarrietAkot 
5698cb33240Swren romano   /// Factory method. Constructs a sparse tensor storage scheme with
5708cb33240Swren romano   /// the given dimensions, permutation, and per-dimension dense/sparse
5718cb33240Swren romano   /// annotations, using the sparse tensor for the initial contents.
5728cb33240Swren romano   ///
5738cb33240Swren romano   /// Preconditions:
5748cb33240Swren romano   /// * `shape`, `perm`, and `sparsity` must be valid for `rank`.
5758cb33240Swren romano   /// * The `tensor` must have the same value type `V`.
5768cb33240Swren romano   static SparseTensorStorage<P, I, V> *
5778cb33240Swren romano   newSparseTensor(uint64_t rank, const uint64_t *shape, const uint64_t *perm,
5788cb33240Swren romano                   const DimLevelType *sparsity,
5798cb33240Swren romano                   const SparseTensorStorageBase *source) {
5808cb33240Swren romano     assert(source && "Got nullptr for source");
5818cb33240Swren romano     SparseTensorEnumeratorBase<V> *enumerator;
5828cb33240Swren romano     source->newEnumerator(&enumerator, rank, perm);
5838cb33240Swren romano     const auto &permsz = enumerator->permutedSizes();
5848cb33240Swren romano     assertPermutedSizesMatchShape(permsz, rank, perm, shape);
5858cb33240Swren romano     auto *tensor =
5868cb33240Swren romano         new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, *source);
5878cb33240Swren romano     delete enumerator;
5888cb33240Swren romano     return tensor;
5898cb33240Swren romano   }
5908cb33240Swren romano 
5918a91bc7bSHarrietAkot private:
59272ec2f76Swren romano   /// Appends an arbitrary new position to `pointers[d]`.  This method
59372ec2f76Swren romano   /// checks that `pos` is representable in the `P` type; however, it
59472ec2f76Swren romano   /// does not check that `pos` is semantically valid (i.e., larger than
59572ec2f76Swren romano   /// the previous position and smaller than `indices[d].capacity()`).
5968d8b566fSwren romano   void appendPointer(uint64_t d, uint64_t pos, uint64_t count = 1) {
59772ec2f76Swren romano     assert(isCompressedDim(d));
59872ec2f76Swren romano     assert(pos <= std::numeric_limits<P>::max() &&
5994d0a18d0Swren romano            "Pointer value is too large for the P-type");
60072ec2f76Swren romano     pointers[d].insert(pointers[d].end(), count, static_cast<P>(pos));
6014d0a18d0Swren romano   }
6024d0a18d0Swren romano 
60372ec2f76Swren romano   /// Appends index `i` to dimension `d`, in the semantically general
60472ec2f76Swren romano   /// sense.  For non-dense dimensions, that means appending to the
60572ec2f76Swren romano   /// `indices[d]` array, checking that `i` is representable in the `I`
60672ec2f76Swren romano   /// type; however, we do not verify other semantic requirements (e.g.,
607fa6aed2aSwren romano   /// that `i` is in bounds for `dimSizes[d]`, and not previously occurring
60872ec2f76Swren romano   /// in the same segment).  For dense dimensions, this method instead
60972ec2f76Swren romano   /// appends the appropriate number of zeros to the `values` array,
61072ec2f76Swren romano   /// where `full` is the number of "entries" already written to `values`
61172ec2f76Swren romano   /// for this segment (aka one after the highest index previously appended).
61272ec2f76Swren romano   void appendIndex(uint64_t d, uint64_t full, uint64_t i) {
61372ec2f76Swren romano     if (isCompressedDim(d)) {
6144d0a18d0Swren romano       assert(i <= std::numeric_limits<I>::max() &&
6154d0a18d0Swren romano              "Index value is too large for the I-type");
61672ec2f76Swren romano       indices[d].push_back(static_cast<I>(i));
61772ec2f76Swren romano     } else { // Dense dimension.
61872ec2f76Swren romano       assert(i >= full && "Index was already filled");
61972ec2f76Swren romano       if (i == full)
62072ec2f76Swren romano         return; // Short-circuit, since it'll be a nop.
62172ec2f76Swren romano       if (d + 1 == getRank())
62272ec2f76Swren romano         values.insert(values.end(), i - full, 0);
62372ec2f76Swren romano       else
62472ec2f76Swren romano         finalizeSegment(d + 1, 0, i - full);
62572ec2f76Swren romano     }
6264d0a18d0Swren romano   }
6274d0a18d0Swren romano 
6288cb33240Swren romano   /// Writes the given coordinate to `indices[d][pos]`.  This method
6298cb33240Swren romano   /// checks that `i` is representable in the `I` type; however, it
6308cb33240Swren romano   /// does not check that `i` is semantically valid (i.e., in bounds
631fa6aed2aSwren romano   /// for `dimSizes[d]` and not elsewhere occurring in the same segment).
6328cb33240Swren romano   void writeIndex(uint64_t d, uint64_t pos, uint64_t i) {
6338cb33240Swren romano     assert(isCompressedDim(d));
6348cb33240Swren romano     // Subscript assignment to `std::vector` requires that the `pos`-th
6358cb33240Swren romano     // entry has been initialized; thus we must be sure to check `size()`
6368cb33240Swren romano     // here, instead of `capacity()` as would be ideal.
6378cb33240Swren romano     assert(pos < indices[d].size() && "Index position is out of bounds");
6388cb33240Swren romano     assert(i <= std::numeric_limits<I>::max() &&
6398cb33240Swren romano            "Index value is too large for the I-type");
6408cb33240Swren romano     indices[d][pos] = static_cast<I>(i);
6418cb33240Swren romano   }
6428cb33240Swren romano 
6438cb33240Swren romano   /// Computes the assembled-size associated with the `d`-th dimension,
6448cb33240Swren romano   /// given the assembled-size associated with the `(d-1)`-th dimension.
6458cb33240Swren romano   /// "Assembled-sizes" correspond to the (nominal) sizes of overhead
6468cb33240Swren romano   /// storage, as opposed to "dimension-sizes" which are the cardinality
6478cb33240Swren romano   /// of coordinates for that dimension.
6488cb33240Swren romano   ///
6498cb33240Swren romano   /// Precondition: the `pointers[d]` array must be fully initialized
6508cb33240Swren romano   /// before calling this method.
6518cb33240Swren romano   uint64_t assembledSize(uint64_t parentSz, uint64_t d) const {
6528cb33240Swren romano     if (isCompressedDim(d))
6538cb33240Swren romano       return pointers[d][parentSz];
6548cb33240Swren romano     // else if dense:
6558cb33240Swren romano     return parentSz * getDimSizes()[d];
6568cb33240Swren romano   }
6578cb33240Swren romano 
6588a91bc7bSHarrietAkot   /// Initializes sparse tensor storage scheme from a memory-resident sparse
6598a91bc7bSHarrietAkot   /// tensor in coordinate scheme. This method prepares the pointers and
6608a91bc7bSHarrietAkot   /// indices arrays under the given per-dimension dense/sparse annotations.
6614d0a18d0Swren romano   ///
6624d0a18d0Swren romano   /// Preconditions:
6634d0a18d0Swren romano   /// (1) the `elements` must be lexicographically sorted.
664fa6aed2aSwren romano   /// (2) the indices of every element are valid for `dimSizes` (equal rank
6654d0a18d0Swren romano   ///     and pointwise less-than).
666ceda1ae9Swren romano   void fromCOO(const std::vector<Element<V>> &elements, uint64_t lo,
667ceda1ae9Swren romano                uint64_t hi, uint64_t d) {
668753fe330Swren romano     uint64_t rank = getRank();
669753fe330Swren romano     assert(d <= rank && hi <= elements.size());
6708a91bc7bSHarrietAkot     // Once dimensions are exhausted, insert the numerical values.
671753fe330Swren romano     if (d == rank) {
672c4017f9dSwren romano       assert(lo < hi);
6731ce77b56SAart Bik       values.push_back(elements[lo].value);
6748a91bc7bSHarrietAkot       return;
6758a91bc7bSHarrietAkot     }
6768a91bc7bSHarrietAkot     // Visit all elements in this interval.
6778a91bc7bSHarrietAkot     uint64_t full = 0;
678c4017f9dSwren romano     while (lo < hi) { // If `hi` is unchanged, then `lo < elements.size()`.
6798a91bc7bSHarrietAkot       // Find segment in interval with same index elements in this dimension.
680f66e5769SAart Bik       uint64_t i = elements[lo].indices[d];
6818a91bc7bSHarrietAkot       uint64_t seg = lo + 1;
682f66e5769SAart Bik       while (seg < hi && elements[seg].indices[d] == i)
6838a91bc7bSHarrietAkot         seg++;
6848a91bc7bSHarrietAkot       // Handle segment in interval for sparse or dense dimension.
68572ec2f76Swren romano       appendIndex(d, full, i);
68672ec2f76Swren romano       full = i + 1;
687ceda1ae9Swren romano       fromCOO(elements, lo, seg, d + 1);
6888a91bc7bSHarrietAkot       // And move on to next segment in interval.
6898a91bc7bSHarrietAkot       lo = seg;
6908a91bc7bSHarrietAkot     }
6918a91bc7bSHarrietAkot     // Finalize the sparse pointer structure at this dimension.
69272ec2f76Swren romano     finalizeSegment(d, full);
6938a91bc7bSHarrietAkot   }
6948a91bc7bSHarrietAkot 
69572ec2f76Swren romano   /// Finalize the sparse pointer structure at this dimension.
69672ec2f76Swren romano   void finalizeSegment(uint64_t d, uint64_t full = 0, uint64_t count = 1) {
69772ec2f76Swren romano     if (count == 0)
69872ec2f76Swren romano       return; // Short-circuit, since it'll be a nop.
69972ec2f76Swren romano     if (isCompressedDim(d)) {
70072ec2f76Swren romano       appendPointer(d, indices[d].size(), count);
70172ec2f76Swren romano     } else { // Dense dimension.
7028d8b566fSwren romano       const uint64_t sz = getDimSizes()[d];
70372ec2f76Swren romano       assert(sz >= full && "Segment is overfull");
7048d8b566fSwren romano       count = checkedMul(count, sz - full);
70572ec2f76Swren romano       // For dense storage we must enumerate all the remaining coordinates
70672ec2f76Swren romano       // in this dimension (i.e., coordinates after the last non-zero
70772ec2f76Swren romano       // element), and either fill in their zero values or else recurse
70872ec2f76Swren romano       // to finalize some deeper dimension.
70972ec2f76Swren romano       if (d + 1 == getRank())
71072ec2f76Swren romano         values.insert(values.end(), count, 0);
71172ec2f76Swren romano       else
71272ec2f76Swren romano         finalizeSegment(d + 1, 0, count);
7131ce77b56SAart Bik     }
7141ce77b56SAart Bik   }
7151ce77b56SAart Bik 
7161ce77b56SAart Bik   /// Wraps up a single insertion path, inner to outer.
7171ce77b56SAart Bik   void endPath(uint64_t diff) {
7181ce77b56SAart Bik     uint64_t rank = getRank();
7191ce77b56SAart Bik     assert(diff <= rank);
7201ce77b56SAart Bik     for (uint64_t i = 0; i < rank - diff; i++) {
72172ec2f76Swren romano       const uint64_t d = rank - i - 1;
72272ec2f76Swren romano       finalizeSegment(d, idx[d] + 1);
7231ce77b56SAart Bik     }
7241ce77b56SAart Bik   }
7251ce77b56SAart Bik 
7261ce77b56SAart Bik   /// Continues a single insertion path, outer to inner.
727c03fd1e6Swren romano   void insPath(const uint64_t *cursor, uint64_t diff, uint64_t top, V val) {
7281ce77b56SAart Bik     uint64_t rank = getRank();
7291ce77b56SAart Bik     assert(diff < rank);
7301ce77b56SAart Bik     for (uint64_t d = diff; d < rank; d++) {
7311ce77b56SAart Bik       uint64_t i = cursor[d];
73272ec2f76Swren romano       appendIndex(d, top, i);
7331ce77b56SAart Bik       top = 0;
7341ce77b56SAart Bik       idx[d] = i;
7351ce77b56SAart Bik     }
7361ce77b56SAart Bik     values.push_back(val);
7371ce77b56SAart Bik   }
7381ce77b56SAart Bik 
7391ce77b56SAart Bik   /// Finds the lexicographic differing dimension.
74046bdacaaSwren romano   uint64_t lexDiff(const uint64_t *cursor) const {
7411ce77b56SAart Bik     for (uint64_t r = 0, rank = getRank(); r < rank; r++)
7421ce77b56SAart Bik       if (cursor[r] > idx[r])
7431ce77b56SAart Bik         return r;
7441ce77b56SAart Bik       else
7451ce77b56SAart Bik         assert(cursor[r] == idx[r] && "non-lexicographic insertion");
7461ce77b56SAart Bik     assert(0 && "duplication insertion");
7471ce77b56SAart Bik     return -1u;
7481ce77b56SAart Bik   }
7491ce77b56SAart Bik 
750753fe330Swren romano   // Allow `SparseTensorEnumerator` to access the data-members (to avoid
751753fe330Swren romano   // the cost of virtual-function dispatch in inner loops), without
752753fe330Swren romano   // making them public to other client code.
753753fe330Swren romano   friend class SparseTensorEnumerator<P, I, V>;
754753fe330Swren romano 
7558a91bc7bSHarrietAkot   std::vector<std::vector<P>> pointers;
7568a91bc7bSHarrietAkot   std::vector<std::vector<I>> indices;
7578a91bc7bSHarrietAkot   std::vector<V> values;
7588d8b566fSwren romano   std::vector<uint64_t> idx; // index cursor for lexicographic insertion.
7598a91bc7bSHarrietAkot };
7608a91bc7bSHarrietAkot 
761753fe330Swren romano /// A (higher-order) function object for enumerating the elements of some
762753fe330Swren romano /// `SparseTensorStorage` under a permutation.  That is, the `forallElements`
763753fe330Swren romano /// method encapsulates the loop-nest for enumerating the elements of
764753fe330Swren romano /// the source tensor (in whatever order is best for the source tensor),
765753fe330Swren romano /// and applies a permutation to the coordinates/indices before handing
766753fe330Swren romano /// each element to the callback.  A single enumerator object can be
767753fe330Swren romano /// freely reused for several calls to `forallElements`, just so long
768753fe330Swren romano /// as each call is sequential with respect to one another.
769753fe330Swren romano ///
770753fe330Swren romano /// N.B., this class stores a reference to the `SparseTensorStorageBase`
771753fe330Swren romano /// passed to the constructor; thus, objects of this class must not
772753fe330Swren romano /// outlive the sparse tensor they depend on.
773753fe330Swren romano ///
774753fe330Swren romano /// Design Note: The reason we define this class instead of simply using
775753fe330Swren romano /// `SparseTensorEnumerator<P,I,V>` is because we need to hide/generalize
776753fe330Swren romano /// the `<P,I>` template parameters from MLIR client code (to simplify the
777753fe330Swren romano /// type parameters used for direct sparse-to-sparse conversion).  And the
778753fe330Swren romano /// reason we define the `SparseTensorEnumerator<P,I,V>` subclasses rather
779753fe330Swren romano /// than simply using this class, is to avoid the cost of virtual-method
780753fe330Swren romano /// dispatch within the loop-nest.
781753fe330Swren romano template <typename V>
782753fe330Swren romano class SparseTensorEnumeratorBase {
783753fe330Swren romano public:
784753fe330Swren romano   /// Constructs an enumerator with the given permutation for mapping
785753fe330Swren romano   /// the semantic-ordering of dimensions to the desired target-ordering.
786753fe330Swren romano   ///
787753fe330Swren romano   /// Preconditions:
788753fe330Swren romano   /// * the `tensor` must have the same `V` value type.
789753fe330Swren romano   /// * `perm` must be valid for `rank`.
790753fe330Swren romano   SparseTensorEnumeratorBase(const SparseTensorStorageBase &tensor,
791753fe330Swren romano                              uint64_t rank, const uint64_t *perm)
792753fe330Swren romano       : src(tensor), permsz(src.getRev().size()), reord(getRank()),
793753fe330Swren romano         cursor(getRank()) {
794753fe330Swren romano     assert(perm && "Received nullptr for permutation");
795753fe330Swren romano     assert(rank == getRank() && "Permutation rank mismatch");
796fa6aed2aSwren romano     const auto &rev = src.getRev();           // source-order -> semantic-order
797fa6aed2aSwren romano     const auto &dimSizes = src.getDimSizes(); // in source storage-order
798753fe330Swren romano     for (uint64_t s = 0; s < rank; s++) {     // `s` source storage-order
799753fe330Swren romano       uint64_t t = perm[rev[s]];              // `t` target-order
800753fe330Swren romano       reord[s] = t;
801fa6aed2aSwren romano       permsz[t] = dimSizes[s];
802753fe330Swren romano     }
803753fe330Swren romano   }
804753fe330Swren romano 
805753fe330Swren romano   virtual ~SparseTensorEnumeratorBase() = default;
806753fe330Swren romano 
807753fe330Swren romano   // We disallow copying to help avoid leaking the `src` reference.
808753fe330Swren romano   // (In addition to avoiding the problem of slicing.)
809753fe330Swren romano   SparseTensorEnumeratorBase(const SparseTensorEnumeratorBase &) = delete;
810753fe330Swren romano   SparseTensorEnumeratorBase &
811753fe330Swren romano   operator=(const SparseTensorEnumeratorBase &) = delete;
812753fe330Swren romano 
813753fe330Swren romano   /// Returns the source/target tensor's rank.  (The source-rank and
814753fe330Swren romano   /// target-rank are always equal since we only support permutations.
815753fe330Swren romano   /// Though once we add support for other dimension mappings, this
816753fe330Swren romano   /// method will have to be split in two.)
817753fe330Swren romano   uint64_t getRank() const { return permsz.size(); }
818753fe330Swren romano 
819753fe330Swren romano   /// Returns the target tensor's dimension sizes.
820753fe330Swren romano   const std::vector<uint64_t> &permutedSizes() const { return permsz; }
821753fe330Swren romano 
822753fe330Swren romano   /// Enumerates all elements of the source tensor, permutes their
823753fe330Swren romano   /// indices, and passes the permuted element to the callback.
824753fe330Swren romano   /// The callback must not store the cursor reference directly,
825753fe330Swren romano   /// since this function reuses the storage.  Instead, the callback
826753fe330Swren romano   /// must copy it if they want to keep it.
827753fe330Swren romano   virtual void forallElements(ElementConsumer<V> yield) = 0;
828753fe330Swren romano 
829753fe330Swren romano protected:
830753fe330Swren romano   const SparseTensorStorageBase &src;
831753fe330Swren romano   std::vector<uint64_t> permsz; // in target order.
832753fe330Swren romano   std::vector<uint64_t> reord;  // source storage-order -> target order.
833753fe330Swren romano   std::vector<uint64_t> cursor; // in target order.
834753fe330Swren romano };
835753fe330Swren romano 
836753fe330Swren romano template <typename P, typename I, typename V>
837753fe330Swren romano class SparseTensorEnumerator final : public SparseTensorEnumeratorBase<V> {
838753fe330Swren romano   using Base = SparseTensorEnumeratorBase<V>;
839753fe330Swren romano 
840753fe330Swren romano public:
841753fe330Swren romano   /// Constructs an enumerator with the given permutation for mapping
842753fe330Swren romano   /// the semantic-ordering of dimensions to the desired target-ordering.
843753fe330Swren romano   ///
844753fe330Swren romano   /// Precondition: `perm` must be valid for `rank`.
845753fe330Swren romano   SparseTensorEnumerator(const SparseTensorStorage<P, I, V> &tensor,
846753fe330Swren romano                          uint64_t rank, const uint64_t *perm)
847753fe330Swren romano       : Base(tensor, rank, perm) {}
848753fe330Swren romano 
849f38765a8SMehdi Amini   ~SparseTensorEnumerator() final = default;
850753fe330Swren romano 
851f38765a8SMehdi Amini   void forallElements(ElementConsumer<V> yield) final {
852753fe330Swren romano     forallElements(yield, 0, 0);
853753fe330Swren romano   }
854753fe330Swren romano 
855753fe330Swren romano private:
856753fe330Swren romano   /// The recursive component of the public `forallElements`.
857753fe330Swren romano   void forallElements(ElementConsumer<V> yield, uint64_t parentPos,
858753fe330Swren romano                       uint64_t d) {
859753fe330Swren romano     // Recover the `<P,I,V>` type parameters of `src`.
860753fe330Swren romano     const auto &src =
861753fe330Swren romano         static_cast<const SparseTensorStorage<P, I, V> &>(this->src);
862753fe330Swren romano     if (d == Base::getRank()) {
863753fe330Swren romano       assert(parentPos < src.values.size() &&
864753fe330Swren romano              "Value position is out of bounds");
865753fe330Swren romano       // TODO: <https://github.com/llvm/llvm-project/issues/54179>
866753fe330Swren romano       yield(this->cursor, src.values[parentPos]);
867753fe330Swren romano     } else if (src.isCompressedDim(d)) {
868753fe330Swren romano       // Look up the bounds of the `d`-level segment determined by the
869753fe330Swren romano       // `d-1`-level position `parentPos`.
870753fe330Swren romano       const std::vector<P> &pointers_d = src.pointers[d];
871753fe330Swren romano       assert(parentPos + 1 < pointers_d.size() &&
872753fe330Swren romano              "Parent pointer position is out of bounds");
873753fe330Swren romano       const uint64_t pstart = static_cast<uint64_t>(pointers_d[parentPos]);
874753fe330Swren romano       const uint64_t pstop = static_cast<uint64_t>(pointers_d[parentPos + 1]);
875753fe330Swren romano       // Loop-invariant code for looking up the `d`-level coordinates/indices.
876753fe330Swren romano       const std::vector<I> &indices_d = src.indices[d];
8773b13f880SAart Bik       assert(pstop <= indices_d.size() && "Index position is out of bounds");
878753fe330Swren romano       uint64_t &cursor_reord_d = this->cursor[this->reord[d]];
879753fe330Swren romano       for (uint64_t pos = pstart; pos < pstop; pos++) {
880753fe330Swren romano         cursor_reord_d = static_cast<uint64_t>(indices_d[pos]);
881753fe330Swren romano         forallElements(yield, pos, d + 1);
882753fe330Swren romano       }
883753fe330Swren romano     } else { // Dense dimension.
884753fe330Swren romano       const uint64_t sz = src.getDimSizes()[d];
885753fe330Swren romano       const uint64_t pstart = parentPos * sz;
886753fe330Swren romano       uint64_t &cursor_reord_d = this->cursor[this->reord[d]];
887753fe330Swren romano       for (uint64_t i = 0; i < sz; i++) {
888753fe330Swren romano         cursor_reord_d = i;
889753fe330Swren romano         forallElements(yield, pstart + i, d + 1);
890753fe330Swren romano       }
891753fe330Swren romano     }
892753fe330Swren romano   }
893753fe330Swren romano };
894753fe330Swren romano 
8958cb33240Swren romano /// Statistics regarding the number of nonzero subtensors in
8968cb33240Swren romano /// a source tensor, for direct sparse=>sparse conversion a la
8978cb33240Swren romano /// <https://arxiv.org/abs/2001.02609>.
8988cb33240Swren romano ///
8998cb33240Swren romano /// N.B., this class stores references to the parameters passed to
9008cb33240Swren romano /// the constructor; thus, objects of this class must not outlive
9018cb33240Swren romano /// those parameters.
90276944420Swren romano class SparseTensorNNZ final {
9038cb33240Swren romano public:
9048cb33240Swren romano   /// Allocate the statistics structure for the desired sizes and
9058cb33240Swren romano   /// sparsity (in the target tensor's storage-order).  This constructor
9068cb33240Swren romano   /// does not actually populate the statistics, however; for that see
9078cb33240Swren romano   /// `initialize`.
9088cb33240Swren romano   ///
909fa6aed2aSwren romano   /// Precondition: `dimSizes` must not contain zeros.
910fa6aed2aSwren romano   SparseTensorNNZ(const std::vector<uint64_t> &dimSizes,
9118cb33240Swren romano                   const std::vector<DimLevelType> &sparsity)
912fa6aed2aSwren romano       : dimSizes(dimSizes), dimTypes(sparsity), nnz(getRank()) {
9138cb33240Swren romano     assert(dimSizes.size() == dimTypes.size() && "Rank mismatch");
9148cb33240Swren romano     bool uncompressed = true;
9158cb33240Swren romano     uint64_t sz = 1; // the product of all `dimSizes` strictly less than `r`.
9168cb33240Swren romano     for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
9178cb33240Swren romano       switch (dimTypes[r]) {
9188cb33240Swren romano       case DimLevelType::kCompressed:
9198cb33240Swren romano         assert(uncompressed &&
9208cb33240Swren romano                "Multiple compressed layers not currently supported");
9218cb33240Swren romano         uncompressed = false;
9228cb33240Swren romano         nnz[r].resize(sz, 0); // Both allocate and zero-initialize.
9238cb33240Swren romano         break;
9248cb33240Swren romano       case DimLevelType::kDense:
9258cb33240Swren romano         assert(uncompressed &&
9268cb33240Swren romano                "Dense after compressed not currently supported");
9278cb33240Swren romano         break;
9288cb33240Swren romano       case DimLevelType::kSingleton:
9298cb33240Swren romano         // Singleton after Compressed causes no problems for allocating
9308cb33240Swren romano         // `nnz` nor for the yieldPos loop.  This remains true even
9318cb33240Swren romano         // when adding support for multiple compressed dimensions or
9328cb33240Swren romano         // for dense-after-compressed.
9338cb33240Swren romano         break;
9348cb33240Swren romano       }
9358cb33240Swren romano       sz = checkedMul(sz, dimSizes[r]);
9368cb33240Swren romano     }
9378cb33240Swren romano   }
9388cb33240Swren romano 
9398cb33240Swren romano   // We disallow copying to help avoid leaking the stored references.
9408cb33240Swren romano   SparseTensorNNZ(const SparseTensorNNZ &) = delete;
9418cb33240Swren romano   SparseTensorNNZ &operator=(const SparseTensorNNZ &) = delete;
9428cb33240Swren romano 
9438cb33240Swren romano   /// Returns the rank of the target tensor.
9448cb33240Swren romano   uint64_t getRank() const { return dimSizes.size(); }
9458cb33240Swren romano 
9468cb33240Swren romano   /// Enumerate the source tensor to fill in the statistics.  The
9478cb33240Swren romano   /// enumerator should already incorporate the permutation (from
9488cb33240Swren romano   /// semantic-order to the target storage-order).
9498cb33240Swren romano   template <typename V>
9508cb33240Swren romano   void initialize(SparseTensorEnumeratorBase<V> &enumerator) {
9518cb33240Swren romano     assert(enumerator.getRank() == getRank() && "Tensor rank mismatch");
9528cb33240Swren romano     assert(enumerator.permutedSizes() == dimSizes && "Tensor size mismatch");
9538cb33240Swren romano     enumerator.forallElements(
9548cb33240Swren romano         [this](const std::vector<uint64_t> &ind, V) { add(ind); });
9558cb33240Swren romano   }
9568cb33240Swren romano 
9578cb33240Swren romano   /// The type of callback functions which receive an nnz-statistic.
9588cb33240Swren romano   using NNZConsumer = const std::function<void(uint64_t)> &;
9598cb33240Swren romano 
9608cb33240Swren romano   /// Lexicographically enumerates all indicies for dimensions strictly
9618cb33240Swren romano   /// less than `stopDim`, and passes their nnz statistic to the callback.
9628cb33240Swren romano   /// Since our use-case only requires the statistic not the coordinates
9638cb33240Swren romano   /// themselves, we do not bother to construct those coordinates.
9648cb33240Swren romano   void forallIndices(uint64_t stopDim, NNZConsumer yield) const {
9658cb33240Swren romano     assert(stopDim < getRank() && "Stopping-dimension is out of bounds");
9668cb33240Swren romano     assert(dimTypes[stopDim] == DimLevelType::kCompressed &&
9678cb33240Swren romano            "Cannot look up non-compressed dimensions");
9688cb33240Swren romano     forallIndices(yield, stopDim, 0, 0);
9698cb33240Swren romano   }
9708cb33240Swren romano 
9718cb33240Swren romano private:
9728cb33240Swren romano   /// Adds a new element (i.e., increment its statistics).  We use
9738cb33240Swren romano   /// a method rather than inlining into the lambda in `initialize`,
9748cb33240Swren romano   /// to avoid spurious templating over `V`.  And this method is private
9758cb33240Swren romano   /// to avoid needing to re-assert validity of `ind` (which is guaranteed
9768cb33240Swren romano   /// by `forallElements`).
9778cb33240Swren romano   void add(const std::vector<uint64_t> &ind) {
9788cb33240Swren romano     uint64_t parentPos = 0;
9798cb33240Swren romano     for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
9808cb33240Swren romano       if (dimTypes[r] == DimLevelType::kCompressed)
9818cb33240Swren romano         nnz[r][parentPos]++;
9828cb33240Swren romano       parentPos = parentPos * dimSizes[r] + ind[r];
9838cb33240Swren romano     }
9848cb33240Swren romano   }
9858cb33240Swren romano 
9868cb33240Swren romano   /// Recursive component of the public `forallIndices`.
9878cb33240Swren romano   void forallIndices(NNZConsumer yield, uint64_t stopDim, uint64_t parentPos,
9888cb33240Swren romano                      uint64_t d) const {
9898cb33240Swren romano     assert(d <= stopDim);
9908cb33240Swren romano     if (d == stopDim) {
9918cb33240Swren romano       assert(parentPos < nnz[d].size() && "Cursor is out of range");
9928cb33240Swren romano       yield(nnz[d][parentPos]);
9938cb33240Swren romano     } else {
9948cb33240Swren romano       const uint64_t sz = dimSizes[d];
9958cb33240Swren romano       const uint64_t pstart = parentPos * sz;
9968cb33240Swren romano       for (uint64_t i = 0; i < sz; i++)
9978cb33240Swren romano         forallIndices(yield, stopDim, pstart + i, d + 1);
9988cb33240Swren romano     }
9998cb33240Swren romano   }
10008cb33240Swren romano 
10018cb33240Swren romano   // All of these are in the target storage-order.
10028cb33240Swren romano   const std::vector<uint64_t> &dimSizes;
10038cb33240Swren romano   const std::vector<DimLevelType> &dimTypes;
10048cb33240Swren romano   std::vector<std::vector<uint64_t>> nnz;
10058cb33240Swren romano };
10068cb33240Swren romano 
10078cb33240Swren romano template <typename P, typename I, typename V>
10088cb33240Swren romano SparseTensorStorage<P, I, V>::SparseTensorStorage(
1009fa6aed2aSwren romano     const std::vector<uint64_t> &dimSizes, const uint64_t *perm,
10108cb33240Swren romano     const DimLevelType *sparsity, const SparseTensorStorageBase &tensor)
1011fa6aed2aSwren romano     : SparseTensorStorage(dimSizes, perm, sparsity) {
10128cb33240Swren romano   SparseTensorEnumeratorBase<V> *enumerator;
10138cb33240Swren romano   tensor.newEnumerator(&enumerator, getRank(), perm);
10148cb33240Swren romano   {
10158cb33240Swren romano     // Initialize the statistics structure.
10168cb33240Swren romano     SparseTensorNNZ nnz(getDimSizes(), getDimTypes());
10178cb33240Swren romano     nnz.initialize(*enumerator);
10188cb33240Swren romano     // Initialize "pointers" overhead (and allocate "indices", "values").
10198cb33240Swren romano     uint64_t parentSz = 1; // assembled-size (not dimension-size) of `r-1`.
10208cb33240Swren romano     for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
10218cb33240Swren romano       if (isCompressedDim(r)) {
10228cb33240Swren romano         pointers[r].reserve(parentSz + 1);
10238cb33240Swren romano         pointers[r].push_back(0);
10248cb33240Swren romano         uint64_t currentPos = 0;
10258cb33240Swren romano         nnz.forallIndices(r, [this, &currentPos, r](uint64_t n) {
10268cb33240Swren romano           currentPos += n;
10278cb33240Swren romano           appendPointer(r, currentPos);
10288cb33240Swren romano         });
10298cb33240Swren romano         assert(pointers[r].size() == parentSz + 1 &&
10308cb33240Swren romano                "Final pointers size doesn't match allocated size");
10318cb33240Swren romano         // That assertion entails `assembledSize(parentSz, r)`
10328cb33240Swren romano         // is now in a valid state.  That is, `pointers[r][parentSz]`
10338cb33240Swren romano         // equals the present value of `currentPos`, which is the
10348cb33240Swren romano         // correct assembled-size for `indices[r]`.
10358cb33240Swren romano       }
10368cb33240Swren romano       // Update assembled-size for the next iteration.
10378cb33240Swren romano       parentSz = assembledSize(parentSz, r);
10388cb33240Swren romano       // Ideally we need only `indices[r].reserve(parentSz)`, however
10398cb33240Swren romano       // the `std::vector` implementation forces us to initialize it too.
10408cb33240Swren romano       // That is, in the yieldPos loop we need random-access assignment
10418cb33240Swren romano       // to `indices[r]`; however, `std::vector`'s subscript-assignment
10428cb33240Swren romano       // only allows assigning to already-initialized positions.
10438cb33240Swren romano       if (isCompressedDim(r))
10448cb33240Swren romano         indices[r].resize(parentSz, 0);
10458cb33240Swren romano     }
10468cb33240Swren romano     values.resize(parentSz, 0); // Both allocate and zero-initialize.
10478cb33240Swren romano   }
10488cb33240Swren romano   // The yieldPos loop
10498cb33240Swren romano   enumerator->forallElements([this](const std::vector<uint64_t> &ind, V val) {
10508cb33240Swren romano     uint64_t parentSz = 1, parentPos = 0;
10518cb33240Swren romano     for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
10528cb33240Swren romano       if (isCompressedDim(r)) {
10538cb33240Swren romano         // If `parentPos == parentSz` then it's valid as an array-lookup;
10548cb33240Swren romano         // however, it's semantically invalid here since that entry
10558cb33240Swren romano         // does not represent a segment of `indices[r]`.  Moreover, that
10568cb33240Swren romano         // entry must be immutable for `assembledSize` to remain valid.
10578cb33240Swren romano         assert(parentPos < parentSz && "Pointers position is out of bounds");
10588cb33240Swren romano         const uint64_t currentPos = pointers[r][parentPos];
10598cb33240Swren romano         // This increment won't overflow the `P` type, since it can't
10608cb33240Swren romano         // exceed the original value of `pointers[r][parentPos+1]`
10618cb33240Swren romano         // which was already verified to be within bounds for `P`
10628cb33240Swren romano         // when it was written to the array.
10638cb33240Swren romano         pointers[r][parentPos]++;
10648cb33240Swren romano         writeIndex(r, currentPos, ind[r]);
10658cb33240Swren romano         parentPos = currentPos;
10668cb33240Swren romano       } else { // Dense dimension.
10678cb33240Swren romano         parentPos = parentPos * getDimSizes()[r] + ind[r];
10688cb33240Swren romano       }
10698cb33240Swren romano       parentSz = assembledSize(parentSz, r);
10708cb33240Swren romano     }
10718cb33240Swren romano     assert(parentPos < values.size() && "Value position is out of bounds");
10728cb33240Swren romano     values[parentPos] = val;
10738cb33240Swren romano   });
10748cb33240Swren romano   // No longer need the enumerator, so we'll delete it ASAP.
10758cb33240Swren romano   delete enumerator;
10768cb33240Swren romano   // The finalizeYieldPos loop
10778cb33240Swren romano   for (uint64_t parentSz = 1, rank = getRank(), r = 0; r < rank; r++) {
10788cb33240Swren romano     if (isCompressedDim(r)) {
10798cb33240Swren romano       assert(parentSz == pointers[r].size() - 1 &&
10808cb33240Swren romano              "Actual pointers size doesn't match the expected size");
10818cb33240Swren romano       // Can't check all of them, but at least we can check the last one.
10828cb33240Swren romano       assert(pointers[r][parentSz - 1] == pointers[r][parentSz] &&
10838cb33240Swren romano              "Pointers got corrupted");
10848cb33240Swren romano       // TODO: optimize this by using `memmove` or similar.
10858cb33240Swren romano       for (uint64_t n = 0; n < parentSz; n++) {
10868cb33240Swren romano         const uint64_t parentPos = parentSz - n;
10878cb33240Swren romano         pointers[r][parentPos] = pointers[r][parentPos - 1];
10888cb33240Swren romano       }
10898cb33240Swren romano       pointers[r][0] = 0;
10908cb33240Swren romano     }
10918cb33240Swren romano     parentSz = assembledSize(parentSz, r);
10928cb33240Swren romano   }
10938cb33240Swren romano }
10948cb33240Swren romano 
10958a91bc7bSHarrietAkot /// Helper to convert string to lower case.
10968a91bc7bSHarrietAkot static char *toLower(char *token) {
10978a91bc7bSHarrietAkot   for (char *c = token; *c; c++)
10988a91bc7bSHarrietAkot     *c = tolower(*c);
10998a91bc7bSHarrietAkot   return token;
11008a91bc7bSHarrietAkot }
11018a91bc7bSHarrietAkot 
11028a91bc7bSHarrietAkot /// Read the MME header of a general sparse matrix of type real.
110303fe15ceSAart Bik static void readMMEHeader(FILE *file, char *filename, char *line,
110433e8ab8eSAart Bik                           uint64_t *idata, bool *isPattern, bool *isSymmetric) {
11058a91bc7bSHarrietAkot   char header[64];
11068a91bc7bSHarrietAkot   char object[64];
11078a91bc7bSHarrietAkot   char format[64];
11088a91bc7bSHarrietAkot   char field[64];
11098a91bc7bSHarrietAkot   char symmetry[64];
11108a91bc7bSHarrietAkot   // Read header line.
11118a91bc7bSHarrietAkot   if (fscanf(file, "%63s %63s %63s %63s %63s\n", header, object, format, field,
1112774674ceSwren romano              symmetry) != 5)
1113774674ceSwren romano     FATAL("Corrupt header in %s\n", filename);
111433e8ab8eSAart Bik   // Set properties
111533e8ab8eSAart Bik   *isPattern = (strcmp(toLower(field), "pattern") == 0);
1116bb56c2b3SMehdi Amini   *isSymmetric = (strcmp(toLower(symmetry), "symmetric") == 0);
11178a91bc7bSHarrietAkot   // Make sure this is a general sparse matrix.
11188a91bc7bSHarrietAkot   if (strcmp(toLower(header), "%%matrixmarket") ||
11198a91bc7bSHarrietAkot       strcmp(toLower(object), "matrix") ||
112033e8ab8eSAart Bik       strcmp(toLower(format), "coordinate") ||
112133e8ab8eSAart Bik       (strcmp(toLower(field), "real") && !(*isPattern)) ||
1122774674ceSwren romano       (strcmp(toLower(symmetry), "general") && !(*isSymmetric)))
1123774674ceSwren romano     FATAL("Cannot find a general sparse matrix in %s\n", filename);
11248a91bc7bSHarrietAkot   // Skip comments.
1125e5639b3fSMehdi Amini   while (true) {
1126774674ceSwren romano     if (!fgets(line, kColWidth, file))
1127774674ceSwren romano       FATAL("Cannot find data in %s\n", filename);
11288a91bc7bSHarrietAkot     if (line[0] != '%')
11298a91bc7bSHarrietAkot       break;
11308a91bc7bSHarrietAkot   }
11318a91bc7bSHarrietAkot   // Next line contains M N NNZ.
11328a91bc7bSHarrietAkot   idata[0] = 2; // rank
11338a91bc7bSHarrietAkot   if (sscanf(line, "%" PRIu64 "%" PRIu64 "%" PRIu64 "\n", idata + 2, idata + 3,
1134774674ceSwren romano              idata + 1) != 3)
1135774674ceSwren romano     FATAL("Cannot find size in %s\n", filename);
11368a91bc7bSHarrietAkot }
11378a91bc7bSHarrietAkot 
11388a91bc7bSHarrietAkot /// Read the "extended" FROSTT header. Although not part of the documented
11398a91bc7bSHarrietAkot /// format, we assume that the file starts with optional comments followed
11408a91bc7bSHarrietAkot /// by two lines that define the rank, the number of nonzeros, and the
11418a91bc7bSHarrietAkot /// dimensions sizes (one per rank) of the sparse tensor.
114203fe15ceSAart Bik static void readExtFROSTTHeader(FILE *file, char *filename, char *line,
114303fe15ceSAart Bik                                 uint64_t *idata) {
11448a91bc7bSHarrietAkot   // Skip comments.
1145e5639b3fSMehdi Amini   while (true) {
1146774674ceSwren romano     if (!fgets(line, kColWidth, file))
1147774674ceSwren romano       FATAL("Cannot find data in %s\n", filename);
11488a91bc7bSHarrietAkot     if (line[0] != '#')
11498a91bc7bSHarrietAkot       break;
11508a91bc7bSHarrietAkot   }
11518a91bc7bSHarrietAkot   // Next line contains RANK and NNZ.
1152774674ceSwren romano   if (sscanf(line, "%" PRIu64 "%" PRIu64 "\n", idata, idata + 1) != 2)
1153774674ceSwren romano     FATAL("Cannot find metadata in %s\n", filename);
11548a91bc7bSHarrietAkot   // Followed by a line with the dimension sizes (one per rank).
1155774674ceSwren romano   for (uint64_t r = 0; r < idata[0]; r++)
1156774674ceSwren romano     if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1)
1157774674ceSwren romano       FATAL("Cannot find dimension size %s\n", filename);
115803fe15ceSAart Bik   fgets(line, kColWidth, file); // end of line
11598a91bc7bSHarrietAkot }
11608a91bc7bSHarrietAkot 
11618a91bc7bSHarrietAkot /// Reads a sparse tensor with the given filename into a memory-resident
11628a91bc7bSHarrietAkot /// sparse tensor in coordinate scheme.
11638a91bc7bSHarrietAkot template <typename V>
11648a91bc7bSHarrietAkot static SparseTensorCOO<V> *openSparseTensorCOO(char *filename, uint64_t rank,
1165d83a7068Swren romano                                                const uint64_t *shape,
11668a91bc7bSHarrietAkot                                                const uint64_t *perm) {
11678a91bc7bSHarrietAkot   // Open the file.
11683734c078Swren romano   assert(filename && "Received nullptr for filename");
1169774674ceSwren romano   FILE *file = fopen(filename, "r");
1170774674ceSwren romano   if (!file)
1171774674ceSwren romano     FATAL("Cannot find file %s\n", filename);
11728a91bc7bSHarrietAkot   // Perform some file format dependent set up.
117303fe15ceSAart Bik   char line[kColWidth];
11748a91bc7bSHarrietAkot   uint64_t idata[512];
117533e8ab8eSAart Bik   bool isPattern = false;
1176bb56c2b3SMehdi Amini   bool isSymmetric = false;
11778a91bc7bSHarrietAkot   if (strstr(filename, ".mtx")) {
117833e8ab8eSAart Bik     readMMEHeader(file, filename, line, idata, &isPattern, &isSymmetric);
11798a91bc7bSHarrietAkot   } else if (strstr(filename, ".tns")) {
118003fe15ceSAart Bik     readExtFROSTTHeader(file, filename, line, idata);
11818a91bc7bSHarrietAkot   } else {
1182774674ceSwren romano     FATAL("Unknown format %s\n", filename);
11838a91bc7bSHarrietAkot   }
11848a91bc7bSHarrietAkot   // Prepare sparse tensor object with per-dimension sizes
11858a91bc7bSHarrietAkot   // and the number of nonzeros as initial capacity.
11868a91bc7bSHarrietAkot   assert(rank == idata[0] && "rank mismatch");
11878a91bc7bSHarrietAkot   uint64_t nnz = idata[1];
11888a91bc7bSHarrietAkot   for (uint64_t r = 0; r < rank; r++)
1189d83a7068Swren romano     assert((shape[r] == 0 || shape[r] == idata[2 + r]) &&
11908a91bc7bSHarrietAkot            "dimension size mismatch");
11918a91bc7bSHarrietAkot   SparseTensorCOO<V> *tensor =
11928a91bc7bSHarrietAkot       SparseTensorCOO<V>::newSparseTensorCOO(rank, idata + 2, perm, nnz);
11938a91bc7bSHarrietAkot   // Read all nonzero elements.
11948a91bc7bSHarrietAkot   std::vector<uint64_t> indices(rank);
11958a91bc7bSHarrietAkot   for (uint64_t k = 0; k < nnz; k++) {
1196774674ceSwren romano     if (!fgets(line, kColWidth, file))
1197774674ceSwren romano       FATAL("Cannot find next line of data in %s\n", filename);
119803fe15ceSAart Bik     char *linePtr = line;
119903fe15ceSAart Bik     for (uint64_t r = 0; r < rank; r++) {
120003fe15ceSAart Bik       uint64_t idx = strtoul(linePtr, &linePtr, 10);
12018a91bc7bSHarrietAkot       // Add 0-based index.
12028a91bc7bSHarrietAkot       indices[perm[r]] = idx - 1;
12038a91bc7bSHarrietAkot     }
12048a91bc7bSHarrietAkot     // The external formats always store the numerical values with the type
12058a91bc7bSHarrietAkot     // double, but we cast these values to the sparse tensor object type.
120633e8ab8eSAart Bik     // For a pattern tensor, we arbitrarily pick the value 1 for all entries.
120733e8ab8eSAart Bik     double value = isPattern ? 1.0 : strtod(linePtr, &linePtr);
12088a91bc7bSHarrietAkot     tensor->add(indices, value);
120902710413SBixia Zheng     // We currently chose to deal with symmetric matrices by fully constructing
121002710413SBixia Zheng     // them. In the future, we may want to make symmetry implicit for storage
121102710413SBixia Zheng     // reasons.
1212bb56c2b3SMehdi Amini     if (isSymmetric && indices[0] != indices[1])
121302710413SBixia Zheng       tensor->add({indices[1], indices[0]}, value);
12148a91bc7bSHarrietAkot   }
12158a91bc7bSHarrietAkot   // Close the file and return tensor.
12168a91bc7bSHarrietAkot   fclose(file);
12178a91bc7bSHarrietAkot   return tensor;
12188a91bc7bSHarrietAkot }
12198a91bc7bSHarrietAkot 
12202046e11aSwren romano /// Writes the sparse tensor to `dest` in extended FROSTT format.
1221efa15f41SAart Bik template <typename V>
122246bdacaaSwren romano static void outSparseTensor(void *tensor, void *dest, bool sort) {
12236438783fSAart Bik   assert(tensor && dest);
12246438783fSAart Bik   auto coo = static_cast<SparseTensorCOO<V> *>(tensor);
12256438783fSAart Bik   if (sort)
12266438783fSAart Bik     coo->sort();
12276438783fSAart Bik   char *filename = static_cast<char *>(dest);
1228fa6aed2aSwren romano   auto &dimSizes = coo->getDimSizes();
12296438783fSAart Bik   auto &elements = coo->getElements();
12306438783fSAart Bik   uint64_t rank = coo->getRank();
1231efa15f41SAart Bik   uint64_t nnz = elements.size();
1232efa15f41SAart Bik   std::fstream file;
1233efa15f41SAart Bik   file.open(filename, std::ios_base::out | std::ios_base::trunc);
1234efa15f41SAart Bik   assert(file.is_open());
1235efa15f41SAart Bik   file << "; extended FROSTT format\n" << rank << " " << nnz << std::endl;
1236efa15f41SAart Bik   for (uint64_t r = 0; r < rank - 1; r++)
1237fa6aed2aSwren romano     file << dimSizes[r] << " ";
1238fa6aed2aSwren romano   file << dimSizes[rank - 1] << std::endl;
1239efa15f41SAart Bik   for (uint64_t i = 0; i < nnz; i++) {
1240efa15f41SAart Bik     auto &idx = elements[i].indices;
1241efa15f41SAart Bik     for (uint64_t r = 0; r < rank; r++)
1242efa15f41SAart Bik       file << (idx[r] + 1) << " ";
1243efa15f41SAart Bik     file << elements[i].value << std::endl;
1244efa15f41SAart Bik   }
1245efa15f41SAart Bik   file.flush();
1246efa15f41SAart Bik   file.close();
1247efa15f41SAart Bik   assert(file.good());
12486438783fSAart Bik }
12496438783fSAart Bik 
12506438783fSAart Bik /// Initializes sparse tensor from an external COO-flavored format.
12516438783fSAart Bik template <typename V>
125246bdacaaSwren romano static SparseTensorStorage<uint64_t, uint64_t, V> *
12536438783fSAart Bik toMLIRSparseTensor(uint64_t rank, uint64_t nse, uint64_t *shape, V *values,
125420eaa88fSBixia Zheng                    uint64_t *indices, uint64_t *perm, uint8_t *sparse) {
125520eaa88fSBixia Zheng   const DimLevelType *sparsity = (DimLevelType *)(sparse);
125620eaa88fSBixia Zheng #ifndef NDEBUG
125720eaa88fSBixia Zheng   // Verify that perm is a permutation of 0..(rank-1).
125820eaa88fSBixia Zheng   std::vector<uint64_t> order(perm, perm + rank);
125920eaa88fSBixia Zheng   std::sort(order.begin(), order.end());
1260774674ceSwren romano   for (uint64_t i = 0; i < rank; ++i)
1261774674ceSwren romano     if (i != order[i])
1262774674ceSwren romano       FATAL("Not a permutation of 0..%" PRIu64 "\n", rank);
126320eaa88fSBixia Zheng 
126420eaa88fSBixia Zheng   // Verify that the sparsity values are supported.
1265774674ceSwren romano   for (uint64_t i = 0; i < rank; ++i)
126620eaa88fSBixia Zheng     if (sparsity[i] != DimLevelType::kDense &&
1267774674ceSwren romano         sparsity[i] != DimLevelType::kCompressed)
1268774674ceSwren romano       FATAL("Unsupported sparsity value %d\n", static_cast<int>(sparsity[i]));
126920eaa88fSBixia Zheng #endif
127020eaa88fSBixia Zheng 
12716438783fSAart Bik   // Convert external format to internal COO.
127263bdcaf9Swren romano   auto *coo = SparseTensorCOO<V>::newSparseTensorCOO(rank, shape, perm, nse);
12736438783fSAart Bik   std::vector<uint64_t> idx(rank);
12746438783fSAart Bik   for (uint64_t i = 0, base = 0; i < nse; i++) {
12756438783fSAart Bik     for (uint64_t r = 0; r < rank; r++)
1276d8b229a1SAart Bik       idx[perm[r]] = indices[base + r];
127763bdcaf9Swren romano     coo->add(idx, values[i]);
12786438783fSAart Bik     base += rank;
12796438783fSAart Bik   }
12806438783fSAart Bik   // Return sparse tensor storage format as opaque pointer.
128163bdcaf9Swren romano   auto *tensor = SparseTensorStorage<uint64_t, uint64_t, V>::newSparseTensor(
128263bdcaf9Swren romano       rank, shape, perm, sparsity, coo);
128363bdcaf9Swren romano   delete coo;
128463bdcaf9Swren romano   return tensor;
12856438783fSAart Bik }
12866438783fSAart Bik 
12876438783fSAart Bik /// Converts a sparse tensor to an external COO-flavored format.
12886438783fSAart Bik template <typename V>
128946bdacaaSwren romano static void fromMLIRSparseTensor(void *tensor, uint64_t *pRank, uint64_t *pNse,
129046bdacaaSwren romano                                  uint64_t **pShape, V **pValues,
129146bdacaaSwren romano                                  uint64_t **pIndices) {
1292736c1b66SAart Bik   assert(tensor);
12936438783fSAart Bik   auto sparseTensor =
12946438783fSAart Bik       static_cast<SparseTensorStorage<uint64_t, uint64_t, V> *>(tensor);
12956438783fSAart Bik   uint64_t rank = sparseTensor->getRank();
12966438783fSAart Bik   std::vector<uint64_t> perm(rank);
12976438783fSAart Bik   std::iota(perm.begin(), perm.end(), 0);
12986438783fSAart Bik   SparseTensorCOO<V> *coo = sparseTensor->toCOO(perm.data());
12996438783fSAart Bik 
13006438783fSAart Bik   const std::vector<Element<V>> &elements = coo->getElements();
13016438783fSAart Bik   uint64_t nse = elements.size();
13026438783fSAart Bik 
13036438783fSAart Bik   uint64_t *shape = new uint64_t[rank];
13046438783fSAart Bik   for (uint64_t i = 0; i < rank; i++)
1305fa6aed2aSwren romano     shape[i] = coo->getDimSizes()[i];
13066438783fSAart Bik 
13076438783fSAart Bik   V *values = new V[nse];
13086438783fSAart Bik   uint64_t *indices = new uint64_t[rank * nse];
13096438783fSAart Bik 
13106438783fSAart Bik   for (uint64_t i = 0, base = 0; i < nse; i++) {
13116438783fSAart Bik     values[i] = elements[i].value;
13126438783fSAart Bik     for (uint64_t j = 0; j < rank; j++)
13136438783fSAart Bik       indices[base + j] = elements[i].indices[j];
13146438783fSAart Bik     base += rank;
13156438783fSAart Bik   }
13166438783fSAart Bik 
13176438783fSAart Bik   delete coo;
13186438783fSAart Bik   *pRank = rank;
13196438783fSAart Bik   *pNse = nse;
13206438783fSAart Bik   *pShape = shape;
13216438783fSAart Bik   *pValues = values;
13226438783fSAart Bik   *pIndices = indices;
1323efa15f41SAart Bik }
1324efa15f41SAart Bik 
13252046e11aSwren romano } // anonymous namespace
13268a91bc7bSHarrietAkot 
13278a91bc7bSHarrietAkot extern "C" {
13288a91bc7bSHarrietAkot 
13298a91bc7bSHarrietAkot //===----------------------------------------------------------------------===//
13308a91bc7bSHarrietAkot //
13312046e11aSwren romano // Public functions which operate on MLIR buffers (memrefs) to interact
13322046e11aSwren romano // with sparse tensors (which are only visible as opaque pointers externally).
13338a91bc7bSHarrietAkot //
13348a91bc7bSHarrietAkot //===----------------------------------------------------------------------===//
13358a91bc7bSHarrietAkot 
13368a91bc7bSHarrietAkot #define CASE(p, i, v, P, I, V)                                                 \
13378a91bc7bSHarrietAkot   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
133863bdcaf9Swren romano     SparseTensorCOO<V> *coo = nullptr;                                         \
1339845561ecSwren romano     if (action <= Action::kFromCOO) {                                          \
1340845561ecSwren romano       if (action == Action::kFromFile) {                                       \
13418a91bc7bSHarrietAkot         char *filename = static_cast<char *>(ptr);                             \
134263bdcaf9Swren romano         coo = openSparseTensorCOO<V>(filename, rank, shape, perm);             \
1343845561ecSwren romano       } else if (action == Action::kFromCOO) {                                 \
134463bdcaf9Swren romano         coo = static_cast<SparseTensorCOO<V> *>(ptr);                          \
13458a91bc7bSHarrietAkot       } else {                                                                 \
1346845561ecSwren romano         assert(action == Action::kEmpty);                                      \
13478a91bc7bSHarrietAkot       }                                                                        \
134863bdcaf9Swren romano       auto *tensor = SparseTensorStorage<P, I, V>::newSparseTensor(            \
134963bdcaf9Swren romano           rank, shape, perm, sparsity, coo);                                   \
135063bdcaf9Swren romano       if (action == Action::kFromFile)                                         \
135163bdcaf9Swren romano         delete coo;                                                            \
135263bdcaf9Swren romano       return tensor;                                                           \
1353bb56c2b3SMehdi Amini     }                                                                          \
13548cb33240Swren romano     if (action == Action::kSparseToSparse) {                                   \
13558cb33240Swren romano       auto *tensor = static_cast<SparseTensorStorageBase *>(ptr);              \
13568cb33240Swren romano       return SparseTensorStorage<P, I, V>::newSparseTensor(rank, shape, perm,  \
13578cb33240Swren romano                                                            sparsity, tensor);  \
13588cb33240Swren romano     }                                                                          \
1359bb56c2b3SMehdi Amini     if (action == Action::kEmptyCOO)                                           \
1360d83a7068Swren romano       return SparseTensorCOO<V>::newSparseTensorCOO(rank, shape, perm);        \
136163bdcaf9Swren romano     coo = static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm);       \
1362845561ecSwren romano     if (action == Action::kToIterator) {                                       \
136363bdcaf9Swren romano       coo->startIterator();                                                    \
13648a91bc7bSHarrietAkot     } else {                                                                   \
1365845561ecSwren romano       assert(action == Action::kToCOO);                                        \
13668a91bc7bSHarrietAkot     }                                                                          \
136763bdcaf9Swren romano     return coo;                                                                \
13688a91bc7bSHarrietAkot   }
13698a91bc7bSHarrietAkot 
1370845561ecSwren romano #define CASE_SECSAME(p, v, P, V) CASE(p, p, v, P, P, V)
13714f2ec7f9SAart Bik 
1372d2215e79SRainer Orth // Assume index_type is in fact uint64_t, so that _mlir_ciface_newSparseTensor
1373bc04a470Swren romano // can safely rewrite kIndex to kU64.  We make this assertion to guarantee
1374bc04a470Swren romano // that this file cannot get out of sync with its header.
1375d2215e79SRainer Orth static_assert(std::is_same<index_type, uint64_t>::value,
1376d2215e79SRainer Orth               "Expected index_type == uint64_t");
1377bc04a470Swren romano 
13788a91bc7bSHarrietAkot void *
1379845561ecSwren romano _mlir_ciface_newSparseTensor(StridedMemRefType<DimLevelType, 1> *aref, // NOLINT
1380d2215e79SRainer Orth                              StridedMemRefType<index_type, 1> *sref,
1381d2215e79SRainer Orth                              StridedMemRefType<index_type, 1> *pref,
1382845561ecSwren romano                              OverheadType ptrTp, OverheadType indTp,
1383845561ecSwren romano                              PrimaryType valTp, Action action, void *ptr) {
13848a91bc7bSHarrietAkot   assert(aref && sref && pref);
13858a91bc7bSHarrietAkot   assert(aref->strides[0] == 1 && sref->strides[0] == 1 &&
13868a91bc7bSHarrietAkot          pref->strides[0] == 1);
13878a91bc7bSHarrietAkot   assert(aref->sizes[0] == sref->sizes[0] && sref->sizes[0] == pref->sizes[0]);
1388845561ecSwren romano   const DimLevelType *sparsity = aref->data + aref->offset;
1389d83a7068Swren romano   const index_type *shape = sref->data + sref->offset;
1390d2215e79SRainer Orth   const index_type *perm = pref->data + pref->offset;
13918a91bc7bSHarrietAkot   uint64_t rank = aref->sizes[0];
13928a91bc7bSHarrietAkot 
1393bc04a470Swren romano   // Rewrite kIndex to kU64, to avoid introducing a bunch of new cases.
1394bc04a470Swren romano   // This is safe because of the static_assert above.
1395bc04a470Swren romano   if (ptrTp == OverheadType::kIndex)
1396bc04a470Swren romano     ptrTp = OverheadType::kU64;
1397bc04a470Swren romano   if (indTp == OverheadType::kIndex)
1398bc04a470Swren romano     indTp = OverheadType::kU64;
1399bc04a470Swren romano 
14008a91bc7bSHarrietAkot   // Double matrices with all combinations of overhead storage.
1401845561ecSwren romano   CASE(OverheadType::kU64, OverheadType::kU64, PrimaryType::kF64, uint64_t,
1402845561ecSwren romano        uint64_t, double);
1403845561ecSwren romano   CASE(OverheadType::kU64, OverheadType::kU32, PrimaryType::kF64, uint64_t,
1404845561ecSwren romano        uint32_t, double);
1405845561ecSwren romano   CASE(OverheadType::kU64, OverheadType::kU16, PrimaryType::kF64, uint64_t,
1406845561ecSwren romano        uint16_t, double);
1407845561ecSwren romano   CASE(OverheadType::kU64, OverheadType::kU8, PrimaryType::kF64, uint64_t,
1408845561ecSwren romano        uint8_t, double);
1409845561ecSwren romano   CASE(OverheadType::kU32, OverheadType::kU64, PrimaryType::kF64, uint32_t,
1410845561ecSwren romano        uint64_t, double);
1411845561ecSwren romano   CASE(OverheadType::kU32, OverheadType::kU32, PrimaryType::kF64, uint32_t,
1412845561ecSwren romano        uint32_t, double);
1413845561ecSwren romano   CASE(OverheadType::kU32, OverheadType::kU16, PrimaryType::kF64, uint32_t,
1414845561ecSwren romano        uint16_t, double);
1415845561ecSwren romano   CASE(OverheadType::kU32, OverheadType::kU8, PrimaryType::kF64, uint32_t,
1416845561ecSwren romano        uint8_t, double);
1417845561ecSwren romano   CASE(OverheadType::kU16, OverheadType::kU64, PrimaryType::kF64, uint16_t,
1418845561ecSwren romano        uint64_t, double);
1419845561ecSwren romano   CASE(OverheadType::kU16, OverheadType::kU32, PrimaryType::kF64, uint16_t,
1420845561ecSwren romano        uint32_t, double);
1421845561ecSwren romano   CASE(OverheadType::kU16, OverheadType::kU16, PrimaryType::kF64, uint16_t,
1422845561ecSwren romano        uint16_t, double);
1423845561ecSwren romano   CASE(OverheadType::kU16, OverheadType::kU8, PrimaryType::kF64, uint16_t,
1424845561ecSwren romano        uint8_t, double);
1425845561ecSwren romano   CASE(OverheadType::kU8, OverheadType::kU64, PrimaryType::kF64, uint8_t,
1426845561ecSwren romano        uint64_t, double);
1427845561ecSwren romano   CASE(OverheadType::kU8, OverheadType::kU32, PrimaryType::kF64, uint8_t,
1428845561ecSwren romano        uint32_t, double);
1429845561ecSwren romano   CASE(OverheadType::kU8, OverheadType::kU16, PrimaryType::kF64, uint8_t,
1430845561ecSwren romano        uint16_t, double);
1431845561ecSwren romano   CASE(OverheadType::kU8, OverheadType::kU8, PrimaryType::kF64, uint8_t,
1432845561ecSwren romano        uint8_t, double);
14338a91bc7bSHarrietAkot 
14348a91bc7bSHarrietAkot   // Float matrices with all combinations of overhead storage.
1435845561ecSwren romano   CASE(OverheadType::kU64, OverheadType::kU64, PrimaryType::kF32, uint64_t,
1436845561ecSwren romano        uint64_t, float);
1437845561ecSwren romano   CASE(OverheadType::kU64, OverheadType::kU32, PrimaryType::kF32, uint64_t,
1438845561ecSwren romano        uint32_t, float);
1439845561ecSwren romano   CASE(OverheadType::kU64, OverheadType::kU16, PrimaryType::kF32, uint64_t,
1440845561ecSwren romano        uint16_t, float);
1441845561ecSwren romano   CASE(OverheadType::kU64, OverheadType::kU8, PrimaryType::kF32, uint64_t,
1442845561ecSwren romano        uint8_t, float);
1443845561ecSwren romano   CASE(OverheadType::kU32, OverheadType::kU64, PrimaryType::kF32, uint32_t,
1444845561ecSwren romano        uint64_t, float);
1445845561ecSwren romano   CASE(OverheadType::kU32, OverheadType::kU32, PrimaryType::kF32, uint32_t,
1446845561ecSwren romano        uint32_t, float);
1447845561ecSwren romano   CASE(OverheadType::kU32, OverheadType::kU16, PrimaryType::kF32, uint32_t,
1448845561ecSwren romano        uint16_t, float);
1449845561ecSwren romano   CASE(OverheadType::kU32, OverheadType::kU8, PrimaryType::kF32, uint32_t,
1450845561ecSwren romano        uint8_t, float);
1451845561ecSwren romano   CASE(OverheadType::kU16, OverheadType::kU64, PrimaryType::kF32, uint16_t,
1452845561ecSwren romano        uint64_t, float);
1453845561ecSwren romano   CASE(OverheadType::kU16, OverheadType::kU32, PrimaryType::kF32, uint16_t,
1454845561ecSwren romano        uint32_t, float);
1455845561ecSwren romano   CASE(OverheadType::kU16, OverheadType::kU16, PrimaryType::kF32, uint16_t,
1456845561ecSwren romano        uint16_t, float);
1457845561ecSwren romano   CASE(OverheadType::kU16, OverheadType::kU8, PrimaryType::kF32, uint16_t,
1458845561ecSwren romano        uint8_t, float);
1459845561ecSwren romano   CASE(OverheadType::kU8, OverheadType::kU64, PrimaryType::kF32, uint8_t,
1460845561ecSwren romano        uint64_t, float);
1461845561ecSwren romano   CASE(OverheadType::kU8, OverheadType::kU32, PrimaryType::kF32, uint8_t,
1462845561ecSwren romano        uint32_t, float);
1463845561ecSwren romano   CASE(OverheadType::kU8, OverheadType::kU16, PrimaryType::kF32, uint8_t,
1464845561ecSwren romano        uint16_t, float);
1465845561ecSwren romano   CASE(OverheadType::kU8, OverheadType::kU8, PrimaryType::kF32, uint8_t,
1466845561ecSwren romano        uint8_t, float);
14678a91bc7bSHarrietAkot 
1468845561ecSwren romano   // Integral matrices with both overheads of the same type.
1469845561ecSwren romano   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI64, uint64_t, int64_t);
1470845561ecSwren romano   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI32, uint64_t, int32_t);
1471845561ecSwren romano   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI16, uint64_t, int16_t);
1472845561ecSwren romano   CASE_SECSAME(OverheadType::kU64, PrimaryType::kI8, uint64_t, int8_t);
14732046e11aSwren romano   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI64, uint32_t, int64_t);
1474845561ecSwren romano   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI32, uint32_t, int32_t);
1475845561ecSwren romano   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI16, uint32_t, int16_t);
1476845561ecSwren romano   CASE_SECSAME(OverheadType::kU32, PrimaryType::kI8, uint32_t, int8_t);
14772046e11aSwren romano   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI64, uint16_t, int64_t);
1478845561ecSwren romano   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI32, uint16_t, int32_t);
1479845561ecSwren romano   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI16, uint16_t, int16_t);
1480845561ecSwren romano   CASE_SECSAME(OverheadType::kU16, PrimaryType::kI8, uint16_t, int8_t);
14812046e11aSwren romano   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI64, uint8_t, int64_t);
1482845561ecSwren romano   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI32, uint8_t, int32_t);
1483845561ecSwren romano   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI16, uint8_t, int16_t);
1484845561ecSwren romano   CASE_SECSAME(OverheadType::kU8, PrimaryType::kI8, uint8_t, int8_t);
14858a91bc7bSHarrietAkot 
1486736c1b66SAart Bik   // Complex matrices with wide overhead.
1487736c1b66SAart Bik   CASE_SECSAME(OverheadType::kU64, PrimaryType::kC64, uint64_t, complex64);
1488736c1b66SAart Bik   CASE_SECSAME(OverheadType::kU64, PrimaryType::kC32, uint64_t, complex32);
1489736c1b66SAart Bik 
14908a91bc7bSHarrietAkot   // Unsupported case (add above if needed).
1491774674ceSwren romano   // TODO: better pretty-printing of enum values!
1492774674ceSwren romano   FATAL("unsupported combination of types: <P=%d, I=%d, V=%d>\n",
1493774674ceSwren romano         static_cast<int>(ptrTp), static_cast<int>(indTp),
1494774674ceSwren romano         static_cast<int>(valTp));
14958a91bc7bSHarrietAkot }
14968a91bc7bSHarrietAkot #undef CASE
14971313f5d3Swren romano #undef CASE_SECSAME
14986438783fSAart Bik 
1499bfadd13dSwren romano #define IMPL_SPARSEVALUES(VNAME, V)                                            \
1500bfadd13dSwren romano   void _mlir_ciface_sparseValues##VNAME(StridedMemRefType<V, 1> *ref,          \
1501bfadd13dSwren romano                                         void *tensor) {                        \
1502bfadd13dSwren romano     assert(ref &&tensor);                                                      \
1503bfadd13dSwren romano     std::vector<V> *v;                                                         \
1504bfadd13dSwren romano     static_cast<SparseTensorStorageBase *>(tensor)->getValues(&v);             \
1505bfadd13dSwren romano     ref->basePtr = ref->data = v->data();                                      \
1506bfadd13dSwren romano     ref->offset = 0;                                                           \
1507bfadd13dSwren romano     ref->sizes[0] = v->size();                                                 \
1508bfadd13dSwren romano     ref->strides[0] = 1;                                                       \
1509bfadd13dSwren romano   }
1510bfadd13dSwren romano FOREVERY_V(IMPL_SPARSEVALUES)
1511bfadd13dSwren romano #undef IMPL_SPARSEVALUES
1512bfadd13dSwren romano 
1513bfadd13dSwren romano #define IMPL_GETOVERHEAD(NAME, TYPE, LIB)                                      \
1514bfadd13dSwren romano   void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor,      \
1515bfadd13dSwren romano                            index_type d) {                                     \
1516bfadd13dSwren romano     assert(ref &&tensor);                                                      \
1517bfadd13dSwren romano     std::vector<TYPE> *v;                                                      \
1518bfadd13dSwren romano     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d);                \
1519bfadd13dSwren romano     ref->basePtr = ref->data = v->data();                                      \
1520bfadd13dSwren romano     ref->offset = 0;                                                           \
1521bfadd13dSwren romano     ref->sizes[0] = v->size();                                                 \
1522bfadd13dSwren romano     ref->strides[0] = 1;                                                       \
1523bfadd13dSwren romano   }
1524a9a19f59Swren romano #define IMPL_SPARSEPOINTERS(PNAME, P)                                          \
1525a9a19f59Swren romano   IMPL_GETOVERHEAD(sparsePointers##PNAME, P, getPointers)
1526a9a19f59Swren romano FOREVERY_O(IMPL_SPARSEPOINTERS)
1527a9a19f59Swren romano #undef IMPL_SPARSEPOINTERS
1528bfadd13dSwren romano 
1529a9a19f59Swren romano #define IMPL_SPARSEINDICES(INAME, I)                                           \
1530a9a19f59Swren romano   IMPL_GETOVERHEAD(sparseIndices##INAME, I, getIndices)
1531a9a19f59Swren romano FOREVERY_O(IMPL_SPARSEINDICES)
1532a9a19f59Swren romano #undef IMPL_SPARSEINDICES
1533bfadd13dSwren romano #undef IMPL_GETOVERHEAD
1534bfadd13dSwren romano 
1535bfadd13dSwren romano #define IMPL_ADDELT(VNAME, V)                                                  \
1536bfadd13dSwren romano   void *_mlir_ciface_addElt##VNAME(void *coo, V value,                         \
1537bfadd13dSwren romano                                    StridedMemRefType<index_type, 1> *iref,     \
1538bfadd13dSwren romano                                    StridedMemRefType<index_type, 1> *pref) {   \
1539bfadd13dSwren romano     assert(coo &&iref &&pref);                                                 \
1540bfadd13dSwren romano     assert(iref->strides[0] == 1 && pref->strides[0] == 1);                    \
1541bfadd13dSwren romano     assert(iref->sizes[0] == pref->sizes[0]);                                  \
1542bfadd13dSwren romano     const index_type *indx = iref->data + iref->offset;                        \
1543bfadd13dSwren romano     const index_type *perm = pref->data + pref->offset;                        \
1544bfadd13dSwren romano     uint64_t isize = iref->sizes[0];                                           \
1545bfadd13dSwren romano     std::vector<index_type> indices(isize);                                    \
1546bfadd13dSwren romano     for (uint64_t r = 0; r < isize; r++)                                       \
1547bfadd13dSwren romano       indices[perm[r]] = indx[r];                                              \
1548bfadd13dSwren romano     static_cast<SparseTensorCOO<V> *>(coo)->add(indices, value);               \
1549bfadd13dSwren romano     return coo;                                                                \
1550bfadd13dSwren romano   }
1551bfadd13dSwren romano FOREVERY_SIMPLEX_V(IMPL_ADDELT)
1552bfadd13dSwren romano IMPL_ADDELT(C64, complex64)
15532046e11aSwren romano // Marked static because it's not part of the public API.
1554*0fbe3f3fSwren romano // NOTE: the `static` keyword confuses clang-format here, causing
1555*0fbe3f3fSwren romano // the strange indentation of the `_mlir_ciface_addEltC32` prototype.
1556*0fbe3f3fSwren romano // In C++11 we can add a semicolon after the call to `IMPL_ADDELT`
1557*0fbe3f3fSwren romano // and that will correct clang-format.  Alas, this file is compiled
1558*0fbe3f3fSwren romano // in C++98 mode where that semicolon is illegal (and there's no portable
1559*0fbe3f3fSwren romano // macro magic to license a no-op semicolon at the top level).
1560*0fbe3f3fSwren romano static IMPL_ADDELT(C32ABI, complex32)
15612046e11aSwren romano #undef IMPL_ADDELT
1562bfadd13dSwren romano     void *_mlir_ciface_addEltC32(void *coo, float r, float i,
1563bfadd13dSwren romano                                  StridedMemRefType<index_type, 1> *iref,
1564bfadd13dSwren romano                                  StridedMemRefType<index_type, 1> *pref) {
1565bfadd13dSwren romano   return _mlir_ciface_addEltC32ABI(coo, complex32(r, i), iref, pref);
1566bfadd13dSwren romano }
1567bfadd13dSwren romano 
1568bfadd13dSwren romano #define IMPL_GETNEXT(VNAME, V)                                                 \
1569bfadd13dSwren romano   bool _mlir_ciface_getNext##VNAME(void *coo,                                  \
1570bfadd13dSwren romano                                    StridedMemRefType<index_type, 1> *iref,     \
1571bfadd13dSwren romano                                    StridedMemRefType<V, 0> *vref) {            \
1572bfadd13dSwren romano     assert(coo &&iref &&vref);                                                 \
1573bfadd13dSwren romano     assert(iref->strides[0] == 1);                                             \
1574bfadd13dSwren romano     index_type *indx = iref->data + iref->offset;                              \
1575bfadd13dSwren romano     V *value = vref->data + vref->offset;                                      \
1576bfadd13dSwren romano     const uint64_t isize = iref->sizes[0];                                     \
1577bfadd13dSwren romano     const Element<V> *elem =                                                   \
1578bfadd13dSwren romano         static_cast<SparseTensorCOO<V> *>(coo)->getNext();                     \
1579bfadd13dSwren romano     if (elem == nullptr)                                                       \
1580bfadd13dSwren romano       return false;                                                            \
1581bfadd13dSwren romano     for (uint64_t r = 0; r < isize; r++)                                       \
1582bfadd13dSwren romano       indx[r] = elem->indices[r];                                              \
1583bfadd13dSwren romano     *value = elem->value;                                                      \
1584bfadd13dSwren romano     return true;                                                               \
1585bfadd13dSwren romano   }
1586bfadd13dSwren romano FOREVERY_V(IMPL_GETNEXT)
1587bfadd13dSwren romano #undef IMPL_GETNEXT
1588bfadd13dSwren romano 
1589bfadd13dSwren romano #define IMPL_LEXINSERT(VNAME, V)                                               \
1590bfadd13dSwren romano   void _mlir_ciface_lexInsert##VNAME(                                          \
1591bfadd13dSwren romano       void *tensor, StridedMemRefType<index_type, 1> *cref, V val) {           \
1592bfadd13dSwren romano     assert(tensor &&cref);                                                     \
1593bfadd13dSwren romano     assert(cref->strides[0] == 1);                                             \
1594bfadd13dSwren romano     index_type *cursor = cref->data + cref->offset;                            \
1595bfadd13dSwren romano     assert(cursor);                                                            \
1596bfadd13dSwren romano     static_cast<SparseTensorStorageBase *>(tensor)->lexInsert(cursor, val);    \
1597bfadd13dSwren romano   }
1598bfadd13dSwren romano FOREVERY_SIMPLEX_V(IMPL_LEXINSERT)
1599bfadd13dSwren romano IMPL_LEXINSERT(C64, complex64)
16002046e11aSwren romano // Marked static because it's not part of the public API.
1601*0fbe3f3fSwren romano // NOTE: see the note for `_mlir_ciface_addEltC32ABI`
1602*0fbe3f3fSwren romano static IMPL_LEXINSERT(C32ABI, complex32)
16032046e11aSwren romano #undef IMPL_LEXINSERT
1604bfadd13dSwren romano     void _mlir_ciface_lexInsertC32(void *tensor,
1605*0fbe3f3fSwren romano                                    StridedMemRefType<index_type, 1> *cref,
1606*0fbe3f3fSwren romano                                    float r, float i) {
1607bfadd13dSwren romano   _mlir_ciface_lexInsertC32ABI(tensor, cref, complex32(r, i));
1608bfadd13dSwren romano }
1609bfadd13dSwren romano 
1610bfadd13dSwren romano #define IMPL_EXPINSERT(VNAME, V)                                               \
1611bfadd13dSwren romano   void _mlir_ciface_expInsert##VNAME(                                          \
1612bfadd13dSwren romano       void *tensor, StridedMemRefType<index_type, 1> *cref,                    \
1613bfadd13dSwren romano       StridedMemRefType<V, 1> *vref, StridedMemRefType<bool, 1> *fref,         \
1614bfadd13dSwren romano       StridedMemRefType<index_type, 1> *aref, index_type count) {              \
1615bfadd13dSwren romano     assert(tensor &&cref &&vref &&fref &&aref);                                \
1616bfadd13dSwren romano     assert(cref->strides[0] == 1);                                             \
1617bfadd13dSwren romano     assert(vref->strides[0] == 1);                                             \
1618bfadd13dSwren romano     assert(fref->strides[0] == 1);                                             \
1619bfadd13dSwren romano     assert(aref->strides[0] == 1);                                             \
1620bfadd13dSwren romano     assert(vref->sizes[0] == fref->sizes[0]);                                  \
1621bfadd13dSwren romano     index_type *cursor = cref->data + cref->offset;                            \
1622bfadd13dSwren romano     V *values = vref->data + vref->offset;                                     \
1623bfadd13dSwren romano     bool *filled = fref->data + fref->offset;                                  \
1624bfadd13dSwren romano     index_type *added = aref->data + aref->offset;                             \
1625bfadd13dSwren romano     static_cast<SparseTensorStorageBase *>(tensor)->expInsert(                 \
1626bfadd13dSwren romano         cursor, values, filled, added, count);                                 \
1627bfadd13dSwren romano   }
1628bfadd13dSwren romano FOREVERY_V(IMPL_EXPINSERT)
1629bfadd13dSwren romano #undef IMPL_EXPINSERT
1630bfadd13dSwren romano 
16318a91bc7bSHarrietAkot //===----------------------------------------------------------------------===//
16328a91bc7bSHarrietAkot //
16332046e11aSwren romano // Public functions which accept only C-style data structures to interact
16342046e11aSwren romano // with sparse tensors (which are only visible as opaque pointers externally).
16358a91bc7bSHarrietAkot //
16368a91bc7bSHarrietAkot //===----------------------------------------------------------------------===//
16378a91bc7bSHarrietAkot 
1638d2215e79SRainer Orth index_type sparseDimSize(void *tensor, index_type d) {
16398a91bc7bSHarrietAkot   return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
16408a91bc7bSHarrietAkot }
16418a91bc7bSHarrietAkot 
1642f66e5769SAart Bik void endInsert(void *tensor) {
1643f66e5769SAart Bik   return static_cast<SparseTensorStorageBase *>(tensor)->endInsert();
1644f66e5769SAart Bik }
1645f66e5769SAart Bik 
164605c17bc4Swren romano #define IMPL_OUTSPARSETENSOR(VNAME, V)                                         \
164705c17bc4Swren romano   void outSparseTensor##VNAME(void *coo, void *dest, bool sort) {              \
164805c17bc4Swren romano     return outSparseTensor<V>(coo, dest, sort);                                \
164905c17bc4Swren romano   }
165005c17bc4Swren romano FOREVERY_V(IMPL_OUTSPARSETENSOR)
165105c17bc4Swren romano #undef IMPL_OUTSPARSETENSOR
165205c17bc4Swren romano 
16538a91bc7bSHarrietAkot void delSparseTensor(void *tensor) {
16548a91bc7bSHarrietAkot   delete static_cast<SparseTensorStorageBase *>(tensor);
16558a91bc7bSHarrietAkot }
16568a91bc7bSHarrietAkot 
165763bdcaf9Swren romano #define IMPL_DELCOO(VNAME, V)                                                  \
165863bdcaf9Swren romano   void delSparseTensorCOO##VNAME(void *coo) {                                  \
165963bdcaf9Swren romano     delete static_cast<SparseTensorCOO<V> *>(coo);                             \
166063bdcaf9Swren romano   }
16611313f5d3Swren romano FOREVERY_V(IMPL_DELCOO)
166263bdcaf9Swren romano #undef IMPL_DELCOO
166363bdcaf9Swren romano 
166405c17bc4Swren romano char *getTensorFilename(index_type id) {
166505c17bc4Swren romano   char var[80];
166605c17bc4Swren romano   sprintf(var, "TENSOR%" PRIu64, id);
166705c17bc4Swren romano   char *env = getenv(var);
166805c17bc4Swren romano   if (!env)
166905c17bc4Swren romano     FATAL("Environment variable %s is not set\n", var);
167005c17bc4Swren romano   return env;
167105c17bc4Swren romano }
167205c17bc4Swren romano 
167320eaa88fSBixia Zheng // TODO: generalize beyond 64-bit indices.
16741313f5d3Swren romano #define IMPL_CONVERTTOMLIRSPARSETENSOR(VNAME, V)                               \
16751313f5d3Swren romano   void *convertToMLIRSparseTensor##VNAME(                                      \
16761313f5d3Swren romano       uint64_t rank, uint64_t nse, uint64_t *shape, V *values,                 \
16771313f5d3Swren romano       uint64_t *indices, uint64_t *perm, uint8_t *sparse) {                    \
16781313f5d3Swren romano     return toMLIRSparseTensor<V>(rank, nse, shape, values, indices, perm,      \
16791313f5d3Swren romano                                  sparse);                                      \
16808a91bc7bSHarrietAkot   }
16811313f5d3Swren romano FOREVERY_V(IMPL_CONVERTTOMLIRSPARSETENSOR)
16821313f5d3Swren romano #undef IMPL_CONVERTTOMLIRSPARSETENSOR
16838a91bc7bSHarrietAkot 
16842f49e6b0SBixia Zheng // TODO: Currently, values are copied from SparseTensorStorage to
16852046e11aSwren romano // SparseTensorCOO, then to the output.  We may want to reduce the number
16862046e11aSwren romano // of copies.
16872f49e6b0SBixia Zheng //
16886438783fSAart Bik // TODO: generalize beyond 64-bit indices, no dim ordering, all dimensions
16896438783fSAart Bik // compressed
16901313f5d3Swren romano #define IMPL_CONVERTFROMMLIRSPARSETENSOR(VNAME, V)                             \
16911313f5d3Swren romano   void convertFromMLIRSparseTensor##VNAME(void *tensor, uint64_t *pRank,       \
16921313f5d3Swren romano                                           uint64_t *pNse, uint64_t **pShape,   \
16931313f5d3Swren romano                                           V **pValues, uint64_t **pIndices) {  \
16941313f5d3Swren romano     fromMLIRSparseTensor<V>(tensor, pRank, pNse, pShape, pValues, pIndices);   \
16952f49e6b0SBixia Zheng   }
16961313f5d3Swren romano FOREVERY_V(IMPL_CONVERTFROMMLIRSPARSETENSOR)
16971313f5d3Swren romano #undef IMPL_CONVERTFROMMLIRSPARSETENSOR
1698efa15f41SAart Bik 
16998a91bc7bSHarrietAkot } // extern "C"
17008a91bc7bSHarrietAkot 
17018a91bc7bSHarrietAkot #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
1702