[flang-commits] [flang] 8c3b6bd - [flang][cuda] Do not lower device variables in main program as globals (#102512)

[via flang-commits]

Author: Valentin Clement (バレンタイン クレメン)
Date: 2024-08-08T13:12:59-07:00
New Revision: 8c3b6bd0cb04a6edc3b294f9048cbfb1ea6c200c

URL: https://github.com/llvm/llvm-project/commit/8c3b6bd0cb04a6edc3b294f9048cbfb1ea6c200c
DIFF: https://github.com/llvm/llvm-project/commit/8c3b6bd0cb04a6edc3b294f9048cbfb1ea6c200c.diff

LOG: [flang][cuda] Do not lower device variables in main program as globals (#102512)

Flang considers arrays in main program larger than 32 bytes having the
SAVE attribute and lowers them as globals. In CUDA Fortran, device
variables are not allowed to have the SAVE attribute and should be
allocated dynamically in the main program scope.
This patch updates lowering so CUDA Fortran device variables are not
considered with the SAVE attribute.

Added: 
    

Modified: 
    flang/include/flang/Evaluate/tools.h
    flang/lib/Evaluate/tools.cpp
    flang/test/Lower/CUDA/cuda-program-global.cuf

Removed: 
    


################################################################################
diff  --git a/flang/include/flang/Evaluate/tools.h b/flang/include/flang/Evaluate/tools.h
index 8c6d3b37166a92..de4d415eda6fd6 100644
--- a/flang/include/flang/Evaluate/tools.h
+++ b/flang/include/flang/Evaluate/tools.h
@@ -1243,6 +1243,18 @@ bool CheckForCoindexedObject(parser::ContextualMessages &,
     const std::optional<ActualArgument> &, const std::string &procName,
     const std::string &argName);
 
+inline bool CanCUDASymbolHasSave(const Symbol &sym) {
+  if (const auto *details =
+          sym.GetUltimate().detailsIf<semantics::ObjectEntityDetails>()) {
+    if (details->cudaDataAttr() &&
+        *details->cudaDataAttr() != common::CUDADataAttr::Pinned &&
+        *details->cudaDataAttr() != common::CUDADataAttr::Unified) {
+      return false;
+    }
+  }
+  return true;
+}
+
 inline bool IsCUDADeviceSymbol(const Symbol &sym) {
   if (const auto *details =
           sym.GetUltimate().detailsIf<semantics::ObjectEntityDetails>()) {

diff  --git a/flang/lib/Evaluate/tools.cpp b/flang/lib/Evaluate/tools.cpp
index 34faba39ffd46f..4d78f814f8ef2c 100644
--- a/flang/lib/Evaluate/tools.cpp
+++ b/flang/lib/Evaluate/tools.cpp
@@ -1696,7 +1696,8 @@ bool IsSaved(const Symbol &original) {
       (features.IsEnabled(common::LanguageFeature::SaveMainProgram) ||
           (features.IsEnabled(
                common::LanguageFeature::SaveBigMainProgramVariables) &&
-              symbol.size() > 32))) {
+              symbol.size() > 32)) &&
+      Fortran::evaluate::CanCUDASymbolHasSave(symbol)) {
     // With SaveBigMainProgramVariables, keeping all unsaved main program
     // variables of 32 bytes or less on the stack allows keeping numerical and
     // logical scalars, small scalar characters or derived, small arrays, and

diff  --git a/flang/test/Lower/CUDA/cuda-program-global.cuf b/flang/test/Lower/CUDA/cuda-program-global.cuf
index 97b9927b3082fd..a3c9e1ba8d253c 100644
--- a/flang/test/Lower/CUDA/cuda-program-global.cuf
+++ b/flang/test/Lower/CUDA/cuda-program-global.cuf
@@ -1,19 +1,22 @@
 ! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
 
-! Test lowering of program local variable that are global
+! Test lowering of program local variables. Make sure CUDA device variables are
+! not lowered as global.
 
 program test
   integer, device :: a(10)
+  integer, unified :: u(10)
   integer :: b(10)
   integer :: i
   print*,i
 end
 
 ! CHECK-LABEL: func.func @_QQmain()
-! CHECK: fir.address_of(@_QFEa) : !fir.ref<!fir.array<10xi32>>
+! CHECK: cuf.alloc !fir.array<10xi32> {bindc_name = "a", data_attr = #cuf.cuda<device>, uniq_name = "_QFEa"} -> !fir.ref<!fir.array<10xi32>>
 ! CHECK: fir.address_of(@_QFEb) : !fir.ref<!fir.array<10xi32>>
 ! CHECK: %[[ALLOCA:.*]] = fir.alloca i32 {bindc_name = "i", uniq_name = "_QFEi"}
 ! CHECK: hlfir.declare %[[ALLOCA]] {uniq_name = "_QFEi"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
 
-! CHECK: fir.global internal @_QFEa {data_attr = #cuf.cuda<device>} : !fir.array<10xi32> {{{$}}
+! CHECK-NOT: fir.global internal @_QFEa {data_attr = #cuf.cuda<device>} : !fir.array<10xi32> {{{$}}
 ! CHECK: fir.global internal @_QFEb : !fir.array<10xi32> {{{$}}
+! CHECK: fir.global internal @_QFEu {data_attr = #cuf.cuda<unified>} : !fir.array<10xi32>