[flang] Implement !DIR$ UNROLL_AND_JAM [N] #125046

JDPailleux · 2025-01-30T10:18:05Z

This patch implements support for the UNROLL_AND_JAM directive to enable or disable unrolling and jamming on a DO LOOP.
It must be placed immediately before a DO LOOP and applies only to the loop that follows. N is an integer that specifying the unrolling factor.
This is done by adding an attribute to the branch into the loop in LLVM to indicate that the loop should unrolled and jammed.

llvmbot · 2025-01-30T10:18:40Z

@llvm/pr-subscribers-flang-semantics
@llvm/pr-subscribers-flang-parser

@llvm/pr-subscribers-flang-fir-hlfir

Author: Jean-Didier PAILLEUX (JDPailleux)

Changes

This patch implements support for the UNROLL_AND_JAM directive to enable or disable unrolling and jamming on a DO LOOP.
It must be placed immediately before a DO LOOP and applies only to the loop that follows. N is an integer that specifying the unrolling factor.
This is done by adding an attribute to the branch into the loop in LLVM to indicate that the loop should unrolled and jammed.

Full diff: https://github.com/llvm/llvm-project/pull/125046.diff

10 Files Affected:

(modified) flang/include/flang/Parser/dump-parse-tree.h (+1)
(modified) flang/include/flang/Parser/parse-tree.h (+5-1)
(modified) flang/lib/Lower/Bridge.cpp (+17-2)
(modified) flang/lib/Parser/Fortran-parsers.cpp (+3)
(modified) flang/lib/Parser/unparse.cpp (+4)
(modified) flang/lib/Semantics/canonicalize-directives.cpp (+5-1)
(modified) flang/lib/Semantics/resolve-names.cpp (+2-1)
(added) flang/test/Integration/unroll_and_jam.f90 (+15)
(added) flang/test/Lower/unroll_and_jam.f90 (+34)
(modified) flang/test/Parser/compiler-directives.f90 (+11)

diff --git a/flang/include/flang/Parser/dump-parse-tree.h b/flang/include/flang/Parser/dump-parse-tree.h
index 87bb65fa5c4664..768a078ebc7063 100644
--- a/flang/include/flang/Parser/dump-parse-tree.h
+++ b/flang/include/flang/Parser/dump-parse-tree.h
@@ -209,6 +209,7 @@ class ParseTreeDumper {
   NODE(CompilerDirective, Unrecognized)
   NODE(CompilerDirective, VectorAlways)
   NODE(CompilerDirective, Unroll)
+  NODE(CompilerDirective, UnrollAndJam)
   NODE(parser, ComplexLiteralConstant)
   NODE(parser, ComplexPart)
   NODE(parser, ComponentArraySpec)
diff --git a/flang/include/flang/Parser/parse-tree.h b/flang/include/flang/Parser/parse-tree.h
index be3b1fbde8c3cd..f3cbcd3d502689 100644
--- a/flang/include/flang/Parser/parse-tree.h
+++ b/flang/include/flang/Parser/parse-tree.h
@@ -3349,6 +3349,7 @@ struct StmtFunctionStmt {
 // !DIR$ IGNORE_TKR [ [(tkrdmac...)] name ]...
 // !DIR$ LOOP COUNT (n1[, n2]...)
 // !DIR$ name[=value] [, name[=value]]...    = can be :
+// !DIR$ UNROLL_AND_JAM [N]
 // !DIR$ <anything else>
 struct CompilerDirective {
   UNION_CLASS_BOILERPLATE(CompilerDirective);
@@ -3371,10 +3372,13 @@ struct CompilerDirective {
   struct Unroll {
     WRAPPER_CLASS_BOILERPLATE(Unroll, std::optional<std::uint64_t>);
   };
+  struct UnrollAndJam {
+    WRAPPER_CLASS_BOILERPLATE(UnrollAndJam, std::optional<std::uint64_t>);
+  };
   EMPTY_CLASS(Unrecognized);
   CharBlock source;
   std::variant<std::list<IgnoreTKR>, LoopCount, std::list<AssumeAligned>,
-      VectorAlways, std::list<NameValue>, Unroll, Unrecognized>
+      VectorAlways, std::list<NameValue>, Unroll, UnrollAndJam, Unrecognized>
       u;
 };
 
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index d6bde8ecbbb25e..540434111e1371 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -2177,6 +2177,7 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     mlir::BoolAttr t = mlir::BoolAttr::get(builder->getContext(), true);
     mlir::LLVM::LoopVectorizeAttr va;
     mlir::LLVM::LoopUnrollAttr ua;
+    mlir::LLVM::LoopUnrollAndJamAttr uja;
     bool has_attrs = false;
     for (const auto *dir : dirs) {
       Fortran::common::visit(
@@ -2198,12 +2199,23 @@ class FirConverter : public Fortran::lower::AbstractConverter {
                     {}, /*full*/ u.v.has_value() ? f : t, {}, {}, {});
                 has_attrs = true;
               },
+              [&](const Fortran::parser::CompilerDirective::UnrollAndJam &u) {
+                mlir::IntegerAttr countAttr;
+                if (u.v.has_value()) {
+                  countAttr = builder->getIntegerAttr(builder->getI64Type(),
+                                                      u.v.value());
+                }
+                uja = mlir::LLVM::LoopUnrollAndJamAttr::get(
+                    builder->getContext(), /*disable=*/f, /*count*/ countAttr,
+                    {}, {}, {}, {}, {});
+                has_attrs = true;
+              },
               [&](const auto &) {}},
           dir->u);
     }
     mlir::LLVM::LoopAnnotationAttr la = mlir::LLVM::LoopAnnotationAttr::get(
-        builder->getContext(), {}, /*vectorize=*/va, {}, /*unroll*/ ua, {}, {},
-        {}, {}, {}, {}, {}, {}, {}, {}, {});
+        builder->getContext(), {}, /*vectorize=*/va, {}, /*unroll*/ ua,
+        /*unroll_and_jam*/ uja, {}, {}, {}, {}, {}, {}, {}, {}, {}, {});
     if (has_attrs)
       info.doLoop.setLoopAnnotationAttr(la);
   }
