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

[Valentin Clement バレンタイン クレメン via flang-commits]

https://github.com/clementval updated https://github.com/llvm/llvm-project/pull/102512

>From 8999705495edbb5443b8a35bfc76ff797d2e836c Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 8 Aug 2024 10:56:28 -0700
Subject: [PATCH 1/2] [flang][cuda] Do not lower device variables in main
 program as globals

---
 flang/lib/Evaluate/tools.cpp                  | 3 ++-
 flang/test/Lower/CUDA/cuda-program-global.cuf | 7 ++++---
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/flang/lib/Evaluate/tools.cpp b/flang/lib/Evaluate/tools.cpp
index 34faba39ffd46f..e042d4d5a09d5d 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::IsCUDADeviceSymbol(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..498725cd44905f 100644
--- a/flang/test/Lower/CUDA/cuda-program-global.cuf
+++ b/flang/test/Lower/CUDA/cuda-program-global.cuf
@@ -1,6 +1,7 @@
 ! 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)
@@ -10,10 +11,10 @@ program test
 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> {{{$}}

>From 9a7769622ad5c422bc522f397a8fefcf4e845a7e Mon Sep 17 00:00:00 2001
From: Valentin Clement <clementval at gmail.com>
Date: Thu, 8 Aug 2024 11:40:25 -0700
Subject: [PATCH 2/2] Allow UNIFIED

---
 flang/include/flang/Evaluate/tools.h | 12 ++++++++++++
 flang/lib/Evaluate/tools.cpp         |  6 +++---
 2 files changed, 15 insertions(+), 3 deletions(-)

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 e042d4d5a09d5d..1bfd253acc8e11 100644
--- a/flang/lib/Evaluate/tools.cpp
+++ b/flang/lib/Evaluate/tools.cpp
@@ -1697,7 +1697,7 @@ bool IsSaved(const Symbol &original) {
           (features.IsEnabled(
                common::LanguageFeature::SaveBigMainProgramVariables) &&
               symbol.size() > 32)) &&
-      !Fortran::evaluate::IsCUDADeviceSymbol(symbol)) {
+      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
@@ -1715,7 +1715,7 @@ bool IsSaved(const Symbol &original) {
   } else if (symbol.test(Symbol::Flag::InDataStmt)) {
     return true;
   } else if (const auto *object{symbol.detailsIf<ObjectEntityDetails>()};
-             object && object->init()) {
+      object && object->init()) {
     return true;
   } else if (IsProcedurePointer(symbol) && symbol.has<ProcEntityDetails>() &&
       symbol.get<ProcEntityDetails>().init()) {
@@ -1723,7 +1723,7 @@ bool IsSaved(const Symbol &original) {
   } else if (scope.hasSAVE()) {
     return true; // bare SAVE statement
   } else if (const Symbol * block{FindCommonBlockContaining(symbol)};
-             block && block->attrs().test(Attr::SAVE)) {
+      block && block->attrs().test(Attr::SAVE)) {
     return true; // in COMMON with SAVE
   } else {
     return false;