[flang-commits] [flang] [flang][mlir][fir] Unroll inner loops in presence of user vectorizati… (PR #210820)
Jason Van Beusekom via flang-commits
flang-commits at lists.llvm.org
Mon Jul 20 14:45:58 PDT 2026
https://github.com/Jason-Van-Beusekom created https://github.com/llvm/llvm-project/pull/210820
…on directives
This is a PR for https://discourse.llvm.org/t/rfc-unroll-inner-loops-when-user-requests-outer-loop-vectorization/90742, I originally had a version of this solution as a LLVM pass in https://github.com/llvm/llvm-project/pull/202444. This PR differs by having it be a FIR pass that marks inner loops with the unroll full attribute which force unrolls the inner loops.
A recursive cost model is used to estimate the unroll cost to limit code size explosion.
For the given reproducer this PR results in a 4.5x speed up on an AMD EPYC 9654:
``` fortran
program hot_loop
implicit none
real(8) :: stiffness ! single accumulator (was s(i,j))
real(8) :: array1(3,3,3,3) ! generic 4D tensor
real(8) :: array2(3,3) ! generic 2D matrix
real(8) :: array3(3,3) ! generic 2D matrix
real(8) :: scalar1 ! generic scalar (Gauss weight)
integer :: loop1, loop2, loop3, loop4, loop5, loop6 ! loop indices (all 1:3)
integer :: iter, niter
character(len=32) :: arg
real(8) :: t0, t1
! --- Parse optional iteration count ---
niter = 5000000
if (command_argument_count() >= 1) then
call get_command_argument(1, arg)
read(arg, *) niter
end if
! Initialise with non-trivial data the compiler can't predict
call random_seed()
call random_number(array1)
call random_number(array2)
call random_number(array3)
array1 = array1 * 1.0d4
array2 = array2 * 1.0d-2
array3 = array3 * 1.0d-3
scalar1 = 0.125d0
stiffness = 0.0d0
write(*,'(A,I0,A)') 'Running ', niter, ' iterations ...'
call cpu_time(t0)
do iter = 1, niter
! === THE HOT LOOP NEST (6 levels, all trip-count 3) ===
!dir$ vector always
do loop1 = 1, 3
do loop2 = 1, 3
do loop3 = 1, 3
do loop4 = 1, 3
stiffness = stiffness &
+ array1(loop1,loop3,loop2,loop4) * array2(loop3,loop4) * scalar1
do loop5 = 1, 3
stiffness = stiffness &
+ array1(loop1,loop3,loop5,loop4) * array2(loop3,loop4) &
* array3(loop2,loop5) * scalar1 &
+ array1(loop5,loop3,loop2,loop4) * array2(loop3,loop4) &
* array3(loop1,loop5) * scalar1
do loop6 = 1, 3
stiffness = stiffness &
+ array1(loop5,loop3,loop6,loop4) * array2(loop3,loop4) &
* array3(loop1,loop5) * array3(loop2,loop6) * scalar1
end do ! loop6
end do ! loop5
end do ! loop4
end do ! loop3
end do ! loop2
end do ! loop1
! === END HOT LOOP NEST ===
end do ! iter
call cpu_time(t1)
write(*,'(A,F10.4,A)') 'Time: ', t1 - t0, ' s'
write(*,'(A,ES15.8)') 'Checksum: ', stiffness
end program hot_loop
```
``` bash
flang -O3 -ffast-math hot_loop_orig.f90
./a.out
Running 5000000 iterations ...
Time: 0.3934 s
Checksum: 1.42594131E+09
flang_nightly -O3 -ffast-math hot_loop_orig.f90 && ./a.out
warning: <unknown>:0:0: loop not vectorized: the optimizer was unable to perform the requested transformation; the transformation might be disabled or specified as part ofan unsupported transformation ordering [-Wbackend-plugin]
Running 5000000 iterations ...
Time: 1.7620 s
Checksum: 1.42594131E+09
```
Assisted-by: Opus 4.8
>From cbea80045d0fe7db8e062a9bc0f30b182557253b Mon Sep 17 00:00:00 2001
From: Jason-Van-Beusekom <jason.van-beusekom at hpe.com>
Date: Mon, 20 Jul 2026 16:38:40 -0500
Subject: [PATCH] [flang][mlir][fir] Unroll inner loops in presence of user
vectorization directives
---
.../include/flang/Optimizer/Dialect/FIROps.td | 2 +-
.../flang/Optimizer/Transforms/Passes.td | 21 ++
flang/lib/Optimizer/Dialect/FIROps.cpp | 18 ++
flang/lib/Optimizer/Passes/Pipelines.cpp | 13 ++
flang/lib/Optimizer/Transforms/CMakeLists.txt | 1 +
.../Transforms/VectorAlwaysUnroll.cpp | 191 ++++++++++++++++++
flang/test/Driver/bbc-mlir-pass-pipeline.f90 | 2 +
flang/test/Driver/mlir-pass-pipeline.f90 | 2 +
.../Driver/vector-always-unroll-pipeline.f90 | 23 +++
flang/test/Fir/basic-program.fir | 2 +
.../test/Transforms/vector-always-unroll.fir | 117 +++++++++++
11 files changed, 391 insertions(+), 1 deletion(-)
create mode 100644 flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
create mode 100644 flang/test/Driver/vector-always-unroll-pipeline.f90
create mode 100644 flang/test/Transforms/vector-always-unroll.fir
diff --git a/flang/include/flang/Optimizer/Dialect/FIROps.td b/flang/include/flang/Optimizer/Dialect/FIROps.td
index ae7b8796f8957..0acb2d35d5596 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROps.td
+++ b/flang/include/flang/Optimizer/Dialect/FIROps.td
@@ -2472,7 +2472,7 @@ class region_Op<string mnemonic, list<Trait> traits = []> :
def fir_DoLoopOp : region_Op<"do_loop", [AttrSizedOperandSegments,
AllTypesMatch<["lowerBound", "upperBound", "step"]>,
DeclareOpInterfaceMethods<LoopLikeOpInterface,
- ["getYieldedValuesMutable"]>,
+ ["getStaticTripCount", "getYieldedValuesMutable"]>,
DeclareOpInterfaceMethods<RegionBranchOpInterface,
["getEntrySuccessorOperands", "getSuccessorInputs"]>]> {
let summary = "generalized loop operation";
diff --git a/flang/include/flang/Optimizer/Transforms/Passes.td b/flang/include/flang/Optimizer/Transforms/Passes.td
index d6cbb26882df5..6809b30f15974 100644
--- a/flang/include/flang/Optimizer/Transforms/Passes.td
+++ b/flang/include/flang/Optimizer/Transforms/Passes.td
@@ -401,6 +401,27 @@ def LoopVersioning : Pass<"loop-versioning", "mlir::func::FuncOp"> {
let dependentDialects = [ "fir::FIROpsDialect", "mlir::DLTIDialect" ];
}
+def VectorAlwaysUnroll : Pass<"vector-always-unroll", "mlir::func::FuncOp"> {
+ let summary = "Tag inner loops of a vectorized loop nest for unrolling";
+ let description = [{
+ Lowering encodes the vectorization directives `!dir$ vector always`,
+ `!dir$ vector length`, and `!dir$ simd` as a loop annotation carrying
+ `vectorize.enable`. For each such loop, this pass attaches an
+ `llvm.loop.unroll.full` annotation to every `fir.do_loop` nested within it.
+ This allows llvm's vectorizer to properly handle this loop.
+
+ Because full unrolling of nested loops can cause multiplicative code-size
+ and compile-time growth, tagging is guarded by the cost heuristic:
+ `maxUnrollOps`, Loop nests with non constant trip counts are also skipped.
+ }];
+ let dependentDialects = [ "fir::FIROpsDialect", "mlir::LLVM::LLVMDialect" ];
+ let options = [
+ Option<"maxUnrollOps", "max-unroll-ops", "unsigned", /*default=*/"16384",
+ "Maximum estimated unrolled operation count that will be tagged "
+ "for full unrolling">,
+ ];
+}
+
def VScaleAttr : Pass<"vscale-attr", "mlir::func::FuncOp"> {
let summary = "Add vscale_range attribute to functions";
let description = [{
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index b9968897ade6d..dd656efed8276 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -4078,6 +4078,24 @@ llvm::SmallVector<mlir::Region *> fir::DoLoopOp::getLoopRegions() {
return {&getRegion()};
}
+std::optional<llvm::APInt> fir::DoLoopOp::getStaticTripCount() {
+ auto getConstant = [](mlir::Value v) -> std::optional<std::int64_t> {
+ while (auto cvt =
+ mlir::dyn_cast_or_null<fir::ConvertOp>(v.getDefiningOp()))
+ v = cvt.getValue();
+ return fir::getIntIfConstant(v);
+ };
+ std::optional<std::int64_t> lb = getConstant(getLowerBound());
+ std::optional<std::int64_t> ub = getConstant(getUpperBound());
+ std::optional<std::int64_t> step = getConstant(getStep());
+ if (!lb || !ub || !step || *step == 0)
+ return std::nullopt;
+ std::int64_t count = (*ub - *lb + *step) / *step;
+ if (count < 0)
+ count = 0;
+ return llvm::APInt(64, static_cast<std::uint64_t>(count));
+}
+
/// Translate a value passed as an iter_arg to the corresponding block
/// argument in the body of the loop.
mlir::BlockArgument fir::DoLoopOp::iterArgToBlockArg(mlir::Value iterArg) {
diff --git a/flang/lib/Optimizer/Passes/Pipelines.cpp b/flang/lib/Optimizer/Passes/Pipelines.cpp
index 879a224824119..f2cef2e2b8d82 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -27,6 +27,14 @@ static llvm::cl::opt<bool> disableArgumentFakeUse("disable-argument-fake-use",
llvm::cl::Hidden,
llvm::cl::init(false));
+static bool isVPlanNativePathEnabled() {
+ auto ®isteredOptions = llvm::cl::getRegisteredOptions();
+ auto it = registeredOptions.find("enable-vplan-native-path");
+ if (it == registeredOptions.end())
+ return false;
+ return static_cast<llvm::cl::opt<bool> *>(it->second)->getValue();
+}
+
namespace fir {
template <typename F>
@@ -221,6 +229,11 @@ void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm,
if (pc.LoopVersioning)
pm.addPass(fir::createLoopVersioning());
+ if ((pc.OptLevel == llvm::OptimizationLevel::O2 ||
+ pc.OptLevel == llvm::OptimizationLevel::O3) &&
+ !isVPlanNativePathEnabled())
+ pm.addPass(fir::createVectorAlwaysUnroll());
+
pm.addPass(mlir::createCSEPass());
if (pc.StackArrays)
diff --git a/flang/lib/Optimizer/Transforms/CMakeLists.txt b/flang/lib/Optimizer/Transforms/CMakeLists.txt
index 1a4940af95d3e..5f37d3bc5b4bc 100644
--- a/flang/lib/Optimizer/Transforms/CMakeLists.txt
+++ b/flang/lib/Optimizer/Transforms/CMakeLists.txt
@@ -58,6 +58,7 @@ add_flang_library(FIRTransforms
StackArrays.cpp
StackReclaim.cpp
VScaleAttr.cpp
+ VectorAlwaysUnroll.cpp
DEPENDS
CUFAttrs
diff --git a/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp b/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
new file mode 100644
index 0000000000000..3e70722e6e4df
--- /dev/null
+++ b/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
@@ -0,0 +1,191 @@
+//===- VectorAlwaysUnroll.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
+/// This pass tags inner loops when their outer loop has a user provided
+/// vectorization attribute:(`!dir$ vector always`, `!dir$ vector length`,
+/// and `!dir$ simd`).
+///
+/// For each such loop, this pass attaches an `llvm.loop.unroll.full` annotation
+/// to every `fir.do_loop` nested within it. Fully unrolling those inner loops
+/// later (in LLVM's LoopFullUnrollPass), which allows outer-loop vectorization
+/// of the annotated loop.
+///
+/// Full unrolling of nested loops is multiplicative in code size and compile
+/// time, so tagging is guarded by a cost heuristic:
+/// * only loops with compile-time-constant trip counts are considered
+/// (LoopFullUnrollPass cannot unroll otherwise);
+/// * and the estimated unrolled op count (trip product times per-iteration
+/// op count) stays within `max-unroll-ops`.
+//===----------------------------------------------------------------------===//
+
+#include "flang/Optimizer/Dialect/FIRDialect.h"
+#include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Transforms/Passes.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/ADT/APInt.h"
+#include "llvm/Support/Debug.h"
+#include "llvm/Support/MathExtras.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cstdint>
+
+#include <optional>
+
+namespace fir {
+#define GEN_PASS_DEF_VECTORALWAYSUNROLL
+#include "flang/Optimizer/Transforms/Passes.h.inc"
+} // namespace fir
+
+#define DEBUG_TYPE "flang-vector-always-unroll"
+
+namespace {
+
+static std::optional<std::uint64_t> estimateUnrolledCost(fir::DoLoopOp loop) {
+ std::optional<llvm::APInt> trip = loop.getStaticTripCount();
+ if (!trip)
+ return std::nullopt;
+ std::uint64_t tripCount = trip->getZExtValue();
+
+ std::uint64_t bodyOps = 0;
+ for (mlir::Operation &op : loop.getBody()->without_terminator()) {
+ auto nested = mlir::dyn_cast<fir::DoLoopOp>(&op);
+ if (!nested) {
+ // Count a non-loop operation counts as one op
+ bodyOps = llvm::SaturatingAdd(bodyOps, std::uint64_t{1});
+ continue;
+ }
+ std::optional<std::uint64_t> childCost = estimateUnrolledCost(nested);
+ if (!childCost)
+ return std::nullopt;
+ bodyOps = llvm::SaturatingAdd(bodyOps, *childCost);
+ }
+
+ return llvm::SaturatingMultiply(tripCount, bodyOps);
+}
+
+class VectorAlwaysUnrollPass
+ : public fir::impl::VectorAlwaysUnrollBase<VectorAlwaysUnrollPass> {
+public:
+ using fir::impl::VectorAlwaysUnrollBase<
+ VectorAlwaysUnrollPass>::VectorAlwaysUnrollBase;
+
+ void runOnOperation() override;
+
+private:
+ /// Tag qualifying nested inner loops with `llvm.loop.unroll.full` annotations.
+ void tagNest(fir::DoLoopOp outerLoop,
+ mlir::LLVM::LoopAnnotationAttr unrollAnnotation);
+};
+
+} // namespace
+
+void VectorAlwaysUnrollPass::runOnOperation() {
+ LLVM_DEBUG(llvm::dbgs() << "=== Begin " DEBUG_TYPE " ===\n");
+ mlir::func::FuncOp func = getOperation();
+ mlir::MLIRContext *ctx = &getContext();
+
+ LLVM_DEBUG(llvm::dbgs() << "Func-name:" << func.getSymName() << "\n");
+
+ mlir::BoolAttr trueAttr = mlir::BoolAttr::get(ctx, true);
+ mlir::LLVM::LoopUnrollAttr unrollFull = mlir::LLVM::LoopUnrollAttr::get(
+ ctx, /*disable=*/{}, /*count=*/{}, /*runtimeDisable=*/{},
+ /*full=*/trueAttr, /*followupUnrolled=*/{}, /*followupRemainder=*/{},
+ /*followupAll=*/{});
+ mlir::LLVM::LoopAnnotationAttr unrollAnnotation =
+ mlir::LLVM::LoopAnnotationAttr::get(
+ ctx, /*disableNonforced=*/{}, /*vectorize=*/{}, /*interleave=*/{},
+ /*unroll=*/unrollFull, /*unrollAndJam=*/{}, /*licm=*/{},
+ /*distribute=*/{}, /*pipeline=*/{}, /*peeled=*/{}, /*unswitch=*/{},
+ /*mustProgress=*/{}, /*isVectorized=*/{}, /*startLoc=*/{},
+ /*endLoc=*/{}, /*parallelAccesses=*/{});
+
+
+ func.walk([&](fir::DoLoopOp outerLoop) {
+ // Only act on loops that request vectorization. Lowering encodes
+ // `!dir$ vector always`, `!dir$ vector length`, and `!dir$ simd` as a
+ // loop_annotation with vectorize.enable (disable = false).
+ mlir::LLVM::LoopAnnotationAttr ann = outerLoop.getLoopAnnotationAttr();
+ if (!ann)
+ return;
+ mlir::LLVM::LoopVectorizeAttr vec = ann.getVectorize();
+ if (!vec)
+ return;
+ mlir::BoolAttr disable = vec.getDisable();
+ if (!disable || disable.getValue())
+ return;
+ LLVM_DEBUG(llvm::dbgs()
+ << "VectorAlwaysUnroll: outer loop at " << outerLoop.getLoc()
+ << " (max-unroll-ops=" << maxUnrollOps << ")\n");
+ tagNest(outerLoop, unrollAnnotation);
+ });
+
+ LLVM_DEBUG(llvm::dbgs() << "=== End " DEBUG_TYPE " ===\n");
+}
+
+void VectorAlwaysUnrollPass::tagNest(
+ fir::DoLoopOp outerLoop,
+ mlir::LLVM::LoopAnnotationAttr unrollAnnotation) {
+ std::uint64_t estimatedOps = 0;
+ for (mlir::Operation &op : outerLoop.getBody()->without_terminator()) {
+ auto nested = mlir::dyn_cast<fir::DoLoopOp>(&op);
+ if (!nested)
+ continue;
+ std::optional<std::uint64_t> cost = estimateUnrolledCost(nested);
+ if (!cost) {
+ LLVM_DEBUG(llvm::dbgs()
+ << " abort nest: contains a non-constant trip count loop\n");
+ return;
+ }
+ estimatedOps = llvm::SaturatingAdd(estimatedOps, *cost);
+ }
+
+ LLVM_DEBUG(llvm::dbgs() << " nest cost: estimatedOps=" << estimatedOps
+ << "\n");
+
+ if (estimatedOps > static_cast<std::uint64_t>(maxUnrollOps)) {
+ LLVM_DEBUG(llvm::dbgs()
+ << " estimatedOps exceeds threshold; tagging nothing\n");
+ return;
+ }
+
+ // The nest is small enough: tag every nested loop for full unrolling.
+ outerLoop.walk([&](fir::DoLoopOp innerLoop) {
+ if (innerLoop == outerLoop)
+ return;
+ LLVM_DEBUG(llvm::dbgs() << " tagging loop at " << innerLoop.getLoc()
+ << " with unroll.full\n");
+ mlir::LLVM::LoopAnnotationAttr existing =
+ innerLoop.getLoopAnnotationAttr();
+ if (!existing) {
+ innerLoop.setLoopAnnotationAttr(unrollAnnotation);
+ return;
+ }
+
+ if (existing.getUnroll()) {
+ LLVM_DEBUG(llvm::dbgs()
+ << " keep: loop already has an unroll annotation\n");
+ return;
+ }
+ // Append the unroll.full annotation to the existing loop_annotation
+ mlir::MLIRContext *ctx = innerLoop.getContext();
+ mlir::LLVM::LoopAnnotationAttr merged =
+ mlir::LLVM::LoopAnnotationAttr::get(
+ ctx, existing.getDisableNonforced(), existing.getVectorize(),
+ existing.getInterleave(), /*unroll=*/unrollAnnotation.getUnroll(),
+ existing.getUnrollAndJam(), existing.getLicm(),
+ existing.getDistribute(), existing.getPipeline(),
+ existing.getPeeled(), existing.getUnswitch(),
+ existing.getMustProgress(), existing.getIsVectorized(),
+ existing.getStartLoc(), existing.getEndLoc(),
+ existing.getParallelAccesses());
+ innerLoop.setLoopAnnotationAttr(merged);
+ });
+}
diff --git a/flang/test/Driver/bbc-mlir-pass-pipeline.f90 b/flang/test/Driver/bbc-mlir-pass-pipeline.f90
index 21697485a2a89..416082e2315eb 100644
--- a/flang/test/Driver/bbc-mlir-pass-pipeline.f90
+++ b/flang/test/Driver/bbc-mlir-pass-pipeline.f90
@@ -34,6 +34,8 @@
! CHECK-NEXT: SimplifyRegionLite
! CHECK-NEXT: SimplifyIntrinsics
! CHECK-NEXT: AlgebraicSimplification
+! CHECK-NEXT: 'func.func' Pipeline
+! CHECK-NEXT: VectorAlwaysUnroll
! CHECK-NEXT: CSE
! CHECK-NEXT: (S) 0 num-cse'd - Number of operations CSE'd
! CHECK-NEXT: (S) 0 num-dce'd - Number of operations DCE'd
diff --git a/flang/test/Driver/mlir-pass-pipeline.f90 b/flang/test/Driver/mlir-pass-pipeline.f90
index b679564adff10..3b5db0f6f6acd 100644
--- a/flang/test/Driver/mlir-pass-pipeline.f90
+++ b/flang/test/Driver/mlir-pass-pipeline.f90
@@ -123,6 +123,8 @@
! ALL-NEXT: SimplifyRegionLite
! O2-NEXT: SimplifyIntrinsics
! O2-NEXT: AlgebraicSimplification
+! O2-NEXT: 'func.func' Pipeline
+! O2-NEXT: VectorAlwaysUnroll
! ALL-NEXT: CSE
! ALL-NEXT: (S) 0 num-cse'd - Number of operations CSE'd
! ALL-NEXT: (S) 0 num-dce'd - Number of operations DCE'd
diff --git a/flang/test/Driver/vector-always-unroll-pipeline.f90 b/flang/test/Driver/vector-always-unroll-pipeline.f90
new file mode 100644
index 0000000000000..60b6fd5997dd4
--- /dev/null
+++ b/flang/test/Driver/vector-always-unroll-pipeline.f90
@@ -0,0 +1,23 @@
+! Test that the VectorAlwaysUnroll pass is scheduled in the FIR optimizer
+! pipeline at -O2 by default, and is skipped when LLVM's VPlan-native
+! outer-loop vectorization path is enabled (-enable-vplan-native-path).
+!
+! The pass tags inner loops for full unrolling so the regular loop vectorizer
+! can vectorize the annotated outer loop. When the VPlan-native path is
+! available it can vectorize outer loops directly, so this workaround is
+! unnecessary and must be skipped.
+
+! RUN: %flang_fc1 -S -O2 -mmlir --mlir-pass-statistics -mmlir --mlir-pass-statistics-display=pipeline -o /dev/null %s 2>&1 | FileCheck --check-prefix=DEFAULT %s
+! RUN: %flang_fc1 -S -O2 -mllvm -enable-vplan-native-path -mmlir --mlir-pass-statistics -mmlir --mlir-pass-statistics-display=pipeline -o /dev/null %s 2>&1 | FileCheck --check-prefix=VPLAN %s
+
+! REQUIRES: asserts
+
+end program
+
+! Default (no -enable-vplan-native-path): the pass is scheduled.
+! DEFAULT: Pass statistics report
+! DEFAULT: VectorAlwaysUnroll
+
+! With -enable-vplan-native-path: the pass is skipped.
+! VPLAN: Pass statistics report
+! VPLAN-NOT: VectorAlwaysUnroll
diff --git a/flang/test/Fir/basic-program.fir b/flang/test/Fir/basic-program.fir
index 1fb7ed81f547c..82972215b243f 100644
--- a/flang/test/Fir/basic-program.fir
+++ b/flang/test/Fir/basic-program.fir
@@ -106,6 +106,8 @@ func.func @_QQmain() {
// PASSES-NEXT: SimplifyRegionLite
// PASSES-NEXT: SimplifyIntrinsics
// PASSES-NEXT: AlgebraicSimplification
+// PASSES-NEXT: 'func.func' Pipeline
+// PASSES-NEXT: VectorAlwaysUnroll
// PASSES-NEXT: CSE
// PASSES-NEXT: (S) 0 num-cse'd - Number of operations CSE'd
// PASSES-NEXT: (S) 0 num-dce'd - Number of operations DCE'd
diff --git a/flang/test/Transforms/vector-always-unroll.fir b/flang/test/Transforms/vector-always-unroll.fir
new file mode 100644
index 0000000000000..a80a3374551f7
--- /dev/null
+++ b/flang/test/Transforms/vector-always-unroll.fir
@@ -0,0 +1,117 @@
+// RUN: fir-opt --vector-always-unroll %s | FileCheck %s
+// RUN: fir-opt --vector-always-unroll="max-unroll-ops=1" %s | FileCheck %s --check-prefix=CAPPED
+
+// Vectorization directives (`!dir$ vector always`, `!dir$ vector length`,
+// `!dir$ simd`) lower to a vectorize.enable (disable = false) loop_annotation.
+// For each such loop this pass tags every nested fir.do_loop with
+// llvm.loop.unroll.full (when trip counts are constant and within thresholds),
+// leaving the outer loop untouched.
+
+#vec_enable = #llvm.loop_vectorize<disable = false>
+#anno_vec = #llvm.loop_annotation<vectorize = #vec_enable>
+#vec_width = #llvm.loop_vectorize<disable = false, scalableEnable = false, width = 4 : i64>
+#anno_width = #llvm.loop_annotation<vectorize = #vec_width>
+#vec_disable = #llvm.loop_vectorize<disable = true>
+#anno_novec = #llvm.loop_annotation<vectorize = #vec_disable>
+
+// CHECK: #[[UNROLL:[a-zA-Z0-9_]+]] = #llvm.loop_unroll<full = true>
+// CHECK: #[[UNROLL_ANNO:[a-zA-Z0-9_]+]] = #llvm.loop_annotation<unroll = #[[UNROLL]]>
+
+// CHECK-LABEL: func.func @tag_middle
+// CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #{{.+}}}
+// CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #[[UNROLL_ANNO]]}
+// CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #[[UNROLL_ANNO]]}
+
+// CAPPED-LABEL: func.func @tag_middle
+// CAPPED-NOT: unroll
+func.func @tag_middle() {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ fir.do_loop %i = %c1 to %c10 step %c1 attributes {loopAnnotation = #anno_vec} {
+ fir.do_loop %j = %c1 to %c10 step %c1 {
+ fir.do_loop %k = %c1 to %c10 step %c1 {
+ %0 = arith.addi %i, %j : index
+ }
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @tag_middle_width
+// CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #{{.+}}}
+// CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #[[UNROLL_ANNO]]}
+// CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #[[UNROLL_ANNO]]}
+func.func @tag_middle_width() {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ fir.do_loop %i = %c1 to %c10 step %c1 attributes {loopAnnotation = #anno_width} {
+ fir.do_loop %j = %c1 to %c10 step %c1 {
+ fir.do_loop %k = %c1 to %c10 step %c1 {
+ %0 = arith.addi %i, %j : index
+ }
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @skip_too_many_ops
+// CHECK-NOT: unroll
+func.func @skip_too_many_ops() {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ %c100 = arith.constant 100 : index
+ %c200 = arith.constant 200 : index
+ fir.do_loop %i = %c1 to %c10 step %c1 attributes {loopAnnotation = #anno_vec} {
+ fir.do_loop %j = %c1 to %c100 step %c1 {
+ fir.do_loop %k = %c1 to %c200 step %c1 {
+ %0 = arith.addi %i, %j : index
+ }
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @skip_nonconst
+// CHECK-NOT: unroll
+func.func @skip_nonconst(%n: index) {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ fir.do_loop %i = %c1 to %c10 step %c1 attributes {loopAnnotation = #anno_vec} {
+ fir.do_loop %j = %c1 to %n step %c1 {
+ fir.do_loop %k = %c1 to %c10 step %c1 {
+ %0 = arith.addi %i, %j : index
+ }
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @skip_novector
+// CHECK-NOT: unroll
+func.func @skip_novector() {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ fir.do_loop %i = %c1 to %c10 step %c1 attributes {loopAnnotation = #anno_novec} {
+ fir.do_loop %j = %c1 to %c10 step %c1 {
+ fir.do_loop %k = %c1 to %c10 step %c1 {
+ %0 = arith.addi %i, %j : index
+ }
+ }
+ }
+ return
+}
+
+// CHECK-LABEL: func.func @no_annotation
+// CHECK-NOT: loopAnnotation
+func.func @no_annotation() {
+ %c1 = arith.constant 1 : index
+ %c10 = arith.constant 10 : index
+ fir.do_loop %i = %c1 to %c10 step %c1 {
+ fir.do_loop %j = %c1 to %c10 step %c1 {
+ fir.do_loop %k = %c1 to %c10 step %c1 {
+ %0 = arith.addi %i, %j : index
+ }
+ }
+ }
+ return
+}
More information about the flang-commits
mailing list