[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
Mon Oct 19 06:48:11 PDT 2026
================
@@ -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,
----------------
cenewcombe wrote:
Updated the PR description and Passes.td ✔️
https://github.com/llvm/llvm-project/pull/231205
More information about the flang-commits
mailing list