[flang] Implement !DIR$ UNROLL_AND_JAM [N]

Author: JDPailleux

flang/docs/Directives.md

@@ -39,6 +39,13 @@ A list of non-standard directives supported by Flang
 * `!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].
+* `!dir$ unroll [N]` control how many times a loop should be unrolled. It must
+ be placed immediately before a loop. `N` is an optional integer that specifying the 
+ unrolling factor. If `N` is omitted the unrolling will be full.
+* `!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. 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

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

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

@@ -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{

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}

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