@@ -2859,6 +2871,9 @@ class FirConverter : public Fortran::lower::AbstractConverter {
             [&](const Fortran::parser::CompilerDirective::Unroll &) {
               attachDirectiveToLoop(dir, &eval);
             },
+            [&](const Fortran::parser::CompilerDirective::UnrollAndJam &) {
+              attachDirectiveToLoop(dir, &eval);
+            },
             [&](const auto &) {}},
         dir.u);
   }
diff --git a/flang/lib/Parser/Fortran-parsers.cpp b/flang/lib/Parser/Fortran-parsers.cpp
index b5bcb53a127613..cfe9ecb29b0b72 100644
--- a/flang/lib/Parser/Fortran-parsers.cpp
+++ b/flang/lib/Parser/Fortran-parsers.cpp
@@ -1308,11 +1308,14 @@ constexpr auto vectorAlways{
     "VECTOR ALWAYS" >> construct<CompilerDirective::VectorAlways>()};
 constexpr auto unroll{
     "UNROLL" >> construct<CompilerDirective::Unroll>(maybe(digitString64))};
+constexpr auto unrollAndJam{"UNROLL_AND_JAM" >>
+    construct<CompilerDirective::UnrollAndJam>(maybe(digitString64))};
 TYPE_PARSER(beginDirective >> "DIR$ "_tok >>
     sourced((construct<CompilerDirective>(ignore_tkr) ||
                 construct<CompilerDirective>(loopCount) ||
                 construct<CompilerDirective>(assumeAligned) ||
                 construct<CompilerDirective>(vectorAlways) ||
+                construct<CompilerDirective>(unrollAndJam) ||
                 construct<CompilerDirective>(unroll) ||
                 construct<CompilerDirective>(
                     many(construct<CompilerDirective::NameValue>(
diff --git a/flang/lib/Parser/unparse.cpp b/flang/lib/Parser/unparse.cpp
index 813dd652e1e9f7..d5b97f95fc5c3e 100644
--- a/flang/lib/Parser/unparse.cpp
+++ b/flang/lib/Parser/unparse.cpp
@@ -1851,6 +1851,10 @@ class UnparseVisitor {
               Word("!DIR$ UNROLL");
               Walk(" ", unroll.v);
             },
+            [&](const CompilerDirective::UnrollAndJam &unrollAndJam) {
+              Word("!DIR$ UNROLL_AND_JAM");
+              Walk(" ", unrollAndJam.v);
+            },
             [&](const CompilerDirective::Unrecognized &) {
               Word("!DIR$ ");
               Word(x.source.ToString());
diff --git a/flang/lib/Semantics/canonicalize-directives.cpp b/flang/lib/Semantics/canonicalize-directives.cpp
index b27a27618808bc..1a0a0d145b3e2d 100644
--- a/flang/lib/Semantics/canonicalize-directives.cpp
+++ b/flang/lib/Semantics/canonicalize-directives.cpp
@@ -56,7 +56,8 @@ bool CanonicalizeDirectives(
 static bool IsExecutionDirective(const parser::CompilerDirective &dir) {
   return std::holds_alternative<parser::CompilerDirective::VectorAlways>(
              dir.u) ||
-      std::holds_alternative<parser::CompilerDirective::Unroll>(dir.u);
+      std::holds_alternative<parser::CompilerDirective::Unroll>(dir.u) ||
+      std::holds_alternative<parser::CompilerDirective::UnrollAndJam>(dir.u);
 }
 
 void CanonicalizationOfDirectives::Post(parser::SpecificationPart &spec) {
@@ -115,6 +116,9 @@ void CanonicalizationOfDirectives::Post(parser::Block &block) {
               [&](parser::CompilerDirective::Unroll &) {
                 CheckLoopDirective(*dir, block, it);
               },
+              [&](parser::CompilerDirective::UnrollAndJam &) {
+                CheckLoopDirective(*dir, block, it);
+              },
               [&](auto &) {}},
           dir->u);
     }
diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp
index a586b8e969ec61..ae1eb0362d46c2 100644
--- a/flang/lib/Semantics/resolve-names.cpp
+++ b/flang/lib/Semantics/resolve-names.cpp
@@ -9459,7 +9459,8 @@ void ResolveNamesVisitor::Post(const parser::AssignedGotoStmt &x) {
 
 void ResolveNamesVisitor::Post(const parser::CompilerDirective &x) {
   if (std::holds_alternative<parser::CompilerDirective::VectorAlways>(x.u) ||
-      std::holds_alternative<parser::CompilerDirective::Unroll>(x.u)) {
+      std::holds_alternative<parser::CompilerDirective::Unroll>(x.u) ||
+      std::holds_alternative<parser::CompilerDirective::UnrollAndJam>(x.u)) {
     return;
   }
   if (const auto *tkr{
diff --git a/flang/test/Integration/unroll_and_jam.f90 b/flang/test/Integration/unroll_and_jam.f90
new file mode 100644
index 00000000000000..472acffefbda84
--- /dev/null
+++ b/flang/test/Integration/unroll_and_jam.f90
@@ -0,0 +1,15 @@
+! RUN: %flang_fc1 -emit-llvm -o - %s | FileCheck %s
+
+! CHECK-LABEL: unroll_and_jam_dir
+subroutine unroll_and_jam_dir
+  integer :: a(10)
+  !dir$ unroll_and_jam 4
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION:.*]]
+  do i=1,10
+     a(i)=i
+  end do
+end subroutine unroll_and_jam_dir
+
+! CHECK: ![[ANNOTATION]] = distinct !{![[ANNOTATION]], ![[UNROLL_AND_JAM:.*]], ![[UNROLL_AND_JAM_COUNT:.*]]}
+! CHECK: ![[UNROLL_AND_JAM]] = !{!"llvm.loop.unroll_and_jam.enable"}
+! CHECK: ![[UNROLL_AND_JAM_COUNT]] = !{!"llvm.loop.unroll_and_jam.count", i32 4}
diff --git a/flang/test/Lower/unroll_and_jam.f90 b/flang/test/Lower/unroll_and_jam.f90
new file mode 100644
index 00000000000000..afc5a7b6b271e7
--- /dev/null
+++ b/flang/test/Lower/unroll_and_jam.f90
@@ -0,0 +1,34 @@
+! RUN: %flang_fc1 -emit-hlfir -o - %s | FileCheck %s
+
+! CHECK: #loop_unroll_and_jam = #llvm.loop_unroll_and_jam<disable = false>
+! CHECK: #loop_unroll_and_jam1 = #llvm.loop_unroll_and_jam<disable = false, count = 2 : i64>
+! CHECK: #loop_annotation = #llvm.loop_annotation<unrollAndJam = #loop_unroll_and_jam>
+! CHECK: #loop_annotation1 = #llvm.loop_annotation<unrollAndJam = #loop_unroll_and_jam1>
+
+! CHECK-LABEL: unroll_and_jam_dir
+subroutine unroll_and_jam_dir
+  integer :: a(10)
+  !dir$ unroll_and_jam
+  !CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation}
+  do i=1,10
+     a(i)=i
+  end do
+
+  !dir$ unroll_and_jam 2
+  !CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation1}
+  do i=1,10
+     a(i)=i
+  end do
+end subroutine unroll_and_jam_dir
+
+
+! CHECK-LABEL: intermediate_directive
+subroutine intermediate_directive
+  integer :: a(10)
+  !dir$ unroll_and_jam
+  !dir$ unknown
+  !CHECK: fir.do_loop {{.*}} attributes {loopAnnotation = #loop_annotation}
+  do i=1,10
+     a(i)=i
+  end do
+end subroutine intermediate_directive
diff --git a/flang/test/Parser/compiler-directives.f90 b/flang/test/Parser/compiler-directives.f90
index f372a9f533a355..d1e386a01dd4dd 100644
--- a/flang/test/Parser/compiler-directives.f90
+++ b/flang/test/Parser/compiler-directives.f90
@@ -46,3 +46,14 @@ subroutine unroll
   do i=1,10
   enddo
 end subroutine
+
+subroutine unroll_and_jam
+  !dir$ unroll_and_jam
+  ! CHECK: !DIR$ UNROLL_AND_JAM
+  do i=1,10
+  enddo
+  !dir$ unroll_and_jam 2
+  ! CHECK: !DIR$ UNROLL_AND_JAM 2
+  do i=1,10
+  enddo
+end subroutine

jeanPerier

LGTM, thanks for adding this

flang/include/flang/Parser/parse-tree.h

kiranchandramohan · 2025-01-30T13:36:05Z

Could you add entries for unroll, unroll and jam in https://github.com/llvm/llvm-project/blob/main/flang/docs/Directives.md ?

JDPailleux · 2025-01-30T21:08:49Z

enable or disable unrolling and jamming

Yes, of course

kiranchandramohan · 2025-02-03T13:20:56Z

flang/docs/Directives.md

+* `!dir$ unroll [N]` control how many times a loop should be unrolled. It must
+  be placed immediately before a loop. `N` is an integer that specifying the 
+  unrolling factor.
+* `!dir$ unroll_and_jam [N]` control how many times a loop should be unrolled and
+  jammed. It must be placed immediately before a loop that follows. `N` is an 
+  integer that specifying the unrolling factor.


Can you also say whether N is optional or not? And if it is optional what happens?

Yes of course.

kiranchandramohan

LGTM. Minor comment inline.

JDPailleux · 2025-02-06T09:02:15Z

Hi All, if there is no other review, is it possible for someone to merge this PR ? Thanks in advance :)

ashermancinelli · 2025-02-09T15:51:22Z

What do you think should happen when the user passes an unrolling factor of 0 or 1? After looking at other fortran compilers, I think the unrolling factor should be forced when N>1, and when N is 0 or 1 we should disable unroll (and jam). This is the behavior I intend with #126170. If you agree, would you mind updating this directive handling? Thanks!

JDPailleux · 2025-02-10T13:35:48Z

What do you think should happen when the user passes an unrolling factor of 0 or 1? After looking at other fortran compilers, I think the unrolling factor should be forced when N>1, and when N is 0 or 1 we should disable unroll (and jam). This is the behavior I intend with #126170. If you agree, would you mind updating this directive handling? Thanks!

Yes of course.

JDPailleux · 2025-02-11T10:04:26Z

@ashermancinelli If you are OK, can you merge this PR ?

flang/test/Integration/unroll_and_jam.f90

JDPailleux · 2025-02-19T14:41:07Z

@ashermancinelli Is it possible to merge this PR ?

llvmbot added flang Flang issues not falling into any other category flang:fir-hlfir flang:semantics flang:parser labels Jan 30, 2025

jeanPerier requested review from DavidTruby and ashermancinelli January 30, 2025 10:31

jeanPerier approved these changes Jan 30, 2025

View reviewed changes

flang/include/flang/Parser/parse-tree.h Show resolved Hide resolved

JDPailleux force-pushed the jdp/flang/unrol_and_jam_directive branch from 07abc54 to 470b03d Compare January 30, 2025 12:45

JDPailleux force-pushed the jdp/flang/unrol_and_jam_directive branch from 470b03d to 316f3ca Compare January 30, 2025 21:05

JDPailleux force-pushed the jdp/flang/unrol_and_jam_directive branch from 316f3ca to 835d7fd Compare January 30, 2025 21:08

kiranchandramohan reviewed Feb 3, 2025

View reviewed changes

kiranchandramohan approved these changes Feb 3, 2025

View reviewed changes

JDPailleux force-pushed the jdp/flang/unrol_and_jam_directive branch from 835d7fd to 2536303 Compare February 4, 2025 07:33

JDPailleux force-pushed the jdp/flang/unrol_and_jam_directive branch from 2536303 to 48f619b Compare February 10, 2025 21:27

ashermancinelli reviewed Feb 11, 2025

View reviewed changes

flang/test/Integration/unroll_and_jam.f90 Show resolved Hide resolved

[flang] Implement !DIR$ UNROLL_AND_JAM [N]

b5fbdc1

JDPailleux force-pushed the jdp/flang/unrol_and_jam_directive branch from 48f619b to b5fbdc1 Compare February 12, 2025 08:41

kiranchandramohan merged commit d6c6bde into llvm:main Feb 19, 2025
9 checks passed

joaosaffran mentioned this pull request Feb 26, 2025

[libc++] Set feature-test macro __cpp_lib_atomic_float (#127559) joaosaffran/llvm-project#1

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[flang] Implement !DIR$ UNROLL_AND_JAM [N] #125046

[flang] Implement !DIR$ UNROLL_AND_JAM [N] #125046

Uh oh!

JDPailleux commented Jan 30, 2025

Uh oh!

llvmbot commented Jan 30, 2025 •

edited

Loading

Uh oh!

jeanPerier left a comment

Uh oh!

Uh oh!

kiranchandramohan commented Jan 30, 2025

Uh oh!

JDPailleux commented Jan 30, 2025 •

edited

Loading

Uh oh!

kiranchandramohan Feb 3, 2025

Uh oh!

JDPailleux Feb 4, 2025

Uh oh!

kiranchandramohan left a comment

Uh oh!

JDPailleux commented Feb 6, 2025

Uh oh!

ashermancinelli commented Feb 9, 2025 •

edited

Loading

Uh oh!

JDPailleux commented Feb 10, 2025

Uh oh!

JDPailleux commented Feb 11, 2025

Uh oh!

Uh oh!

JDPailleux commented Feb 19, 2025

Uh oh!

Uh oh!

Uh oh!

[flang] Implement !DIR$ UNROLL_AND_JAM [N] #125046

[flang] Implement !DIR$ UNROLL_AND_JAM [N] #125046

Uh oh!

Conversation

JDPailleux commented Jan 30, 2025

Uh oh!

llvmbot commented Jan 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jeanPerier left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kiranchandramohan commented Jan 30, 2025

Uh oh!

JDPailleux commented Jan 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kiranchandramohan Feb 3, 2025

Choose a reason for hiding this comment

Uh oh!

JDPailleux Feb 4, 2025

Choose a reason for hiding this comment

Uh oh!

kiranchandramohan left a comment

Choose a reason for hiding this comment

Uh oh!

JDPailleux commented Feb 6, 2025

Uh oh!

ashermancinelli commented Feb 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

JDPailleux commented Feb 10, 2025

Uh oh!

JDPailleux commented Feb 11, 2025

Uh oh!

Uh oh!

JDPailleux commented Feb 19, 2025

Uh oh!

Uh oh!

Uh oh!

llvmbot commented Jan 30, 2025 •

edited

Loading

JDPailleux commented Jan 30, 2025 •

edited

Loading

ashermancinelli commented Feb 9, 2025 •

edited

Loading