[Mlir-commits] [mlir] d5802c3 - [mlir] Optimize const values AffineMap::compose (#141005)

[llvmlistbot at llvm.org]

Author: qazwsxedcrfvtg14
Date: 2025-05-23T17:30:48-07:00
New Revision: d5802c30ae6cf296489daf12b36582e9e1d658bb

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

LOG: [mlir] Optimize const values AffineMap::compose (#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.

Added: 
    

Modified: 
    mlir/lib/IR/AffineMap.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 72effb38d614c..2838843e2eba6 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -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;
 }