[flang-commits] [flang] [flang][cuda] Enforce DEVICE attribute when ALLOCATE with STREAM option (PR #89459)

[via flang-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

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

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

<details>
<summary>Changes</summary>

When the STREAM option is specified on an ALLOCATE statement, the object must have the DEVICE attribute.

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


6 Files Affected:

- (modified) flang/lib/Semantics/check-allocate.cpp (+7) 
- (modified) flang/test/Lower/CUDA/cuda-allocatable.cuf (+3-3) 
- (modified) flang/test/Parser/cuf-sanity-common (+1-1) 
- (modified) flang/test/Parser/cuf-sanity-tree.CUF (-2) 
- (modified) flang/test/Parser/cuf-sanity-unparse.CUF (+1-1) 
- (modified) flang/test/Semantics/cuf07.cuf (+8) 


``````````diff
diff --git a/flang/lib/Semantics/check-allocate.cpp b/flang/lib/Semantics/check-allocate.cpp
index a7244e1c58330a..364b5ece591f09 100644
--- a/flang/lib/Semantics/check-allocate.cpp
+++ b/flang/lib/Semantics/check-allocate.cpp
@@ -611,6 +611,13 @@ bool AllocationCheckerHelper::RunChecks(SemanticsContext &context) {
       return false;
     }
   }
+  if (allocateInfo_.gotStream) {
+    std::optional<common::CUDADataAttr> cudaAttr{GetCUDADataAttr(ultimate_)};
+    if (!cudaAttr || *cudaAttr != common::CUDADataAttr::Device) {
+      context.Say(name_.source,
+          "Object in ALLOCATE must have DEVICE attribute when STREAM option is specified"_err_en_US);
+    }
+  }
   return RunCoarrayRelatedChecks(context);
 }
 
