[Mlir-commits] [mlir] [mlir][affine] Add AffineValueMap::combine and sum (PR #228837)

[lonely eagle]

https://github.com/linuxlonelyeagle created https://github.com/llvm/llvm-project/pull/228837

Applying a schedule to loops requires bounds for the new loops. For supported schedules, these bounds can be derived by combining lower or upper bounds from multiple source loops. Add `AffineValueMap::combine` to apply an expression over `d0` and `d1` to corresponding results of two value maps. Compose, canonicalize, and simplify the result. Use it to implement difference and pairwise sum.

>From eb4dae05b8f4e1229eb2e2eb285ffd877aba3e1d Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sun, 4 Oct 2026 06:37:42 +0000
Subject: [PATCH] Add AffineValueMap::combine and sum.

---
 .../mlir/Dialect/Affine/IR/AffineValueMap.h   | 13 +++++
 mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp | 55 ++++++++++++++-----
 2 files changed, 54 insertions(+), 14 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h b/mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h
index ac480d6997ef430..1f2f8e788fbbe66 100644
--- a/mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h
+++ b/mlir/include/mlir/Dialect/Affine/IR/AffineValueMap.h
@@ -49,12 +49,25 @@ class AffineValueMap {
   /// operands, but the result count remains the same.
   void composeSimplifyAndCanonicalize();
 
+  /// Combine corresponding results of 'a' and 'b' using 'expr'. The expression
+  /// may refer only to d0 and d1, which stand for the corresponding results of
+  /// 'a' and 'b'. For example, difference() uses `d0 - d1`, and sum() uses
+  /// `d0 + d1`.
+  static void combine(const AffineValueMap &a, const AffineValueMap &b,
+                      AffineExpr expr, AffineValueMap *res);
+
   /// Return the value map that is the difference of value maps 'a' and 'b',
   /// represented as an affine map and its operands. The output map + operands
   /// are canonicalized and simplified.
   static void difference(const AffineValueMap &a, const AffineValueMap &b,
                          AffineValueMap *res);
 
+  /// Return the value map that is the sum of value maps 'a' and 'b',
+  /// represented as an affine map and its operands. The output map + operands
+  /// are canonicalized and simplified.
+  static void sum(const AffineValueMap &a, const AffineValueMap &b,
+                  AffineValueMap *res);
+
   /// Return true if the idx^th result can be proved to be a multiple of
   /// 'factor', false otherwise.
   inline bool isMultipleOf(unsigned idx, int64_t factor) const;
diff --git a/mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp b/mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp
index 6d2d3212c696b50..3a02b56154aac66 100644
--- a/mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp
+++ b/mlir/lib/Dialect/Affine/IR/AffineValueMap.cpp
@@ -33,10 +33,19 @@ void AffineValueMap::composeSimplifyAndCanonicalize() {
   this->map.reset(sMap);
 }
 
-void AffineValueMap::difference(const AffineValueMap &a,
-                                const AffineValueMap &b, AffineValueMap *res) {
+void AffineValueMap::combine(const AffineValueMap &a, const AffineValueMap &b,
+                             AffineExpr expr, AffineValueMap *res) {
   assert(a.getNumResults() == b.getNumResults() && "invalid inputs");
 
+  bool validExpr = true;
+  expr.walk([&](AffineExpr subExpr) {
+    if (auto dim = llvm::dyn_cast<AffineDimExpr>(subExpr))
+      validExpr &= dim.getPosition() < 2;
+    else if (llvm::isa<AffineSymbolExpr>(subExpr))
+      validExpr = false;
+  });
+  assert(validExpr && "combining expression must use only d0 and d1");
+
   SmallVector<Value, 4> allOperands;
   allOperands.reserve(a.getNumOperands() + b.getNumOperands());
   auto aDims = a.getOperands().take_front(a.getNumDims());
@@ -53,19 +62,37 @@ void AffineValueMap::difference(const AffineValueMap &a,
                   .shiftDims(a.getNumDims())
                   .shiftSymbols(a.getNumSymbols());
 
-  // Construct the difference expressions.
+  // Substitute the corresponding results into the combining expression.
   auto aMap = a.getAffineMap();
-  SmallVector<AffineExpr, 4> diffExprs;
-  diffExprs.reserve(a.getNumResults());
-  for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i)
-    diffExprs.push_back(aMap.getResult(i) - bMap.getResult(i));
-
-  auto diffMap = AffineMap::get(bMap.getNumDims(), bMap.getNumSymbols(),
-                                diffExprs, bMap.getContext());
-  fullyComposeAffineMapAndOperands(&diffMap, &allOperands);
-  canonicalizeMapAndOperands(&diffMap, &allOperands);
-  diffMap = simplifyAffineMap(diffMap);
-  res->reset(diffMap, allOperands);
+  SmallVector<AffineExpr, 4> resultExprs;
+  resultExprs.reserve(a.getNumResults());
+  for (unsigned i = 0, e = bMap.getNumResults(); i < e; ++i) {
+    AffineExpr replacements[] = {aMap.getResult(i), bMap.getResult(i)};
+    resultExprs.push_back(expr.replaceDims(replacements));
+  }
+
+  auto resultMap = AffineMap::get(bMap.getNumDims(), bMap.getNumSymbols(),
+                                  resultExprs, bMap.getContext());
+  fullyComposeAffineMapAndOperands(&resultMap, &allOperands);
+  canonicalizeMapAndOperands(&resultMap, &allOperands);
+  resultMap = simplifyAffineMap(resultMap);
+  res->reset(resultMap, allOperands);
+}
+
+void AffineValueMap::difference(const AffineValueMap &a,
+                                const AffineValueMap &b, AffineValueMap *res) {
+  MLIRContext *context = a.getAffineMap().getContext();
+  AffineExpr lhs = getAffineDimExpr(0, context);
+  AffineExpr rhs = getAffineDimExpr(1, context);
+  combine(a, b, lhs - rhs, res);
+}
+
+void AffineValueMap::sum(const AffineValueMap &a, const AffineValueMap &b,
+                         AffineValueMap *res) {
+  MLIRContext *context = a.getAffineMap().getContext();
+  AffineExpr lhs = getAffineDimExpr(0, context);
+  AffineExpr rhs = getAffineDimExpr(1, context);
+  combine(a, b, lhs + rhs, res);
 }
 
 // Returns true and sets 'indexOfMatch' if 'valueToMatch' is found in