[flang-commits] [flang] [flang][OpenMP] Fold canonical loop trip counts during lowering (PR #215238)
Tom Eccles via flang-commits
flang-commits at lists.llvm.org
Mon Aug 10 03:39:54 PDT 2026
https://github.com/tblah created https://github.com/llvm/llvm-project/pull/215238
Canonical loop lowering currently materializes the complete trip-count arithmetic even when the loop bounds and step are compile-time constants.
Use FirOpBuilder::createOrFold for the comparisons, selects, subtraction, unsigned division, and addition in the trip-count prologue. This folds constant trip counts while preserving the existing dynamic lowering and overflow flags.
Add bare-unroll lowering coverage for ascending, descending, zero-trip, non-unit-step, dynamic, and i64 loops.
I'm doing this because
1) Surprisingly, no canonicalization pass between lowering and LLVM-IR
translation is doing this already. Folding here saves us building
LLVM instructions just for the middle-end to fold it.
2) This could simplify some analysis for #214115
Assisted-by: Codex
>From 41c70f673572754954f08f751b6e0955129d6212 Mon Sep 17 00:00:00 2001
From: Tom Eccles <tom.eccles at arm.com>
Date: Fri, 7 Aug 2026 16:58:54 +0100
Subject: [PATCH] [flang][OpenMP] Fold canonical loop trip counts during
lowering
Canonical loop lowering currently materializes the complete trip-count
arithmetic even when the loop bounds and step are compile-time constants.
Use FirOpBuilder::createOrFold for the comparisons, selects,
subtraction, unsigned division, and addition in the trip-count prologue.
This folds constant trip counts while preserving the existing dynamic
lowering and overflow flags.
Add bare-unroll lowering coverage for ascending, descending, zero-trip,
non-unit-step, dynamic, and i64 loops.
I'm doing this because
1) Surprisingly, no canonicalization pass between lowering and LLVM-IR
translation is doing this already. Folding here saves us building
LLVM instructions just for the middle-end to fold it.
2) This could simplify some analysis for #214115
Assisted-by: Codex
---
flang/lib/Lower/OpenMP/OpenMP.cpp | 49 ++++----
.../OpenMP/canonical-loop-trip-count.f90 | 119 ++++++++++++++++++
2 files changed, 143 insertions(+), 25 deletions(-)
create mode 100644 flang/test/Lower/OpenMP/canonical-loop-trip-count.f90
diff --git a/flang/lib/Lower/OpenMP/OpenMP.cpp b/flang/lib/Lower/OpenMP/OpenMP.cpp
index 3876799b3a081..71df19ff4d5f2 100644
--- a/flang/lib/Lower/OpenMP/OpenMP.cpp
+++ b/flang/lib/Lower/OpenMP/OpenMP.cpp
@@ -3167,6 +3167,11 @@ static void genCanonicalLoopNest(
converter.genExprValue(*semantics::GetExpr(bounds->Lower()), stmtCtx));
mlir::Value loopUBVar = fir::getBase(
converter.genExprValue(*semantics::GetExpr(bounds->Upper()), stmtCtx));
+
+ // Get the integer kind for the loop variable and cast the loop bounds.
+ size_t loopVarTypeSize = bounds->Name().thing.symbol->GetUltimate().size();
+ mlir::Type loopVarType = getLoopVarType(converter, loopVarTypeSize);
+ loopVarTypes.push_back(loopVarType);
mlir::Value loopStepVar = [&]() {
if (auto &step = bounds->Step()) {
return fir::getBase(
@@ -3174,14 +3179,9 @@ static void genCanonicalLoopNest(
}
// If `step` is not present, assume it is `1`.
- auto intTy = firOpBuilder.getI32Type();
- return firOpBuilder.createIntegerConstant(loc, intTy, 1);
+ return firOpBuilder.createIntegerConstant(loc, loopVarType, 1);
}();
- // Get the integer kind for the loop variable and cast the loop bounds
- size_t loopVarTypeSize = bounds->Name().thing.symbol->GetUltimate().size();
- mlir::Type loopVarType = getLoopVarType(converter, loopVarTypeSize);
- loopVarTypes.push_back(loopVarType);
loopLBVar = firOpBuilder.createConvert(loc, loopVarType, loopLBVar);
loopUBVar = firOpBuilder.createConvert(loc, loopVarType, loopUBVar);
loopStepVar = firOpBuilder.createConvert(loc, loopVarType, loopStepVar);
@@ -3191,34 +3191,33 @@ static void genCanonicalLoopNest(
// Start lowering
mlir::Value zero = firOpBuilder.createIntegerConstant(loc, loopVarType, 0);
mlir::Value one = firOpBuilder.createIntegerConstant(loc, loopVarType, 1);
- mlir::Value isDownwards = mlir::arith::CmpIOp::create(
- firOpBuilder, loc, mlir::arith::CmpIPredicate::slt, loopStepVar, zero);
+ mlir::Value isDownwards = firOpBuilder.createOrFold<mlir::arith::CmpIOp>(
+ loc, mlir::arith::CmpIPredicate::slt, loopStepVar, zero);
// Ensure we are counting upwards. If not, negate step and swap lb and ub.
mlir::Value negStep =
- mlir::arith::SubIOp::create(firOpBuilder, loc, zero, loopStepVar);
- mlir::Value incr = mlir::arith::SelectOp::create(
- firOpBuilder, loc, isDownwards, negStep, loopStepVar);
- mlir::Value lb = mlir::arith::SelectOp::create(
- firOpBuilder, loc, isDownwards, loopUBVar, loopLBVar);
- mlir::Value ub = mlir::arith::SelectOp::create(
- firOpBuilder, loc, isDownwards, loopLBVar, loopUBVar);
+ firOpBuilder.createOrFold<mlir::arith::SubIOp>(loc, zero, loopStepVar);
+ mlir::Value incr = firOpBuilder.createOrFold<mlir::arith::SelectOp>(
+ loc, isDownwards, negStep, loopStepVar);
+ mlir::Value lb = firOpBuilder.createOrFold<mlir::arith::SelectOp>(
+ loc, isDownwards, loopUBVar, loopLBVar);
+ mlir::Value ub = firOpBuilder.createOrFold<mlir::arith::SelectOp>(
+ loc, isDownwards, loopLBVar, loopUBVar);
// Compute the trip count assuming lb <= ub. This guarantees that the result
// is non-negative and we can use unsigned arithmetic.
- mlir::Value span = mlir::arith::SubIOp::create(
- firOpBuilder, loc, ub, lb, ::mlir::arith::IntegerOverflowFlags::nuw);
+ mlir::Value span = firOpBuilder.createOrFold<mlir::arith::SubIOp>(
+ loc, ub, lb, ::mlir::arith::IntegerOverflowFlags::nuw);
mlir::Value tcMinusOne =
- mlir::arith::DivUIOp::create(firOpBuilder, loc, span, incr);
- mlir::Value tcIfLooping =
- mlir::arith::AddIOp::create(firOpBuilder, loc, tcMinusOne, one,
- ::mlir::arith::IntegerOverflowFlags::nuw);
+ firOpBuilder.createOrFold<mlir::arith::DivUIOp>(loc, span, incr);
+ mlir::Value tcIfLooping = firOpBuilder.createOrFold<mlir::arith::AddIOp>(
+ loc, tcMinusOne, one, ::mlir::arith::IntegerOverflowFlags::nuw);
// Fall back to 0 if lb > ub
- mlir::Value isZeroTC = mlir::arith::CmpIOp::create(
- firOpBuilder, loc, mlir::arith::CmpIPredicate::slt, ub, lb);
- mlir::Value tripcount = mlir::arith::SelectOp::create(
- firOpBuilder, loc, isZeroTC, zero, tcIfLooping);
+ mlir::Value isZeroTC = firOpBuilder.createOrFold<mlir::arith::CmpIOp>(
+ loc, mlir::arith::CmpIPredicate::slt, ub, lb);
+ mlir::Value tripcount = firOpBuilder.createOrFold<mlir::arith::SelectOp>(
+ loc, isZeroTC, zero, tcIfLooping);
tripcounts.push_back(tripcount);
// Create the CLI handle.
diff --git a/flang/test/Lower/OpenMP/canonical-loop-trip-count.f90 b/flang/test/Lower/OpenMP/canonical-loop-trip-count.f90
new file mode 100644
index 0000000000000..b96fa03f64b5a
--- /dev/null
+++ b/flang/test/Lower/OpenMP/canonical-loop-trip-count.f90
@@ -0,0 +1,119 @@
+! RUN: %flang_fc1 -emit-hlfir -fopenmp -fopenmp-version=51 -o - %s 2>&1 | FileCheck %s
+
+subroutine constant_ascending()
+ integer :: i
+ !$omp unroll
+ do i = 1, 100
+ end do
+ !$omp end unroll
+end subroutine
+
+! CHECK-LABEL: func.func @_QPconstant_ascending()
+! CHECK-NOT: arith.cmpi
+! CHECK-NOT: arith.select
+! CHECK-NOT: arith.subi
+! CHECK-NOT: arith.divui
+! CHECK-NOT: arith.addi
+! CHECK: %[[TRIP_COUNT:.*]] = arith.constant 100 : i32
+! CHECK-NEXT: %[[CLI:.*]] = omp.new_cli
+! CHECK-NEXT: omp.canonical_loop(%[[CLI]]) {{.*}} : i32 in range(%[[TRIP_COUNT]]) {
+
+subroutine constant_descending()
+ integer :: i
+ !$omp unroll
+ do i = 100, 1, -1
+ end do
+ !$omp end unroll
+end subroutine
+
+! CHECK-LABEL: func.func @_QPconstant_descending()
+! CHECK-NOT: arith.cmpi
+! CHECK-NOT: arith.select
+! CHECK-NOT: arith.subi
+! CHECK-NOT: arith.divui
+! CHECK-NOT: arith.addi
+! CHECK: arith.constant -1 : i32
+! CHECK: %[[TRIP_COUNT:.*]] = arith.constant 100 : i32
+! CHECK-NEXT: %[[CLI:.*]] = omp.new_cli
+! CHECK-NEXT: omp.canonical_loop(%[[CLI]]) {{.*}} : i32 in range(%[[TRIP_COUNT]]) {
+
+subroutine constant_zero_trip()
+ integer :: i
+ !$omp unroll
+ do i = 100, 1
+ end do
+ !$omp end unroll
+end subroutine
+
+! CHECK-LABEL: func.func @_QPconstant_zero_trip()
+! CHECK-NOT: arith.cmpi
+! CHECK-NOT: arith.select
+! CHECK-NOT: arith.subi
+! CHECK-NOT: arith.divui
+! CHECK-NOT: arith.addi
+! CHECK: %[[TRIP_COUNT:.*]] = arith.constant 0 : i32
+! CHECK-NEXT: %[[CLI:.*]] = omp.new_cli
+! CHECK-NEXT: omp.canonical_loop(%[[CLI]]) {{.*}} : i32 in range(%[[TRIP_COUNT]]) {
+
+subroutine constant_non_unit_step()
+ integer :: i
+ !$omp unroll
+ do i = 1, 100, 3
+ end do
+ !$omp end unroll
+end subroutine
+
+! CHECK-LABEL: func.func @_QPconstant_non_unit_step()
+! CHECK-NOT: arith.cmpi
+! CHECK-NOT: arith.select
+! CHECK-NOT: arith.subi
+! CHECK-NOT: arith.divui
+! CHECK-NOT: arith.addi
+! CHECK: %[[TRIP_COUNT:.*]] = arith.constant 34 : i32
+! CHECK-NEXT: %[[CLI:.*]] = omp.new_cli
+! CHECK-NEXT: omp.canonical_loop(%[[CLI]]) {{.*}} : i32 in range(%[[TRIP_COUNT]]) {
+
+subroutine runtime_bounds(lb, ub, step)
+ integer :: i, lb, ub, step
+ !$omp unroll
+ do i = lb, ub, step
+ end do
+ !$omp end unroll
+end subroutine
+
+! CHECK-LABEL: func.func @_QPruntime_bounds
+! CHECK: %[[LB:.*]] = fir.load {{.*}} : !fir.ref<i32>
+! CHECK-NEXT: %[[UB:.*]] = fir.load {{.*}} : !fir.ref<i32>
+! CHECK-NEXT: %[[STEP:.*]] = fir.load {{.*}} : !fir.ref<i32>
+! CHECK: %[[ZERO:.*]] = arith.constant 0 : i32
+! CHECK-NEXT: %[[ONE:.*]] = arith.constant 1 : i32
+! CHECK-NEXT: %[[IS_DOWNWARDS:.*]] = arith.cmpi slt, %[[STEP]], %[[ZERO]] : i32
+! CHECK-NEXT: %[[NEG_STEP:.*]] = arith.subi %[[ZERO]], %[[STEP]] : i32
+! CHECK-NEXT: %[[INCR:.*]] = arith.select %[[IS_DOWNWARDS]], %[[NEG_STEP]], %[[STEP]] : i32
+! CHECK-NEXT: %[[LOWER:.*]] = arith.select %[[IS_DOWNWARDS]], %[[UB]], %[[LB]] : i32
+! CHECK-NEXT: %[[UPPER:.*]] = arith.select %[[IS_DOWNWARDS]], %[[LB]], %[[UB]] : i32
+! CHECK-NEXT: %[[SPAN:.*]] = arith.subi %[[UPPER]], %[[LOWER]] overflow<nuw> : i32
+! CHECK-NEXT: %[[TC_MINUS_ONE:.*]] = arith.divui %[[SPAN]], %[[INCR]] : i32
+! CHECK-NEXT: %[[TC_IF_LOOPING:.*]] = arith.addi %[[TC_MINUS_ONE]], %[[ONE]] overflow<nuw> : i32
+! CHECK-NEXT: %[[IS_ZERO_TC:.*]] = arith.cmpi slt, %[[UPPER]], %[[LOWER]] : i32
+! CHECK-NEXT: %[[TRIP_COUNT:.*]] = arith.select %[[IS_ZERO_TC]], %[[ZERO]], %[[TC_IF_LOOPING]] : i32
+! CHECK-NEXT: %[[CLI:.*]] = omp.new_cli
+! CHECK-NEXT: omp.canonical_loop(%[[CLI]]) {{.*}} : i32 in range(%[[TRIP_COUNT]]) {
+
+subroutine constant_i64()
+ integer(kind=8) :: i
+ !$omp unroll
+ do i = 1_8, 100_8
+ end do
+ !$omp end unroll
+end subroutine
+
+! CHECK-LABEL: func.func @_QPconstant_i64()
+! CHECK-NOT: arith.cmpi
+! CHECK-NOT: arith.select
+! CHECK-NOT: arith.subi
+! CHECK-NOT: arith.divui
+! CHECK-NOT: arith.addi
+! CHECK: %[[TRIP_COUNT:.*]] = arith.constant 100 : i64
+! CHECK-NEXT: %[[CLI:.*]] = omp.new_cli
+! CHECK-NEXT: omp.canonical_loop(%[[CLI]]) {{.*}} : i64 in range(%[[TRIP_COUNT]]) {
More information about the flang-commits
mailing list