[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
Thu Jul 23 10:52:28 PDT 2026


https://github.com/Jason-Van-Beusekom updated https://github.com/llvm/llvm-project/pull/210820

>From b6be3bc606794331f773eb2601ac18cbb035f771 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 1/4] [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 8de944277e758..3a76d4377c544 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 798ab130c1d30..0524d3ee73988 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -4080,6 +4080,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 7fe8bccc14d54..48ee70b61d318 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 &registeredOptions = 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 14dea7818f230..1cd9f7bb0031e 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
+}

>From f93eaea65d977bdcb3047f48b1536ef7a179537e Mon Sep 17 00:00:00 2001
From: Jason-Van-Beusekom <jason.van-beusekom at hpe.com>
Date: Mon, 20 Jul 2026 16:52:09 -0500
Subject: [PATCH 2/4] format

---
 flang/lib/Optimizer/Dialect/FIROps.cpp        |  3 +-
 .../Transforms/VectorAlwaysUnroll.cpp         | 40 +++++++++----------
 2 files changed, 19 insertions(+), 24 deletions(-)

diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 0524d3ee73988..448846b141a7a 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -4082,8 +4082,7 @@ llvm::SmallVector<mlir::Region *> fir::DoLoopOp::getLoopRegions() {
 
 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()))
+    while (auto cvt = mlir::dyn_cast_or_null<fir::ConvertOp>(v.getDefiningOp()))
       v = cvt.getValue();
     return fir::getIntIfConstant(v);
   };
diff --git a/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp b/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
index 3e70722e6e4df..4ddf3cf216356 100644
--- a/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
+++ b/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
@@ -8,13 +8,13 @@
 
 //===----------------------------------------------------------------------===//
 /// \file
-/// This pass tags inner loops when their outer loop has a user provided 
-/// vectorization attribute:(`!dir$ vector always`, `!dir$ vector length`, 
+/// 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 
+/// 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
@@ -80,7 +80,8 @@ class VectorAlwaysUnrollPass
   void runOnOperation() override;
 
 private:
-  /// Tag qualifying nested inner loops with `llvm.loop.unroll.full` annotations.
+  /// Tag qualifying nested inner loops with `llvm.loop.unroll.full`
+  /// annotations.
   void tagNest(fir::DoLoopOp outerLoop,
                mlir::LLVM::LoopAnnotationAttr unrollAnnotation);
 };
