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

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.

Author: JDPailleux

flang/docs/Directives.md

@@ -45,6 +45,11 @@ A list of non-standard directives supported by Flang
   times if possible. When `n` is omitted, the compiler should attempt to fully
   unroll the loop. Some compilers accept an optional `=` before the `n` when `n`
   is present in the directive. Flang does not.
+* `!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 optional 
+  integer that specifying the unrolling factor. When `N` is `0` or `1`, the loop 
+  should not be unrolled at all. If `N` is omitted the optimizer will
+  selects the number of times to unroll the loop.
 
 # Directive Details
 

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)

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

@@ -3349,6 +3349,8 @@ struct StmtFunctionStmt {
 // !DIR$ IGNORE_TKR [ [(tkrdmac...)] name ]...
 // !DIR$ LOOP COUNT (n1[, n2]...)
 // !DIR$ name[=value] [, name[=value]]...    = can be :
+// !DIR$ UNROLL [N]
+// !DIR$ UNROLL_AND_JAM [N]
 // !DIR$ <anything else>
 struct CompilerDirective {
   UNION_CLASS_BOILERPLATE(CompilerDirective);
@@ -3371,10 +3373,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;
 };
 

flang/lib/Lower/Bridge.cpp

@@ -2205,11 +2205,39 @@ class FirConverter : public Fortran::lower::AbstractConverter {
         /*full=*/fullUnrollAttr, {}, {}, {});
   }
 
+  // Enabling unroll and jamming directive without a value.
+  // For directives with a value, if the value is greater than 1,
+  // force unrolling with the given factor. Otherwise, disable unrolling and
+  // jamming.
+  mlir::LLVM::LoopUnrollAndJamAttr
+  genLoopUnrollAndJamAttr(std::optional<std::uint64_t> count) {
+    mlir::BoolAttr falseAttr =
+        mlir::BoolAttr::get(builder->getContext(), false);
+    mlir::BoolAttr trueAttr = mlir::BoolAttr::get(builder->getContext(), true);
+    mlir::IntegerAttr countAttr;
+    bool shouldUnroll = true;
+    if (count.has_value()) {
+      auto unrollingFactor = count.value();
+      if (unrollingFactor == 0 || unrollingFactor == 1) {
+        shouldUnroll = false;
+      } else {
+        countAttr =
+            builder->getIntegerAttr(builder->getI64Type(), unrollingFactor);
+      }
+    }
+
+    mlir::BoolAttr disableAttr = shouldUnroll ? falseAttr : trueAttr;
+    return mlir::LLVM::LoopUnrollAndJamAttr::get(
+        builder->getContext(), /*disable=*/disableAttr, /*count*/ countAttr, {},
+        {}, {}, {}, {});
+  }
+
   void addLoopAnnotationAttr(
       IncrementLoopInfo &info,
       llvm::SmallVectorImpl<const Fortran::parser::CompilerDirective *> &dirs) {
     mlir::LLVM::LoopVectorizeAttr va;
     mlir::LLVM::LoopUnrollAttr ua;
+    mlir::LLVM::LoopUnrollAndJamAttr uja;
     bool has_attrs = false;
     for (const auto *dir : dirs) {
       Fortran::common::visit(
@@ -2226,12 +2254,16 @@ class FirConverter : public Fortran::lower::AbstractConverter {
                 ua = genLoopUnrollAttr(u.v);
                 has_attrs = true;
               },
+              [&](const Fortran::parser::CompilerDirective::UnrollAndJam &u) {
+                uja = genLoopUnrollAndJamAttr(u.v);
+                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);
   }
@@ -2887,6 +2919,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);
   }

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>(

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());

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);
     }

flang/lib/Semantics/resolve-names.cpp

@@ -9552,7 +9552,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{

flang/test/Integration/unroll_and_jam.f90

@@ -0,0 +1,48 @@
+! 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-LABEL: unroll_and_jam_dir_0
+subroutine unroll_and_jam_dir_0
+  integer :: a(10)
+  !dir$ unroll_and_jam 0
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE:.*]]
+  do i=1,10
+  a(i)=i
+  end do
+end subroutine unroll_and_jam_dir_0
+
+! CHECK-LABEL: unroll_and_jam_dir_1
+subroutine unroll_and_jam_dir_1
+  integer :: a(10)
+  !dir$ unroll_and_jam 1 
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_DISABLE]]
+  do i=1,10
+  a(i)=i
+  end do
+end subroutine unroll_and_jam_dir_1
+
+! CHECK-LABEL: unroll_and_jam_dir_no_factor
+subroutine unroll_and_jam_dir_no_factor
+  integer :: a(10)
+  !dir$ unroll_and_jam
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION_NO_FACTOR:.*]]
+  do i=1,10
+  a(i)=i
+  end do
+end subroutine unroll_and_jam_dir_no_factor
+
+! 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}
+! CHECK: ![[ANNOTATION_DISABLE]] = distinct !{![[ANNOTATION_DISABLE]], ![[UNROLL_AND_JAM2:.*]]}
+! CHECK: ![[UNROLL_AND_JAM2]] = !{!"llvm.loop.unroll_and_jam.disable"}
+! CHECK: ![[ANNOTATION_NO_FACTOR]] = distinct !{![[ANNOTATION_NO_FACTOR]], ![[UNROLL_AND_JAM]]}

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

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