[flang-commits] [flang] [flang] Add pass to version innermost loops on a likely constant trip count (1/2) (PR #231205)

[Sergey Shcherbinin via flang-commits]

================
@@ -0,0 +1,338 @@
+//===- 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`. 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. 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/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
+#include "flang/Optimizer/Dialect/FIRType.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"
+#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(NumLoopsTooLarge,
+          "Number of candidates skipped: N * body size > max-unrolled-ops");
+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;
+}
+
+/// 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;
+  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.
+///
+/// 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::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()));
+  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;
+  } else if (coor->getBlock() != body) {
----------------
SergeyShch01 wrote:

Thanks, reverting the filter makes sense. It was just a heuristic idea; I don't have strong evidence for that particular rule.
The broader question is whether we're comfortable with the current heuristic largely guessing profitability: it may admit many false positives - loops whose cloning increases code size without improving runtime - and the current checks don't specifically distinguish those cases. As possible follow-up work, we could identify the largest code-size increases with no measurable speedup, then look for IR properties that distinguish them from cases with repeatable gains. That could help reduce unnecessary cloning.

https://github.com/llvm/llvm-project/pull/231205