[flang-commits] [flang] [flang] Correctly handle `!dir$ unroll` with unrolling factors of 0 and 1 (PR #126170)

Asher Mancinelli via flang-commits flang-commits at lists.llvm.org
Thu Feb 6 18:23:52 PST 2025


https://github.com/ashermancinelli updated https://github.com/llvm/llvm-project/pull/126170

>From 132d4f3e19c62402439c649713ea8f8dfb38a381 Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli at gmail.com>
Date: Thu, 6 Feb 2025 17:14:22 -0800
Subject: [PATCH 1/4] Correctly handle unroll directives that disable unrolling

---
 flang/docs/Directives.md          | 12 ++++++++-
 flang/lib/Lower/Bridge.cpp        | 45 ++++++++++++++++++++++---------
 flang/test/Integration/unroll.f90 | 45 ++++++++++++++++++++++++++-----
 3 files changed, 83 insertions(+), 19 deletions(-)

diff --git a/flang/docs/Directives.md b/flang/docs/Directives.md
index f356f762b13a2d2..73967fd53949494 100644
--- a/flang/docs/Directives.md
+++ b/flang/docs/Directives.md
@@ -45,9 +45,10 @@ A list of non-standard directives supported by Flang
 ## Introduction
 Directives are commonly used in Fortran programs to specify additional actions 
 to be performed by the compiler. The directives are always specified with the 
-`!dir$` or `cdir$` prefix. 
+`!dir$` or `cdir$` prefix.
 
 ## Loop Directives
+
 Some directives are associated with the following construct, for example loop
 directives. Directives on loops are used to specify additional transformation to
 be performed by the compiler like enabling vectorisation, unrolling, interchange
@@ -57,6 +58,15 @@ Currently loop directives are not accepted in the presence of OpenMP or OpenACC
 constructs on the loop. This should be implemented as it is used in some
 applications.
 
+### Unrolling Directive `!dir$ unroll [n]`
+
+This directive specifies that the compiler ought to unroll the immediately
+folling 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`
+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.
+
 ### Array Expressions
 It is to be decided whether loop directives should also be able to be associated
 with array expressions.
diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index a31629b17cf295e..5f3f918815308a4 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -63,6 +63,7 @@
 #include "flang/Semantics/tools.h"
 #include "flang/Support/Version.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Parser/Parser.h"
@@ -2170,11 +2171,36 @@ class FirConverter : public Fortran::lower::AbstractConverter {
     return builder->createIntegerConstant(loc, controlType, 1); // step
   }
 
+  // For unroll directives without a value, force full unrolling.
+  // For unroll directives with a value, if the value is greater than 1,
+  // force unrolling with the given factor. Otherwise, disable unrolling.
+  mlir::LLVM::LoopUnrollAttr
+  genLoopUnrollAttr(std::optional<std::uint64_t> directiveArg) {
+    mlir::BoolAttr falseAttr = mlir::BoolAttr::get(builder->getContext(), false);
+    mlir::BoolAttr trueAttr = mlir::BoolAttr::get(builder->getContext(), true);
+    mlir::IntegerAttr countAttr;
+    mlir::BoolAttr fullUnrollAttr;
+    bool shouldUnroll = true;
+    if (directiveArg.has_value()) {
+      auto unrollingFactor = directiveArg.value();
+      if (unrollingFactor == 0 || unrollingFactor == 1) {
+        shouldUnroll = false;
+      } else {
+        countAttr = builder->getIntegerAttr(builder->getI64Type(), unrollingFactor);
+      }
+    } else {
+      fullUnrollAttr = trueAttr;
+    }
+
+    mlir::BoolAttr disableAttr = shouldUnroll ? falseAttr : trueAttr;
+    return mlir::LLVM::LoopUnrollAttr::get(
+        builder->getContext(), /*disable=*/disableAttr, /*count=*/countAttr, {},
+        /*full=*/fullUnrollAttr, {}, {}, {});
+  }
+
   void addLoopAnnotationAttr(
       IncrementLoopInfo &info,
       llvm::SmallVectorImpl<const Fortran::parser::CompilerDirective *> &dirs) {
-    mlir::BoolAttr f = mlir::BoolAttr::get(builder->getContext(), false);
-    mlir::BoolAttr t = mlir::BoolAttr::get(builder->getContext(), true);
     mlir::LLVM::LoopVectorizeAttr va;
     mlir::LLVM::LoopUnrollAttr ua;
     bool has_attrs = false;
@@ -2182,20 +2208,15 @@ class FirConverter : public Fortran::lower::AbstractConverter {
       Fortran::common::visit(
           Fortran::common::visitors{
               [&](const Fortran::parser::CompilerDirective::VectorAlways &) {
+                mlir::BoolAttr falseAttr =
+                    mlir::BoolAttr::get(builder->getContext(), false);
                 va = mlir::LLVM::LoopVectorizeAttr::get(builder->getContext(),
-                                                        /*disable=*/f, {}, {},
-                                                        {}, {}, {}, {});
+                                                        /*disable=*/falseAttr,
+                                                        {}, {}, {}, {}, {}, {});
                 has_attrs = true;
               },
               [&](const Fortran::parser::CompilerDirective::Unroll &u) {
-                mlir::IntegerAttr countAttr;
-                if (u.v.has_value()) {
-                  countAttr = builder->getIntegerAttr(builder->getI64Type(),
-                                                      u.v.value());
-                }
-                ua = mlir::LLVM::LoopUnrollAttr::get(
-                    builder->getContext(), /*disable=*/f, /*count*/ countAttr,
-                    {}, /*full*/ u.v.has_value() ? f : t, {}, {}, {});
+                ua = genLoopUnrollAttr(u.v);
                 has_attrs = true;
               },
               [&](const auto &) {}},
diff --git a/flang/test/Integration/unroll.f90 b/flang/test/Integration/unroll.f90
index 9d69605e10d1b3c..4c35d690d2907cf 100644
--- a/flang/test/Integration/unroll.f90
+++ b/flang/test/Integration/unroll.f90
@@ -3,14 +3,47 @@
 ! CHECK-LABEL: unroll_dir
 subroutine unroll_dir
   integer :: a(10)
-  !dir$ unroll 
-  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[ANNOTATION:.*]]
+  !dir$ unroll
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_ENABLE_FULL_ANNO:.*]]
   do i=1,10
-     a(i)=i
+  a(i)=i
   end do
 end subroutine unroll_dir
 
-! CHECK: ![[ANNOTATION]] = distinct !{![[ANNOTATION]], ![[UNROLL:.*]], ![[UNROLL_FULL:.*]]}
-! CHECK: ![[UNROLL]] = !{!"llvm.loop.unroll.enable"}
-! CHECK: ![[UNROLL_FULL]] = !{!"llvm.loop.unroll.full"}
+! CHECK-LABEL: unroll_dir_0
+subroutine unroll_dir_0
+  integer :: a(10)
+  !dir$ unroll 0
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_DISABLE:.*]]
+  do i=1,10
+  a(i)=i
+  end do
+end subroutine unroll_dir_0
+
+! CHECK-LABEL: unroll_dir_1
+subroutine unroll_dir_1
+  integer :: a(10)
+  !dir$ unroll 1
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_DISABLE_ANNO:.*]]
+  do i=1,10
+  a(i)=i
+  end do
+end subroutine unroll_dir_1
+
+! CHECK-LABEL: unroll_dir_2
+subroutine unroll_dir_2
+  integer :: a(10)
+  !dir$ unroll 2
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_ENABLE_COUNT_2:.*]]
+  do i=1,10
+  a(i)=i
+  end do
+end subroutine unroll_dir_2
 
+! CHECK: ![[UNROLL_ENABLE_FULL_ANNO]] = distinct !{![[UNROLL_ENABLE_FULL_ANNO]], ![[UNROLL_ENABLE:.*]], ![[UNROLL_FULL:.*]]}
+! CHECK: ![[UNROLL_ENABLE:.*]] = !{!"llvm.loop.unroll.enable"}
+! CHECK: ![[UNROLL_FULL:.*]] = !{!"llvm.loop.unroll.full"}
+! CHECK: ![[UNROLL_DISABLE_ANNO]] = distinct !{![[UNROLL_DISABLE_ANNO]], ![[UNROLL_DISABLE:.*]]}
+! CHECK: ![[UNROLL_DISABLE]] = !{!"llvm.loop.unroll.disable"}
+! CHECK: ![[UNROLL_ENABLE_COUNT_2]] = distinct !{![[UNROLL_ENABLE_COUNT_2]], ![[UNROLL_ENABLE]], ![[UNROLL_COUNT_2:.*]]}
+! CHECK: ![[UNROLL_COUNT_2]] = !{!"llvm.loop.unroll.count", i32 2}

>From d7a6ae775adbb916cce90a520e5f13f64920f1e8 Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli at gmail.com>
Date: Thu, 6 Feb 2025 17:25:49 -0800
Subject: [PATCH 2/4] Touch up filecheck test

---
 flang/test/Integration/unroll.f90 | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/flang/test/Integration/unroll.f90 b/flang/test/Integration/unroll.f90
index 4c35d690d2907cf..aa47e465b63fceb 100644
--- a/flang/test/Integration/unroll.f90
+++ b/flang/test/Integration/unroll.f90
@@ -14,7 +14,7 @@ end subroutine unroll_dir
 subroutine unroll_dir_0
   integer :: a(10)
   !dir$ unroll 0
-  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_DISABLE:.*]]
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_DISABLE_ANNO:.*]]
   do i=1,10
   a(i)=i
   end do
@@ -24,7 +24,7 @@ end subroutine unroll_dir_0
 subroutine unroll_dir_1
   integer :: a(10)
   !dir$ unroll 1
-  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_DISABLE_ANNO:.*]]
+  ! CHECK:   br i1 {{.*}}, label {{.*}}, label {{.*}}, !llvm.loop ![[UNROLL_DISABLE_ANNO]]
   do i=1,10
   a(i)=i
   end do

>From 71c5dec3db8e51f6f31098fe9c3b8636f1e6dba5 Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli at gmail.com>
Date: Thu, 6 Feb 2025 17:42:10 -0800
Subject: [PATCH 3/4] Spelling

---
 flang/docs/Directives.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/flang/docs/Directives.md b/flang/docs/Directives.md
index 73967fd53949494..b27ab8bb36effe5 100644
--- a/flang/docs/Directives.md
+++ b/flang/docs/Directives.md
@@ -61,7 +61,7 @@ applications.
 ### Unrolling Directive `!dir$ unroll [n]`
 
 This directive specifies that the compiler ought to unroll the immediately
-folling loop `n` times. When `n` is `0` or `1`, the loop should not be unrolled
+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`
 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`

>From 84d0a76c7cc4bc5fa45f3401d0bc443b9bd12e66 Mon Sep 17 00:00:00 2001
From: Asher Mancinelli <ashermancinelli at gmail.com>
Date: Thu, 6 Feb 2025 18:23:27 -0800
Subject: [PATCH 4/4] Formatting

---
 flang/lib/Lower/Bridge.cpp | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/flang/lib/Lower/Bridge.cpp b/flang/lib/Lower/Bridge.cpp
index 5f3f918815308a4..36e58e456dea350 100644
--- a/flang/lib/Lower/Bridge.cpp
+++ b/flang/lib/Lower/Bridge.cpp
@@ -2176,7 +2176,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
   // force unrolling with the given factor. Otherwise, disable unrolling.
   mlir::LLVM::LoopUnrollAttr
   genLoopUnrollAttr(std::optional<std::uint64_t> directiveArg) {
-    mlir::BoolAttr falseAttr = mlir::BoolAttr::get(builder->getContext(), false);
+    mlir::BoolAttr falseAttr =
+        mlir::BoolAttr::get(builder->getContext(), false);
     mlir::BoolAttr trueAttr = mlir::BoolAttr::get(builder->getContext(), true);
     mlir::IntegerAttr countAttr;
     mlir::BoolAttr fullUnrollAttr;
@@ -2186,7 +2187,8 @@ class FirConverter : public Fortran::lower::AbstractConverter {
       if (unrollingFactor == 0 || unrollingFactor == 1) {
         shouldUnroll = false;
       } else {
-        countAttr = builder->getIntegerAttr(builder->getI64Type(), unrollingFactor);
+        countAttr =
+            builder->getIntegerAttr(builder->getI64Type(), unrollingFactor);
       }
     } else {
       fullUnrollAttr = trueAttr;



More information about the flang-commits mailing list