[mlir][sparse] Reducing computational complexity

[wrengr]

Summary

This is a followup to D128847. The AffineMap::getPermutedPosition method performs a linear scan of the map, thus the previous implementation had asymptotic complexity of O(|topSort| * |m|). This change reduces that to O(|topSort| + |m|).

[aartbik]

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

119

why do you compute the invSz and permSz twice?
just use the one below, and assert on it if you are worried about a mismatch
(although that is done at many other places too already)

122

Perhaps, in the long run, we should push such logic into AffineMap for permutations, so that getPermutedPosition is O(1) again, or we get a direct inverse back, because the original code is easier to read than the optimized code (and it would benefit all clients). Maybe add a TODO for that (bonus points for also making it happen ;-)

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

119

Each one is only computed once? I use them separately because of the mismatch worry. I considered asserting they were the same, but wasn't sure if that was actually an invariant or not (it's been too ling since I looked at this code). I'll add the assertion

122

+1. I'll write up an RFC for that

[wrengr]

addressing comments

[aartbik]

Ship it!

mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp

 //===----------------------------------------------------------------------===//
 // Sparse compiler analysis methods.
 //===----------------------------------------------------------------------===//
 
 /// Helper method to construct a permuted dimension ordering
 /// that adheres to the given topological sort.
 static AffineMap permute(MLIRContext *context, AffineMap m,
                          std::vector<unsigned> &topSort) {
-  unsigned sz = topSort.size();
-  SmallVector<unsigned, 4> perm(sz);
-  for (unsigned i = 0; i < sz; i++)
-    perm[i] = m.getPermutedPosition(topSort[i]);
+  // Construct the inverse of `m`; to avoid the asymptotic complexity
+  // of calling `m.getPermutedPosition` repeatedly.
+  unsigned invSz = m.getNumResults();
+  SmallVector<unsigned, 4> inv(invSz);
+  for (unsigned i = 0; i < invSz; i++)
+    inv[i] = m.getDimPosition(i);
+  // Construct the permutation.
+  unsigned permSz = topSort.size();
+  SmallVector<unsigned, 4> perm(permSz);
+  for (unsigned i = 0; i < permSz; i++)
+    perm[i] = inv[topSort[i]];
   return AffineMap::getPermutationMap(perm, context);
 }
 
 /// Helper method to apply dimension ordering permutation.
 static unsigned perm(const SparseTensorEncodingAttr &enc, unsigned d) {
   if (enc) {
     auto order = enc.getDimOrdering();
     if (order) {