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

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

Author: clementval

flang/lib/Semantics/check-allocate.cpp

@@ -618,6 +618,13 @@ bool AllocationCheckerHelper::RunChecks(SemanticsContext &context) {
           "Object in ALLOCATE must have PINNED attribute when PINNED option is specified"_err_en_US);
     }
   }
+  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);
 }
 

flang/test/Lower/CUDA/cuda-allocatable.cuf

@@ -68,21 +68,21 @@ end subroutine
 ! CHECK: }
 
 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
 ! CHECK: fir.if %{{.*}} {
-! CHECK:   %{{.*}} = fir.cuda_deallocate %[[BOX_DECL]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {cuda_attr = #fir.cuda<unified>} -> i32
+! CHECK:   %{{.*}} = fir.cuda_deallocate %[[BOX_DECL]]#1 : !fir.ref<!fir.box<!fir.heap<!fir.array<?xf32>>>> {cuda_attr = #fir.cuda<device>} -> i32
 ! CHECK: }
 
 subroutine sub5()

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

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 -> 

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

flang/test/Semantics/cuf07.cuf

@@ -31,4 +31,12 @@ module m
     !ERROR: Object in ALLOCATE must have PINNED attribute when PINNED option is specified
     allocate(ia(100), pinned = plog)
   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