[mlir][sparse] Replace Element with ElementId for sorting.

[bixia]

Summary

Previously, we represented a COO element with a pointer to the coordinates and a
value. We use this representing for sorting. Such a representation
unnecessarily involves the element in the compare function as well as in the
swap function for sorting.

We now replace this Element representation with ElementId, which is the index
of the element value in the values-array.

This improves read-CSC-arabic2005-lib by more than 10%.

[wrengr]

Nit: in the CL summary it should be "Previously, we represented a"

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

31

"at indices"

49–50

I think it would be nicer to keep the Element class around, since it's still helpful to have a struct that gives better names for things than std::pair<ElementId, V> does. Although the new Element class can't give the user the uint64_t const* coords pointer, without also storing a reference to the enclosing SparseTensorCOO, that's not too bad of a problem since the old Element class already had issues because it did not store the rank. (And since the ElementConsumer already unpacks the Element class for its arguments, it would be easy enough to change the new Element class to store that reference.)

53–54

This should be named coordinates since it's not just the set of coordinates for a single element, which is why that's the named used by SparseTensorCOO as well. (This is different than the coords pointer of the Element class, which is indeed the set of coordinates for that single element.)

59–61

Much as I love consting all the things, unfortunately MLIR style says that scalar function parameters shouldn't use const (only pointers, references, etc may use it).

63–65

Since e1, e2, and rank are all constant, it would be better to define uint64_t const* const coords1 = coordinates + e1 * rank; uint64_t const* const coords2 = coordinates + e2 * rank; and then use if (coords1[d] == coords2[d]) etc. Although this may help performance due to LICM, the real reason I suggest it is that it aids code clarity by removing redundancies. Moreover, it helps ensure correctness despite making the parameters non-const as per my other comment.

