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

[Tom Eccles via flang-commits]

================
@@ -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>())
----------------
tblah wrote:

nit: why does `!shapeLike` return true? Is that just for convenience?

It doesn't matter too much because I think a null shape argument to fir::ArrayCoorOp wouldn't get past the verifier.

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