diff --git a/flang-rt/lib/runtime/CMakeLists.txt b/flang-rt/lib/runtime/CMakeLists.txt index 332c0872e065f..5200b2b710a5e 100644 --- a/flang-rt/lib/runtime/CMakeLists.txt +++ b/flang-rt/lib/runtime/CMakeLists.txt @@ -21,6 +21,7 @@ set(supported_sources allocatable.cpp array-constructor.cpp assign.cpp + assign_omp.cpp buffer.cpp character.cpp connection.cpp @@ -99,6 +100,7 @@ set(gpu_sources allocatable.cpp array-constructor.cpp assign.cpp + assign_omp.cpp buffer.cpp character.cpp connection.cpp diff --git a/flang-rt/lib/runtime/assign_omp.cpp b/flang-rt/lib/runtime/assign_omp.cpp new file mode 100644 index 0000000000000..10eb22a2650e1 --- /dev/null +++ b/flang-rt/lib/runtime/assign_omp.cpp @@ -0,0 +1,77 @@ +//===-- lib/runtime/assign_omp.cpp ----------------------------------*- C++ +//-*-===// +// +// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. +// See https://llvm.org/LICENSE.txt for license information. +// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +// +//===----------------------------------------------------------------------===// + +#include "flang-rt/runtime/assign-impl.h" +#include "flang-rt/runtime/derived.h" +#include "flang-rt/runtime/descriptor.h" +#include "flang-rt/runtime/stat.h" +#include "flang-rt/runtime/terminator.h" +#include "flang-rt/runtime/tools.h" +#include "flang-rt/runtime/type-info.h" +#include "flang/Runtime/assign.h" + +#include + +namespace Fortran::runtime { +namespace omp { + +typedef int32_t OMPDeviceTy; + +template static T *getDevicePtr(T *anyPtr, OMPDeviceTy ompDevice) { + auto voidAnyPtr = reinterpret_cast(anyPtr); + // If not present on the device it should already be a device ptr + if (!omp_target_is_present(voidAnyPtr, ompDevice)) + return anyPtr; + T *device_ptr = omp_get_mapped_ptr(anyPtr, ompDevice); + return device_ptr; +} + +RT_API_ATTRS static void Assign(Descriptor &to, const Descriptor &from, + Terminator &terminator, int flags, OMPDeviceTy omp_device) { + std::size_t toElementBytes{to.ElementBytes()}; + std::size_t fromElementBytes{from.ElementBytes()}; + std::size_t toElements{to.Elements()}; + std::size_t fromElements{from.Elements()}; + + if (toElementBytes != fromElementBytes) + terminator.Crash("Assign: toElementBytes != fromElementBytes"); + if (toElements != fromElements) + terminator.Crash("Assign: toElements != fromElements"); + + // Get base addresses and calculate length + void *to_base = to.raw().base_addr; + void *from_base = from.raw().base_addr; + size_t length = toElements * toElementBytes; + + // Get device pointers after ensuring data is on device + void *to_ptr = getDevicePtr(to_base, omp_device); + void *from_ptr = getDevicePtr(from_base, omp_device); + + // Perform copy between device pointers + int result = omp_target_memcpy(to_ptr, from_ptr, length, + /*dst_offset*/ 0, /*src_offset*/ 0, omp_device, omp_device); + + if (result != 0) + terminator.Crash("Assign: omp_target_memcpy failed"); + return; +} + +extern "C" { +RT_EXT_API_GROUP_BEGIN +void RTDEF(Assign_omp)(Descriptor &to, const Descriptor &from, + const char *sourceFile, int sourceLine, omp::OMPDeviceTy omp_device) { + Terminator terminator{sourceFile, sourceLine}; + Fortran::runtime::omp::Assign(to, from, terminator, + MaybeReallocate | NeedFinalization | ComponentCanBeDefinedAssignment, + omp_device); +} + +} // extern "C" +} // namespace omp +} // namespace Fortran::runtime diff --git a/flang/include/flang/Runtime/assign.h b/flang/include/flang/Runtime/assign.h index 7d198bdcc9e89..b510d2bbfcc8b 100644 --- a/flang/include/flang/Runtime/assign.h +++ b/flang/include/flang/Runtime/assign.h @@ -56,6 +56,9 @@ extern "C" { // API for lowering assignment void RTDECL(Assign)(Descriptor &to, const Descriptor &from, const char *sourceFile = nullptr, int sourceLine = 0); +void RTDECL(Assign_omp)(Descriptor &to, const Descriptor &from, + const char *sourceFile = nullptr, int sourceLine = 0, + int32_t omp_device = 0); // This variant has no finalization, defined assignment, or allocatable // reallocation. void RTDECL(AssignTemporary)(Descriptor &to, const Descriptor &from,