[Mlir-commits] [mlir] [mlir][affine] Add affine-loop-constantize-bounds transfrom pass to affine dialect (PR #211989)
lonely eagle
llvmlistbot at llvm.org
Mon Jul 27 01:46:08 PDT 2026
https://github.com/linuxlonelyeagle updated https://github.com/llvm/llvm-project/pull/211989
>From 559a1cebe5d06c1b21e84ea5f40cfbebca7a219d Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Sat, 25 Jul 2026 05:56:24 +0000
Subject: [PATCH 1/2] add affine-loop-constantize-bounds pass.
---
.../mlir/Dialect/Affine/Transforms/Passes.td | 13 +++
.../AffineLoopConstantizeBounds.cpp | 104 ++++++++++++++++++
.../Dialect/Affine/Transforms/CMakeLists.txt | 1 +
.../affine-loop-constantize-bounds.mlir | 45 ++++++++
4 files changed, 163 insertions(+)
create mode 100644 mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
create mode 100644 mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
diff --git a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
index 7cecbcac5b932..44cb733c525b9 100644
--- a/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
+++ b/mlir/include/mlir/Dialect/Affine/Transforms/Passes.td
@@ -465,4 +465,17 @@ def AffineFoldMemRefAliasOps : Pass<"affine-fold-memref-alias-ops"> {
let dependentDialects = ["memref::MemRefDialect"];
}
+def AffineLoopConstantizeBounds : Pass<"affine-loop-constantize-bounds"> {
+ let summary = "Constantize affine loop bounds using range analysis";
+ let description = [{
+ The pass leverages value range analysis (e.g., via the Presburger library)
+ to transform affine loops with dynamic symbolic bounds into a static-trip-count
+ main loop and a residual tail loop.
+
+ By rewriting the main loop to have strict compile-time constant bounds
+ (typically normalized to `0` to `TC_safe` with step `1`), it exposes the loop
+ body to downstream affine optimization pipelines.
+ }];
+}
+
#endif // MLIR_DIALECT_AFFINE_PASSES
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
new file mode 100644
index 0000000000000..529d4043f5693
--- /dev/null
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
@@ -0,0 +1,104 @@
+//===- AffineLoopConstantizeBounds.cpp - Constantize loop bounds pass ---===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a pass to transform affine loops with symbolic bounds
+// into a static-trip-count main loop and a residual tail loop using range
+// analysis.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Transforms/Passes.h"
+#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
+#include <cstdint>
+
+namespace mlir {
+namespace affine {
+#define GEN_PASS_DEF_AFFINELOOPCONSTANTIZEBOUNDS
+#include "mlir/Dialect/Affine/Transforms/Passes.h.inc"
+} // namespace affine
+} // namespace mlir
+
+#define DEBUG_TYPE "affine-loop-constantize-bounds"
+
+using namespace mlir;
+using namespace mlir::affine;
+
+namespace {
+
+/// Computes the constant upper or lower bound for a given affine map expression
+/// \p map and its operands \p operands, constrained by the specified \p type.
+static FailureOr<int64_t> computeConstantBound(AffineMap map,
+ ValueRange operands,
+ presburger::BoundType type) {
+ ValueBoundsConstraintSet::Variable var(map, operands);
+ return ValueBoundsConstraintSet::computeConstantBound(type, var, nullptr,
+ {true, true});
+}
+
+static LogicalResult inferAffineLoopUpperConstantBound(AffineForOp forOp) {
+ // Ensure the loop is normalized (lower bound is strictly 0 and step is 1).
+ if (!forOp.hasConstantLowerBound() || forOp.getConstantLowerBound() != 0)
+ return failure();
+ if (forOp.getStepAsInt() != 1)
+ return failure();
+
+ // Infer the range [upperMin, upperMax] for the upper bound. We require a
+ // strictly positive minimum bound (upperMin > 0) to guarantee a safe,
+ // non-empty static trip count for the main loop.
+ FailureOr<int64_t> upperMin = computeConstantBound(
+ forOp.getUpperBoundMap(), forOp.getUpperBoundOperands(),
+ presburger::BoundType::LB);
+ FailureOr<int64_t> upperMax = computeConstantBound(
+ forOp.getUpperBoundMap(), forOp.getUpperBoundOperands(),
+ presburger::BoundType::UB);
+ if (failed(upperMin) || *upperMin <= 0)
+ return failure();
+
+ IRRewriter b(forOp->getContext());
+
+ // The upper bound is dynamic within [upperMin, upperMax]. Split the loop into
+ // a static main loop (0 to upperMin) and a residual tail loop (upperMin to
+ // dynamic bound).
+ if (failed(upperMax) || *upperMax > *upperMin) {
+ b.setInsertionPoint(forOp);
+ AffineForOp clonedForOp = cast<AffineForOp>(b.clone(*forOp));
+ clonedForOp.setConstantUpperBound(*upperMin);
+ forOp.setConstantLowerBound(*upperMin);
+ forOp.getInitsMutable().assign(clonedForOp->getResults());
+ return success();
+ }
+
+ // If upperMin == upperMax. The upper bound is proven to be a strict constant
+ // at compile time. Directly constantize the bound without peeling a tail
+ // loop.
+ forOp.setConstantUpperBound(*upperMin);
+ return success();
+}
+
+struct AffineLoopConstantizeBounds
+ : public affine::impl::AffineLoopConstantizeBoundsBase<
+ AffineLoopConstantizeBounds> {
+ void runOnOperation() override;
+};
+} // namespace
+
+void AffineLoopConstantizeBounds::runOnOperation() {
+ SmallVector<AffineForOp, 4> loops;
+ getOperation()->walk([&](AffineForOp forOp) {
+ // First collect eligible loops and normalize them so that their lower bound
+ // is locked to 0 and step to 1.
+ if (succeeded(normalizeAffineFor(forOp, false)))
+ loops.push_back(forOp);
+ });
+ // Infer and rewrite the upper bound into a compile-time constant for each
+ // loop.
+ for (AffineForOp loop : loops)
+ (void)inferAffineLoopUpperConstantBound(loop);
+}
diff --git a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
index 9d912139810b2..3a854e1c7dc36 100644
--- a/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
+++ b/mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt
@@ -2,6 +2,7 @@ add_mlir_dialect_library(MLIRAffineTransforms
AffineDataCopyGeneration.cpp
AffineExpandIndexOps.cpp
AffineExpandIndexOpsAsAffine.cpp
+ AffineLoopConstantizeBounds.cpp
AffineLoopInvariantCodeMotion.cpp
AffineLoopNormalize.cpp
AffineParallelize.cpp
diff --git a/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir b/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
new file mode 100644
index 0000000000000..6bb9aed223e83
--- /dev/null
+++ b/mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir
@@ -0,0 +1,45 @@
+// RUN: mlir-opt %s -affine-loop-constantize-bounds -split-input-file | FileCheck %s
+
+// CHECK-DAG: #[[$MAP_APPLY:.+]] = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
+// CHECK-DAG: #[[$MAP_UB:.+]] = affine_map<()[s0] -> ((-s0 + 9) ceildiv 2)>
+
+// CHECK-LABEL: func @peeling_main_loop
+
+func.func @peeling_main_loop() {
+ %c0 = arith.constant 0 :index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ affine.for %iv = %bound to 9 step 2 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %bound : index
+ affine.yield %sum : index
+ }
+ return
+}
+
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[BOUND:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+// CHECK: %[[MAIN_RES:.*]] = affine.for %[[IV_MAIN:.*]] = 0 to 4 iter_args(%[[ARG_MAIN:.*]] = %[[C0]]) -> (index) {
+// CHECK: %{{.*}} = affine.apply #map(%[[IV_MAIN]])[%[[BOUND]]]
+// CHECK: }
+// CHECK: %[[TAIL_RES:.*]] = affine.for %[[IV_TAIL:.*]] = 4 to #map1()[%[[BOUND]]] iter_args(%[[ARG_TAIL:.*]] = %[[MAIN_RES]]) -> (index) {
+// CHECK: %{{.*}} = affine.apply #map(%[[IV_TAIL]])[%[[BOUND]]]
+// CHECK: }
+
+// -----
+// CHECK: #map = affine_map<(d0)[s0] -> (d0 * 2 + s0)>
+// CHECK-LABEL: func @fully_constantized_no_peeling
+
+func.func @fully_constantized_no_peeling() {
+ %c0 = arith.constant 0 :index
+ %bound = test.value_with_bounds { min = 0 : index, max = 1 : index}
+ affine.for %iv = %bound to 6 step 2 iter_args(%arg = %c0) -> index {
+ %sum = arith.addi %arg, %bound : index
+ affine.yield %sum : index
+ }
+ return
+}
+
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[BOUND:.*]] = test.value_with_bounds {max = 1 : index, min = 0 : index}
+// CHECK: %{{.*}} = affine.for %[[IV:.*]] = 0 to 3 iter_args(%{{.*}} = %[[C0]]) -> (index) {
+// CHECK: %{{.*}} = affine.apply #map(%[[IV]])[%[[BOUND]]]
+// CHECK: }
>From dbe3e12c7d80edb8585bfc6794ebaa204a622569 Mon Sep 17 00:00:00 2001
From: linuxlonelyeagle <2020382038 at qq.com>
Date: Mon, 27 Jul 2026 08:45:50 +0000
Subject: [PATCH 2/2] fix nit.
---
.../AffineLoopConstantizeBounds.cpp | 23 +++++++++++++------
1 file changed, 16 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
index 529d4043f5693..ed424d2638631 100644
--- a/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp
@@ -15,6 +15,7 @@
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Affine/Transforms/Passes.h"
#include "mlir/Dialect/Affine/Utils.h"
+#include "mlir/IR/PatternMatch.h"
#include "mlir/Interfaces/ValueBoundsOpInterface.h"
#include <cstdint>
@@ -33,16 +34,24 @@ using namespace mlir::affine;
namespace {
/// Computes the constant upper or lower bound for a given affine map expression
-/// \p map and its operands \p operands, constrained by the specified \p type.
+/// and its operands, constrained by the specified type.
static FailureOr<int64_t> computeConstantBound(AffineMap map,
ValueRange operands,
presburger::BoundType type) {
ValueBoundsConstraintSet::Variable var(map, operands);
- return ValueBoundsConstraintSet::computeConstantBound(type, var, nullptr,
- {true, true});
+ return ValueBoundsConstraintSet::computeConstantBound(
+ type, var, nullptr, {/*closedUb=*/true, /*allowIntegerType=*/true});
}
-static LogicalResult inferAffineLoopUpperConstantBound(AffineForOp forOp) {
+/// Attempts to infer a static constant upper bound for the given normalized
+/// `affine.for` loop using Value Bounds Analysis. If the dynamic upper bound's
+/// range [upperMin, upperMax] is proven to be a single constant value (upperMin
+/// == upperMax), the upper bound is directly replaced with this constant.
+/// Otherwise, if upperMin > 0, the loop is split (peeled) into a static main
+/// loop with a constant upper bound (`upperMin`) and a residual tail loop
+/// iterating from `upperMin` to the original dynamic bound.
+static LogicalResult inferAffineLoopUpperConstantBound(AffineForOp forOp,
+ RewriterBase &b) {
// Ensure the loop is normalized (lower bound is strictly 0 and step is 1).
if (!forOp.hasConstantLowerBound() || forOp.getConstantLowerBound() != 0)
return failure();
@@ -61,8 +70,6 @@ static LogicalResult inferAffineLoopUpperConstantBound(AffineForOp forOp) {
if (failed(upperMin) || *upperMin <= 0)
return failure();
- IRRewriter b(forOp->getContext());
-
// The upper bound is dynamic within [upperMin, upperMax]. Split the loop into
// a static main loop (0 to upperMin) and a residual tail loop (upperMin to
// dynamic bound).
@@ -97,8 +104,10 @@ void AffineLoopConstantizeBounds::runOnOperation() {
if (succeeded(normalizeAffineFor(forOp, false)))
loops.push_back(forOp);
});
+
// Infer and rewrite the upper bound into a compile-time constant for each
// loop.
+ IRRewriter b(&getContext());
for (AffineForOp loop : loops)
- (void)inferAffineLoopUpperConstantBound(loop);
+ (void)inferAffineLoopUpperConstantBound(loop, b);
}
More information about the Mlir-commits
mailing list