shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[mlir] Optimize const values AffineMap::compose (#141005)

Browse Source 
The original implementation will create two intermediate AffineMap in
the context, calling this compose function with different values
multiple times will occupy a lot of memory.

To improve the performance, we can call the AffineExpr::replace
directly, so we don't need to store all combinations of values in the
context.
This commit is contained in:
qazwsxedcrfvtg14 2025-05-24 08:30:48 +08:00 committed by GitHub
parent e9dbf31be5
commit d5802c30ae
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 5 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
mlir/lib/IR/AffineMap.cpp
View File

@ -580,15 +580,13 @@ AffineMap AffineMap::compose(AffineMap map) const {
SmallVector<int64_t, 4> AffineMap::compose(ArrayRef<int64_t> values) const {
assert(getNumSymbols() == 0 && "Expected symbol-less map");
SmallVector<AffineExpr, 4> exprs;
exprs.reserve(values.size());
MLIRContext *ctx = getContext();
for (auto v : values)
exprs.push_back(getAffineConstantExpr(v, ctx));
auto resMap = compose(AffineMap::get(0, 0, exprs, ctx));
for (int64_t value : values)
exprs.push_back(getAffineConstantExpr(value, ctx));
SmallVector<int64_t, 4> res;
res.reserve(resMap.getNumResults());
for (auto e : resMap.getResults())
res.push_back(cast<AffineConstantExpr>(e).getValue());
res.reserve(getNumResults());
for (auto e : getResults())
res.push_back(cast<AffineConstantExpr>(e.replaceDims(exprs)).getValue());
return res;
}