[Mlir-commits] [mlir] [mlir][sparse] complete migration to dim2lvl/lvl2dim in library (PR #69268)

[Aart Bik]

================
@@ -411,23 +377,18 @@ class SparseTensorStorage final : public SparseTensorStorageBase {
       endPath(0);
   }
 
-  /// Allocates a new COO object and initializes it with the contents
-  /// of this tensor under the given mapping from the `getDimSizes()`
-  /// coordinate-space to the `trgSizes` coordinate-space. Callers must
-  /// make sure to delete the COO when they're done with it.
-  SparseTensorCOO<V> *toCOO(uint64_t trgRank, const uint64_t *trgSizes,
-                            uint64_t srcRank,
-                            const uint64_t *src2trg, // FIXME: dim2lvl
-                            const uint64_t *lvl2dim) const {
-    // TODO: use MapRef here too for the translation
-    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.
-    // <https://github.com/llvm/llvm-project/issues/54179>
+  /// Allocates a new COO object and initializes it with the contents.
+  /// Callers must make sure to delete the COO when they're done with it.
+  SparseTensorCOO<V> *toCOO() const {
+    std::vector<uint64_t> dimCoords(getDimRank());
+    std::vector<uint64_t> lvlCoords(getLvlRank());
+    auto *coo = new SparseTensorCOO<V>(getDimSizes(), values.size());
+    forallElements(
+        [this, &dimCoords, &lvlCoords, &coo](V val) {
----------------
aartbik wrote:

I actually made the whole recursion a bit simpler.

https://github.com/llvm/llvm-project/pull/69268