[Mlir-commits] [mlir] [mlir][NFC] Move `foldAttributesIntoMap` to `IR` build unit (PR #70155)

[Matthias Springer]

https://github.com/matthias-springer created https://github.com/llvm/llvm-project/pull/70155

`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.

>From 218bf2aa82ca6a177e6aca66d9b6327bd0383f49 Mon Sep 17 00:00:00 2001
From: Matthias Springer <springerm at google.com>
Date: Wed, 25 Oct 2023 11:16:20 +0900
Subject: [PATCH] [mlir][NFC] Move `foldAttributesIntoMap` to `IR` build unit

`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.
---
 mlir/include/mlir/IR/AffineMap.h         |  9 ++++++
 mlir/lib/Dialect/Affine/IR/AffineOps.cpp | 38 +++---------------------
 mlir/lib/IR/AffineMap.cpp                | 29 ++++++++++++++++++
 3 files changed, 42 insertions(+), 34 deletions(-)

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