diff --git a/mlir/lib/ExecutionEngine/SparseUtils.cpp b/mlir/lib/ExecutionEngine/SparseUtils.cpp --- a/mlir/lib/ExecutionEngine/SparseUtils.cpp +++ b/mlir/lib/ExecutionEngine/SparseUtils.cpp @@ -34,10 +34,11 @@ // The following memory-resident sparse storage schemes are supported: // // (a) A coordinate scheme for temporarily storing and lexicographically -// sorting a sparse tensor by index. +// sorting a sparse tensor by index (SparseTensorCOO). // // (b) A "one-size-fits-all" sparse tensor storage scheme defined by per-rank -// sparse/dense annnotations to be used by generated MLIR code. +// sparse/dense annnotations together with a dimension ordering to be +// used by MLIR compiler-generated code (SparseTensorStorage). // // The following external formats are supported: // @@ -47,6 +48,19 @@ // (2) Formidable Repository of Open Sparse Tensors and Tools (FROSTT): *.tns // http://frostt.io/tensors/file-formats.html // +// Two public APIs are supported: +// +// (I) Methods operating on MLIR buffers (memrefs) to interact with sparse +// tensors. These methods should be used exclusively by MLIR +// compiler-generated code. +// +// (II) Methods that accept C-style data structures to interact with sparse +// tensors. These methods can be used by any external runtime that wants +// to interact with MLIR compiler-generated code. +// +// In both cases (I) and (II), the SparseTensorStorage format is externally +// only visible as an opague pointer. +// //===----------------------------------------------------------------------===// namespace { @@ -69,9 +83,9 @@ /// by indices before passing it back to the client (most packed storage /// formats require the elements to appear in lexicographic index order). template -struct SparseTensor { +struct SparseTensorCOO { public: - SparseTensor(const std::vector &szs, uint64_t capacity) + SparseTensorCOO(const std::vector &szs, uint64_t capacity) : sizes(szs), pos(0) { if (capacity) elements.reserve(capacity); @@ -98,13 +112,13 @@ /// the given ordering and expects subsequent add() calls to honor /// that same ordering for the given indices. The result is a /// fully permuted coordinate scheme. - static SparseTensor *newSparseTensor(uint64_t size, uint64_t *sizes, - uint64_t *perm, - uint64_t capacity = 0) { + static SparseTensorCOO *newSparseTensorCOO(uint64_t size, uint64_t *sizes, + uint64_t *perm, + uint64_t capacity = 0) { std::vector permsz(size); for (uint64_t r = 0; r < size; r++) permsz[perm[r]] = sizes[r]; - return new SparseTensor(permsz, capacity); + return new SparseTensorCOO(permsz, capacity); } private: @@ -160,8 +174,7 @@ /// A memory-resident sparse tensor using a storage scheme based on per-rank /// annotations on dense/sparse. This data structure provides a bufferized -/// form of an imaginary SparseTensorType, until such a type becomes a -/// first-class citizen of MLIR. In contrast to generating setup methods for +/// form of a sparse tensor type. In contrast to generating setup methods for /// each differently annotated sparse tensor, this method provides a convenient /// "one-size-fits-all" solution that simply takes an input tensor and /// annotations to implement all required setup in a general manner. @@ -171,7 +184,7 @@ /// Constructs a sparse tensor storage scheme from the given sparse /// tensor in coordinate scheme following the given per-rank dimension /// dense/sparse annotations. - SparseTensorStorage(SparseTensor *tensor, uint8_t *sparsity, + SparseTensorStorage(SparseTensorCOO *tensor, uint8_t *sparsity, uint64_t *perm) : sizes(tensor->getSizes()), rev(getRank()), pointers(getRank()), indices(getRank()) { @@ -217,14 +230,14 @@ /// Returns this sparse tensor storage scheme as a new memory-resident /// sparse tensor in coordinate scheme with the given dimension order. - SparseTensor *toCOO(uint64_t *perm) { + SparseTensorCOO *toCOO(uint64_t *perm) { // Restore original order of the dimension sizes and allocate coordinate // scheme with desired new ordering specified in perm. uint64_t size = getRank(); std::vector orgsz(size); for (uint64_t r = 0; r < size; r++) orgsz[rev[r]] = sizes[r]; - SparseTensor *tensor = SparseTensor::newSparseTensor( + SparseTensorCOO *tensor = SparseTensorCOO::newSparseTensorCOO( size, orgsz.data(), perm, values.size()); // Populate coordinate scheme restored from old ordering and changed with // new ordering. Rather than applying both reorderings during the recursion, @@ -240,7 +253,7 @@ /// Factory method. Expects a coordinate scheme that respects the same /// permutation as is desired for the new sparse storage scheme. static SparseTensorStorage * - newSparseTensor(SparseTensor *t, uint8_t *sparsity, uint64_t *perm) { + newSparseTensor(SparseTensorCOO *t, uint8_t *sparsity, uint64_t *perm) { t->sort(); // sort lexicographically SparseTensorStorage *n = new SparseTensorStorage(t, sparsity, perm); @@ -252,7 +265,7 @@ /// Initializes sparse tensor storage scheme from a memory-resident sparse /// tensor in coordinate scheme. This method prepares the pointers and indices /// arrays under the given per-rank dimension dense/sparse annotations. - void fromCOO(SparseTensor *tensor, uint8_t *sparsity, uint64_t lo, + void fromCOO(SparseTensorCOO *tensor, uint8_t *sparsity, uint64_t lo, uint64_t hi, uint64_t d) { const std::vector> &elements = tensor->getElements(); // Once dimensions are exhausted, insert the numerical values. @@ -296,7 +309,7 @@ /// Stores the sparse tensor storage scheme into a memory-resident sparse /// tensor in coordinate scheme. - void toCOO(SparseTensor *tensor, std::vector &reord, + void toCOO(SparseTensorCOO *tensor, std::vector &reord, std::vector &idx, uint64_t pos, uint64_t d) { if (d == getRank()) { tensor->add(idx, values[pos]); @@ -403,8 +416,9 @@ /// Reads a sparse tensor with the given filename into a memory-resident /// sparse tensor in coordinate scheme. template -static SparseTensor *openTensor(char *filename, uint64_t size, - uint64_t *sizes, uint64_t *perm) { +static SparseTensorCOO *openSparseTensorCOO(char *filename, uint64_t size, + uint64_t *sizes, + uint64_t *perm) { // Open the file. FILE *file = fopen(filename, "r"); if (!file) { @@ -428,8 +442,8 @@ for (uint64_t r = 0; r < size; r++) assert((sizes[r] == 0 || sizes[r] == idata[2 + r]) && "dimension size mismatch"); - SparseTensor *tensor = - SparseTensor::newSparseTensor(size, idata + 2, perm, nnz); + SparseTensorCOO *tensor = + SparseTensorCOO::newSparseTensorCOO(size, idata + 2, perm, nnz); // Read all nonzero elements. std::vector indices(size); for (uint64_t k = 0; k < nnz; k++) { @@ -471,13 +485,13 @@ //===----------------------------------------------------------------------===// // -// Public API of the sparse runtime support library that support an opaque -// implementation of a bufferized SparseTensor in MLIR. This could be replaced -// by actual codegen in MLIR. +// Public API with methods that operate on MLIR buffers (memrefs) to interact +// with sparse tensors, which are only visible as opaque pointers externally. +// These methods should be used exclusively by MLIR compiler-generated code. // // Because we cannot use C++ templates with C linkage, some macro magic is used // to generate implementations for all required type combinations that can be -// called from MLIR generated code. +// called from MLIR compiler-generated code. // //===----------------------------------------------------------------------===// @@ -492,13 +506,14 @@ #define CASE(p, i, v, P, I, V) \ if (ptrTp == (p) && indTp == (i) && valTp == (v)) { \ - SparseTensor *tensor; \ + SparseTensorCOO *tensor; \ if (action == 0) \ - tensor = openTensor(static_cast(ptr), asize, sizes, perm); \ + tensor = openSparseTensorCOO(static_cast(ptr), asize, sizes, \ + perm); \ else if (action == 1) \ - tensor = static_cast *>(ptr); \ + tensor = static_cast *>(ptr); \ else if (action == 2) \ - return SparseTensor::newSparseTensor(asize, sizes, perm); \ + return SparseTensorCOO::newSparseTensorCOO(asize, sizes, perm); \ else \ return static_cast *>(ptr)->toCOO(perm); \ return SparseTensorStorage::newSparseTensor(tensor, sparsity, \ @@ -532,7 +547,7 @@ std::vector indices(isize); \ for (uint64_t r = 0; r < isize; r++) \ indices[perm[r]] = indx[r]; \ - static_cast *>(tensor)->add(indices, value); \ + static_cast *>(tensor)->add(indices, value); \ return tensor; \ } @@ -634,11 +649,6 @@ exit(1); } -/// Returns size of sparse tensor in given dimension. -uint64_t sparseDimSize(void *tensor, uint64_t d) { - return static_cast(tensor)->getDimSize(d); -} - /// Methods that provide direct access to pointers, indices, and values. IMPL2(MemRef1DU64, sparsePointers, uint64_t, getPointers) IMPL2(MemRef1DU64, sparsePointers64, uint64_t, getPointers) @@ -657,11 +667,6 @@ IMPL1(MemRef1DI16, sparseValuesI16, int16_t, getValues) IMPL1(MemRef1DI8, sparseValuesI8, int8_t, getValues) -/// Releases sparse tensor storage. -void delSparseTensor(void *tensor) { - delete static_cast(tensor); -} - /// Helper to add value to coordinate scheme, one per value type. IMPL3(addEltF64, double) IMPL3(addEltF32, float) @@ -676,6 +681,27 @@ #undef IMPL2 #undef IMPL3 +//===----------------------------------------------------------------------===// +// +// Public API with methods that accept C-style data structures to interact +// with sparse tensors, which are only visible as opaque pointers externally. +// These methods can be used by any external runtime that wants to interact +// with MLIR compiler-generated code. +// +//===----------------------------------------------------------------------===// + +// TODO: setup sparse tensor from C style format + +/// Returns size of sparse tensor in given dimension. +uint64_t sparseDimSize(void *tensor, uint64_t d) { + return static_cast(tensor)->getDimSize(d); +} + +/// Releases sparse tensor storage. +void delSparseTensor(void *tensor) { + delete static_cast(tensor); +} + } // extern "C" #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS