[mlir][NFC] Move foldAttributesIntoMap to IR build unit

[matthias-springer]

foldAttributesIntoMap is a helper function that folds constant OpFoldResult into an affine map. This commit moves the function from the affine dialect to AffineMap.h, so that it can be used without depending on the affine dialect.

[llvmbot]

@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-affine

@llvm/pr-subscribers-mlir-core

Author: Matthias Springer (matthias-springer)

Changes

foldAttributesIntoMap is a helper function that folds constant OpFoldResult into an affine map. This commit moves the function from the affine dialect to AffineMap.h, so that it can be used without depending on the affine dialect.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/IR/AffineMap.h (+9)
• (modified) mlir/lib/Dialect/Affine/IR/AffineOps.cpp (+4-34)
• (modified) mlir/lib/IR/AffineMap.cpp (+29)

diff --git a/mlir/include/mlir/IR/AffineMap.h b/mlir/include/mlir/IR/AffineMap.h
index 3430db2b99c3f2e..5af7835258f6bd2 100644
--- a/mlir/include/mlir/IR/AffineMap.h
+++ b/mlir/include/mlir/IR/AffineMap.h
@@ -15,6 +15,7 @@
 #define MLIR_IR_AFFINEMAP_H
 
 #include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/Value.h"
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMapInfo.h"
@@ -33,7 +34,9 @@ struct AffineMapStorage;
 } // namespace detail
 
 class Attribute;
+class Builder;
 struct LogicalResult;
+class OpFoldResult;
 class MLIRContext;
 
 /// A multi-dimensional affine map
@@ -447,6 +450,12 @@ AffineMap compressUnusedSymbols(AffineMap map);
 /// dims and symbols.
 SmallVector<AffineMap> compressUnusedSymbols(ArrayRef<AffineMap> maps);
 
+/// Fold all attributes among the given operands into the affine map. Return the
+/// folded affine map. Return all remaining values via `remainingValues`.
+AffineMap foldAttributesIntoMap(Builder &b, AffineMap map,
+                                ArrayRef<OpFoldResult> operands,
+                                SmallVector<Value> &remainingValues);
+
 /// Returns a map with the same dimension and symbol count as `map`, but whose
 /// results are the unique affine expressions of `map`.
 AffineMap removeDuplicateExprs(AffineMap map);
diff --git a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
index 85d16088c43fb1e..4d79c458889d2a9 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineOps.cpp
@@ -1193,42 +1193,11 @@ void mlir::affine::fullyComposeAffineMapAndOperands(
   }
 }
 
-/// Fold all attributes among the given operands into the affine map and return
-/// all remaining values. The affine map is modified in-place.
-static SmallVector<Value>
-foldAttributesIntoMap(Builder &b, AffineMap *map,
-                      ArrayRef<OpFoldResult> operands) {
-  SmallVector<AffineExpr> dimReplacements, symReplacements;
-  SmallVector<Value> valueOperands;
-  int64_t numDims = 0;
-  for (int64_t i = 0; i < map->getNumDims(); ++i) {
-    if (auto attr = operands[i].dyn_cast<Attribute>()) {
-      dimReplacements.push_back(
-          b.getAffineConstantExpr(attr.cast<IntegerAttr>().getInt()));
-    } else {
-      dimReplacements.push_back(b.getAffineDimExpr(numDims++));
-      valueOperands.push_back(operands[i].get<Value>());
-    }
-  }
-  int64_t numSymbols = 0;
-  for (int64_t i = 0; i < map->getNumSymbols(); ++i) {
-    if (auto attr = operands[i + map->getNumDims()].dyn_cast<Attribute>()) {
-      symReplacements.push_back(
-          b.getAffineConstantExpr(attr.cast<IntegerAttr>().getInt()));
-    } else {
-      symReplacements.push_back(b.getAffineSymbolExpr(numSymbols++));
-      valueOperands.push_back(operands[i + map->getNumDims()].get<Value>());
-    }
-  }
-  *map = map->replaceDimsAndSymbols(dimReplacements, symReplacements, numDims,
-                                    numSymbols);
-  return valueOperands;
-}
-
 AffineApplyOp
 mlir::affine::makeComposedAffineApply(OpBuilder &b, Location loc, AffineMap map,
                                       ArrayRef<OpFoldResult> operands) {
-  SmallVector<Value> valueOperands = foldAttributesIntoMap(b, &map, operands);
+  SmallVector<Value> valueOperands;
+  map = foldAttributesIntoMap(b, map, operands, valueOperands);
   composeAffineMapAndOperands(&map, &valueOperands);
   assert(map);
   return b.create<AffineApplyOp>(loc, map, valueOperands);
@@ -1331,7 +1300,8 @@ mlir::affine::makeComposedFoldedMultiResultAffineApply(
 template <typename OpTy>
 static OpTy makeComposedMinMax(OpBuilder &b, Location loc, AffineMap map,
                                ArrayRef<OpFoldResult> operands) {
-  SmallVector<Value> valueOperands = foldAttributesIntoMap(b, &map, operands);
+  SmallVector<Value> valueOperands;
+  map = foldAttributesIntoMap(b, map, operands, valueOperands);
   composeMultiResultAffineMap(map, valueOperands);
   return b.create<OpTy>(loc, b.getIndexType(), map, valueOperands);
 }
diff --git a/mlir/lib/IR/AffineMap.cpp b/mlir/lib/IR/AffineMap.cpp
index 9cdac964710ca86..3bd1181b6c7bbd8 100644
--- a/mlir/lib/IR/AffineMap.cpp
+++ b/mlir/lib/IR/AffineMap.cpp
@@ -9,6 +9,7 @@
 #include "mlir/IR/AffineMap.h"
 #include "AffineMapDetail.h"
 #include "mlir/IR/AffineExpr.h"
+#include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Support/LogicalResult.h"
@@ -658,6 +659,34 @@ SmallVector<AffineMap> mlir::compressUnusedSymbols(ArrayRef<AffineMap> maps) {
       maps, [](AffineMap m) { return compressUnusedSymbols(m); });
 }
 
+AffineMap mlir::foldAttributesIntoMap(Builder &b, AffineMap map,
+                                      ArrayRef<OpFoldResult> operands,
+                                      SmallVector<Value> &remainingValues) {
+  SmallVector<AffineExpr> dimReplacements, symReplacements;
+  int64_t numDims = 0;
+  for (int64_t i = 0; i < map.getNumDims(); ++i) {
+    if (auto attr = operands[i].dyn_cast<Attribute>()) {
+      dimReplacements.push_back(
+          b.getAffineConstantExpr(attr.cast<IntegerAttr>().getInt()));
+    } else {
+      dimReplacements.push_back(b.getAffineDimExpr(numDims++));
+      remainingValues.push_back(operands[i].get<Value>());
+    }
+  }
+  int64_t numSymbols = 0;
+  for (int64_t i = 0; i < map.getNumSymbols(); ++i) {
+    if (auto attr = operands[i + map.getNumDims()].dyn_cast<Attribute>()) {
+      symReplacements.push_back(
+          b.getAffineConstantExpr(attr.cast<IntegerAttr>().getInt()));
+    } else {
+      symReplacements.push_back(b.getAffineSymbolExpr(numSymbols++));
+      remainingValues.push_back(operands[i + map.getNumDims()].get<Value>());
+    }
+  }
+  return map.replaceDimsAndSymbols(dimReplacements, symReplacements, numDims,
+                                   numSymbols);
+}
+
 AffineMap mlir::simplifyAffineMap(AffineMap map) {
   SmallVector<AffineExpr, 8> exprs;
   for (auto e : map.getResults()) {