[Mlir-commits] [mlir] [mlir][affine] Add affine-loop-constantize-bounds transfrom pass to affine dialect (PR #211989)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 24 23:04:19 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: lonely eagle (linuxlonelyeagle)
<details>
<summary>Changes</summary>
This PR introduces the `affine-loop-constantize-bounds` pass, which leverages `ValueBoundsConstraintSet` analysis to refine dynamic/symbolic affine.for loop bounds into compile-time constant bounds.This transformation effectively bridges the gap between constrained symbolic bounds and the rigid constant requirements of the existing affine infrastructure. RFC: https://discourse.llvm.org/t/rfc-mlir-enable-dynamic-and-tighter-affine-unrolling-via-valueboundsconstraintset/91055/2
---
Full diff: https://github.com/llvm/llvm-project/pull/211989.diff
4 Files Affected:
- (modified) mlir/include/mlir/Dialect/Affine/Transforms/Passes.td (+13)
- (added) mlir/lib/Dialect/Affine/Transforms/AffineLoopConstantizeBounds.cpp (+104)
- (modified) mlir/lib/Dialect/Affine/Transforms/CMakeLists.txt (+1)
- (added) mlir/test/Dialect/Affine/affine-loop-constantize-bounds.mlir (+45)
``````````diff
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: }
``````````
</details>
https://github.com/llvm/llvm-project/pull/211989
More information about the Mlir-commits
mailing list