diff --git a/flang/test/Lower/CUDA/cuda-allocatable.cuf b/flang/test/Lower/CUDA/cuda-allocatable.cuf
index 5b10334ecdbc14..251ff16a56c797 100644
--- a/flang/test/Lower/CUDA/cuda-allocatable.cuf
+++ b/flang/test/Lower/CUDA/cuda-allocatable.cuf
@@ -52,19 +52,19 @@ end subroutine
 ! CHECK: %{{.*}} = fir.cuda_allocate %[[BOX_DECL]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?x?xi32>>>> pinned(%[[PLOG_DECL]]#1 : !fir.ref<!fir.logical<4>>) {cuda_attr = #fir.cuda<pinned>} -> i32
 
 subroutine sub4()
-  real, allocatable, unified :: a(:)
+  real, allocatable, device :: a(:)
   integer :: istream
   allocate(a(10), stream=istream)
 end subroutine
 
 ! CHECK-LABEL: func.func @_QPsub4()
 ! CHECK: %[[BOX:.*]] = fir.alloca !fir.box<!fir.heap<!fir.array<?xf32>>> {bindc_name = "a", uniq_name = "_QFsub4Ea"}
-! CHECK: %[[BOX_DECL:.*]]:2 = hlfir.declare %0 {cuda_attr = #fir.cuda<unified>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFsub4Ea"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
+! CHECK: %[[BOX_DECL:.*]]:2 = hlfir.declare %0 {cuda_attr = #fir.cuda<device>, fortran_attrs = #fir.var_attrs<allocatable>, uniq_name = "_QFsub4Ea"} : (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>) -> (!fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>, !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>>)
 ! CHECK: %[[ISTREAM:.*]] = fir.alloca i32 {bindc_name = "istream", uniq_name = "_QFsub4Eistream"}
 ! CHECK: %[[ISTREAM_DECL:.*]]:2 = hlfir.declare %[[ISTREAM]] {uniq_name = "_QFsub4Eistream"} : (!fir.ref<i32>) -> (!fir.ref<i32>, !fir.ref<i32>)
 ! CHECK: fir.call @_FortranAAllocatableSetBounds
 ! CHECK: %[[STREAM:.*]] = fir.load %[[ISTREAM_DECL]]#0 : !fir.ref<i32>
-! CHECK: %{{.*}} = fir.cuda_allocate %[[BOX_DECL]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> stream(%[[STREAM]] : i32) {cuda_attr = #fir.cuda<unified>} -> i32
+! CHECK: %{{.*}} = fir.cuda_allocate %[[BOX_DECL]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> stream(%[[STREAM]] : i32) {cuda_attr = #fir.cuda<device>} -> i32
 
 subroutine sub5()
   real, allocatable, device :: a(:)
diff --git a/flang/test/Parser/cuf-sanity-common b/flang/test/Parser/cuf-sanity-common
index 7f4217fb58355d..b097a6aa300458 100644
--- a/flang/test/Parser/cuf-sanity-common
+++ b/flang/test/Parser/cuf-sanity-common
@@ -32,6 +32,6 @@ module m
     call globalsub<<<1, 2>>>
     call globalsub<<<1, 2, 3>>>
     call globalsub<<<1, 2, 3, 4>>>
-    allocate(pa(32), stream = 1, pinned = isPinned)
+    allocate(pa(32), pinned = isPinned)
   end subroutine
 end module
diff --git a/flang/test/Parser/cuf-sanity-tree.CUF b/flang/test/Parser/cuf-sanity-tree.CUF
index dc12759d3ce52f..2820441d5b5f0a 100644
--- a/flang/test/Parser/cuf-sanity-tree.CUF
+++ b/flang/test/Parser/cuf-sanity-tree.CUF
@@ -199,8 +199,6 @@ include "cuf-sanity-common"
 !CHECK: | | | | | | AllocateShapeSpec
 !CHECK: | | | | | | | Scalar -> Integer -> Expr = '32_4'
 !CHECK: | | | | | | | | LiteralConstant -> IntLiteralConstant = '32'
-!CHECK: | | | | | AllocOpt -> Stream -> Scalar -> Integer -> Expr = '1_4'
-!CHECK: | | | | | | LiteralConstant -> IntLiteralConstant = '1'
 !CHECK: | | | | | AllocOpt -> Pinned -> Scalar -> Logical -> Variable = 'ispinned'
 !CHECK: | | | | | | Designator -> DataRef -> Name = 'ispinned'
 !CHECK: | | | EndSubroutineStmt -> 
diff --git a/flang/test/Parser/cuf-sanity-unparse.CUF b/flang/test/Parser/cuf-sanity-unparse.CUF
index 7ac39448d7bd45..b6921e74fc05ae 100644
--- a/flang/test/Parser/cuf-sanity-unparse.CUF
+++ b/flang/test/Parser/cuf-sanity-unparse.CUF
@@ -37,6 +37,6 @@ include "cuf-sanity-common"
 !CHECK:    CALL globalsub<<<1_4,2_4>>>()
 !CHECK:    CALL globalsub<<<1_4,2_4,3_4>>>()
 !CHECK:    CALL globalsub<<<1_4,2_4,3_4,4_4>>>()
-!CHECK:   ALLOCATE(pa(32_4), STREAM=1_4, PINNED=ispinned)
+!CHECK:   ALLOCATE(pa(32_4), PINNED=ispinned)
 !CHECK:  END SUBROUTINE
 !CHECK: END MODULE
diff --git a/flang/test/Semantics/cuf07.cuf b/flang/test/Semantics/cuf07.cuf
index b520b5da51264b..7f9918f907bcb7 100644
--- a/flang/test/Semantics/cuf07.cuf
+++ b/flang/test/Semantics/cuf07.cuf
@@ -23,4 +23,12 @@ module m
     !BECAUSE: 'ma' is a host-associated allocatable and is not definable in a device subprogram
     deallocate(ma)
   end subroutine
+
+  subroutine host2()
+    integer, allocatable, pinned :: ia(:)
+    integer :: istream
+
+    !ERROR: Object in ALLOCATE must have DEVICE attribute when STREAM option is specified
+    allocate(ia(100), stream = istream)
+  end subroutine
 end module

``````````

</details>


https://github.com/llvm/llvm-project/pull/89459