maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[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
This commit is contained in:
wren romano 2022-07-01 12:41:01 -07:00
parent b19cbda45a
commit 875ee0ed1c
1 changed files with 8 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
mlir/lib/Dialect/SparseTensor/Transforms/Sparsification.cpp
View File

@ -115,9 +115,16 @@ struct CodeGen {
static AffineMap permute(MLIRContext *context, AffineMap m, static AffineMap permute(MLIRContext *context, AffineMap m,
std::vector<unsigned> &topSort) { std::vector<unsigned> &topSort) {
unsigned sz = topSort.size(); 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); SmallVector<unsigned, 4> perm(sz);
for (unsigned i = 0; i < sz; i++) for (unsigned i = 0; i < sz; i++)
perm[i] = m.getPermutedPosition(topSort[i]); perm[i] = inv[topSort[i]];
return AffineMap::getPermutationMap(perm, context); return AffineMap::getPermutationMap(perm, context);
} }