diff --git a/flang/include/flang/Runtime/CUDA/pointer.h b/flang/include/flang/Runtime/CUDA/pointer.h index db5242696303f..2197d85f4b93e 100644 --- a/flang/include/flang/Runtime/CUDA/pointer.h +++ b/flang/include/flang/Runtime/CUDA/pointer.h @@ -21,6 +21,13 @@ int RTDECL(CUFPointerAllocate)(Descriptor &, int64_t stream = -1, bool hasStat = false, const Descriptor *errMsg = nullptr, const char *sourceFile = nullptr, int sourceLine = 0); +/// Perform allocation of the descriptor without synchronization. Assign data +/// from source. +int RTDEF(CUFPointerAllocateSource)(Descriptor &pointer, + const Descriptor &source, int64_t stream = -1, bool hasStat = false, + const Descriptor *errMsg = nullptr, const char *sourceFile = nullptr, + int sourceLine = 0); + } // extern "C" } // namespace Fortran::runtime::cuda diff --git a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp index 23248f6d12622..b0d6b0f0993a6 100644 --- a/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp +++ b/flang/lib/Optimizer/Transforms/CUFOpConversion.cpp @@ -189,11 +189,12 @@ struct CUFAllocateOpConversion mlir::func::FuncOp func; if (op.getSource()) { - if (isPointer) - TODO(loc, "pointer allocation with source"); func = - fir::runtime::getRuntimeFunc( - loc, builder); + isPointer + ? fir::runtime::getRuntimeFunc( + loc, builder) + : fir::runtime::getRuntimeFunc(loc, builder); } else { func = isPointer diff --git a/flang/runtime/CUDA/pointer.cpp b/flang/runtime/CUDA/pointer.cpp index 0c5d3a5a6297d..35f373b0a56c3 100644 --- a/flang/runtime/CUDA/pointer.cpp +++ b/flang/runtime/CUDA/pointer.cpp @@ -7,8 +7,10 @@ //===----------------------------------------------------------------------===// #include "flang/Runtime/CUDA/pointer.h" +#include "../assign-impl.h" #include "../stat.h" #include "../terminator.h" +#include "flang/Runtime/CUDA/memmove-function.h" #include "flang/Runtime/pointer.h" #include "cuda_runtime.h" @@ -33,6 +35,19 @@ int RTDEF(CUFPointerAllocate)(Descriptor &desc, int64_t stream, bool hasStat, return stat; } +int RTDEF(CUFPointerAllocateSource)(Descriptor &pointer, + const Descriptor &source, int64_t stream, bool hasStat, + const Descriptor *errMsg, const char *sourceFile, int sourceLine) { + int stat{RTNAME(CUFPointerAllocate)( + pointer, stream, hasStat, errMsg, sourceFile, sourceLine)}; + if (stat == StatOk) { + Terminator terminator{sourceFile, sourceLine}; + Fortran::runtime::DoFromSourceAssign( + pointer, source, terminator, &MemmoveHostToDevice); + } + return stat; +} + RT_EXT_API_GROUP_END } // extern "C" diff --git a/flang/test/Fir/CUDA/cuda-allocate.fir b/flang/test/Fir/CUDA/cuda-allocate.fir index 2ac9498d35541..804bb8636685d 100644 --- a/flang/test/Fir/CUDA/cuda-allocate.fir +++ b/flang/test/Fir/CUDA/cuda-allocate.fir @@ -192,4 +192,22 @@ func.func @_QPp_alloc() { // CHECK-LABEL: func.func @_QPp_alloc() // CHECK: fir.call @_FortranACUFPointerAllocate +func.func @_QPpointer_source() { + %c0_i64 = arith.constant 0 : i64 + %c1_i32 = arith.constant 1 : i32 + %c0_i32 = arith.constant 0 : i32 + %c1 = arith.constant 1 : index + %c0 = arith.constant 0 : index + %0 = fir.alloca !fir.box>> {bindc_name = "a", uniq_name = "_QFpointer_sourceEa"} + %4 = fir.declare %0 {fortran_attrs = #fir.var_attrs, uniq_name = "_QFpointer_sourceEa"} : (!fir.ref>>>) -> !fir.ref>>> + %5 = cuf.alloc !fir.box>> {bindc_name = "a_d", data_attr = #cuf.cuda, uniq_name = "_QFpointer_sourceEa_d"} -> !fir.ref>>> + %7 = fir.declare %5 {data_attr = #cuf.cuda, fortran_attrs = #fir.var_attrs, uniq_name = "_QFpointer_sourceEa_d"} : (!fir.ref>>>) -> !fir.ref>>> + %8 = fir.load %4 : !fir.ref>>> + %22 = cuf.allocate %7 : !fir.ref>>> source(%8 : !fir.box>>) {data_attr = #cuf.cuda} -> i32 + return +} + +// CHECK-LABEL: func.func @_QPpointer_source() +// CHECK: _FortranACUFPointerAllocateSource + } // end of module