@@ -107,7 +108,6 @@ void VectorAlwaysUnrollPass::runOnOperation() {
           /*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
@@ -131,8 +131,7 @@ void VectorAlwaysUnrollPass::runOnOperation() {
 }
 
 void VectorAlwaysUnrollPass::tagNest(
-    fir::DoLoopOp outerLoop,
-    mlir::LLVM::LoopAnnotationAttr unrollAnnotation) {
+    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);
@@ -162,13 +161,12 @@ void VectorAlwaysUnrollPass::tagNest(
       return;
     LLVM_DEBUG(llvm::dbgs() << "    tagging loop at " << innerLoop.getLoc()
                             << " with unroll.full\n");
-    mlir::LLVM::LoopAnnotationAttr existing =
-        innerLoop.getLoopAnnotationAttr();
+    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");
@@ -176,16 +174,14 @@ void VectorAlwaysUnrollPass::tagNest(
     }
     // 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());
+    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);
   });
 }

>From 17a442bdcb90e406604d96cb11955c0d3e2533b1 Mon Sep 17 00:00:00 2001
From: Jason-Van-Beusekom <jason.van-beusekom at hpe.com>
Date: Thu, 23 Jul 2026 12:18:24 -0500
Subject: [PATCH 3/4] feedback

---
 flang/docs/Directives.md                      |  7 +-
 .../flang/Optimizer/Dialect/FIROpsSupport.h   |  8 ++
 flang/lib/Optimizer/Dialect/FIROps.cpp        | 50 ++++++++---
 flang/lib/Optimizer/Passes/Pipelines.cpp      | 13 +--
 .../Transforms/VectorAlwaysUnroll.cpp         | 90 ++++++++++++-------
 .../vplan-outer-loop-vectorization.f90        | 43 +++++++++
 .../test/Transforms/vector-always-unroll.fir  | 46 ++++++++++
 7 files changed, 199 insertions(+), 58 deletions(-)
 create mode 100644 flang/test/Integration/vplan-outer-loop-vectorization.f90

diff --git a/flang/docs/Directives.md b/flang/docs/Directives.md
index 45080acb778e3..4021f0b33392b 100644
--- a/flang/docs/Directives.md
+++ b/flang/docs/Directives.md
@@ -109,6 +109,10 @@ contains
 * `!dir$ vector always` forces vectorization on the following loop regardless
   of cost model decisions. The loop must still be vectorizable.
   [This directive currently only works on plain do loops without labels].
+  When the annotated loop encloses other loops, flang also fully unrolls the
+  inner loops so the loop vectorizer, which only handles innermost loops, can
+  vectorize the annotated loop. Inner loops are only unrolled when they have
+  constant trip counts and the nest stays within an internal size budget.
 * `!dir$ simd` works the same as `vector always` above, but provides an alternative
   spelling and support for projects which would have used the classic-flang frontend
   previously.
@@ -120,7 +124,8 @@ contains
   vectorization, though it can choose to use fixed length vectorization or not
   at all. `<num>` means that the compiler should consider using this specific
   vectorization factor, which should be an integer literal. This directive
-  currently has the same limitations as `!dir$ vector always`.
+  currently has the same limitations and inner-loop unrolling behavior as
+  `!dir$ vector always`.
 * `!dir$ unroll [n]` specifies that the compiler ought to unroll the immediately
   following loop `n` times. When `n` is `0` or `1`, the loop should not be unrolled
   at all. When `n` is `2` or greater, the loop should be unrolled exactly `n`
diff --git a/flang/include/flang/Optimizer/Dialect/FIROpsSupport.h b/flang/include/flang/Optimizer/Dialect/FIROpsSupport.h
index a354f7aef511b..f5ed04c20bad3 100644
--- a/flang/include/flang/Optimizer/Dialect/FIROpsSupport.h
+++ b/flang/include/flang/Optimizer/Dialect/FIROpsSupport.h
@@ -192,6 +192,14 @@ bool anyFuncArgsHaveAttr(mlir::func::FuncOp func, llvm::StringRef attr);
 /// Unwrap an integer constant from an mlir::Value as an APInt.
 std::optional<llvm::APInt> getIntIfConstant(mlir::Value value);
 
+/// Compute the static trip count of a loop with the given bounds and stride.
+/// Returns std::nullopt if any of \p lb, \p ub or \p step is not a compile-time
+/// integer constant, or if \p step is zero. \p inclusive selects whether \p ub
+/// belongs to the iteration space (as in a Fortran do loop) or is excluded (a
+/// half-open range, as in a default omp.loop_nest).
+std::optional<llvm::APInt> computeTripCount(mlir::Value lb, mlir::Value ub,
+                                            mlir::Value step, bool inclusive);
+
 static constexpr llvm::StringRef getAdaptToByRefAttrName() {
   return "adapt.valuebyref";
 }
diff --git a/flang/lib/Optimizer/Dialect/FIROps.cpp b/flang/lib/Optimizer/Dialect/FIROps.cpp
index 448846b141a7a..651f2eae7ea80 100644
--- a/flang/lib/Optimizer/Dialect/FIROps.cpp
+++ b/flang/lib/Optimizer/Dialect/FIROps.cpp
@@ -4081,20 +4081,8 @@ llvm::SmallVector<mlir::Region *> fir::DoLoopOp::getLoopRegions() {
 }
 
 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));
+  return fir::computeTripCount(getLowerBound(), getUpperBound(), getStep(),
+                               /*inclusive=*/true);
 }
 
 /// Translate a value passed as an iter_arg to the corresponding block
@@ -6152,6 +6140,40 @@ std::optional<llvm::APInt> fir::getIntIfConstant(mlir::Value value) {
   return {};
 }
 
