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, ¤tPos, 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