Skip to content

[flang] Correctly handle !dir$ unroll with unrolling factors of 0 and 1 #126170

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Feb 10, 2025

Conversation

ashermancinelli
Copy link
Contributor

#123331 added support for the unrolling directive. In the presence of an explicit unrolling factor, that unrolling factor would be unconditionally passed into the metadata even when it was 1 or 0. These special cases should instead disable unrolling. Adding an explicit unrolling factor of 0 triggered this assertion which is fixed by this patch:

unsigned int unrollCountPragmaValue(const llvm::Loop*):
Assertion `Count >= 1 && "Unroll count must be positive."' failed.

Updated tests and documentation.

@ashermancinelli ashermancinelli self-assigned this Feb 7, 2025
@llvmbot llvmbot added flang Flang issues not falling into any other category flang:fir-hlfir labels Feb 7, 2025
@llvmbot
Copy link
Member

llvmbot commented Feb 7, 2025

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

Author: Asher Mancinelli (ashermancinelli)

Changes

#123331 added support for the unrolling directive. In the presence of an explicit unrolling factor, that unrolling factor would be unconditionally passed into the metadata even when it was 1 or 0. These special cases should instead disable unrolling. Adding an explicit unrolling factor of 0 triggered this assertion which is fixed by this patch:

unsigned int unrollCountPragmaValue(const llvm::Loop*):
Assertion `Count >= 1 && "Unroll count must be positive."' failed.

Updated tests and documentation.


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

3 Files Affected:

  • (modified) flang/docs/Directives.md (+11-1)
  • (modified) flang/lib/Lower/Bridge.cpp (+33-12)
  • (modified) flang/test/Integration/unroll.f90 (+39-6)
diff --git a/flang/docs/Directives.md b/flang/docs/Directives.md
index f356f762b13a2d2..b27ab8bb36effe5 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
+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`
+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..aa47e465b63fceb 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_ANNO:.*]]
+  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}

Copy link

github-actions bot commented Feb 7, 2025

✅ With the latest revision this PR passed the C/C++ code formatter.

Copy link
Contributor

@vzakhari vzakhari left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you, Asher!

@ashermancinelli
Copy link
Contributor Author

Thanks for reviewing! I'll give David a little more time to respond. I'd like to merge this before #126082 since this fixes a crash and the other PR fixes semantics.

When an explicit `N` is passed to the loop unroll directive, the
unrolling factor should N if N > 1, and unrolling should be disabled
when N is 0 or 1. Update docs and add test cases.
Copy link
Contributor

@JDPailleux JDPailleux left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@ashermancinelli ashermancinelli merged commit 6b52fb2 into llvm:main Feb 10, 2025
9 checks passed
ashermancinelli added a commit to ashermancinelli/llvm-project that referenced this pull request Feb 10, 2025
Icohedron pushed a commit to Icohedron/llvm-project that referenced this pull request Feb 11, 2025
…nd 1 (llvm#126170)

llvm#123331 added support for the
unrolling directive. In the presence of an explicit unrolling factor,
that unrolling factor would be unconditionally passed into the metadata
even when it was 1 or 0. These special cases should instead disable
unrolling. Adding an explicit unrolling factor of 0 triggered this
assertion which is fixed by this patch:

```
unsigned int unrollCountPragmaValue(const llvm::Loop*):
Assertion `Count >= 1 && "Unroll count must be positive."' failed.
```

Updated tests and documentation.
sivan-shani pushed a commit to sivan-shani/llvm-project that referenced this pull request Feb 24, 2025
…nd 1 (llvm#126170)

llvm#123331 added support for the
unrolling directive. In the presence of an explicit unrolling factor,
that unrolling factor would be unconditionally passed into the metadata
even when it was 1 or 0. These special cases should instead disable
unrolling. Adding an explicit unrolling factor of 0 triggered this
assertion which is fixed by this patch:

```
unsigned int unrollCountPragmaValue(const llvm::Loop*):
Assertion `Count >= 1 && "Unroll count must be positive."' failed.
```

Updated tests and documentation.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
flang:fir-hlfir flang Flang issues not falling into any other category
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants