[mlir] Optimize const values AffineMap::compose

[qazwsxedcrfvtg14]

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.

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: qazwsxedcrfvtg14 (qazwsxedcrfvtg14)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/141005.diff

1 Files Affected:

• (modified) mlir/lib/IR/AffineMap.cpp (+8-8)

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;
 }
 

[qazwsxedcrfvtg14]

It looks like I don't have the ability to merge the PR.
Could you merge the PR on my behalf?

Thanks!