[flang-commits] [flang] [flang] [cuda] implicitly set DEVICE attribute to scalars in device routines (PR #140834)

[Zhen Wang via flang-commits]

https://github.com/wangzpgi updated https://github.com/llvm/llvm-project/pull/140834

>From 564ff8f169f7807dedd95fe2d3eb995c0472f277 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Tue, 20 May 2025 19:47:59 -0700
Subject: [PATCH 1/2] implicitly set DEVICE attribute to scalars in device
 routines

---
 flang/lib/Semantics/resolve-names.cpp |  2 +-
 flang/test/Lower/CUDA/cuda-shared.cuf |  1 +
 flang/test/Semantics/cuf21.cuf        | 38 +++++++++++++++++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 flang/test/Semantics/cuf21.cuf

diff --git a/flang/lib/Semantics/resolve-names.cpp b/flang/lib/Semantics/resolve-names.cpp
index 92a3277191ae0..3f4a06444c4f3 100644
--- a/flang/lib/Semantics/resolve-names.cpp
+++ b/flang/lib/Semantics/resolve-names.cpp
@@ -9376,7 +9376,7 @@ static void SetImplicitCUDADevice(bool inDeviceSubprogram, Symbol &symbol) {
   if (inDeviceSubprogram && symbol.has<ObjectEntityDetails>()) {
     auto *object{symbol.detailsIf<ObjectEntityDetails>()};
     if (!object->cudaDataAttr() && !IsValue(symbol) &&
-        (IsDummy(symbol) || object->IsArray())) {
+        !IsFunctionResult(symbol)) {
       // Implicitly set device attribute if none is set in device context.
       object->set_cudaDataAttr(common::CUDADataAttr::Device);
     }
diff --git a/flang/test/Lower/CUDA/cuda-shared.cuf b/flang/test/Lower/CUDA/cuda-shared.cuf
index f41011df06ae7..565857f01bdb8 100644
--- a/flang/test/Lower/CUDA/cuda-shared.cuf
+++ b/flang/test/Lower/CUDA/cuda-shared.cuf
@@ -9,4 +9,5 @@ end subroutine
 
 ! CHECK-LABEL: func.func @_QPsharedmem() attributes {cuf.proc_attr = #cuf.cuda_proc<global>}
 ! CHECK: %{{.*}} = cuf.shared_memory !fir.array<32xf32> {bindc_name = "s", uniq_name = "_QFsharedmemEs"} -> !fir.ref<!fir.array<32xf32>>
+! CHECK: cuf.free %{{.*}}#0 : !fir.ref<i32> {data_attr = #cuf.cuda<device>}
 ! CHECK-NOT: cuf.free
diff --git a/flang/test/Semantics/cuf21.cuf b/flang/test/Semantics/cuf21.cuf
new file mode 100644
index 0000000000000..52343daaf66f1
--- /dev/null
+++ b/flang/test/Semantics/cuf21.cuf
@@ -0,0 +1,38 @@
+! RUN: bbc -emit-hlfir -fcuda %s -o - | FileCheck %s
+
+module mlocModule
+  interface maxlocUpdate
+     module procedure :: &
+          maxlocUpdateR_32F, &
+          maxlocUpdateR_64F, &
+          maxlocUpdateR_32I, &
+          maxlocUpdateR_64I
+  end interface maxlocUpdate
+contains
+
+  attributes(global) subroutine maxlocPartialMaskR_32F1D()
+    implicit none
+    real(4) :: mval
+
+    call maxlocUpdate(mval)
+
+  end subroutine maxlocPartialMaskR_32F1D
+
+  attributes(device) subroutine maxlocUpdateR_32F(mval)
+    real(4) :: mval
+  end subroutine maxlocUpdateR_32F
+
+  attributes(device) subroutine maxlocUpdateR_64F(mval)
+    real(8) :: mval
+  end subroutine maxlocUpdateR_64F
+
+  attributes(device) subroutine maxlocUpdateR_32I(mval)
+    integer(4) :: mval
+  end subroutine maxlocUpdateR_32I
+
+  attributes(device) subroutine maxlocUpdateR_64I(mval)
+    integer(8) :: mval
+  end subroutine maxlocUpdateR_64I
+end module
+
+! CHECK-LABEL: func.func @_QMmlocmodulePmaxlocpartialmaskr_32f1d()

>From ab3d3013ff32778652add6c1fced10af28b772b5 Mon Sep 17 00:00:00 2001
From: Zhen Wang <zhenw at nvidia.com>
Date: Wed, 21 May 2025 09:12:24 -0700
Subject: [PATCH 2/2] shorten test case

---
 flang/test/Semantics/cuf21.cuf | 12 +-----------
 1 file changed, 1 insertion(+), 11 deletions(-)

diff --git a/flang/test/Semantics/cuf21.cuf b/flang/test/Semantics/cuf21.cuf
index 52343daaf66f1..8f9c4f7064bb2 100644
--- a/flang/test/Semantics/cuf21.cuf
+++ b/flang/test/Semantics/cuf21.cuf
@@ -4,9 +4,7 @@ module mlocModule
   interface maxlocUpdate
      module procedure :: &
           maxlocUpdateR_32F, &
-          maxlocUpdateR_64F, &
-          maxlocUpdateR_32I, &
-          maxlocUpdateR_64I
+          maxlocUpdateR_64F
   end interface maxlocUpdate
 contains
 
@@ -25,14 +23,6 @@ contains
   attributes(device) subroutine maxlocUpdateR_64F(mval)
     real(8) :: mval
   end subroutine maxlocUpdateR_64F
-
-  attributes(device) subroutine maxlocUpdateR_32I(mval)
-    integer(4) :: mval
-  end subroutine maxlocUpdateR_32I
-
-  attributes(device) subroutine maxlocUpdateR_64I(mval)
-    integer(8) :: mval
-  end subroutine maxlocUpdateR_64I
 end module
 
 ! CHECK-LABEL: func.func @_QMmlocmodulePmaxlocpartialmaskr_32f1d()