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

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

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

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()) {