diff --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp --- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp +++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp @@ -246,6 +246,19 @@ : sizes(szs), rev(getRank()), idx(getRank()), pointers(getRank()), indices(getRank()) { uint64_t rank = getRank(); + for (uint64_t r = 0; r < rank; r++) { + // Tensors with a size-zero dimension have trivial storage, so they + // should be handled by a separate `SparseTensorStorageBase` subclass. + assert(sizes[r] > 0 && "Dimension size is zero"); + // In the worst case, the `I` type must be large enough to store + // every possible index on any dimension (aka every `sizes[r]-1`). + // The library doesn't require this, since that would conservatively + // disallow situations where the largest populated index is much + // smaller than the largest possible index. However, for codegen: + // if `I` is indeed large enough, then assertions of `addIndex` + // may be avoided (discussed at the call-sites); and users may + // want to be warned if `I` is too small. + } // Store "reverse" permutation. for (uint64_t r = 0; r < rank; r++) rev[perm[r]] = r; @@ -256,6 +269,12 @@ for (uint64_t r = 0; r < rank; r++) { sz *= sizes[r]; if (sparsity[r] == DimLevelType::kCompressed) { + // In the worst case, the `P` type must to large enough to store + // the position after the end of every `indices[r]` (aka `sz`). + // We do not require this, since that worst case only occurs + // when the tensor is fully populated! Nevertheless, for codegen: + // if `P` happens to be large enough, then assertions of + // `addPointer` may be avoided. pointers[r].reserve(sz + 1); indices[r].reserve(sz); sz = 1; @@ -266,13 +285,18 @@ } } // Prepare sparse pointer structures for all dimensions. + // (We cannot use `addPointer` here, because `isCompressedDim` + // won't work until after this loop. But this is safe: `r` is + // in-bounds and compressed, and 0 fits into every supported `P`.) for (uint64_t r = 0; r < rank; r++) if (sparsity[r] == DimLevelType::kCompressed) pointers[r].push_back(0); // Then assign contents from coordinate scheme tensor if provided. if (tensor) { - // Lexicographically sort the tensor, to ensure precondition of `fromCOO`. + // Ensure both preconditions of `fromCOO`. + assert(tensor->getSizes() == sizes && "Tensor size mismatch"); tensor->sort(); + // Now actually insert the `elements`. const std::vector> &elements = tensor->getElements(); uint64_t nnz = elements.size(); values.reserve(nnz); @@ -403,10 +427,36 @@ } private: + /// Appends the next free position of `indices[d]` to `pointers[d]`. + /// Thus, when called after inserting the last element of a segment, + /// it will append the position where the next segment begins. + inline void addPointer(uint64_t d) { + assert(isCompressedDim(d)); // Entails `d < getRank()`. + uint64_t p = indices[d].size(); + assert(p <= std::numeric_limits

::max() && + "Pointer value is too large for the P-type"); + pointers[d].push_back(p); // Here is where we convert the `uint64_t` to `P`. + } + + /// Appends the given index to `indices[d]`. + inline void addIndex(uint64_t d, uint64_t i) { + assert(isCompressedDim(d)); // Entails `d < getRank()`. + assert(i <= std::numeric_limits::max() && + "Index value is too large for the I-type"); + indices[d].push_back(i); // Here is where we convert the `uint64_t` to `I`. + } + /// 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-dimension dense/sparse annotations. - /// Precondition: the `elements` must be lexicographically sorted. + /// + /// Preconditions: + /// (1) the `elements` must be lexicographically sorted. + /// (2) the indices of every element are valid for `sizes` (equal rank + /// and pointwise less-than). N.B., if `elements = coo->getElements()` + /// for some `SparseTensorCOO coo`, then it is sufficient to + /// check that `coo->getSizes() == sizes`; since `SparseTensorCOO::add` + /// ensures that every element is valid for `coo->getSizes()`. void fromCOO(const std::vector> &elements, uint64_t lo, uint64_t hi, uint64_t d) { // Once dimensions are exhausted, insert the numerical values. @@ -427,7 +477,10 @@ seg++; // Handle segment in interval for sparse or dense dimension. if (isCompressedDim(d)) { - indices[d].push_back(i); + // Note for codegen: If we know `sizes[d]-1` is valid for `I`, + // then we can avoid assertions here: the second precondition + // of `fromCOO` entails that this `i` is valid for `I` too. + addIndex(d, i); } else { // For dense storage we must fill in all the zero values between // the previous element (when last we ran this for-loop) and the @@ -442,7 +495,7 @@ } // Finalize the sparse pointer structure at this dimension. if (isCompressedDim(d)) { - pointers[d].push_back(indices[d].size()); + addPointer(d) } else { // For dense storage we must fill in all the zero values after // the last element. @@ -480,7 +533,7 @@ if (d == getRank()) { values.push_back(0); } else if (isCompressedDim(d)) { - pointers[d].push_back(indices[d].size()); + addPointer(d); } else { for (uint64_t full = 0, sz = sizes[d]; full < sz; full++) endDim(d + 1); @@ -494,7 +547,7 @@ for (uint64_t i = 0; i < rank - diff; i++) { uint64_t d = rank - i - 1; if (isCompressedDim(d)) { - pointers[d].push_back(indices[d].size()); + addPointer(d); } else { for (uint64_t full = idx[d] + 1, sz = sizes[d]; full < sz; full++) endDim(d + 1); @@ -509,7 +562,10 @@ for (uint64_t d = diff; d < rank; d++) { uint64_t i = cursor[d]; if (isCompressedDim(d)) { - indices[d].push_back(i); + // Note for codegen: knowing `sizes[d]-1 is valid for `I` is not + // enough to avoid the assertions here; we also need to know the + // `cursor` is valid for `sizes`. + addIndex(d, i); } else { for (uint64_t full = top; full < i; full++) endDim(d + 1); @@ -533,6 +589,7 @@ /// Returns true if dimension is compressed. inline bool isCompressedDim(uint64_t d) const { + assert(d < getRank()); return (!pointers[d].empty()); }