[flang][acc] allow and ignore DIR between ACC and loops

[jeanPerier]

The current pattern was failing OpenACC semantics in acc parse tree canonicalization:

!acc loop
!dir vector aligned
do i=1,n
...

Fix it by moving the directive before the OpenACC construct node.

Note that I think it could make sense to propagate the $dir info to the acc.loop, at least with classic flang, the $dir seems to make a difference. This is not done here since few directives are supported anyway.

[llvmbot]

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

@llvm/pr-subscribers-openacc

Author: None (jeanPerier)

Changes

The current pattern was failing OpenACC semantics in acc parse tree canonicalization:

!acc loop
!dir vector aligned
do i=1,n
...

Fix it by moving the directive before the OpenACC construct node.

Note that I think it could make sense to propagate the $dir info to the acc.loop, at least with classic flang, the $dir seems to make a difference. This is not done here since few directives are supported anyway.

---

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

5 Files Affected:

• (modified) flang/lib/Lower/OpenACC.cpp (+7)
• (modified) flang/lib/Semantics/canonicalize-acc.cpp (+16)
• (modified) flang/lib/Semantics/canonicalize-directives.cpp (+7-12)
• (added) flang/test/Lower/OpenACC/acc-loop-and-cpu-dir.f90 (+75)
• (modified) flang/test/Semantics/loop-directives.f90 (+11)

diff --git a/flang/lib/Lower/OpenACC.cpp b/flang/lib/Lower/OpenACC.cpp
index be184aeead6ee5..431fab52872d33 100644
--- a/flang/lib/Lower/OpenACC.cpp
+++ b/flang/lib/Lower/OpenACC.cpp
@@ -2080,6 +2080,13 @@ static mlir::acc::LoopOp createLoopOp(
     loopOp.setCombinedAttr(mlir::acc::CombinedConstructsTypeAttr::get(
         builder.getContext(), *combinedConstructs));
 
+  // TODO: retrieve directives from NonLabelDoStmt pft::Evaluation, and add them
+  // as attribute to the acc.loop as an extra attribute. It is not quite clear
+  // how useful these $dir are in acc contexts, but they could still provide
+  // more information about the loop acc codegen. They can be obtained by
+  // looking for the first lexicalSuccessor of eval that is a NonLabelDoStmt,
+  // and using the related `dirs` member.
+
   return loopOp;
 }
 
diff --git a/flang/lib/Semantics/canonicalize-acc.cpp b/flang/lib/Semantics/canonicalize-acc.cpp
index f9e44b9540dbb9..108fd33c2ed949 100644
--- a/flang/lib/Semantics/canonicalize-acc.cpp
+++ b/flang/lib/Semantics/canonicalize-acc.cpp
@@ -107,6 +107,20 @@ class CanonicalizationOfAcc {
     }
   }
 
