[flang-commits] [flang] [flang] Add pass to version innermost loops on a likely constant trip count (1/2) (PR #231205)
Caroline Newcombe via flang-commits
flang-commits at lists.llvm.org
Wed Oct 14 08:55:27 PDT 2026
https://github.com/cenewcombe updated https://github.com/llvm/llvm-project/pull/231205
>From 6b0a7a9bf62ab5c3646016932f50c075eae85a38 Mon Sep 17 00:00:00 2001
From: Caroline Newcombe <caroline.newcombe at hpe.com>
Date: Thu, 8 Oct 2026 13:44:41 -0500
Subject: [PATCH 1/4] [flang] Add pass to version innermost loops on a likely
constant trip count
---
.../flang/Optimizer/Transforms/Passes.td | 25 +
flang/lib/Optimizer/Transforms/CMakeLists.txt | 1 +
.../Transforms/LoopTripCountVersioning.cpp | 297 +++++++++
.../Transforms/loop-tripcount-versioning.fir | 603 ++++++++++++++++++
4 files changed, 926 insertions(+)
create mode 100644 flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
create mode 100644 flang/test/Transforms/loop-tripcount-versioning.fir
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index 07d3296a17be96..de04a6c41581eb 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -481,6 +481,31 @@ def LoopVersioning : Pass<"loop-versioning", "mlir::func::FuncOp"> {
];
}
+def LoopTripCountVersioning
+ : Pass<"loop-tripcount-versioning", "mlir::func::FuncOp"> {
+ let summary = "Version innermost loops on a likely constant trip count";
+ let description = [{
+ Adds a runtime check and two variants of an innermost loop whose trip count
+ is likely to equal a compile-time constant, so that one variant has a
+ constant trip count that LLVM can fully unroll, or at least vectorize
+ without a minimum-iteration check and remainder loop. Both variants are the
+ original loop with only the bound changed, so the transformation is always
+ legal; the heuristics decide only where it is worthwhile.
+
+ The first heuristic takes the constant from the declared extent of an
+ explicit-shape array whose leading dimension the loop indexes: for
+ `a(1:n, :)` on `a(4, :)` the innermost loop runs `1..n`, and versioning on
+ `n == 4` gives the fast variant a trip count of 4. Further heuristics may
+ be added in the future.
+ }];
+ let dependentDialects = [ "fir::FIROpsDialect", "mlir::arith::ArithDialect" ];
+ let options = [
+ Option<"maxSectionExtent", "max-section-extent", "unsigned",
+ /*default=*/"32",
+ "Do not version on declared extents larger than this; 0 disables">
+ ];
+}
+
def VScaleAttr : Pass<"vscale-attr", "mlir::func::FuncOp"> {
let summary = "Add vscale_range attribute to functions";
let description = [{
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index b7526e2cec289f..4d17dcdc97bb4b 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -45,6 +45,7 @@ add_flang_library(FIRTransforms
FunctionAttr.cpp
GenRuntimeCallsForTest.cpp
LoopInvariantCodeMotion.cpp
+ LoopTripCountVersioning.cpp
LoopVersioning.cpp
SelectOpsConversion.cpp
MIFOpConversion.cpp
diff --git a/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp b/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
new file mode 100644
index 00000000000000..2e6c6d3cdaac56
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
@@ -0,0 +1,297 @@
+//===- LoopTripCountVersioning.cpp ----------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+/// \file
+/// Versions an innermost loop whose runtime trip count is likely to equal a
+/// compile-time constant N:
+///
+/// if (ub == N) { copy of the loop with ub := N } else { original loop }
+///
+/// Both arms are the original loop with only the bound changed, so the
+/// transformation is always legal. The constant trip count lets LLVM fully
+/// unroll the fast arm, or vectorize it without a remainder loop.
+///
+/// Heuristic: the loop variable is the leading subscript of an explicit-shape
+/// array whose declared leading extent N is at most max-section-extent. For
+/// `a(1:m, :) = 0` with `real :: a(4, n)`, the inner loop is versioned on
+/// `m == 4`. This runs on FIR, the only level where both the loop and the
+/// array's declared type exist.
+//
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Builder/FIRBuilder.h"
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Dialect/FIRType.h"
+#include "flang/Optimizer/Dialect/Support/KindMapping.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/OpenACC/OpenACC.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/Interfaces/LoopLikeInterface.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/Support/Debug.h"
+
+namespace fir {
+#define GEN_PASS_DEF_LOOPTRIPCOUNTVERSIONING
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "flang-loop-tripcount-versioning"
+
+STATISTIC(NumLoopsExamined, "Number of innermost loops examined");
+STATISTIC(NumLoopsVersioned, "Number of loops versioned on trip count");
+
+namespace {
+
+/// One specialized copy of a loop: when \p condition holds at runtime, the
+/// loop runs exactly \p tripCount iterations. Both values must dominate the
+/// loop.
+struct LoopVersion {
+ mlir::Value condition;
+ mlir::Value tripCount;
+};
+
+/// An innermost `do i = 1, ub` whose induction variable is the leading
+/// subscript of an array with constant declared leading extent \p extent.
+struct SectionExtentCandidate {
+ fir::DoLoopOp loop;
+ int64_t extent;
+};
+
+class LoopTripCountVersioningPass
+ : public fir::impl::LoopTripCountVersioningBase<
+ LoopTripCountVersioningPass> {
+public:
+ LoopTripCountVersioningPass() = default;
+ LoopTripCountVersioningPass(
+ const fir::LoopTripCountVersioningOptions &options)
+ : Base(options) {}
+ void runOnOperation() override;
+};
+
+} // namespace
+
+/// Look through fir.convert operations.
+static mlir::Value stripConverts(mlir::Value val) {
+ while (auto cvt = val.getDefiningOp<fir::ConvertOp>())
+ val = cvt.getValue();
+ return val;
+}
+
+/// Is \p val the integer constant \p expected, looking through converts?
+static bool isConstantInt(mlir::Value val, int64_t expected) {
+ std::optional<llvm::APInt> cst = fir::getIntIfConstant(stripConverts(val));
+ return cst && cst->trySExtValue() == expected;
+}
+
+/// An absent shape or a plain fir.shape has lower bound one in every
+/// dimension; a fir.shape_shift must say so explicitly.
+static bool hasUnitLeadingLowerBound(mlir::Value shapeLike) {
+ if (!shapeLike || shapeLike.getDefiningOp<fir::ShapeOp>())
+ return true;
+ auto shapeShift = shapeLike.getDefiningOp<fir::ShapeShiftOp>();
+ return shapeShift && isConstantInt(shapeShift.getOrigins().front(), 1);
+}
+
+/// If the leading subscript of \p coor is \p iv, and the array's leading
+/// dimension has a constant declared extent onto which \p iv maps one-to-one
+/// (unit lower bound and, if sliced, a triplet `(1, sectionUb, 1)`), return
+/// the extent.
+///
+/// Only the leading dimension is considered: it becomes the innermost loop,
+/// and unrolling produces contiguous memory accesses. An array section that
+/// reaches the loop through a temporary box (fir.box_addr of a fir.embox) is
+/// not recognized. Both are possible future extensions.
+static std::optional<int64_t> getLeadingExtentIndexedBy(fir::ArrayCoorOp coor,
+ mlir::Value iv,
+ mlir::Value sectionUb) {
+ // unwrapRefType leaves a box type as is, so this also rejects boxed arrays.
+ auto seqTy = mlir::dyn_cast<fir::SequenceType>(
+ fir::unwrapRefType(coor.getMemref().getType()));
+ if (!seqTy || seqTy.getShape().empty())
+ return std::nullopt;
+ int64_t extent = seqTy.getShape().front();
+ if (extent == fir::SequenceType::getUnknownExtent())
+ return std::nullopt;
+ if (!hasUnitLeadingLowerBound(coor.getShape()))
+ return std::nullopt;
+
+ mlir::ValueRange indices = coor.getIndices();
+ if (indices.empty() || stripConverts(indices.front()) != iv)
+ return std::nullopt;
+ if (mlir::Value slice = coor.getSlice()) {
+ auto sliceOp = slice.getDefiningOp<fir::SliceOp>();
+ if (!sliceOp)
+ return std::nullopt;
+ mlir::ValueRange triples = sliceOp.getTriples();
+ if (!isConstantInt(triples[0], 1) || !isConstantInt(triples[2], 1) ||
+ stripConverts(triples[1]) != sectionUb)
+ return std::nullopt;
+ }
+ return extent;
+}
+
+/// If \p loop is an innermost `do i = 1, ub step 1` with a non-constant upper
+/// bound whose induction variable is the leading subscript of at least one
+/// array with a constant declared leading extent, return the smallest such
+/// extent N, provided 1 < N <= \p maxExtent. In a conforming program an
+/// unconditional access bounds ub by its extent, so the smallest extent is the
+/// most likely value; if the smaller array is accessed only conditionally, ub
+/// may be larger and the original loop runs.
+///
+/// Extent 1 is skipped: legacy code declares dummies as `a(1)` and indexes
+/// past them, so `ub == 1` would almost never hold.
+static std::optional<int64_t> getSectionExtent(fir::DoLoopOp loop,
+ unsigned maxExtent) {
+ // A constant upper bound needs no versioning.
+ if (!isConstantInt(loop.getLowerBound(), 1) ||
+ !isConstantInt(loop.getStep(), 1) ||
+ fir::getIntIfConstant(stripConverts(loop.getUpperBound())))
+ return std::nullopt;
+
+ // Array-section lowering clamps the loop bound to max(ub, 0) while the slice
+ // triplet carries the unclamped ub; compare against the unclamped value.
+ mlir::Value ub = loop.getUpperBound();
+ mlir::Value sectionUb = ub;
+ if (auto maxOp = ub.getDefiningOp<mlir::arith::MaxSIOp>()) {
+ if (isConstantInt(maxOp.getRhs(), 0))
+ sectionUb = maxOp.getLhs();
+ else if (isConstantInt(maxOp.getLhs(), 0))
+ sectionUb = maxOp.getRhs();
+ }
+ sectionUb = stripConverts(sectionUb);
+
+ mlir::Value iv = loop.getInductionVar();
+ std::optional<int64_t> extent;
+ loop.getBody()->walk([&](fir::ArrayCoorOp coor) {
+ if (std::optional<int64_t> coorExtent =
+ getLeadingExtentIndexedBy(coor, iv, sectionUb))
+ extent = extent ? std::min(*extent, *coorExtent) : *coorExtent;
+ });
+ if (!extent || *extent <= 1 || *extent > maxExtent)
+ return std::nullopt;
+ return extent;
+}
+
+/// Replace \p loop with a chain of guarded copies, one per entry of
+/// \p versions, each with its upper bound replaced by the constant trip
+/// count, and the original loop as the final else arm:
+///
+/// if (c0) { loop with ub := tc0 }
+/// else if (c1) { loop with ub := tc1 }
+/// else { original loop }
+///
+/// The conditions and trip counts must dominate \p loop. Loop results are
+/// forwarded through fir.result.
+static void versionLoop(fir::FirOpBuilder &builder, fir::DoLoopOp loop,
+ llvm::ArrayRef<LoopVersion> versions) {
+ assert(!versions.empty() && "nothing to version on");
+ mlir::Location loc = loop.getLoc();
+ mlir::TypeRange resultTypes = loop->getResultTypes();
+ bool hasResults = !resultTypes.empty();
+
+ builder.setInsertionPoint(loop);
+ fir::IfOp outerIf;
+ for (const LoopVersion &version : versions) {
+ auto ifOp = fir::IfOp::create(builder, loc, resultTypes, version.condition,
+ /*withElse=*/true);
+ if (!outerIf) {
+ outerIf = ifOp;
+ } else if (hasResults) {
+ // Terminate the enclosing else arm with this nested if's results.
+ fir::ResultOp::create(builder, loc, ifOp.getResults());
+ }
+
+ builder.setInsertionPointToStart(&ifOp.getThenRegion().front());
+ auto fastLoop = mlir::cast<fir::DoLoopOp>(builder.clone(*loop));
+ fastLoop.getUpperBoundMutable().assign(version.tripCount);
+ if (hasResults)
+ fir::ResultOp::create(builder, loc, fastLoop.getResults());
+
+ builder.setInsertionPointToStart(&ifOp.getElseRegion().front());
+ }
+
+ // The insertion point is now at the start of the innermost else arm. Move
+ // the original loop there.
+ loop->replaceAllUsesWith(outerIf);
+ loop->remove();
+ builder.insert(loop);
+ if (hasResults)
+ fir::ResultOp::create(builder, loc, loop.getResults());
+}
+
+/// Innermost means no nested loop of any kind. Loops containing OpenMP or
+/// OpenACC operations are also excluded.
+static bool isVersionableInnermostLoop(fir::DoLoopOp loop) {
+ return !loop.getBody()
+ ->walk([](mlir::Operation *op) {
+ if (mlir::isa<mlir::LoopLikeOpInterface>(op) ||
+ mlir::isa_and_present<mlir::omp::OpenMPDialect,
+ mlir::acc::OpenACCDialect>(
+ op->getDialect()))
+ return mlir::WalkResult::interrupt();
+ return mlir::WalkResult::advance();
+ })
+ .wasInterrupted();
+}
+
+void LoopTripCountVersioningPass::runOnOperation() {
+ if (maxSectionExtent == 0)
+ return;
+ LLVM_DEBUG(llvm::dbgs() << "=== Begin " DEBUG_TYPE " ===\n");
+ mlir::func::FuncOp func = getOperation();
+ LLVM_DEBUG(llvm::dbgs() << "Func-name:" << func.getSymName() << "\n");
+
+ // Collect candidates before transforming anything. Candidates are innermost
+ // loops, so no candidate contains another and the transformations are
+ // independent.
+ llvm::SmallVector<SectionExtentCandidate> candidates;
+ func.walk([&](fir::DoLoopOp loop) {
+ if (!isVersionableInnermostLoop(loop))
+ return;
+ ++NumLoopsExamined;
+ if (std::optional<int64_t> extent =
+ getSectionExtent(loop, maxSectionExtent)) {
+ LLVM_DEBUG(llvm::dbgs() << "Section extent candidate (N=" << *extent
+ << ") at " << loop.getLoc() << '\n');
+ candidates.push_back({loop, *extent});
+ }
+ });
+
+ if (candidates.empty()) {
+ LLVM_DEBUG(llvm::dbgs()
+ << "No loops to transform.\n=== End " DEBUG_TYPE " ===\n");
+ return;
+ }
+
+ mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
+ fir::FirOpBuilder builder{module, fir::getKindMapping(module)};
+
+ for (const SectionExtentCandidate &candidate : candidates) {
+ fir::DoLoopOp loop = candidate.loop;
+ mlir::Location loc = loop.getLoc();
+ mlir::Value ub = loop.getUpperBound();
+
+ // Define the guard values before the loop so they dominate the fir.if.
+ builder.setInsertionPoint(loop);
+ mlir::Value extentCst =
+ builder.createIntegerConstant(loc, ub.getType(), candidate.extent);
+ mlir::Value isDeclaredExtent = mlir::arith::CmpIOp::create(
+ builder, loc, mlir::arith::CmpIPredicate::eq, ub, extentCst);
+
+ versionLoop(builder, loop, {{isDeclaredExtent, extentCst}});
+ ++NumLoopsVersioned;
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << "Func after transformation:\n");
+ LLVM_DEBUG(func->dump());
+ LLVM_DEBUG(llvm::dbgs() << "=== End " DEBUG_TYPE " ===\n");
+}
diff --git a/flang/test/Transforms/loop-tripcount-versioning.fir b/flang/test/Transforms/loop-tripcount-versioning.fir
new file mode 100644
index 00000000000000..37c816fd746691
--- /dev/null
+++ b/flang/test/Transforms/loop-tripcount-versioning.fir
@@ -0,0 +1,603 @@
+// RUN: fir-opt --loop-tripcount-versioning %s | FileCheck %s
+// Every versioned loop below has a declared extent of 4: a bound of 4 must
+// accept it and a bound of 3 must not. A bound of 0 disables the pass.
+// RUN: fir-opt --loop-tripcount-versioning="max-section-extent=4" %s | FileCheck %s
+// RUN: fir-opt --loop-tripcount-versioning="max-section-extent=3" %s | FileCheck %s --check-prefix=OFF
+// RUN: fir-opt --loop-tripcount-versioning="max-section-extent=0" %s | FileCheck %s --check-prefix=OFF
+
+// Array sections with a runtime extent on an explicit-shape array. The pass
+// versions the innermost loop on `extent == declared extent` so that one arm
+// has a constant trip count.
+
+// subroutine zero(tau, ncol)
+// real(8), intent(out) :: tau(4, 26)
+// integer, intent(in) :: ncol
+// tau(1:ncol, 1:26) = 0.0d0
+func.func @zero(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("tau") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c1, %3, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg2 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg3 = %c1 to %4 step %c1 unordered {
+ %6 = fir.array_coor %1(%0) [%5] %arg3, %arg2 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %cst to %6 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @zero(
+// CHECK: %[[UB:.*]] = arith.maxsi %{{.*}}, %{{.*}} : index
+// CHECK-NOT: fir.if
+// CHECK: fir.do_loop %[[J:.*]] = %{{.*}} to %{{.*}} step %{{.*}} unordered {
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK: fir.if %[[CMP]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4]] step %{{.*}} unordered {
+// CHECK: fir.array_coor
+// CHECK: fir.store
+// CHECK: }
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step %{{.*}} unordered {
+// CHECK: fir.array_coor
+// CHECK: fir.store
+// CHECK: }
+// CHECK: }
+// CHECK: }
+
+// OFF-LABEL: func.func @zero(
+// OFF-NOT: fir.if
+// OFF: return
+
+// An unsliced access: the induction variable indexes the array directly, so
+// there is no triplet to check. (The assignment conforms only when n == 4,
+// which is exactly the case the fast arm is for.)
+// subroutine unsliced(a, b, n)
+// real(8) :: a(4, 26), b(8, 26)
+// a = b(1:n, :)
+func.func @unsliced(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<!fir.array<8x26xf64>>, %arg2: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.shape %c8, %c26 : (index, index) -> !fir.shape<2>
+ %3 = fir.declare %arg1(%2) uniq_name("b") : (!fir.ref<!fir.array<8x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<8x26xf64>>
+ %4 = fir.load %arg2 : !fir.ref<i32>
+ %5 = fir.convert %4 : (i32) -> index
+ %6 = arith.maxsi %5, %c0 : index
+ %7 = fir.slice %c1, %5, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg3 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg4 = %c1 to %6 step %c1 unordered {
+ %8 = fir.array_coor %3(%2) [%7] %arg4, %arg3 : (!fir.ref<!fir.array<8x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ %9 = fir.load %8 : !fir.ref<f64>
+ %10 = fir.array_coor %1(%0) %arg4, %arg3 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, index, index) -> !fir.ref<f64>
+ fir.store %9 to %10 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @unsliced(
+// CHECK: %[[UB:.*]] = arith.maxsi %{{.*}}, %{{.*}} : index
+// CHECK: fir.do_loop
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK: fir.if %[[CMP]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4]] step
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
+// CHECK: }
+
+// A constant section bound already gives LLVM a constant trip count; there is
+// nothing to version on.
+// subroutine constant_ub(tau)
+// real(8) :: tau(4, 26)
+// tau(1:4, :) = 0.0d0
+func.func @constant_ub(%arg0: !fir.ref<!fir.array<4x26xf64>>) {
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("tau") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.slice %c1, %c4, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg1 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg2 = %c1 to %c4 step %c1 unordered {
+ %3 = fir.array_coor %1(%0) [%2] %arg2, %arg1 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %cst to %3 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @constant_ub(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// Declared extent 1 is not versioned on: legacy code declares dummies as
+// `a(1)` and indexes past them, so `n == 1` would almost never hold.
+// subroutine extent_one(a, n)
+// real(8) :: a(1)
+// a(1:n) = 0.0d0
+func.func @extent_one(%arg0: !fir.ref<!fir.array<1xf64>>, %arg1: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %0 = fir.shape %c1 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<1xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<1xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c1, %3, %c1 : (index, index, index) -> !fir.slice<1>
+ fir.do_loop %arg2 = %c1 to %4 step %c1 unordered {
+ %6 = fir.array_coor %1(%0) [%5] %arg2 : (!fir.ref<!fir.array<1xf64>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<f64>
+ fir.store %cst to %6 : !fir.ref<f64>
+ }
+ return
+}
+// CHECK-LABEL: func.func @extent_one(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// A section that does not start at 1. Hand-written so that the loop and the
+// slice share the same bound; only the slice lower bound differs from the
+// accepted form, so this is the check being exercised.
+func.func @slice_lb(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("tau") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c2, %3, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg2 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg3 = %c1 to %4 step %c1 unordered {
+ %6 = fir.array_coor %1(%0) [%5] %arg3, %arg2 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %cst to %6 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @slice_lb(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// A strided section. Hand-written as above; only the slice stride differs
+// from the accepted form.
+func.func @slice_stride(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("tau") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c1, %3, %c2, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg2 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg3 = %c1 to %4 step %c1 unordered {
+ %6 = fir.array_coor %1(%0) [%5] %arg3, %arg2 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %cst to %6 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @slice_stride(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// The two sides of an assignment have different section bounds. The loop runs
+// over the left-hand side's bound, so only arrays whose slice shares that bound
+// take part: `b` is left out and the loop is versioned on 4, not 2. (The
+// program conforms only when n == m.)
+// subroutine slice_ub_mismatch(a, b, n, m)
+// real(8) :: a(4, 26), b(2, 26)
+// a(1:n, :) = b(1:m, :)
+func.func @slice_ub_mismatch(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<!fir.array<2x26xf64>>, %arg2: !fir.ref<i32>, %arg3: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.shape %c2, %c26 : (index, index) -> !fir.shape<2>
+ %3 = fir.declare %arg1(%2) uniq_name("b") : (!fir.ref<!fir.array<2x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<2x26xf64>>
+ %4 = fir.load %arg2 : !fir.ref<i32>
+ %5 = fir.convert %4 : (i32) -> index
+ %6 = arith.maxsi %5, %c0 : index
+ %7 = fir.slice %c1, %5, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ %8 = fir.load %arg3 : !fir.ref<i32>
+ %9 = fir.convert %8 : (i32) -> index
+ %10 = fir.slice %c1, %9, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg4 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg5 = %c1 to %6 step %c1 unordered {
+ %11 = fir.array_coor %3(%2) [%10] %arg5, %arg4 : (!fir.ref<!fir.array<2x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ %12 = fir.load %11 : !fir.ref<f64>
+ %13 = fir.array_coor %1(%0) [%7] %arg5, %arg4 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %12 to %13 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @slice_ub_mismatch(
+// CHECK: %[[UB:.*]] = arith.maxsi %{{.*}}, %{{.*}} : index
+// CHECK: fir.do_loop
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK: fir.if %[[CMP]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4]] step
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
+// CHECK: }
+
+// The loop bound is not clamped with arith.maxsi; the slice bound is compared
+// against the loop bound directly.
+func.func @unclamped(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<i32>) {
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("tau") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = fir.slice %c1, %3, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg2 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg3 = %c1 to %3 step %c1 unordered {
+ %5 = fir.array_coor %1(%0) [%4] %arg3, %arg2 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %cst to %5 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @unclamped(
+// CHECK: %[[UB:.*]] = fir.convert %{{.*}} : (i32) -> index
+// CHECK: fir.do_loop
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK: fir.if %[[CMP]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4]] step
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
+// CHECK: }
+
+// The loop does not start at 1, so `ub == 4` does not mean four iterations
+// over the whole dimension.
+// subroutine loop_lb(a, n)
+// real(8) :: a(4)
+// do i = 2, n
+// a(i) = 0.0d0
+// end do
+func.func @loop_lb(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<i32>) {
+ %c1 = arith.constant 1 : index
+ %c2 = arith.constant 2 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ fir.do_loop %arg2 = %c2 to %3 step %c1 {
+ %4 = fir.array_coor %1(%0) %arg2 : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, index) -> !fir.ref<f64>
+ fir.store %cst to %4 : !fir.ref<f64>
+ }
+ return
+}
+// CHECK-LABEL: func.func @loop_lb(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// Non-unit declared lower bound: the slice triplet is the accepted
+// (1, ncol, 1), but position i of the section is element i of a dimension
+// that starts at 0, so `ncol == 4` would not mean the whole dimension.
+// subroutine nonunit_lb(tau, ncol)
+// real(8), intent(out) :: tau(0:3, 26)
+// tau(1:ncol, :) = 0.0d0
+func.func @nonunit_lb(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape_shift %c0, %c4, %c1, %c26 : (index, index, index, index) -> !fir.shapeshift<2>
+ %1 = fir.declare %arg0(%0) uniq_name("tau") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shapeshift<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c1, %3, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg2 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg3 = %c1 to %4 step %c1 unordered {
+ %6 = fir.array_coor %1(%0) [%5] %arg3, %arg2 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shapeshift<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %cst to %6 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @nonunit_lb(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// Two arrays with different declared extents in the same dimension: version
+// on the smaller one, the only extent the section bound can equal in a
+// conforming program.
+// subroutine mixed_extents(a, b, n)
+// real(8) :: a(4, 26), b(8, 26)
+// a(1:n, :) = b(1:n, :)
+func.func @mixed_extents(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<!fir.array<8x26xf64>>, %arg2: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %c8 = arith.constant 8 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.shape %c8, %c26 : (index, index) -> !fir.shape<2>
+ %3 = fir.declare %arg1(%2) uniq_name("b") : (!fir.ref<!fir.array<8x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<8x26xf64>>
+ %4 = fir.load %arg2 : !fir.ref<i32>
+ %5 = fir.convert %4 : (i32) -> index
+ %6 = arith.maxsi %5, %c0 : index
+ %7 = fir.slice %c1, %5, %c1, %c1, %c26, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg3 = %c1 to %c26 step %c1 unordered {
+ fir.do_loop %arg4 = %c1 to %6 step %c1 unordered {
+ %8 = fir.array_coor %3(%2) [%7] %arg4, %arg3 : (!fir.ref<!fir.array<8x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ %9 = fir.load %8 : !fir.ref<f64>
+ %10 = fir.array_coor %1(%0) [%7] %arg4, %arg3 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %9 to %10 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @mixed_extents(
+// CHECK: %[[UB:.*]] = arith.maxsi %{{.*}}, %{{.*}} : index
+// CHECK: fir.do_loop
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK: fir.if %[[CMP]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4]] step
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
+// CHECK: }
+
+// The loop variable is not the leading subscript: only the leading dimension
+// is versioned.
+// subroutine nonleading(a, k, n)
+// real(8) :: a(26, 4)
+// a(k, 1:n) = 0.0d0
+func.func @nonleading(%arg0: !fir.ref<!fir.array<26x4xf64>>, %arg1: !fir.ref<i32>, %arg2: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c26, %c4 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<26x4xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<26x4xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = fir.load %arg2 : !fir.ref<i32>
+ %5 = fir.convert %4 : (i32) -> index
+ %6 = arith.maxsi %5, %c0 : index
+ %7 = fir.slice %3, %3, %c1, %c1, %5, %c1 : (index, index, index, index, index, index) -> !fir.slice<2>
+ fir.do_loop %arg3 = %c1 to %6 step %c1 unordered {
+ %8 = fir.array_coor %1(%0) [%7] %c1, %arg3 : (!fir.ref<!fir.array<26x4xf64>>, !fir.shape<2>, !fir.slice<2>, index, index) -> !fir.ref<f64>
+ fir.store %cst to %8 : !fir.ref<f64>
+ }
+ return
+}
+// CHECK-LABEL: func.func @nonleading(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// A section that reaches the loop through a temporary box (fir.box_addr of a
+// fir.embox) is not recognized: the declared shape is read from the memref
+// type only, and here that type has an unknown extent.
+// subroutine boxed(fsds, ncol)
+// real(8) :: fsds(4)
+// fsds(1:ncol) = 0.0d0
+func.func @boxed(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("fsds") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c1, %3, %c1 : (index, index, index) -> !fir.slice<1>
+ %6 = fir.embox %1(%0) [%5] : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, !fir.slice<1>) -> !fir.box<!fir.array<?xf64>>
+ %7 = fir.box_addr %6 : (!fir.box<!fir.array<?xf64>>) -> !fir.ref<!fir.array<?xf64>>
+ %8 = fir.shape %4 : (index) -> !fir.shape<1>
+ fir.do_loop %arg2 = %c1 to %4 step %c1 unordered {
+ %9 = fir.array_coor %7(%8) %arg2 : (!fir.ref<!fir.array<?xf64>>, !fir.shape<1>, index) -> !fir.ref<f64>
+ fir.store %cst to %9 : !fir.ref<f64>
+ }
+ return
+}
+// CHECK-LABEL: func.func @boxed(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// A loop whose body contains an OpenMP region is not versioned: the region is
+// outlined, so a constant trip count cannot help and cloning would only
+// duplicate the outlined code.
+func.func @contains_omp(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c1, %3, %c1 : (index, index, index) -> !fir.slice<1>
+ fir.do_loop %arg2 = %c1 to %4 step %c1 unordered {
+ %6 = fir.array_coor %1(%0) [%5] %arg2 : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<f64>
+ omp.parallel {
+ fir.store %cst to %6 : !fir.ref<f64>
+ omp.terminator
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @contains_omp(
+// CHECK-NOT: fir.if
+// CHECK: omp.parallel
+// CHECK: return
+
+// A loop containing a loop that is not a fir.do_loop is not innermost and is
+// not versioned.
+func.func @contains_iterate_while(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<i32>, %arg2: i1) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c1, %3, %c1 : (index, index, index) -> !fir.slice<1>
+ fir.do_loop %arg3 = %c1 to %4 step %c1 unordered {
+ %6 = fir.array_coor %1(%0) [%5] %arg3 : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<f64>
+ %7 = fir.iterate_while (%arg4 = %c1 to %c4 step %c1) and (%arg5 = %arg2) {
+ fir.store %cst to %6 : !fir.ref<f64>
+ fir.result %arg5 : i1
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @contains_iterate_while(
+// CHECK-NOT: fir.if
+// CHECK: fir.iterate_while
+// CHECK: return
+
+// An explicit DO nest in the opposite order to array-section lowering: the
+// outer loop variable is the leading subscript of a 4-extent array, but the
+// loop contains another loop, so it is not versioned; the inner loop indexes
+// the second dimension, so neither is it.
+// subroutine outer_qualifies(tau, n)
+// real(8) :: tau(4, 26)
+// do i = 1, n
+// do j = 1, 26
+// tau(i, j) = 0.0d0
+// end do
+// end do
+func.func @outer_qualifies(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<i32>) {
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c26 = arith.constant 26 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4, %c26 : (index, index) -> !fir.shape<2>
+ %1 = fir.declare %arg0(%0) uniq_name("tau") : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>) -> !fir.ref<!fir.array<4x26xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ fir.do_loop %arg2 = %c1 to %3 step %c1 {
+ fir.do_loop %arg3 = %c1 to %c26 step %c1 {
+ %4 = fir.array_coor %1(%0) %arg2, %arg3 : (!fir.ref<!fir.array<4x26xf64>>, !fir.shape<2>, index, index) -> !fir.ref<f64>
+ fir.store %cst to %4 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @outer_qualifies(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+// Two candidates in one function: both are versioned.
+// subroutine two_loops(a, b, n)
+// real(8) :: a(4), b(4)
+// a(1:n) = 0.0d0
+// b(1:n) = 1.0d0
+func.func @two_loops(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<!fir.array<4xf64>>, %arg2: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst0 = arith.constant 0.000000e+00 : f64
+ %cst1 = arith.constant 1.000000e+00 : f64
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %2 = fir.declare %arg1(%0) uniq_name("b") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %3 = fir.load %arg2 : !fir.ref<i32>
+ %4 = fir.convert %3 : (i32) -> index
+ %5 = arith.maxsi %4, %c0 : index
+ %6 = fir.slice %c1, %4, %c1 : (index, index, index) -> !fir.slice<1>
+ fir.do_loop %arg3 = %c1 to %5 step %c1 unordered {
+ %7 = fir.array_coor %1(%0) [%6] %arg3 : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<f64>
+ fir.store %cst0 to %7 : !fir.ref<f64>
+ }
+ fir.do_loop %arg3 = %c1 to %5 step %c1 unordered {
+ %7 = fir.array_coor %2(%0) [%6] %arg3 : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<f64>
+ fir.store %cst1 to %7 : !fir.ref<f64>
+ }
+ return
+}
+// CHECK-LABEL: func.func @two_loops(
+// CHECK: %[[UB:.*]] = arith.maxsi %{{.*}}, %{{.*}} : index
+// CHECK: %[[C4A:.*]] = arith.constant 4 : index
+// CHECK: %[[CMPA:.*]] = arith.cmpi eq, %[[UB]], %[[C4A]] : index
+// CHECK: fir.if %[[CMPA]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4A]] step
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
+// CHECK: }
+// CHECK: %[[C4B:.*]] = arith.constant 4 : index
+// CHECK: %[[CMPB:.*]] = arith.cmpi eq, %[[UB]], %[[C4B]] : index
+// CHECK: fir.if %[[CMPB]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4B]] step
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
+// CHECK: }
+// CHECK: return
+
+// A loop carrying results: fir.result must forward them through both arms.
+// Hand-written; array-section lowering does not produce iter_args, but the
+// rewrite must handle them.
+func.func @with_results(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<i32>) -> f64 {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ %4 = arith.maxsi %3, %c0 : index
+ %5 = fir.slice %c1, %3, %c1 : (index, index, index) -> !fir.slice<1>
+ %6 = fir.do_loop %arg2 = %c1 to %4 step %c1 iter_args(%arg3 = %cst) -> (f64) {
+ %7 = fir.array_coor %1(%0) [%5] %arg2 : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<f64>
+ %8 = fir.load %7 : !fir.ref<f64>
+ %9 = arith.addf %arg3, %8 : f64
+ fir.result %9 : f64
+ }
+ return %6 : f64
+}
+// CHECK-LABEL: func.func @with_results(
+// CHECK: %[[UB:.*]] = arith.maxsi
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK: %[[RES:.*]] = fir.if %[[CMP]] -> (f64) {
+// CHECK: %[[FAST:.*]] = fir.do_loop {{.*}} to %[[C4]] step %{{.*}} iter_args({{.*}}) -> (f64) {
+// CHECK: fir.result %[[FAST]] : f64
+// CHECK: } else {
+// CHECK: %[[SLOW:.*]] = fir.do_loop {{.*}} to %[[UB]] step %{{.*}} iter_args({{.*}}) -> (f64) {
+// CHECK: fir.result %[[SLOW]] : f64
+// CHECK: }
+// CHECK: return %[[RES]] : f64
>From 33d94c2f541920d0311a9c4ff10a2ee3ed747c96 Mon Sep 17 00:00:00 2001
From: Caroline Newcombe <caroline.newcombe at hpe.com>
Date: Tue, 13 Oct 2026 15:09:34 -0500
Subject: [PATCH 2/4] [flang] Add max-unrolled-ops limit and skip conditional
unsliced accesses
---
.../flang/Optimizer/Transforms/Passes.td | 11 ++-
.../Transforms/LoopTripCountVersioning.cpp | 51 ++++++++---
.../Transforms/loop-tripcount-versioning.fir | 84 +++++++++++++++++++
3 files changed, 133 insertions(+), 13 deletions(-)
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index de04a6c41581eb..f1f018e5bad8dc 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -497,12 +497,21 @@ def LoopTripCountVersioning
`a(1:n, :)` on `a(4, :)` the innermost loop runs `1..n`, and versioning on
`n == 4` gives the fast variant a trip count of 4. Further heuristics may
be added in the future.
+
+ Candidates are limited by the declared extent (`max-section-extent`) and
+ by the size of the fully unrolled fast variant, extent times loop body
+ operations (`max-unrolled-ops`, matching LLVM's -O3 full-unroll budget),
+ so that loops are only duplicated where LLVM is expected to unroll them.
}];
let dependentDialects = [ "fir::FIROpsDialect", "mlir::arith::ArithDialect" ];
let options = [
Option<"maxSectionExtent", "max-section-extent", "unsigned",
/*default=*/"32",
- "Do not version on declared extents larger than this; 0 disables">
+ "Do not version on declared extents larger than this; 0 disables">,
+ Option<"maxUnrolledOps", "max-unrolled-ops", "unsigned",
+ /*default=*/"300",
+ "Do not version if the declared extent times the loop body size "
+ "exceeds this; 0 removes the limit">
];
}
diff --git a/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp b/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
index 2e6c6d3cdaac56..7746c17357ccf4 100644
--- a/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
+++ b/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
@@ -19,7 +19,9 @@
/// Heuristic: the loop variable is the leading subscript of an explicit-shape
/// array whose declared leading extent N is at most max-section-extent. For
/// `a(1:m, :) = 0` with `real :: a(4, n)`, the inner loop is versioned on
-/// `m == 4`. This runs on FIR, the only level where both the loop and the
+/// `m == 4`. Loops are skipped when N times the body size exceeds
+/// max-unrolled-ops (LLVM's -O3 full-unroll budget), since the fast arm would
+/// not unroll. This runs on FIR, the only level where both the loop and the
/// array's declared type exist.
//
//===----------------------------------------------------------------------===//
@@ -47,6 +49,8 @@ namespace fir {
#define DEBUG_TYPE "flang-loop-tripcount-versioning"
STATISTIC(NumLoopsExamined, "Number of innermost loops examined");
+STATISTIC(NumLoopsTooLarge,
+ "Number of candidates skipped: N * body size > max-unrolled-ops");
STATISTIC(NumLoopsVersioned, "Number of loops versioned on trip count");
namespace {
@@ -106,13 +110,19 @@ static bool hasUnitLeadingLowerBound(mlir::Value shapeLike) {
/// (unit lower bound and, if sliced, a triplet `(1, sectionUb, 1)`), return
/// the extent.
///
+/// A sliced access `a(1:ub)` requires ub <= extent even if some elements are
+/// skipped, so it may sit under a fir.if (as in WHERE). An unsliced access
+/// `a(i)` only bounds ub if it runs on every iteration, so it must be
+/// directly in \p body.
+///
/// Only the leading dimension is considered: it becomes the innermost loop,
/// and unrolling produces contiguous memory accesses. An array section that
/// reaches the loop through a temporary box (fir.box_addr of a fir.embox) is
/// not recognized. Both are possible future extensions.
static std::optional<int64_t> getLeadingExtentIndexedBy(fir::ArrayCoorOp coor,
mlir::Value iv,
- mlir::Value sectionUb) {
+ mlir::Value sectionUb,
+ mlir::Block *body) {
// unwrapRefType leaves a box type as is, so this also rejects boxed arrays.
auto seqTy = mlir::dyn_cast<fir::SequenceType>(
fir::unwrapRefType(coor.getMemref().getType()));
@@ -135,6 +145,8 @@ static std::optional<int64_t> getLeadingExtentIndexedBy(fir::ArrayCoorOp coor,
if (!isConstantInt(triples[0], 1) || !isConstantInt(triples[2], 1) ||
stripConverts(triples[1]) != sectionUb)
return std::nullopt;
+ } else if (coor->getBlock() != body) {
+ return std::nullopt;
}
return extent;
}
@@ -142,10 +154,10 @@ static std::optional<int64_t> getLeadingExtentIndexedBy(fir::ArrayCoorOp coor,
/// If \p loop is an innermost `do i = 1, ub step 1` with a non-constant upper
/// bound whose induction variable is the leading subscript of at least one
/// array with a constant declared leading extent, return the smallest such
-/// extent N, provided 1 < N <= \p maxExtent. In a conforming program an
-/// unconditional access bounds ub by its extent, so the smallest extent is the
-/// most likely value; if the smaller array is accessed only conditionally, ub
-/// may be larger and the original loop runs.
+/// extent N, provided 1 < N <= \p maxExtent. Only accesses that bound ub in a
+/// conforming program take part (see getLeadingExtentIndexedBy), so the
+/// smallest extent is the value ub must equal when all of them are at full
+/// extent.
///
/// Extent 1 is skipped: legacy code declares dummies as `a(1)` and indexes
/// past them, so `ub == 1` would almost never hold.
@@ -173,7 +185,7 @@ static std::optional<int64_t> getSectionExtent(fir::DoLoopOp loop,
std::optional<int64_t> extent;
loop.getBody()->walk([&](fir::ArrayCoorOp coor) {
if (std::optional<int64_t> coorExtent =
- getLeadingExtentIndexedBy(coor, iv, sectionUb))
+ getLeadingExtentIndexedBy(coor, iv, sectionUb, loop.getBody()))
extent = extent ? std::min(*extent, *coorExtent) : *coorExtent;
});
if (!extent || *extent <= 1 || *extent > maxExtent)
@@ -258,12 +270,27 @@ void LoopTripCountVersioningPass::runOnOperation() {
if (!isVersionableInnermostLoop(loop))
return;
++NumLoopsExamined;
- if (std::optional<int64_t> extent =
- getSectionExtent(loop, maxSectionExtent)) {
- LLVM_DEBUG(llvm::dbgs() << "Section extent candidate (N=" << *extent
- << ") at " << loop.getLoc() << '\n');
- candidates.push_back({loop, *extent});
+ std::optional<int64_t> extent = getSectionExtent(loop, maxSectionExtent);
+ if (!extent)
+ return;
+ unsigned bodyOps = 0;
+ loop.getBody()->walk([&](mlir::Operation *) { ++bodyOps; });
+ uint64_t unrolledOps = static_cast<uint64_t>(*extent) * bodyOps;
+ bool tooLarge = maxUnrolledOps != 0 && unrolledOps > maxUnrolledOps;
+ LLVM_DEBUG({
+ // One write per line so parallel builds do not interleave mid-line.
+ std::string msg;
+ llvm::raw_string_ostream os(msg);
+ os << "Section extent candidate (N=" << *extent
+ << ", body ops=" << bodyOps << ") at " << loop.getLoc()
+ << (tooLarge ? " skipped: too large to unroll\n" : "\n");
+ llvm::dbgs() << msg;
+ });
+ if (tooLarge) {
+ ++NumLoopsTooLarge;
+ return;
}
+ candidates.push_back({loop, *extent});
});
if (candidates.empty()) {
diff --git a/flang/test/Transforms/loop-tripcount-versioning.fir b/flang/test/Transforms/loop-tripcount-versioning.fir
index 37c816fd746691..6307098fbabc28 100644
--- a/flang/test/Transforms/loop-tripcount-versioning.fir
+++ b/flang/test/Transforms/loop-tripcount-versioning.fir
@@ -4,6 +4,11 @@
// RUN: fir-opt --loop-tripcount-versioning="max-section-extent=4" %s | FileCheck %s
// RUN: fir-opt --loop-tripcount-versioning="max-section-extent=3" %s | FileCheck %s --check-prefix=OFF
// RUN: fir-opt --loop-tripcount-versioning="max-section-extent=0" %s | FileCheck %s --check-prefix=OFF
+// The unrolled-size limit is extent times body operations. @zero's body is
+// three operations, so 4 * 3 = 12 is accepted at the default and rejected
+// at 11; 0 removes the limit.
+// RUN: fir-opt --loop-tripcount-versioning="max-unrolled-ops=11" %s | FileCheck %s --check-prefix=OFF
+// RUN: fir-opt --loop-tripcount-versioning="max-unrolled-ops=0" %s | FileCheck %s
// Array sections with a runtime extent on an explicit-shape array. The pass
// versions the innermost loop on `extent == declared extent` so that one arm
@@ -375,6 +380,85 @@ func.func @mixed_extents(%arg0: !fir.ref<!fir.array<4x26xf64>>, %arg1: !fir.ref<
// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
// CHECK: }
+// An unsliced access under a fir.if does not execute on every iteration, so
+// it does not bound the loop and takes no part: here `b(2)` is conditional,
+// `a(4)` is not, and the loop is versioned on 4 rather than the minimum 2.
+// Hand-written: lowering reloads the loop variable in explicit DO loops.
+// do i = 1, n
+// a(i) = 0
+// if (c) b(i) = 0
+// end do
+func.func @conditional_unsliced(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<!fir.array<2xf64>>, %arg2: !fir.ref<i32>, %arg3: i1) {
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c2 = arith.constant 2 : index
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %2 = fir.shape %c2 : (index) -> !fir.shape<1>
+ %3 = fir.declare %arg1(%2) uniq_name("b") : (!fir.ref<!fir.array<2xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<2xf64>>
+ %4 = fir.load %arg2 : !fir.ref<i32>
+ %5 = fir.convert %4 : (i32) -> index
+ fir.do_loop %arg4 = %c1 to %5 step %c1 {
+ %6 = fir.array_coor %1(%0) %arg4 : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, index) -> !fir.ref<f64>
+ fir.store %cst to %6 : !fir.ref<f64>
+ fir.if %arg3 {
+ %7 = fir.array_coor %3(%2) %arg4 : (!fir.ref<!fir.array<2xf64>>, !fir.shape<1>, index) -> !fir.ref<f64>
+ fir.store %cst to %7 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @conditional_unsliced(
+// CHECK: %[[UB:.*]] = fir.convert %{{.*}} : (i32) -> index
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK: fir.if %[[CMP]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4]] step
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
+// CHECK: }
+
+// A sliced access under a fir.if still bounds the loop: the section
+// designator `a(1:n)` must be in bounds whether or not element i is stored.
+// This is how WHERE lowers.
+// subroutine where_sliced(a, m, n)
+// real(8) :: a(4)
+// logical :: m(4)
+// where (m(1:n)) a(1:n) = 0.0d0
+func.func @where_sliced(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<!fir.array<4x!fir.logical<4>>>, %arg2: !fir.ref<i32>) {
+ %c0 = arith.constant 0 : index
+ %c1 = arith.constant 1 : index
+ %cst = arith.constant 0.000000e+00 : f64
+ %c4 = arith.constant 4 : index
+ %0 = fir.shape %c4 : (index) -> !fir.shape<1>
+ %1 = fir.declare %arg0(%0) uniq_name("a") : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>) -> !fir.ref<!fir.array<4xf64>>
+ %2 = fir.declare %arg1(%0) uniq_name("m") : (!fir.ref<!fir.array<4x!fir.logical<4>>>, !fir.shape<1>) -> !fir.ref<!fir.array<4x!fir.logical<4>>>
+ %3 = fir.load %arg2 : !fir.ref<i32>
+ %4 = fir.convert %3 : (i32) -> index
+ %5 = arith.maxsi %4, %c0 : index
+ %6 = fir.slice %c1, %4, %c1 : (index, index, index) -> !fir.slice<1>
+ fir.do_loop %arg3 = %c1 to %5 step %c1 unordered {
+ %7 = fir.array_coor %2(%0) [%6] %arg3 : (!fir.ref<!fir.array<4x!fir.logical<4>>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<!fir.logical<4>>
+ %8 = fir.load %7 : !fir.ref<!fir.logical<4>>
+ %9 = fir.convert %8 : (!fir.logical<4>) -> i1
+ fir.if %9 {
+ %10 = fir.array_coor %1(%0) [%6] %arg3 : (!fir.ref<!fir.array<4xf64>>, !fir.shape<1>, !fir.slice<1>, index) -> !fir.ref<f64>
+ fir.store %cst to %10 : !fir.ref<f64>
+ }
+ }
+ return
+}
+// CHECK-LABEL: func.func @where_sliced(
+// CHECK: %[[UB:.*]] = arith.maxsi %{{.*}}, %{{.*}} : index
+// CHECK: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[CMP:.*]] = arith.cmpi eq, %[[UB]], %[[C4]] : index
+// CHECK: fir.if %[[CMP]] {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[C4]] step
+// CHECK: } else {
+// CHECK: fir.do_loop %{{.*}} = %{{.*}} to %[[UB]] step
+// CHECK: }
+
// The loop variable is not the leading subscript: only the leading dimension
// is versioned.
// subroutine nonleading(a, k, n)
>From a623ae3436c1f23abbe2068a18b06a9f9c70a9b0 Mon Sep 17 00:00:00 2001
From: Caroline Newcombe <caroline.newcombe at hpe.com>
Date: Wed, 14 Oct 2026 08:44:01 -0500
Subject: [PATCH 3/4] [flang] Skip loops in device code in
LoopTripCountVersioning
---
.../Transforms/LoopTripCountVersioning.cpp | 20 ++++-
.../Transforms/loop-tripcount-versioning.fir | 87 +++++++++++++++++++
2 files changed, 104 insertions(+), 3 deletions(-)
diff --git a/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp b/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
index 7746c17357ccf4..04380d8f288128 100644
--- a/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
+++ b/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
@@ -21,11 +21,13 @@
/// `a(1:m, :) = 0` with `real :: a(4, n)`, the inner loop is versioned on
/// `m == 4`. Loops are skipped when N times the body size exceeds
/// max-unrolled-ops (LLVM's -O3 full-unroll budget), since the fast arm would
-/// not unroll. This runs on FIR, the only level where both the loop and the
-/// array's declared type exist.
+/// not unroll. Loops that run on a GPU are skipped: the guard could split
+/// threads across two copies of a barrier. This runs on FIR, the only level
+/// where both the loop and the array's declared type exist.
//
//===----------------------------------------------------------------------===//
+#include "flang/Optimizer/Builder/CUFCommon.h"
#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Dialect/FIRDialect.h"
#include "flang/Optimizer/Dialect/FIROps.h"
@@ -240,6 +242,18 @@ static void versionLoop(fir::FirOpBuilder &builder, fir::DoLoopOp loop,
fir::ResultOp::create(builder, loc, loop.getResults());
}
+/// Device code may contain convergent operations which must not be duplicated
+/// under a condition that can differ between threads. attributes(host,device)
+/// procedures are included: their device copy is split off after this pass
+/// runs, so there is only one func.func here and it would be versioned for
+/// both. omp.target is not covered by the CUF helpers so we add it here.
+static bool runsOnDevice(fir::DoLoopOp loop) {
+ auto func = loop->getParentOfType<mlir::func::FuncOp>();
+ return cuf::isExecutingOnDevice(loop) ||
+ (func && cuf::isDeviceProcedure(func)) ||
+ loop->getParentOfType<mlir::omp::TargetOp>();
+}
+
/// Innermost means no nested loop of any kind. Loops containing OpenMP or
/// OpenACC operations are also excluded.
static bool isVersionableInnermostLoop(fir::DoLoopOp loop) {
@@ -267,7 +281,7 @@ void LoopTripCountVersioningPass::runOnOperation() {
// independent.
llvm::SmallVector<SectionExtentCandidate> candidates;
func.walk([&](fir::DoLoopOp loop) {
- if (!isVersionableInnermostLoop(loop))
+ if (runsOnDevice(loop) || !isVersionableInnermostLoop(loop))
return;
++NumLoopsExamined;
std::optional<int64_t> extent = getSectionExtent(loop, maxSectionExtent);
diff --git a/flang/test/Transforms/loop-tripcount-versioning.fir b/flang/test/Transforms/loop-tripcount-versioning.fir
index 6307098fbabc28..8a2fc59a87e49f 100644
--- a/flang/test/Transforms/loop-tripcount-versioning.fir
+++ b/flang/test/Transforms/loop-tripcount-versioning.fir
@@ -546,6 +546,93 @@ func.func @contains_omp(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<i32>
// CHECK: omp.parallel
// CHECK: return
+// Loops that run on a GPU are not versioned. The host loop below has the same
+// body as the device cases and is versioned, which shows that the device cases
+// are skipped only because of where they run.
+func.func @host_control(%a: !fir.ref<!fir.array<4xi32>>, %n: index) {
+ %c1 = arith.constant 1 : index
+ %c0_i32 = arith.constant 0 : i32
+ fir.do_loop %i = %c1 to %n step %c1 {
+ %p = fir.array_coor %a %i : (!fir.ref<!fir.array<4xi32>>, index) -> !fir.ref<i32>
+ fir.store %c0_i32 to %p : !fir.ref<i32>
+ }
+ return
+}
+// CHECK-LABEL: func.func @host_control(
+// CHECK: fir.if
+// CHECK: return
+
+// Versioning would turn the one barrier into two, and threads whose %n differ
+// would reach different ones.
+func.func @device_proc(%n: index) attributes {cuf.proc_attr = #cuf.cuda_proc<device>} {
+ %a = fir.alloca !fir.array<4xi32>
+ %c1 = arith.constant 1 : index
+ %c0_i32 = arith.constant 0 : i32
+ fir.do_loop %i = %c1 to %n step %c1 {
+ %first = arith.cmpi eq, %i, %c1 : index
+ fir.if %first { nvvm.barrier }
+ %p = fir.array_coor %a %i : (!fir.ref<!fir.array<4xi32>>, index) -> !fir.ref<i32>
+ fir.store %c0_i32 to %p : !fir.ref<i32>
+ }
+ return
+}
+// CHECK-LABEL: func.func @device_proc(
+// CHECK-NOT: arith.constant 4
+// CHECK: nvvm.barrier
+// CHECK-NOT: nvvm.barrier
+// CHECK: return
+
+// attributes(host,device) compiles for both CPU and GPU, so it is not versioned.
+func.func @host_device_proc(%a: !fir.ref<!fir.array<4xi32>>, %n: index) attributes {cuf.proc_attr = #cuf.cuda_proc<host_device>} {
+ %c1 = arith.constant 1 : index
+ %c0_i32 = arith.constant 0 : i32
+ fir.do_loop %i = %c1 to %n step %c1 {
+ %p = fir.array_coor %a %i : (!fir.ref<!fir.array<4xi32>>, index) -> !fir.ref<i32>
+ fir.store %c0_i32 to %p : !fir.ref<i32>
+ }
+ return
+}
+// CHECK-LABEL: func.func @host_device_proc(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+func.func @cuf_kernel(%a: !fir.ref<!fir.array<4xi32>>, %n: index) {
+ %c1 = arith.constant 1 : index
+ %c1_i32 = arith.constant 1 : i32
+ %c0_i32 = arith.constant 0 : i32
+ cuf.kernel<<<%c1_i32, %c1_i32>>> (%j : index) = (%c1 : index) to (%c1 : index) step (%c1 : index) {
+ fir.do_loop %i = %c1 to %n step %c1 {
+ %p = fir.array_coor %a %i : (!fir.ref<!fir.array<4xi32>>, index) -> !fir.ref<i32>
+ fir.store %c0_i32 to %p : !fir.ref<i32>
+ }
+ "fir.end"() : () -> ()
+ }
+ return
+}
+// CHECK-LABEL: func.func @cuf_kernel(
+// CHECK-NOT: fir.if
+// CHECK: return
+
+func.func @omp_target(%a: !fir.ref<!fir.array<4xi32>>, %n: !fir.ref<i32>) {
+ %0 = omp.map.info var_ptr(%a : !fir.ref<!fir.array<4xi32>>, !fir.array<4xi32>) map_clauses(tofrom) capture(ByRef) name("a") -> !fir.ref<!fir.array<4xi32>>
+ %1 = omp.map.info var_ptr(%n : !fir.ref<i32>, i32) map_clauses(to) capture(ByRef) name("n") -> !fir.ref<i32>
+ omp.target kernel_type(generic) map_entries(%0 -> %arg0, %1 -> %arg1 : !fir.ref<!fir.array<4xi32>>, !fir.ref<i32>) {
+ %c1 = arith.constant 1 : index
+ %c0_i32 = arith.constant 0 : i32
+ %2 = fir.load %arg1 : !fir.ref<i32>
+ %3 = fir.convert %2 : (i32) -> index
+ fir.do_loop %i = %c1 to %3 step %c1 {
+ %p = fir.array_coor %arg0 %i : (!fir.ref<!fir.array<4xi32>>, index) -> !fir.ref<i32>
+ fir.store %c0_i32 to %p : !fir.ref<i32>
+ }
+ omp.terminator
+ }
+ return
+}
+// CHECK-LABEL: func.func @omp_target(
+// CHECK-NOT: fir.if
+// CHECK: return
+
// A loop containing a loop that is not a fir.do_loop is not innermost and is
// not versioned.
func.func @contains_iterate_while(%arg0: !fir.ref<!fir.array<4xf64>>, %arg1: !fir.ref<i32>, %arg2: i1) {
>From 3abd6b604860f9e45be774747aab30bf9dee25bd Mon Sep 17 00:00:00 2001
From: Caroline Newcombe <caroline.newcombe at hpe.com>
Date: Wed, 14 Oct 2026 09:32:45 -0500
Subject: [PATCH 4/4] [flang] Use mlir::OpBuilder in LoopTripCountVersioning,
no module needed
---
.../Transforms/LoopTripCountVersioning.cpp | 18 +++++++++---------
1 file changed, 9 insertions(+), 9 deletions(-)
diff --git a/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp b/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
index 04380d8f288128..2f248478faba18 100644
--- a/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
+++ b/flang/lib/Optimizer/Transforms/LoopTripCountVersioning.cpp
@@ -28,16 +28,15 @@
//===----------------------------------------------------------------------===//
#include "flang/Optimizer/Builder/CUFCommon.h"
-#include "flang/Optimizer/Builder/FIRBuilder.h"
#include "flang/Optimizer/Dialect/FIRDialect.h"
#include "flang/Optimizer/Dialect/FIROps.h"
#include "flang/Optimizer/Dialect/FIROpsSupport.h"
#include "flang/Optimizer/Dialect/FIRType.h"
-#include "flang/Optimizer/Dialect/Support/KindMapping.h"
#include "flang/Optimizer/Transforms/Passes.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/OpenACC/OpenACC.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/IR/Builders.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Pass/Pass.h"
#include "llvm/ADT/Statistic.h"
@@ -98,8 +97,9 @@ static bool isConstantInt(mlir::Value val, int64_t expected) {
return cst && cst->trySExtValue() == expected;
}
-/// An absent shape or a plain fir.shape has lower bound one in every
-/// dimension; a fir.shape_shift must say so explicitly.
+/// The shape operand is optional on fir.array_coor. An absent shape or a plain
+/// fir.shape carries only extents, so the lower bound is one in every
+/// dimension; only a fir.shape_shift can say otherwise.
static bool hasUnitLeadingLowerBound(mlir::Value shapeLike) {
if (!shapeLike || shapeLike.getDefiningOp<fir::ShapeOp>())
return true;
@@ -205,7 +205,7 @@ static std::optional<int64_t> getSectionExtent(fir::DoLoopOp loop,
///
/// The conditions and trip counts must dominate \p loop. Loop results are
/// forwarded through fir.result.
-static void versionLoop(fir::FirOpBuilder &builder, fir::DoLoopOp loop,
+static void versionLoop(mlir::OpBuilder &builder, fir::DoLoopOp loop,
llvm::ArrayRef<LoopVersion> versions) {
assert(!versions.empty() && "nothing to version on");
mlir::Location loc = loop.getLoc();
@@ -313,8 +313,7 @@ void LoopTripCountVersioningPass::runOnOperation() {
return;
}
- mlir::ModuleOp module = func->getParentOfType<mlir::ModuleOp>();
- fir::FirOpBuilder builder{module, fir::getKindMapping(module)};
+ mlir::OpBuilder builder(func.getContext());
for (const SectionExtentCandidate &candidate : candidates) {
fir::DoLoopOp loop = candidate.loop;
@@ -323,8 +322,9 @@ void LoopTripCountVersioningPass::runOnOperation() {
// Define the guard values before the loop so they dominate the fir.if.
builder.setInsertionPoint(loop);
- mlir::Value extentCst =
- builder.createIntegerConstant(loc, ub.getType(), candidate.extent);
+ mlir::Value extentCst = mlir::arith::ConstantOp::create(
+ builder, loc, ub.getType(),
+ builder.getIntegerAttr(ub.getType(), candidate.extent));
mlir::Value isDeclaredExtent = mlir::arith::CmpIOp::create(
builder, loc, mlir::arith::CmpIPredicate::eq, ub, extentCst);
More information about the flang-commits
mailing list