[mlir][sparse] Eliminating intermediate arrays/copies in toCOO

[wrengr]

Summary

Superseded by D122060

[aartbik]

Even though this potentially saves an array copy, I find it much easier to reason about invariants in the original structure with the factory method that permutes from original order.

[wrengr]

I think the logic will become a lot more transparent once I finish the differential for introducing the Permutation type, since then the code becomes just: Permutation reord = perm.compose(rev); Permutation permsz = reord.dirImg(sizes); (or if we just store the original permutation directly rather than inverting it in the constructor: Permutation reord = oldPerm.dirImg(newPerm);).

[wrengr]

I'm fine postponing this differential until after the Permutation type is introduced, I only factored it out to keep the logic change separate from the type change

[aartbik]

Not trying to be difficult, since I am all in favor of better code, but before, the code had the nice invariant of having the original sizes in the factory method, so I could set a breakpoint there and check it against the sizes in the MLIR tensor for correctness. We lose that invariant now, and I doubt this runs noticeable faster than the original for real-life tensors. So, keep it for now and send the me next differential to give me a better idea where this is going....

[wrengr]

Superseded by D122060

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

if (values.empty())
       endDim(0);
     else
       endPath(0);
   }
 
   /// Returns this sparse tensor storage scheme as a new memory-resident
   /// sparse tensor in coordinate scheme with the given dimension order.
   SparseTensorCOO<V> *toCOO(const uint64_t *perm) {
-    // Restore original order of the dimension sizes and allocate coordinate
-    // scheme with desired new ordering specified in perm.
     uint64_t rank = getRank();
-    std::vector<uint64_t> orgsz(rank);
-    for (uint64_t r = 0; r < rank; r++)
-      orgsz[rev[r]] = sizes[r];
-    SparseTensorCOO<V> *tensor = SparseTensorCOO<V>::newSparseTensorCOO(
-        rank, 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,
-    // we compute the combine permutation in advance.
+    // The permutation of dimensions from this tensor's storage-order to
+    // the new COO's storage-order.  This is just the composition of `rev`
+    // (to map this tensor's storage-order back to the original/semantic
+    // order) with `perm` (to map that original order to the new COO's
+    // storage-order); but we precompute it to avoid needing to compute
+    // it during the recursion.
     std::vector<uint64_t> reord(rank);
-    for (uint64_t r = 0; r < rank; r++)
-      reord[r] = perm[rev[r]];
+    // The new coordinate scheme's sizes array.
+    std::vector<uint64_t> permsz(rank);
+    for (uint64_t r = 0; r < rank; r++) {
+      uint64_t s = perm[rev[r]];
+      assert(s < rank && "Permutation index is out of range");
+      reord[r] = s;
+      permsz[s] = sizes[r];
+    }
+    SparseTensorCOO<V> *tensor = new SparseTensorCOO<V>(permsz, values.size());
+    // Populate the coordinate scheme.
     toCOO(*tensor, reord, 0, 0);
     assert(tensor->getElements().size() == values.size());
     return tensor;
   }
 
   /// Factory method. Constructs a sparse tensor storage scheme with the given
   /// dimensions, permutation, and per-dimension dense/sparse annotations,
   /// using the coordinate scheme tensor for the initial contents if provided.