+std::optional<llvm::APInt> fir::computeTripCount(mlir::Value lb, mlir::Value ub,
+                                                 mlir::Value step,
+                                                 bool inclusive) {
+  auto getConstant = [](mlir::Value v) -> std::optional<llvm::APInt> {
+    while (auto cvt = mlir::dyn_cast_or_null<fir::ConvertOp>(v.getDefiningOp()))
+      v = cvt.getValue();
+    return fir::getIntIfConstant(v);
+  };
+  std::optional<llvm::APInt> lbv = getConstant(lb);
+  std::optional<llvm::APInt> ubv = getConstant(ub);
+  std::optional<llvm::APInt> stepv = getConstant(step);
+  if (!lbv || !ubv || !stepv || stepv->isZero())
+    return std::nullopt;
+
+  unsigned width =
+      std::max({lbv->getBitWidth(), ubv->getBitWidth(), stepv->getBitWidth()}) +
+      2;
+  llvm::APInt lbi = lbv->sext(width);
+  llvm::APInt ubi = ubv->sext(width);
+  llvm::APInt stepi = stepv->sext(width);
+  llvm::APInt span = ubi - lbi + stepi;
+  if (!inclusive) {
+    llvm::APInt one(width, 1);
+    if (stepi.isNegative())
+      span += one;
+    else
+      span -= one;
+  }
+  llvm::APInt count = span.sdiv(stepi);
+  if (count.isNegative())
+    return llvm::APInt(width, 0);
+  return count;
+}
+
 bool fir::isDummyArgument(mlir::Value v) {
   auto blockArg{mlir::dyn_cast<mlir::BlockArgument>(v)};
   if (!blockArg) {
diff --git a/flang/lib/Optimizer/Passes/Pipelines.cpp b/flang/lib/Optimizer/Passes/Pipelines.cpp
index 48ee70b61d318..2c036fea36734 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -27,14 +27,6 @@ static llvm::cl::opt<bool> disableArgumentFakeUse("disable-argument-fake-use",
                                                   llvm::cl::Hidden,
                                                   llvm::cl::init(false));
 
-static bool isVPlanNativePathEnabled() {
-  auto &registeredOptions = 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>
@@ -229,9 +221,8 @@ void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm,
   if (pc.LoopVersioning)
     pm.addPass(fir::createLoopVersioning());
 
-  if ((pc.OptLevel == llvm::OptimizationLevel::O2 ||
-       pc.OptLevel == llvm::OptimizationLevel::O3) &&
-      !isVPlanNativePathEnabled())
+  if (pc.OptLevel == llvm::OptimizationLevel::O2 ||
+      pc.OptLevel == llvm::OptimizationLevel::O3)
     pm.addPass(fir::createVectorAlwaysUnroll());
 
   pm.addPass(mlir::createCSEPass());
diff --git a/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp b/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
index 4ddf3cf216356..7aaf1b221b931 100644
--- a/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
+++ b/flang/lib/Optimizer/Transforms/VectorAlwaysUnroll.cpp
@@ -27,8 +27,10 @@
 
 #include "flang/Optimizer/Dialect/FIRDialect.h"
 #include "flang/Optimizer/Dialect/FIROps.h"
+#include "flang/Optimizer/Dialect/FIROpsSupport.h"
 #include "flang/Optimizer/Transforms/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/Pass/Pass.h"
 #include "llvm/ADT/APInt.h"
 #include "llvm/Support/Debug.h"
@@ -48,27 +50,57 @@ namespace fir {
 
 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)
+static std::optional<llvm::APInt>
+computeLoopNestTripCount(mlir::omp::LoopNestOp loopNest) {
+  mlir::OperandRange lbs = loopNest.getLoopLowerBounds();
+  mlir::OperandRange ubs = loopNest.getLoopUpperBounds();
+  mlir::OperandRange steps = loopNest.getLoopSteps();
+  bool inclusive = loopNest.getLoopInclusive();
+
+  std::uint64_t product = 1;
+  for (unsigned i = 0, e = steps.size(); i < e; ++i) {
+    std::optional<llvm::APInt> count =
+        fir::computeTripCount(lbs[i], ubs[i], steps[i], inclusive);
+    if (!count)
       return std::nullopt;
-    bodyOps = llvm::SaturatingAdd(bodyOps, *childCost);
+    product = llvm::SaturatingMultiply(product, count->getZExtValue());
   }
+  return llvm::APInt(64, product);
+}
+
+static std::optional<std::uint64_t> estimateBlockCost(mlir::Block &block) {
+  std::uint64_t ops = 0;
+  for (mlir::Operation &op : block.without_terminator()) {
+    std::optional<llvm::APInt> trip;
+    mlir::Block *body = nullptr;
+    if (auto loop = mlir::dyn_cast<fir::DoLoopOp>(&op)) {
+      trip = loop.getStaticTripCount();
+      body = loop.getBody();
+    } else if (auto loopNest = mlir::dyn_cast<mlir::omp::LoopNestOp>(&op)) {
+      trip = computeLoopNestTripCount(loopNest);
+      body = &loopNest.getRegion().front();
+    }
+
+    if (body) {
+      std::optional<std::uint64_t> bodyCost = estimateBlockCost(*body);
+      if (!trip || !bodyCost)
+        return std::nullopt;
+      ops = llvm::SaturatingAdd(
+          ops, llvm::SaturatingMultiply(trip->getZExtValue(), *bodyCost));
+      continue;
+    }
 
-  return llvm::SaturatingMultiply(tripCount, bodyOps);
+    // Count a non-loop operation as one op
+    ops = llvm::SaturatingAdd(ops, std::uint64_t{1});
+    for (mlir::Region &region : op.getRegions())
+      for (mlir::Block &nested : region) {
+        std::optional<std::uint64_t> cost = estimateBlockCost(nested);
+        if (!cost)
+          return std::nullopt;
+        ops = llvm::SaturatingAdd(ops, *cost);
+      }
+  }
+  return ops;
 }
 
 class VectorAlwaysUnrollPass
@@ -132,24 +164,18 @@ void VectorAlwaysUnrollPass::runOnOperation() {
 
 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);
+  std::optional<std::uint64_t> estimatedOps =
+      estimateBlockCost(*outerLoop.getBody());
+  if (!estimatedOps) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "  abort nest: contains a non-constant trip count loop\n");
+    return;
   }
 
-  LLVM_DEBUG(llvm::dbgs() << "  nest cost: estimatedOps=" << estimatedOps
+  LLVM_DEBUG(llvm::dbgs() << "  nest cost: estimatedOps=" << *estimatedOps
                           << "\n");
 
-  if (estimatedOps > static_cast<std::uint64_t>(maxUnrollOps)) {
+  if (*estimatedOps > static_cast<std::uint64_t>(maxUnrollOps)) {
     LLVM_DEBUG(llvm::dbgs()
                << "  estimatedOps exceeds threshold; tagging nothing\n");
     return;
diff --git a/flang/test/Integration/vplan-outer-loop-vectorization.f90 b/flang/test/Integration/vplan-outer-loop-vectorization.f90
new file mode 100644
index 0000000000000..1fb0e8bd34490
--- /dev/null
+++ b/flang/test/Integration/vplan-outer-loop-vectorization.f90
@@ -0,0 +1,43 @@
+! When a `!dir$ vector always` loop encloses another loop, flang relies on the
+! VectorAlwaysUnroll pass to fully unroll the inner loop(s) so that the ordinary
+! (inner-loop) vectorizer can vectorize the annotated loop. That workaround only
+! applies when the inner loops have compile-time-constant trip counts.
+!
+! The loop below deliberately has a runtime inner trip count and a loop-carried
+! dependence on the inner loop, so the VectorAlwaysUnroll workaround does not
+! apply (even though it now runs regardless of `-enable-vplan-native-path`). The
+! annotated outer loop can therefore only be vectorized by VPlan-native
+! outer-loop vectorization. That path does not yet handle this pattern, so the
+! forced vectorization request fails and LLVM emits the "unable to perform the
+! requested transformation" warning checked below.
+!
+! This warning is a stable, loop-specific signal (only a forced `vector always`
+! loop produces it), so the test asserts it directly rather than relying on the
+! absence of an incidental remark via XFAIL. When VPlan-native outer-loop
+! vectorization learns to vectorize the loop below, this warning disappears and
+! the CHECK starts failing. At that point:
+!   * update this test to check for the vectorization instead, and
+!   * the VectorAlwaysUnroll workaround and its pipeline scheduling in
+!     flang/lib/Optimizer/Passes/Pipelines.cpp can be removed.
+
+! REQUIRES: x86-registered-target
+
+! RUN: %flang_fc1 -emit-llvm -O2 -triple x86_64-unknown-linux-gnu \
+! RUN:   -mllvm -enable-vplan-native-path -Rpass=loop-vectorize \
+! RUN:   -o /dev/null %s 2>&1 | FileCheck %s
+
+subroutine outer_vec(a, b, n)
+  integer :: n, i, j
+  real :: a(n, n), b(n, n)
+  ! The inner loop over j carries a dependence (a(j,i) reads a(j-1,i)), so it
+  ! is not a legal inner-loop vectorization candidate. Only the outer loop over
+  ! i (independent columns) can be vectorized, and it is the one annotated.
+  !dir$ vector always
+  do i = 1, n
+     do j = 2, n
+        a(j, i) = a(j - 1, i) + b(j, i)
+     end do
+  end do
+end subroutine outer_vec
+
+! CHECK: loop not vectorized: the optimizer was unable to perform the requested transformation
diff --git a/flang/test/Transforms/vector-always-unroll.fir b/flang/test/Transforms/vector-always-unroll.fir
index a80a3374551f7..14c9f97536b0c 100644
--- a/flang/test/Transforms/vector-always-unroll.fir
+++ b/flang/test/Transforms/vector-always-unroll.fir
@@ -1,5 +1,6 @@
 // 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
+// RUN: fir-opt --vector-always-unroll="max-unroll-ops=50" %s | FileCheck %s --check-prefix=CAP50
 
 // Vectorization directives (`!dir$ vector always`, `!dir$ vector length`,
 // `!dir$ simd`) lower to a vectorize.enable (disable = false) loop_annotation.
@@ -115,3 +116,48 @@ func.func @no_annotation() {
   }
   return
 }
+
+// CHECK-LABEL: func.func @omp_collapse_counted
+// CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #{{.+}}}
+// CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #[[UNROLL_ANNO]]}
+
+// With max-unroll-ops=50: ignoring the collapsed loop the estimate would be 4
+// (<= 50, so the loop would be tagged). Counting the 100-iteration collapsed
+// omp.loop_nest raises it to 404 (> 50), so nothing is tagged.
+// CAP50-LABEL: func.func @omp_collapse_counted
+// CAP50-NOT: unroll
+func.func @omp_collapse_counted() {
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c1_i32 = arith.constant 1 : i32
+  %c10_i32 = arith.constant 10 : i32
+  fir.do_loop %i = %c1 to %c4 step %c1 attributes {loopAnnotation = #anno_vec} {
+    fir.do_loop %j = %c1 to %c4 step %c1 {
+      omp.simd {
+        omp.loop_nest (%k, %l) : i32 = (%c1_i32, %c1_i32) to (%c10_i32, %c10_i32) inclusive step (%c1_i32, %c1_i32) collapse(2) {
+          %0 = arith.addi %k, %l : i32
+          omp.yield
+        }
+      }
+    }
+  }
+  return
+}
+// CHECK-LABEL: func.func @omp_collapse_nonconst
+// CHECK-NOT: unroll
+func.func @omp_collapse_nonconst(%n: i32) {
+  %c1 = arith.constant 1 : index
+  %c4 = arith.constant 4 : index
+  %c1_i32 = arith.constant 1 : i32
+  fir.do_loop %i = %c1 to %c4 step %c1 attributes {loopAnnotation = #anno_vec} {
+    fir.do_loop %j = %c1 to %c4 step %c1 {
+      omp.simd {
+        omp.loop_nest (%k, %l) : i32 = (%c1_i32, %c1_i32) to (%n, %n) inclusive step (%c1_i32, %c1_i32) collapse(2) {
+          %0 = arith.addi %k, %l : i32
+          omp.yield
+        }
+      }
+    }
+  }
+  return
+}

>From 06c052815d8d066d21127d7cedd2dd65b4203ae6 Mon Sep 17 00:00:00 2001
From: Jason-Van-Beusekom <jason.van-beusekom at hpe.com>
Date: Thu, 23 Jul 2026 12:52:11 -0500
Subject: [PATCH 4/4] add XFail test and disable flag

---
 .../flang/Optimizer/Passes/CommandLineOpts.h  |  2 +-
 .../lib/Optimizer/Passes/CommandLineOpts.cpp  |  3 ++
 flang/lib/Optimizer/Passes/Pipelines.cpp      |  3 +-
 .../Driver/vector-always-unroll-pipeline.f90  | 23 --------
 .../vplan-outer-loop-vectorization.f90        | 52 ++++++++-----------
 5 files changed, 29 insertions(+), 54 deletions(-)
 delete mode 100644 flang/test/Driver/vector-always-unroll-pipeline.f90

diff --git a/flang/include/flang/Optimizer/Passes/CommandLineOpts.h b/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
index 882f02032a3b8..8d9f75d305ed9 100644
--- a/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
+++ b/flang/include/flang/Optimizer/Passes/CommandLineOpts.h
@@ -70,5 +70,5 @@ extern llvm::cl::opt<bool> enableSafeTrampoline;
 extern llvm::cl::opt<bool> disableExternalNameConversion;
 extern llvm::cl::opt<bool> enableConstantArgumentGlobalisation;
 extern llvm::cl::opt<bool> disableCompilerGeneratedNamesConversion;
-
+extern llvm::cl::opt<bool> disableVectorAlwaysUnroll;
 #endif // FORTRAN_OPTIMIZER_PASSES_COMMANDLINE_OPTS_H
diff --git a/flang/lib/Optimizer/Passes/CommandLineOpts.cpp b/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
index d461c1b9757b5..6cfbfd1ac6e8f 100644
--- a/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
+++ b/flang/lib/Optimizer/Passes/CommandLineOpts.cpp
@@ -80,3 +80,6 @@ EnableOption(ConstantArgumentGlobalisation, "constant-argument-globalisation",
              "the local constant argument to global constant conversion");
 DisableOption(CompilerGeneratedNamesConversion, "compiler-generated-names",
               "replace special symbols in compiler generated names");
+
+DisableOption(VectorAlwaysUnroll, "vector-always-unroll",
+              "unroll inner loops under user vectorization directives");
diff --git a/flang/lib/Optimizer/Passes/Pipelines.cpp b/flang/lib/Optimizer/Passes/Pipelines.cpp
index 2c036fea36734..d6dbcfcfc09c1 100644
--- a/flang/lib/Optimizer/Passes/Pipelines.cpp
+++ b/flang/lib/Optimizer/Passes/Pipelines.cpp
@@ -223,7 +223,8 @@ void createDefaultFIROptimizerPassPipeline(mlir::PassManager &pm,
 
   if (pc.OptLevel == llvm::OptimizationLevel::O2 ||
       pc.OptLevel == llvm::OptimizationLevel::O3)
-    pm.addPass(fir::createVectorAlwaysUnroll());
+    addPassConditionally(pm, disableVectorAlwaysUnroll,
+                         [&]() { return fir::createVectorAlwaysUnroll(); });
 
   pm.addPass(mlir::createCSEPass());
 
diff --git a/flang/test/Driver/vector-always-unroll-pipeline.f90 b/flang/test/Driver/vector-always-unroll-pipeline.f90
deleted file mode 100644
index 60b6fd5997dd4..0000000000000
--- a/flang/test/Driver/vector-always-unroll-pipeline.f90
+++ /dev/null
@@ -1,23 +0,0 @@
-! 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/Integration/vplan-outer-loop-vectorization.f90 b/flang/test/Integration/vplan-outer-loop-vectorization.f90
index 1fb0e8bd34490..248e7ce24a413 100644
--- a/flang/test/Integration/vplan-outer-loop-vectorization.f90
+++ b/flang/test/Integration/vplan-outer-loop-vectorization.f90
@@ -1,43 +1,37 @@
-! When a `!dir$ vector always` loop encloses another loop, flang relies on the
-! VectorAlwaysUnroll pass to fully unroll the inner loop(s) so that the ordinary
-! (inner-loop) vectorizer can vectorize the annotated loop. That workaround only
-! applies when the inner loops have compile-time-constant trip counts.
+! When a `!dir$ vector always` loop encloses an inner loop with a loop-carried
+! dependence, only the outer loop can be vectorized. flang normally handles this
+! with the VectorAlwaysUnroll pass, which fully unrolls the (constant-trip)
+! inner loop so the ordinary (inner-loop) vectorizer can then vectorize the
+! annotated outer loop.
 !
-! The loop below deliberately has a runtime inner trip count and a loop-carried
-! dependence on the inner loop, so the VectorAlwaysUnroll workaround does not
-! apply (even though it now runs regardless of `-enable-vplan-native-path`). The
-! annotated outer loop can therefore only be vectorized by VPlan-native
-! outer-loop vectorization. That path does not yet handle this pattern, so the
-! forced vectorization request fails and LLVM emits the "unable to perform the
-! requested transformation" warning checked below.
+! This test disables that workaround (-disable-vector-always-unroll) to check
+! whether LLVM's VPlan-native outer-loop vectorization can vectorize the outer
+! loop directly.
 !
-! This warning is a stable, loop-specific signal (only a forced `vector always`
-! loop produces it), so the test asserts it directly rather than relying on the
-! absence of an incidental remark via XFAIL. When VPlan-native outer-loop
-! vectorization learns to vectorize the loop below, this warning disappears and
-! the CHECK starts failing. At that point:
-!   * update this test to check for the vectorization instead, and
-!   * the VectorAlwaysUnroll workaround and its pipeline scheduling in
+! VPlan-native outer-loop vectorization does not yet handle this pattern, so the
+! expected "vectorized loop" remark is not produced and the test is marked
+! XFAIL. When that path learns to vectorize this loop, the remark appears and
+! the test starts passing (XPASS). At that point:
+!   * drop the XFAIL below, and
+!   * the VectorAlwaysUnroll workaround and its scheduling in
 !     flang/lib/Optimizer/Passes/Pipelines.cpp can be removed.
 
 ! REQUIRES: x86-registered-target
+! XFAIL: *
 
 ! RUN: %flang_fc1 -emit-llvm -O2 -triple x86_64-unknown-linux-gnu \
-! RUN:   -mllvm -enable-vplan-native-path -Rpass=loop-vectorize \
-! RUN:   -o /dev/null %s 2>&1 | FileCheck %s
+! RUN:   -mllvm -enable-vplan-native-path -mmlir -disable-vector-always-unroll \
+! RUN:   -Rpass=loop-vectorize -o /dev/null %s 2>&1 | FileCheck %s
 
-subroutine outer_vec(a, b, n)
-  integer :: n, i, j
-  real :: a(n, n), b(n, n)
-  ! The inner loop over j carries a dependence (a(j,i) reads a(j-1,i)), so it
-  ! is not a legal inner-loop vectorization candidate. Only the outer loop over
-  ! i (independent columns) can be vectorized, and it is the one annotated.
+subroutine outer_vec(a, b)
+  real :: a(8, 8), b(8, 8)
+  integer :: i, j
   !dir$ vector always
-  do i = 1, n
-     do j = 2, n
+  do i = 1, 8
+     do j = 2, 8
         a(j, i) = a(j - 1, i) + b(j, i)
      end do
   end do
 end subroutine outer_vec
 
-! CHECK: loop not vectorized: the optimizer was unable to perform the requested transformation
+! CHECK: remark: {{.*}}vectorized loop



More information about the flang-commits mailing list