[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

129

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)

132

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

129

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

132

+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();
+  assert(m.getNumResults() == sz && "TopoSort/AffineMap size mismatch");
+  // Construct the inverse of `m`; to avoid the asymptotic complexity
+  // of calling `m.getPermutedPosition` repeatedly.
+  SmallVector<unsigned, 4> inv(sz);
+  for (unsigned i = 0; i < sz; i++)
+    inv[i] = m.getDimPosition(i);
+  // Construct the permutation.
   SmallVector<unsigned, 4> perm(sz);
   for (unsigned i = 0; i < sz; i++)
-    perm[i] = m.getPermutedPosition(topSort[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) {
       assert(order.isPermutation());
       return order.getDimPosition(d);
     }
   }
   return d;