[flang] Implement !DIR$ UNROLL_AND_JAM [N]

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

@@ -2200,11 +2200,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(
@@ -2221,12 +2249,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);
   }
@@ -2882,6 +2914,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

@@ -9535,7 +9535,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