[Mlir-commits] [mlir] [mlir][sparse] cleanup of COO (PR #69239)

[Aart Bik]

https://github.com/aartbik created https://github.com/llvm/llvm-project/pull/69239

Moves typedef to only file where it is used.
Removes some deadcode. Some minor doc changes.

>From 0d80455c0e1f90ddee10e5bb5215405e673a71ae Mon Sep 17 00:00:00 2001
From: Aart Bik <ajcbik at google.com>
Date: Mon, 16 Oct 2023 12:12:55 -0700
Subject: [PATCH] [mlir][sparse] cleanup of COO

Moves typedef to only file where it is used.
Removes some deadcode. Some minor doc changes.
---
 .../mlir/ExecutionEngine/SparseTensor/COO.h   | 21 +++++--------------
 .../ExecutionEngine/SparseTensor/Storage.h    |  5 +++++
 2 files changed, 10 insertions(+), 16 deletions(-)

diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h
index f6eb45defcc1cfd..721e9bc69adac20 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/COO.h
@@ -58,21 +58,13 @@ struct ElementLT final {
   const uint64_t rank;
 };
 
-/// The type of callback functions which receive an element.
-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.
+/// (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.
 template <typename V>
 class SparseTensorCOO final {
 public:
-  using const_iterator = typename std::vector<Element<V>>::const_iterator;
-
   /// Constructs a new coordinate-scheme sparse tensor with the given
   /// sizes and an optional initial storage capacity.
   explicit SparseTensorCOO(const std::vector<uint64_t> &dimSizes,
@@ -106,7 +98,7 @@ class SparseTensorCOO final {
   /// Returns the `operator<` closure object for the COO's element type.
   ElementLT<V> getElementLT() const { return ElementLT<V>(getRank()); }
 
-  /// Adds an element to the tensor. This method invalidates all iterators.
+  /// Adds an element to the tensor.
   void add(const std::vector<uint64_t> &dimCoords, V val) {
     const uint64_t *base = coordinates.data();
     const uint64_t size = coordinates.size();
@@ -135,12 +127,9 @@ class SparseTensorCOO final {
     elements.push_back(addedElem);
   }
 
-  const_iterator begin() const { return elements.cbegin(); }
-  const_iterator end() const { return elements.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.
+  /// values is unspecified.
   void sort() {
     if (isSorted)
       return;
diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
index 5e57facaf2376f6..c5be3d1acc33783 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h
@@ -37,6 +37,11 @@
 namespace mlir {
 namespace sparse_tensor {
 
+/// The type of callback functions which receive an element.
+template <typename V>
+using ElementConsumer =
+    const std::function<void(const std::vector<uint64_t> &, V)> &;
+
 // Forward references.
 template <typename V>
 class SparseTensorEnumeratorBase;