[Mlir-commits] [mlir] [mlir] Optimize const values AffineMap::compose (PR #141005)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed May 21 21:48:48 PDT 2025
https://github.com/qazwsxedcrfvtg14 created https://github.com/llvm/llvm-project/pull/141005
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.
>From e11eb6f9de57eebcaac9fdca1922161bb25cccdc Mon Sep 17 00:00:00 2001
From: Yi Chou <yich at google.com>
Date: Thu, 22 May 2025 04:31:49 +0000
Subject: [PATCH] [mlir] Optimize const values AffineMap::compose
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.
---
mlir/lib/IR/AffineMap.cpp | 16 ++++++++--------
1 file changed, 8 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 72effb38d614c..db8f76dec1e01 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -579,16 +579,16 @@ 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());
+ unsigned numValues = values.size();
+ DenseMap<AffineExpr, AffineExpr> mapping;
MLIRContext *ctx = getContext();
- for (auto v : values)
- exprs.push_back(getAffineConstantExpr(v, ctx));
- auto resMap = compose(AffineMap::get(0, 0, exprs, ctx));
+ for (unsigned idx = 0; idx < numValues; ++idx)
+ mapping[getAffineDimExpr(idx, ctx)] =
+ getAffineConstantExpr(values[idx], 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.replace(mapping)).getValue());
return res;
}
More information about the Mlir-commits
mailing list