(N.B., for this change, the two local variables are called "coords" because they're the collection of coordinates for a single element.)

(N.B., the two consts are important there. The the const* one is required by the fact that ElementLT::coordinates is const*. Whereas the const coordsN one is what ensures correctness: by making the coordsN local variable immutable itself. Personally I prefer the "East const style" https://isocpp.org/wiki/faq/const-correctness#const-ref-alt because it's more consistant/logical to parse; but since MLIR uses "West const style" everywhere else, feel free to use "const uint64_t * const coordsN" instead if you prefer.)

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

70

as mentioned earlier, this field should be named "coordinates"

111

This change needs to be reverted.

The value_type alias is part of the "Container" concept which is used by the C++-style iterator interface requested by Peter Gavin. The SparseTensorCOO class is not a container of ids, it's a container of elements, and therefore should be iterated as such.

113–114

This change needs to be reverted for the same reason.

116–167

The blank line should remain.

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

49–50

See my other followup comment, re preserving the Element<V> class for the iterator interface of SparseTensorCOO.

111

As a followup to this, since the Element<V> class is needed by the iterator interface and since clients can't convert an ElementId into a uint64_t const* coords pointer, that means we can't have the new Element<V> class just store an ElementId in lieu of the old pointer. I see two obvious solutions:

(1) Keep the Element<V> class exactly as it was before, and adjust the iterator interface to construct the uint64_t const* coords pointers at the time the Element<V> needs to be returned to the client. The main benefit of this approach is that it avoids any breaking changes to client code. The downside is that it might be tricky to construct the pointer when advancing from one element to the next. (I haven't looked that far down this CL to see if it actually would be tricky or not.)

(2) Update the Element<V> class to store all three of { SparseTensorCOO const& coo; const ElementId elemId; V value} and provide a method uint64_t const* coords() const which uses the coo to compute the correct pointer from the elemId. The main benefit of this approach is that it helps the Element<V> be more stable to mutations of the arrays stored in the coo. The downside is that it makes the struct larger.

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

209

This should use the ElementId type. (I'm pretty sure)

Since element identifiers are the thing that's stable under sorting etc, they're the thing that client code will store and pass around as the primary handle on an element. Therefore, this method should be named getValue since it's the primary method; whereas the method currently named that should be renamed to something else like getValueAtPosition or getNthValue.

210

Ditto my comment at getValueForId: the parameter should use the ElementId type, and this method should take the name getCoords.

214

To avoid confusion with the method currently named getValueForId, you should define a second typedef named something like ElementPosition or NthElement, and use that typedef for the argument. Also, you should use different variable names to help distinguish when something is supposed to be an ElementId vs an NthElement (e.g., using i for the former but n for the latter).

Since element identifiers are the thing that's stable under sorting etc, they're the thing that client code will store and pass around as the primary handle on an element. Therefore, this method should not take the name getValue since it's not the primary method client code will want to use. Instead this method should be renamed to something else like getValueAtPosition or getNthValue, so that the method currently named getValueForId can take the name getValue.

215

Ditto my comment at getValue: this method should instead be getNthCoords(NthElement n) or similar.

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

211

Why are you calling .data() before the operator[] instead of just using the operator[] directly on coordinates?

250

you should use ElementId elemId instead, since that's what the variable actually means in this context.

251

Nit: you can/should use elemId != 0 here instead, to avoid the cost of the additional method call.

256–262

As mentioned earlier, the iterator interface really does need to return Element<V> rather than ElementId. Since we no longer store a vector_type i.e. std::vector<Element<V>>, that means you'll need to define a new SparseTensorCOO<V>::Iterator class and use that in lieu of the vector_type typedef when defining the {iterator, const_iterator, difference_type, size_type} typedefs.

Defining the iterator class shouldn't be difficult, though there are a lot of fiddly details to making everything work right. Alas, unfortunately you can't use the llvm::iterator_facade_base base class like I do in D146691. So if you're not familiar with all the details of how to define an iterator from scratch, just let me know and I can define one in a separate CL.

277–279

delete this "COO"

277–279

"stores"

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

94

Why do you need to store two different iterators?

105

I think it would be better to have this method return Element<V>. Once the SparseTensorCOO<V>::Iterator class is defined, then that class will handle everything that's needed, for constructing the new element. If that class is defined correctly then this method can still be implemented via (it < end ? &*it++ : nullptr) or some slight variation thereof.

In particular, the variation I have in mind is return it < end ? std::make_optional(*it++) : std::nullopt; since that handles the liveness issues associated with SparseTensorCOO<V>::Iterator::operator* constructing a new Element<V> to return. Since the SparseTensorIterator::getNext method is not exported directly but rather is only used within the CAPI macro definitions below, there's no problem with using std::option in lieu of const*

108

Should use pre-increment here, since that's what I use throughout the runtime library. The postincrement of the previous implementation was just so that we could squish everything into the succinct "&*it++"

114–119

It feels wrong to me for this class to need these methods. The redundancy of needing SparseTensorIterator::{getValue,getCoords} in addition to SparseTensorCOO::{getValue,getCoords} should belie the fact that the C++ iterator interface for SparseTensorCOO really ought to be yielding Element<V> rather than ElementId.

[bixia]

Address review comments.

[bixia]

In D147011#4228972, @wrengr wrote:

Nit: in the CL summary it should be "Previously, we represented a"

done.

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

49–50

Add Element<V> back and use it for the iterator

59–61

Removed const for ElementId

63–65

Thanks! done.

111

Most of the typedefs in this code block aren't really needed anymore, mostly because the iterator is not a std::iterator, I added comment to the new iterator class to explain this.

111

Add a iterator class and use it to enumerate the COO elements.

113–114

See my comment above, this typedef isn't needed anymore.

209

Fixed the type, but didn't rename the function, per offline discussion

210

Fixed the type, but didn't rename the function, per offline discussion.

210

Similar to the above.

211

Fixed.

214

Added ElementPosition type and used it here and other places. Also use different parameter names.

256–262

Fixed this based on the new iterator class.

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

94

These fields were modified to use the new iterator class.

105

Change to use std::optional<Element<V>>, which is different from std::iterator, but serves our purpose.

108

This code is removed.

114–119

This is replaced by returning Element<V>.

[aartbik]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

29–43

The only downside is good english, right?
Why did you change this to downsize?!

31

I find (ID+1)*rank a bit more intuitive, but that is subjective

116–164

An iterator?

119

constructs

[bixia]

Replace typedef ElementId with a struct and simplify getValueForId/getCoordsforId with getValue/getCoords.

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

29

please add a comma here

29–30

"An element ID is used for retrieving the element's value and coordinates from the appropriate arrays."

29–43

"downside"

30

"pointer"

30

"underlying"

31

"an element"

33

"or" (and dropping the second "the")

46–47

"the"

111

Most of the typedefs in this code block aren't really needed anymore, mostly because the iterator is not a std::iterator, I added comment to the new iterator class to explain this.

This has nothing to do with std::iterator; these aliases are for the std::iterator_traits template, which is used ubiquitously throughout LLVM/MLIR in order to derive new iterators from old ones. (cf., https://en.cppreference.com/w/cpp/iterator/iterator_traits)

To see an example of why it's important to support that template, take a look at the definition of tensor_loop_id::Iterator in D146691 (line 207 at https://reviews.llvm.org/D146691?id=508266#change-rkcdqIP8Zcjv), and then take a look at the definition of llvm::iterator_adaptor_base and how it uses std::iterator_traits. Although this particular case is just a minor inconvenience, failing to provide the typedefs needed by std::iterator_traits can easily snowball into much bigger issues for client code. Therefore, usability dictates that the iterator class provides those typedefs.

116–164

"An iterator over the elements of the COO in the order given by..."

120

Should make this final.

Also, should use the class keyword since (1) all the data members are private, and (2) it provides a bunch of methods. I say that for stylistic reasons, but apparently on Windows the struct-vs-class keywords affect linkage (according to the -Wmismatched-tags warnings emitted by the phabricator buildbot on Debian).

121

I think this typedef should be named Impl instead, since it's the underlying implementation and that's the name used elsewhere for such things

122–123

I don't think this constructor should be public, since that can easily lead to bugs from clients passing some other arbitrary thing as the second argument.

Instead, you should have the iterator class declare SparseTensorCOO<V> to be a friend. However, there is some trickiness about the ordering of definitions vs declarations in order to get things to compile correctly. For a complete worked example, see the {TensorId, TensorId::Iterator, TensorId::Range} classes in D146693. Or the tldr is:

1. Have SparseTensorCOO<V> forward declare the iterator class and the begin/end methods.
2. After the closing "};" of the COO class, then you can define the iterator class.
3. After the closing "};" of the iterator class, then you can define the SparseTensorCOO<V>::{begin,end} methods.

141

You should also provide == and != since those are what's used when desugaring "for (auto elem : coo)"

147

I think it'd be clearer to name this it or iter

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

31

I find that more intuitive as well. I'd go even further and suggest: [rank*ID .. rank*(ID+1))

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

28–44

Previously, this block of text was documentation for the Element<V> type, hence using the "///" so that the doxygen documentation tooling picks it up. Whereas the present CL changes that to a more general commentary on the whole file rather than something specifically attached to any one type. So I'd suggest adjusting this text in one of the following ways. (I'm not sure which is best, so you may want to take a look at the generated doxygen files and/or see what Aart's preference is)

1. Leave the text here but change the "///" to a plain "//" comment instead, and add a blank line between the commentary and the ElementId definition.
   • This approach is best if we think these details are only relevant to developers of this library itself, and hence should not be shown to users of the library.
2. Move the text to the whole-file documentation section (currently at lines 8–14) and keep the "///".
   • This approach is best if we think users of the library should see them, and we think they should see them before seeing the documentation for the various things defined in this file.
3. Move the text to the documentation for the SparseTensorCOO class and keep the "///".
   • This approach is best if we think users of the library should see them, but we think they should only see them when looking at the documentation for the COO class.
   • If going with this approach, then this text should come after the stuff about the Container concept.

[bixia]

Address review comments.

[bixia]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

28–44

Thanks! I took approach 3. and move it to document SparseTensorCOO.

29–43

downside

111

The iterator we defined here isn't a proper iterator that can have all the normal iterator traits, such as how would you define iter::pointer as this iterator doesn't actually return a pointer to the storage but constructs a temp object and returns it?
On the other hand, we only use this iterator to implement SparseTensorIterator::getNext and we don't need most of t he iterator traits to support this.

120

Add final and change to class.

121

Add Impl to name and also make it private.

122–123

Move the constructor to private and make the iterator declare SparseTensorCOO as a friend.

147

change to iter.

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

111

On the other hand, we only use this iterator to implement SparseTensorIterator::getNext and we don't need most of t he iterator traits to support this.

As I have already said several times: We only use this for implementing SparseTensorIterator::getNext, but this iterator interface is not for us. The C++ iterator interface was added because Peter Gavin specifically requested it. (Although it was something on my todo list for a long time.) The runtime library does not merely serve our own internal use. The runtime library was factored out into a standalone C++ library for the express purpose of being reused by other clients, including the SparseCore team. Because the library is intended to be used by external clients working in C++ rather than through the MLIR dialect, the library must therefore support the needs of those external clients in addition to our own internal needs.

By focusing on the getNext method you're putting the cart before the horse. The C++ iterator interface is the principal interface and should therefore take primacy in matters of design. The only reason the SparseTensorIterator class exists is as a stop-gap to paper over the fact that the MLIR bindings cannot use the C++ iterator directly. More specifically, the only purpose of the SparseTensorIterator class is to pair an iterator with its end-iterator so that _mlir_ciface_getNext##VNAME knows when the end has been reached. If we redesigned the MLIR bindings to avoid that need to pair iterators with their end-iterators, then we could remove the SparseTensorIterator class entirely.

[bixia]

Add definitions for iterator trait.

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

111

Add the using for iterator trait back per offline discussion.

[wrengr]

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

30

"coordinates"

31

This class should be final

34–35

For hygiene, the id field shouldn't be public.

Making it private doesn't loose any functionality: since you have operator uint64_t for reading it out, and have the default copy-assignment operator for mutating it.

38

"`ElementId`"

47

This should be "A pointer into the..." (since it's one of several pointers, all of which pointing into a single shared-pool)

95–96

You should either: have the "(2)" start a new line; or, not have the "(1)" start a new line.

(I don't have a strong preference which; though since the "(1)" is so short and the "(2)" is so long, I have a slight preference towards not having them start new lines.)

110

I'm totally fine with this forward declaration, but since you're defining the Iterator class within the SparseTensorCOO<V> definition (i.e., rather than defining the iterator class after closing the definition of the COO class), it would be cleaner to give the Iterator definition first and then give the typedefs afterwards. (Unless there's some style guide thing saying to do typedefs before class definitions?)

111–117

👍

125–126

You should also provide the other typedefs required by https://en.cppreference.com/w/cpp/iterator/iterator_traits. I.e., using value_type = Element<V>; using reference = value_type &; using pointer = void; using iterator_category = std::forward_iterator_tag;.

(N.B., you're only allowed to elide the using pointer = void; definition in C++20; but since MLIR uses C++17, it's required)

128

Once you provide all the typedefs, you can change this to value_type (if you want).

129

Since the ctor isn't marked as explicit, you can change this to return {coo.getCoords(*iter), coo.getValue(*iter)}; (if you want).

157

What I meant was that the entire name should be just "Impl", since the rest of the name is redundant

217–224

Nit: overloads should be grouped together. I.e., the two getValue overloads together, followed by the two getCoords overloads together

mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h

 #include <algorithm>
 #include <cassert>
 #include <cinttypes>
 #include <functional>
 #include <vector>
 
 namespace mlir {
 namespace sparse_tensor {
 
-/// An element of a sparse tensor in coordinate-scheme representation
-/// (i.e., a pair of coordinates and value).  For example, a rank-1
-/// vector element would look like
-///   ({i}, a[i])
-/// and a rank-5 tensor element would look like
-///   ({i,j,k,l,m}, a[i,j,k,l,m])
-///
-/// The coordinates are represented as a (non-owning) pointer into
-/// a shared pool of coordinates, rather than being stored directly in
-/// this object.  This significantly improves performance because it:
-/// (1) reduces the per-element memory footprint, and (2) centralizes
-/// the memory management for coordinates.  The only downside is that
-/// the coordinates themselves cannot be retrieved without knowing the
-/// rank of the tensor to which this element belongs (and that rank is
-/// not stored in this object).
+/// An element ID is used for retrieving the element's value and coordinates
+/// from the value-array and coordinares-array.
+struct ElementId {
+  ElementId() : id(0) {}
+  explicit ElementId(uint64_t id) : id(id) {}
+  operator uint64_t() const { return id; }
+  uint64_t id;
+};
+
+/// The position of an element is the index of the element-ID in the
+/// element-ID-array.
+using ElementPosition = uint64_t;
+
+/// An element returned by a COO iterator is represented by a pointer to the
+/// coordinates and a value.
 template <typename V>
 struct Element final {
-  Element(const uint64_t *coords, V val) : coords(coords), value(val){};
-  const uint64_t *coords; // pointer into shared coordinates pool
+  Element(const uint64_t *coords, V val) : coords(coords), value(val) {}
+  const uint64_t *coords; // The pointer into a shared coordinates pool.
   V value;
 };
 
 /// Closure object for `operator<` on `Element` with a given rank.
-template <typename V>
 struct ElementLT final {
-  ElementLT(uint64_t rank) : rank(rank) {}
+  ElementLT(const uint64_t *coordinates, uint64_t rank)
+      : coordinates(coordinates), rank(rank) {}
 
   /// Compares two elements a la `operator<`.
   ///
   /// Precondition: the elements must both be valid for `rank`.
-  bool operator()(const Element<V> &e1, const Element<V> &e2) const {
+  bool operator()(ElementId e1, ElementId e2) const {
+    uint64_t const *const coords1 = coordinates + e1 * rank;
+    uint64_t const *const coords2 = coordinates + e2 * rank;
     for (uint64_t d = 0; d < rank; ++d) {
-      if (e1.coords[d] == e2.coords[d])
+      if (coords1[d] == coords2[d])
         continue;
-      return e1.coords[d] < e2.coords[d];
+      return coords1[d] < coords2[d];
     }
     return false;
   }
 
+  const uint64_t *coordinates; // A pointer to the coordinates array.
   const uint64_t rank;
 };
 
 /// The type of callback functions which receive an element.  We avoid
 /// packaging the coordinates and value together as an `Element` object
 /// because this helps keep code somewhat cleaner.
 template <typename V>
 using ElementConsumer =
     const std::function<void(const std::vector<uint64_t> &, V)> &;
 
 /// A memory-resident sparse tensor in coordinate-scheme representation
 /// (a collection of `Element`s).  This data structure is used as
 /// an intermediate representation; e.g., for reading sparse tensors
 /// from external formats into memory, or for certain conversions between
 /// different `SparseTensorStorage` formats.
 ///
+/// Elements of a sparse tensor are stored in three arrays: an array for
+/// coordinates, an array for values, and an array for the element IDs. An
+/// element ID is used for retrieving the element's value and coordinates from
+/// the appropriate arrays. The coordinates of an element are stored at indices
+/// [rank*ID .. rank*(ID+1)) in the coordinates-array. The element-ID-array
+/// allows us to sort the elements without having to move the coordinates or
+/// values.
+///
+/// This significantly improves performance because it:
+/// (1) reduces the per-element memory footprint, and (2) centralizes
+/// the memory management for coordinates and values. The only downside is that
+/// the iterator can't simply return a pointer to the underlying storage for the
+/// elements.
+///
 /// This class provides all the typedefs required by the "Container"
 /// concept (<https://en.cppreference.com/w/cpp/named_req/Container>);
 /// however, beware that it cannot fully implement that concept since
 /// it cannot have a default ctor (because the `dimSizes` field is const).
 /// Thus these typedefs are provided for familiarity reasons, rather
 /// than as a proper implementation of the concept.
 template <typename V>
 class SparseTensorCOO final {
 public:
+  class Iterator;
   using value_type = const Element<V>;
   using reference = value_type &;
   using const_reference = reference;
-  // The types associated with `std::vector` differ significantly between
-  // C++11/17 vs C++20; so we explicitly defer to whatever `std::vector`
-  // says the types should be.
-  using vector_type = std::vector<Element<V>>;
-  using iterator = typename vector_type::const_iterator;
+  using iterator = Iterator;
   using const_iterator = iterator;
-  using difference_type = typename vector_type::difference_type;
-  using size_type = typename vector_type::size_type;
+  using difference_type = typename iterator::difference_type;
+  using size_type = typename iterator::size_type;
+
+  // An iterator over the elements of the COO in the order given by the
+  // element-ID-array. This can't be an std::iterator because the value type
+  // Element<V> the iterator returns is not a type that is used in the actual
+  // storage. As such, the operator* constructs an object of type Element<V>.
+  class Iterator final {
+  public:
+    using size_type = size_t;
+    using difference_type = ptrdiff_t;
+
+    Element<V> operator*() const {
+      return Element<V>(coo.getCoords(*iter), coo.getValue(*iter));
+    }
+
+    Iterator &operator++() {
+      ++iter;
+      return *this;
+    }
+
+    Iterator operator++(int) {
+      Iterator tmp = *this;
+      ++iter;
+      return tmp;
+    }
+
+    friend bool operator<(const Iterator &a, const Iterator &b) {
+      return a.iter < b.iter;
+    }
+
+    friend bool operator==(const Iterator &a, const Iterator &b) {
+      return a.iter == b.iter;
+    }
+
+    friend bool operator!=(const Iterator &a, const Iterator &b) {
+      return a.iter != b.iter;
+    }
+
+  private:
+    friend class SparseTensorCOO<V>;
+    using ElementIDIterImpl = typename std::vector<ElementId>::const_iterator;
+
+    Iterator(const SparseTensorCOO<V> &coo, ElementIDIterImpl iter)
+        : coo(coo), iter(iter) {}
+
+    const SparseTensorCOO<V> &coo;
+    ElementIDIterImpl iter;
+  };
 
   /// Constructs a new coordinate-scheme sparse tensor with the given
   /// sizes and initial storage capacity.
   ///
   /// Asserts:
   /// * `dimSizes` has nonzero size.
   /// * the elements of `dimSizes` are nonzero.
   explicit SparseTensorCOO(const std::vector<uint64_t> &dimSizes,
                            uint64_t capacity = 0)
       : SparseTensorCOO(dimSizes.size(), dimSizes.data(), capacity) {}
 
public:
   /// * the elements of `dimSizes` are nonzero.
   explicit SparseTensorCOO(uint64_t dimRank, const uint64_t *dimSizes,
                            uint64_t capacity = 0)
       : dimSizes(dimSizes, dimSizes + dimRank), isSorted(true) {
     assert(dimRank > 0 && "Trivial shape is not supported");
     for (uint64_t d = 0; d < dimRank; ++d)
       assert(dimSizes[d] > 0 && "Dimension size zero has trivial storage");
     if (capacity) {
-      elements.reserve(capacity);
+      elementIds.reserve(capacity);
       coordinates.reserve(capacity * dimRank);
+      values.reserve(capacity);
     }
   }
 
   /// Gets the dimension-rank of the tensor.
   uint64_t getRank() const { return dimSizes.size(); }
 
   /// Gets the dimension-sizes array.
   const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
 
-  /// Gets the elements array.
-  const std::vector<Element<V>> &getElements() const { return elements; }
+  /// Returns the number of stored elements.
+  uint64_t getNse() const { return elementIds.size(); }
+
+  /// Returns the value for the element with the given ID.
+  V getValue(ElementId i) const { return values[i]; }
+
+  /// Returns the coordinates for the element with the given ID.
+  const uint64_t *getCoords(ElementId i) const {
+    return &coordinates[i * getRank()];
+  }
+
+  /// Returns the value for the element at the given position.
+  V getValue(ElementPosition n) const { return getValue(elementIds[n]); }
+
+  /// Returns the coordinates for the element at the given position.
+  const uint64_t *getCoords(ElementPosition n) const {
+    return getCoords(elementIds[n]);
+  }
 
   /// Returns the `operator<` closure object for the COO's element type.
-  ElementLT<V> getElementLT() const { return ElementLT<V>(getRank()); }
+  ElementLT getElementLT() const {
+    return ElementLT(coordinates.data(), getRank());
+  }
 
   /// Adds an element to the tensor.  This method does not check whether
   /// `dimCoords` is already associated with a value, it adds it regardless.
   /// Resolving such conflicts is left up to clients of the iterator
   /// interface.
   ///
   /// This method invalidates all iterators.
   ///
   /// Asserts:
   /// * the `dimCoords` is valid for `getRank`.
   /// * the components of `dimCoords` are valid for `getDimSizes`.
   void add(const std::vector<uint64_t> &dimCoords, V val) {
-    const uint64_t *base = coordinates.data();
-    const uint64_t size = coordinates.size();
     const uint64_t dimRank = getRank();
     assert(dimCoords.size() == dimRank && "Element rank mismatch");
     for (uint64_t d = 0; d < dimRank; ++d) {
       assert(dimCoords[d] < dimSizes[d] &&
              "Coordinate is too large for the dimension");
       coordinates.push_back(dimCoords[d]);
     }
-    // This base only changes if `coordinates` was reallocated.  In which
-    // case, we need to correct all previous pointers into the vector.
-    // Note that this only happens if we did not set the initial capacity
-    // right, and then only for every internal vector reallocation (which
-    // with the doubling rule should only incur an amortized linear overhead).
-    const uint64_t *const newBase = coordinates.data();
-    if (newBase != base) {
-      for (uint64_t i = 0, n = elements.size(); i < n; ++i)
-        elements[i].coords = newBase + (elements[i].coords - base);
-      base = newBase;
-    }
-    // Add the new element and update the sorted bit.
-    const Element<V> addedElem(base + size, val);
-    if (!elements.empty() && isSorted)
-      isSorted = getElementLT()(elements.back(), addedElem);
-    elements.push_back(addedElem);
+    values.push_back(val);
+    const ElementId elemId = ElementId(elementIds.size());
+    if (elemId != 0 && isSorted)
+      isSorted = getElementLT()(elementIds.back(), elemId);
+    elementIds.push_back(elemId);
   }
 
-  const_iterator begin() const { return elements.cbegin(); }
-  const_iterator end() const { return elements.cend(); }
+  const_iterator begin() const {
+    return const_iterator(*this, elementIds.cbegin());
+  }
+  const_iterator end() const {
+    return const_iterator(*this, elementIds.cend());
+  }
 
   /// Sorts elements lexicographically by coordinates.  If a coordinate
   /// is mapped to multiple values, then the relative order of those
   /// values is unspecified.
   ///
   /// This method invalidates all iterators.
   void sort() {
     if (isSorted)
       return;
-    std::sort(elements.begin(), elements.end(), getElementLT());
+    std::sort(elementIds.begin(), elementIds.end(), getElementLT());
     isSorted = true;
   }
 
 private:
   const std::vector<uint64_t> dimSizes; // per-dimension sizes
-  std::vector<Element<V>> elements;     // all COO elements
+  // All element IDs. When isSorted == true, this array stores the IDs in the
+  // order of the sorted elements.
+  std::vector<ElementId> elementIds;
   std::vector<uint64_t> coordinates;    // shared coordinate pool
+  std::vector<V> values;
   bool isSorted;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_COO_H

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

 }
 
 /// Writes the sparse tensor to `filename` in extended FROSTT format.
 template <typename V>
 inline void writeExtFROSTT(const SparseTensorCOO<V> &coo,
                            const char *filename) {
   assert(filename && "Got nullptr for filename");
   const auto &dimSizes = coo.getDimSizes();
-  const auto &elements = coo.getElements();
   const uint64_t dimRank = coo.getRank();
-  const uint64_t nse = elements.size();
+  const uint64_t nse = coo.getNse();
   std::fstream file;
   file.open(filename, std::ios_base::out | std::ios_base::trunc);
   assert(file.is_open());
   file << "; extended FROSTT format\n" << dimRank << " " << nse << std::endl;
   for (uint64_t d = 0; d < dimRank - 1; ++d)
     file << dimSizes[d] << " ";
   file << dimSizes[dimRank - 1] << std::endl;
-  for (uint64_t i = 0; i < nse; ++i) {
-    const auto &coords = elements[i].coords;
+  for (ElementPosition i = 0; i < nse; ++i) {
+    const auto &coords = coo.getCoords(i);
     for (uint64_t d = 0; d < dimRank; ++d)
       file << (coords[d] + 1) << " ";
-    file << elements[i].value << std::endl;
+    file << coo.getValue(i) << std::endl;
   }
   file.flush();
   file.close();
   assert(file.good());
 }
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_SPARSETENSOR_FILE_H

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

SparseTensorCOO<V> *toCOO(uint64_t trgRank, const uint64_t *trgSizes,
     SparseTensorEnumerator<P, C, V> enumerator(*this, trgRank, trgSizes,
                                                srcRank, src2trg);
     auto *coo = new SparseTensorCOO<V>(trgRank, trgSizes, values.size());
     enumerator.forallElements(
         [&coo](const auto &trgCoords, V val) { coo->add(trgCoords, val); });
     // TODO: This assertion assumes there are no stored zeros,
     // or if there are then that we don't filter them out.
     // Cf., <https://github.com/llvm/llvm-project/issues/54179>
-    assert(coo->getElements().size() == values.size());
+    assert(coo->getNse() == values.size());
     return coo;
   }
 
 private:
   /// Appends an arbitrary new position to `positions[lvl]`.  This method
   /// checks that `pos` is representable in the `P` type; however, it
   /// does not check that `pos` is semantically valid (i.e., larger than
   /// the previous position and smaller than `coordinates[lvl].capacity()`).
MLIR_SPARSETENSOR_FATAL("unsupported level type: %d\n",
                             static_cast<uint8_t>(dlt));
   }
 
   /// Initializes sparse tensor storage scheme from a memory-resident sparse
   /// tensor in coordinate scheme. This method prepares the positions and
   /// coordinates arrays under the given per-level dense/sparse annotations.
   ///
   /// Preconditions:
-  /// * the `lvlElements` must be lexicographically sorted.
+  /// * the elements in `lvlCOO` must be lexicographically sorted.
   /// * the coordinates of every element are valid for `getLvlSizes()`
   ///   (i.e., equal rank and pointwise less-than).
-  void fromCOO(const std::vector<Element<V>> &lvlElements, uint64_t lo,
-               uint64_t hi, uint64_t l) {
+  void fromCOO(const SparseTensorCOO<V> &lvlCOO, ElementPosition lo,
+               ElementPosition hi, uint64_t l) {
     const uint64_t lvlRank = getLvlRank();
-    assert(l <= lvlRank && hi <= lvlElements.size());
+    assert(l <= lvlRank && hi <= lvlCOO.getNse());
     // Once levels are exhausted, insert the numerical values.
     if (l == lvlRank) {
       assert(lo < hi);
-      values.push_back(lvlElements[lo].value);
+      values.push_back(lvlCOO.getValue(lo));
       return;
     }
     // Visit all elements in this interval.
     uint64_t full = 0;
     while (lo < hi) { // If `hi` is unchanged, then `lo < lvlElements.size()`.
       // Find segment in interval with same coordinate at this level.
-      const uint64_t c = lvlElements[lo].coords[l];
+      const uint64_t c = lvlCOO.getCoords(lo)[l];
       uint64_t seg = lo + 1;
       if (isUniqueLvl(l))
-        while (seg < hi && lvlElements[seg].coords[l] == c)
+        while (seg < hi && lvlCOO.getCoords(seg)[l] == c)
           ++seg;
       // Handle segment in interval for sparse or dense level.
       appendCrd(l, full, c);
       full = c + 1;
-      fromCOO(lvlElements, lo, seg, l + 1);
+      fromCOO(lvlCOO, lo, seg, l + 1);
       // And move on to next segment in interval.
       lo = seg;
     }
     // Finalize the sparse position structure at this level.
     finalizeSegment(l, full);
   }
 
   /// Finalizes the sparse position structure at this level.
SparseTensorStorage<P, C, V>::SparseTensorStorage( // NOLINT
     : SparseTensorStorage(dimRank, dimSizes, lvlRank,
                           lvlCOO.getDimSizes().data(), lvlTypes, lvl2dim,
                           false) {
   assert(lvlRank == lvlCOO.getDimSizes().size() && "Level-rank mismatch");
   // Ensure the preconditions of `fromCOO`.  (One is already ensured by
   // using `lvlSizes = lvlCOO.getDimSizes()` in the ctor above.)
   lvlCOO.sort();
   // Now actually insert the `elements`.
-  const auto &elements = lvlCOO.getElements();
-  const uint64_t nse = elements.size();
+  const uint64_t nse = lvlCOO.getNse();
   values.reserve(nse);
-  fromCOO(elements, 0, nse, 0);
+  fromCOO(lvlCOO, 0, nse, 0);
 }
 
 template <typename P, typename C, typename V>
 SparseTensorStorage<P, C, V>::SparseTensorStorage(
     uint64_t dimRank, const uint64_t *dimSizes, uint64_t lvlRank,
     const DimLevelType *lvlTypes, const uint64_t *lvl2dim,
     SparseTensorEnumeratorBase<V> &lvlEnumerator)
     : SparseTensorStorage(dimRank, dimSizes, lvlRank,

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

 template <typename V>
 class SparseTensorIterator final {
 public:
   /// This ctor requires `coo` to be a non-null pointer to a dynamically
   /// allocated object, and takes ownership of that object.  Therefore,
   /// callers must not free the underlying COO object, since the iterator's
   /// dtor will do so.
   explicit SparseTensorIterator(const SparseTensorCOO<V> *coo)
       : coo(coo), it(coo->begin()), end(coo->end()) {}
 
   ~SparseTensorIterator() { delete coo; }
 
   // Disable copy-ctor and copy-assignment, to prevent double-free.
   SparseTensorIterator(const SparseTensorIterator<V> &) = delete;
   SparseTensorIterator<V> &operator=(const SparseTensorIterator<V> &) = delete;
 
   /// Gets the next element.  If there are no remaining elements, then
-  /// returns nullptr.
-  const Element<V> *getNext() { return it < end ? &*it++ : nullptr; }
+  /// returns std::nullopt.
+  std::optional<Element<V>> getNext() {
+    return it < end ? std::optional<Element<V>>(*it++) : std::nullopt;
+  }
 
 private:
   const SparseTensorCOO<V> *const coo; // Owning pointer.
   typename SparseTensorCOO<V>::const_iterator it;
   const typename SparseTensorCOO<V>::const_iterator end;
 };
 
 // TODO: When using this library from MLIR, the `toMLIRSparseTensor`/
 // `IMPL_CONVERTTOMLIRSPARSETENSOR` and `fromMLIRSparseTensor`/
 // `IMPL_CONVERTFROMMLIRSPARSETENSOR` constructs will be codegened away;
 // therefore, these functions are only used by PyTACO, one place in the
 // Python integration tests, and possibly by out-of-tree projects.
 // This is notable because neither function can be easily generalized
 // to handle non-permutations.  In particular, while we could adjust
 // the functions to take all the arguments they'd need, that would just
 // push the problem into client code.  So if we want to generalize these
 // functions to support non-permutations, we'll need to figure out how
 // to do so without putting undue burden on clients.
 
 /// Initializes sparse tensor from an external COO-flavored format.
 /// The `rank` argument is both dimension-rank and level-rank, and the
fromMLIRSparseTensor(const SparseTensorStorage<uint64_t, uint64_t, V> *tensor,
   assert(tensor && "Received nullptr for tensor");
   const uint64_t dimRank = tensor->getDimRank();
   const auto &dimSizes = tensor->getDimSizes();
   std::vector<uint64_t> identityPerm(dimRank);
   std::iota(identityPerm.begin(), identityPerm.end(), 0);
   SparseTensorCOO<V> *coo =
       tensor->toCOO(dimRank, dimSizes.data(), dimRank, identityPerm.data());
 
-  const std::vector<Element<V>> &elements = coo->getElements();
-  const uint64_t nse = elements.size();
+  const uint64_t nse = coo->getNse();
 
   const auto &cooSizes = coo->getDimSizes();
   assert(cooSizes.size() == dimRank && "Rank mismatch");
   uint64_t *dimShape = new uint64_t[dimRank];
   std::memcpy(static_cast<void *>(dimShape),
               static_cast<const void *>(cooSizes.data()),
               sizeof(uint64_t) * dimRank);
 
   V *values = new V[nse];
   uint64_t *coordinates = new uint64_t[dimRank * nse];
 
-  for (uint64_t i = 0, base = 0; i < nse; ++i) {
-    values[i] = elements[i].value;
+  for (ElementPosition i = 0, base = 0; i < nse; ++i) {
+    values[i] = coo->getValue(i);
     for (uint64_t d = 0; d < dimRank; ++d)
-      coordinates[base + d] = elements[i].coords[d];
+      coordinates[base + d] = coo->getCoords(i)[d];
     base += dimRank;
   }
 
   delete coo;
   *pRank = dimRank;
   *pNse = nse;
   *pShape = dimShape;
   *pValues = values;
#define IMPL_GETNEXT(VNAME, V)                                                 \
   bool _mlir_ciface_getNext##VNAME(void *iter,                                 \
                                    StridedMemRefType<index_type, 1> *cref,     \
                                    StridedMemRefType<V, 0> *vref) {            \
     assert(iter &&vref);                                                       \
     ASSERT_NO_STRIDE(cref);                                                    \
     index_type *coords = MEMREF_GET_PAYLOAD(cref);                             \
     V *value = MEMREF_GET_PAYLOAD(vref);                                       \
     const uint64_t rank = MEMREF_GET_USIZE(cref);                              \
-    const Element<V> *elem =                                                   \
-        static_cast<SparseTensorIterator<V> *>(iter)->getNext();               \
-    if (elem == nullptr)                                                       \
+    const auto elem =                                                          \
+        static_cast<SparseTensorIterator<V> *>(iter) -> getNext();             \
+    if (!elem.has_value())                                                     \
       return false;                                                            \
     for (uint64_t d = 0; d < rank; d++)                                        \
       coords[d] = elem->coords[d];                                             \
     *value = elem->value;                                                      \
     return true;                                                               \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
 #undef IMPL_GETNEXT