[Mlir-commits] [mlir] MLIR][Affine] Support bound computation for local variables in getBoundForAffineExpr (PR #212167)

[lonely eagle]

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

This PR extends `getBoundForAffineExpr` to recursively compute upper and lower bounds for local variables (`flattener.localExprs`).

>From f4daf40a051826d8d1fc935938e572895c0f0f20 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sun, 26 Jul 2026 15:37:39 +0000
Subject: [PATCH] add compute local var logic in the getBoundForAffineExpr.

---
 mlir/lib/IR/AffineExpr.cpp                 | 56 +++++++++++++++-------
 mlir/test/Dialect/Affine/canonicalize.mlir | 18 +++++++
 2 files changed, 57 insertions(+), 17 deletions(-)

diff --git a/mlir/lib/IR/AffineExpr.cpp b/mlir/lib/IR/AffineExpr.cpp
index da91066815ca4..fbc218f149de0 100644
--- a/mlir/lib/IR/AffineExpr.cpp
+++ b/mlir/lib/IR/AffineExpr.cpp
@@ -17,6 +17,7 @@
 #include "mlir/IR/IntegerSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVectorExtras.h"
+#include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/MathExtras.h"
 #include <numeric>
 #include <optional>
@@ -1710,25 +1711,46 @@ std::optional<int64_t> mlir::getBoundForAffineExpr(
   if (failed(simpleResult))
     return std::nullopt;
   ArrayRef<int64_t> flattenedExpr = flattener.operandExprStack.back();
-  // TODO: Handle local variables. We can get hold of flattener.localExprs and
-  // get bound on the local expr recursively.
-  if (flattener.numLocals > 0)
-    return std::nullopt;
   int64_t bound = 0;
-  // Substitute the constant lower or upper bound for the dimensional or
-  // symbolic input depending on `isUpper` to determine the bound.
-  for (unsigned i = 0, e = numDims + numSymbols; i < e; ++i) {
-    if (flattenedExpr[i] > 0) {
-      auto &constBound = isUpper ? constUpperBounds[i] : constLowerBounds[i];
-      if (!constBound)
-        return std::nullopt;
-      bound += *constBound * flattenedExpr[i];
-    } else if (flattenedExpr[i] < 0) {
-      auto &constBound = isUpper ? constLowerBounds[i] : constUpperBounds[i];
-      if (!constBound)
-        return std::nullopt;
-      bound += *constBound * flattenedExpr[i];
+
+  // Helper to retrieve the bound for a dimension, symbol, or local variable.
+  // Returns std::nullopt if the required bound is not available.
+  auto getBoundForIndex = [&](unsigned index,
+                              bool wantUpper) -> std::optional<int64_t> {
+    unsigned numDimsAndSymbols = numDims + numSymbols;
+
+    // Dimensions and symbols: retrieve directly from the input bounds arrays.
+    if (index < numDimsAndSymbols) {
+      const std::optional<int64_t> &bound =
+          wantUpper ? constUpperBounds[index] : constLowerBounds[index];
+      return bound;
     }
+
+    // Local variables, recursively compute the bound on the underlying local
+    // expression.
+    unsigned localIndex = index - numDimsAndSymbols;
+    return getBoundForAffineExpr(flattener.localExprs[localIndex], numDims,
+                                 numSymbols, constLowerBounds, constUpperBounds,
+                                 wantUpper);
+  };
+
+  // Accumulate bounds for dimensions, symbols, and local variables.
+  unsigned totalVars = numDims + numSymbols + flattener.numLocals;
+  for (unsigned i = 0, e = totalVars; i < e; ++i) {
+    int64_t flattenedValue = flattenedExpr[i];
+    if (flattenedValue == 0)
+      continue;
+
+    // Determine the bound direction needed for the term. If coefficient > 0, we
+    // need upper bound for max (isUpper=true) and lower bound for min
+    // (isUpper=false). If coefficient < 0, the direction flips.
+    bool wantUpper = (flattenedValue > 0) == isUpper;
+
+    std::optional<int64_t> constBound = getBoundForIndex(i, wantUpper);
+    if (!constBound)
+      return std::nullopt;
+
+    bound += *constBound * flattenedValue;
   }
   // Constant term.
   bound += flattenedExpr.back();
diff --git a/mlir/test/Dialect/Affine/canonicalize.mlir b/mlir/test/Dialect/Affine/canonicalize.mlir
index 7d236ef3c2421..95f51a51f695d 100644
--- a/mlir/test/Dialect/Affine/canonicalize.mlir
+++ b/mlir/test/Dialect/Affine/canonicalize.mlir
@@ -2614,3 +2614,21 @@ func.func @split_delinearize_spanning_final_part_vector(
   %1:4 = affine.delinearize_index %0 into (2, 3, 8, 4) : vector<4xindex>, vector<4xindex>, vector<4xindex>, vector<4xindex>
   return %1#0, %1#1, %1#2, %1#3 : vector<4xindex>, vector<4xindex>, vector<4xindex>, vector<4xindex>
 }
+
+// -----
+
+// CHECK-DAG: #[[$MAP:.*]] = affine_map<(d0, d1) -> (d0 + (d0 + d1) floordiv 4 + 2)>
+// CHECK-LABEL: func @simplify_loop_ub_with_local_var
+// CHECK:         affine.for %[[VAL_0:.*]] = 0 to 4 {
+// CHECK:           affine.for %[[VAL_1:.*]] = 0 to 4 {
+// CHECK:             affine.for %[[VAL_2:.*]] = 0 to #[[$MAP]](%[[VAL_0]], %[[VAL_1]]) {
+func.func @simplify_loop_ub_with_local_var(%arg0: index) {
+  affine.for %arg1 = 0 to 4 {
+    affine.for %arg2 = 0 to 4 {
+      affine.for %arg3 = 0 to min affine_map<(d0, d1) -> (10, (d0 + d1) floordiv 4 + d0 + 2)>(%arg1, %arg2) {
+        "test.foo"(%arg3) : (index) -> ()
+      }
+    }
+  }
+  return
+}