[flang][cuda] Implicitly load cudadevice in host,device and grid_global procedures

[clementval]

No description provided.

[llvmbot]

@llvm/pr-subscribers-flang-semantics

Author: Valentin Clement (バレンタイン クレメン) (clementval)

Changes

---

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

2 Files Affected:

• (modified) flang/lib/Semantics/resolve-names.cpp (+3-1)
• (modified) flang/test/Semantics/cuf-device-procedures02.cuf (+12)

diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp
index 50ca58ea01429..74367b5229548 100644
--- a/flang/lib/Semantics/resolve-names.cpp
+++ b/flang/lib/Semantics/resolve-names.cpp
@@ -4343,7 +4343,9 @@ bool SubprogramVisitor::Pre(const parser::PrefixSpec::Attributes &attrs) {
     }
     if (auto attrs{subp->cudaSubprogramAttrs()}) {
       if (*attrs == common::CUDASubprogramAttrs::Global ||
-          *attrs == common::CUDASubprogramAttrs::Device) {
+          *attrs == common::CUDASubprogramAttrs::Grid_Global ||
+          *attrs == common::CUDASubprogramAttrs::Device ||
+          *attrs == common::CUDASubprogramAttrs::HostDevice) {
         const Scope &scope{currScope()};
         const Scope *mod{FindModuleContaining(scope)};
         if (mod &&
diff --git a/flang/test/Semantics/cuf-device-procedures02.cuf b/flang/test/Semantics/cuf-device-procedures02.cuf
index c93fc4033b8f0..5945abb013eed 100644
--- a/flang/test/Semantics/cuf-device-procedures02.cuf
+++ b/flang/test/Semantics/cuf-device-procedures02.cuf
@@ -13,5 +13,17 @@ contains
 !ERROR: 'threadfence' is use-associated from module 'cudadevice' and cannot be re-declared
     integer :: threadfence
   end subroutine
+
+  attributes(host,device) subroutine sub3()
+    if (on_device()) then
+      print*, 'on device'
+    else
+      print*, 'on host'
+    end if
+  end subroutine
+
+  attributes(grid_global) subroutine sub4()
+    call threadfence()
+  end subroutine
 end module