+  // Utility to move all parser::CompilerDirective right after it to right
+  // before it.  This allows preserving loop directives $DIR that may lie
+  // between an $acc directive and loop and leave lowering decide if it should
+  // ignore them or lower/apply them to the acc loops.
+  void moveCompilerDirectivesBefore(
+      parser::Block &block, parser::Block::iterator it) {
+    parser::Block::iterator nextIt = std::next(it);
+    while (nextIt != block.end() &&
+        parser::Unwrap<parser::CompilerDirective>(*nextIt)) {
+      block.emplace(it, std::move(*nextIt));
+      nextIt = block.erase(nextIt);
+    }
+  }
+
   void RewriteOpenACCLoopConstruct(parser::OpenACCLoopConstruct &x,
       parser::Block &block, parser::Block::iterator it) {
     parser::Block::iterator nextIt;
@@ -115,6 +129,7 @@ class CanonicalizationOfAcc {
     auto &nestedDo{std::get<std::optional<parser::DoConstruct>>(x.t)};
 
     if (!nestedDo) {
+      moveCompilerDirectivesBefore(block, it);
       nextIt = it;
       if (++nextIt != block.end()) {
         if (auto *doCons{parser::Unwrap<parser::DoConstruct>(*nextIt)}) {
@@ -151,6 +166,7 @@ class CanonicalizationOfAcc {
     auto &nestedDo{std::get<std::optional<parser::DoConstruct>>(x.t)};
 
     if (!nestedDo) {
+      moveCompilerDirectivesBefore(block, it);
       nextIt = it;
       if (++nextIt != block.end()) {
         if (auto *doCons{parser::Unwrap<parser::DoConstruct>(*nextIt)}) {
diff --git a/flang/lib/Semantics/canonicalize-directives.cpp b/flang/lib/Semantics/canonicalize-directives.cpp
index ae691ab612ba27..739bc3c1992ba6 100644
--- a/flang/lib/Semantics/canonicalize-directives.cpp
+++ b/flang/lib/Semantics/canonicalize-directives.cpp
@@ -8,6 +8,7 @@
 
 #include "canonicalize-directives.h"
 #include "flang/Parser/parse-tree-visitor.h"
+#include "flang/Semantics/tools.h"
 
 namespace Fortran::semantics {
 
@@ -82,25 +83,19 @@ bool CanonicalizationOfDirectives::Pre(parser::ExecutionPart &x) {
   return true;
 }
 
-template <typename T> T *GetConstructIf(parser::ExecutionPartConstruct &x) {
-  if (auto *y{std::get_if<parser::ExecutableConstruct>(&x.u)}) {
-    if (auto *z{std::get_if<common::Indirection<T>>(&y->u)}) {
-      return &z->value();
-    }
-  }
-  return nullptr;
-}
-
 void CanonicalizationOfDirectives::CheckLoopDirective(
     parser::CompilerDirective &dir, parser::Block &block,
     std::list<parser::ExecutionPartConstruct>::iterator it) {
 
   // Skip over this and other compiler directives
-  while (GetConstructIf<parser::CompilerDirective>(*it)) {
+  while (it != block.end() && parser::Unwrap<parser::CompilerDirective>(*it)) {
     ++it;
   }
 
-  if (it == block.end() || !GetConstructIf<parser::DoConstruct>(*it)) {
+  if (it == block.end() ||
+      (!parser::Unwrap<parser::DoConstruct>(*it) &&
+          !parser::Unwrap<parser::OpenACCLoopConstruct>(*it) &&
+          !parser::Unwrap<parser::OpenACCCombinedConstruct>(*it))) {
     std::string s{parser::ToUpperCaseLetters(dir.source.ToString())};
     s.pop_back(); // Remove trailing newline from source string
     messages_.Say(
@@ -110,7 +105,7 @@ void CanonicalizationOfDirectives::CheckLoopDirective(
 
 void CanonicalizationOfDirectives::Post(parser::Block &block) {
   for (auto it{block.begin()}; it != block.end(); ++it) {
-    if (auto *dir{GetConstructIf<parser::CompilerDirective>(*it)}) {
+    if (auto *dir{parser::Unwrap<parser::CompilerDirective>(*it)}) {
       std::visit(
           common::visitors{[&](parser::CompilerDirective::VectorAlways &) {
                              CheckLoopDirective(*dir, block, it);
diff --git a/flang/test/Lower/OpenACC/acc-loop-and-cpu-dir.f90 b/flang/test/Lower/OpenACC/acc-loop-and-cpu-dir.f90
new file mode 100644
index 00000000000000..51c6c367d653e0
--- /dev/null
+++ b/flang/test/Lower/OpenACC/acc-loop-and-cpu-dir.f90
@@ -0,0 +1,75 @@
+! Test that $dir loop directives (known or unknown) are not clashing
+! with $acc lowering.
+
+! RUN: %flang_fc1 -fopenacc -emit-hlfir %s -o - | FileCheck %s
+
+subroutine test_before_acc_loop(a, b, c)
+  real, dimension(10) :: a,b,c
+  !dir$ myloop_directive_1
+  !dir$ myloop_directive_2
+  !$acc loop
+  do i=1,N
+    a(i) = b(i) + c(i)
+  enddo
+end subroutine
+! CHECK-LABEL: test_before_acc_loop
+! CHECK: acc.loop
+
+subroutine test_after_acc_loop(a, b, c)
+  real, dimension(10) :: a,b,c
+  !$acc loop
+  !dir$ myloop_directive_1
+  !dir$ myloop_directive_2
+  do i=1,N
+    a(i) = b(i) + c(i)
+  enddo
+end subroutine
+! CHECK-LABEL: test_after_acc_loop
+! CHECK: acc.loop
+
+subroutine test_before_acc_combined(a, b, c)
+  real, dimension(10) :: a,b,c
+  !dir$ myloop_directive_1
+  !dir$ myloop_directive_2
+  !$acc parallel loop
+  do i=1,N
+    a(i) = b(i) + c(i)
+  enddo
+end subroutine
+! CHECK-LABEL: test_before_acc_combined
+! CHECK: acc.parallel combined(loop)
+
+subroutine test_after_acc_combined(a, b, c)
+  real, dimension(10) :: a,b,c
+  !$acc parallel loop
+  !dir$ myloop_directive_1
+  !dir$ myloop_directive_2
+  do i=1,N
+    a(i) = b(i) + c(i)
+  enddo
+end subroutine
+! CHECK-LABEL: test_after_acc_combined
+! CHECK: acc.parallel combined(loop)
+
+
+subroutine test_vector_always_after_acc(a, b, c)
+  real, dimension(10) :: a,b,c
+  !$acc loop
+  !dir$ vector always
+  do i=1,N
+    a(i) = b(i) + c(i)
+  enddo
+end subroutine
+! CHECK-LABEL: test_vector_always_after_acc
+! CHECK: acc.loop
+
+subroutine test_vector_always_before_acc(a, b, c)
+  real, dimension(10) :: a,b,c
+  !dir$ vector always
+  !$acc loop
+  do i=1,N
+    a(i) = b(i) + c(i)
+  enddo
+end subroutine
+! CHECK-LABEL: test_vector_always_before_acc
+! CHECK: acc.loop
diff --git a/flang/test/Semantics/loop-directives.f90 b/flang/test/Semantics/loop-directives.f90
index 6f994c767d45fe..58fb9b8082bc1a 100644
--- a/flang/test/Semantics/loop-directives.f90
+++ b/flang/test/Semantics/loop-directives.f90
@@ -1,4 +1,5 @@
 ! RUN: %python %S/test_errors.py %s %flang_fc1 -Werror
+! RUN: %python %S/test_errors.py %s %flang_fc1 -fopenacc -Werror
 
 subroutine empty
   ! WARNING: A DO loop must follow the VECTOR ALWAYS directive
@@ -17,3 +18,13 @@ subroutine execution_part
   !dir$ vector always
   end do
 end subroutine execution_part
+
+! OK
+subroutine test_vector_always_before_acc(a, b, c)
+  real, dimension(10) :: a,b,c
+  !dir$ vector always
+  !$acc loop
+  do i=1,N
+    a(i) = b(i) + c(i)
+  enddo
+end subroutine

[clementval]

LGTM

[razvanlupusoru]

This is great! Thank you for adding the TODO - I agree that once a scheme is made to preserve these directives, they should be able to be placed on acc.loop.