[Mlir-commits] [mlir] 875ee0e - [mlir][sparse] Reducing computational complexity

[wren romano]

Author: wren romano
Date: 2022-07-01T12:55:09-07:00
New Revision: 875ee0ed1c5af58cb4909f239093e25a35d7a21a

URL: https://github.com/llvm/llvm-project/commit/875ee0ed1c5af58cb4909f239093e25a35d7a21a
DIFF: https://github.com/llvm/llvm-project/commit/875ee0ed1c5af58cb4909f239093e25a35d7a21a.diff

LOG: [mlir][sparse] Reducing computational complexity

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|)`.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D129011

Added: 
    

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

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
index 614700aba9a16..167155008f8e5 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
@@ -115,9 +115,16 @@ struct CodeGen {
 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);
 }