177 lines · cpp
1//===-- lib/cuda/memory.cpp -------------------------------------*- C++ -*-===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8 9#include "flang/Runtime/CUDA/memory.h"10#include "flang-rt/runtime/assign-impl.h"11#include "flang-rt/runtime/descriptor.h"12#include "flang-rt/runtime/environment.h"13#include "flang-rt/runtime/terminator.h"14#include "flang/Runtime/CUDA/common.h"15#include "flang/Runtime/CUDA/descriptor.h"16#include "flang/Runtime/CUDA/memmove-function.h"17#include "flang/Runtime/assign.h"18 19#include "cuda_runtime.h"20 21namespace Fortran::runtime::cuda {22 23extern "C" {24 25void *RTDEF(CUFMemAlloc)(26 std::size_t bytes, unsigned type, const char *sourceFile, int sourceLine) {27 void *ptr = nullptr;28 bytes = bytes ? bytes : 1;29 if (type == kMemTypeDevice) {30 if (Fortran::runtime::executionEnvironment.cudaDeviceIsManaged) {31 CUDA_REPORT_IF_ERROR(32 cudaMallocManaged((void **)&ptr, bytes, cudaMemAttachGlobal));33 } else {34 CUDA_REPORT_IF_ERROR(cudaMalloc((void **)&ptr, bytes));35 }36 } else if (type == kMemTypeManaged || type == kMemTypeUnified) {37 CUDA_REPORT_IF_ERROR(38 cudaMallocManaged((void **)&ptr, bytes, cudaMemAttachGlobal));39 } else if (type == kMemTypePinned) {40 CUDA_REPORT_IF_ERROR(cudaMallocHost((void **)&ptr, bytes));41 } else {42 Terminator terminator{sourceFile, sourceLine};43 terminator.Crash("unsupported memory type");44 }45 return ptr;46}47 48void RTDEF(CUFMemFree)(49 void *ptr, unsigned type, const char *sourceFile, int sourceLine) {50 if (!ptr)51 return;52 if (type == kMemTypeDevice || type == kMemTypeManaged ||53 type == kMemTypeUnified) {54 CUDA_REPORT_IF_ERROR(cudaFree(ptr));55 } else if (type == kMemTypePinned) {56 CUDA_REPORT_IF_ERROR(cudaFreeHost(ptr));57 } else {58 Terminator terminator{sourceFile, sourceLine};59 terminator.Crash("unsupported memory type");60 }61}62 63void RTDEF(CUFMemsetDescriptor)(64 Descriptor *desc, void *value, const char *sourceFile, int sourceLine) {65 Terminator terminator{sourceFile, sourceLine};66 terminator.Crash("not yet implemented: CUDA data transfer from a scalar "67 "value to a descriptor");68}69 70void RTDEF(CUFDataTransferPtrPtr)(void *dst, void *src, std::size_t bytes,71 unsigned mode, const char *sourceFile, int sourceLine) {72 cudaMemcpyKind kind;73 if (mode == kHostToDevice) {74 kind = cudaMemcpyHostToDevice;75 } else if (mode == kDeviceToHost) {76 kind = cudaMemcpyDeviceToHost;77 } else if (mode == kDeviceToDevice) {78 kind = cudaMemcpyDeviceToDevice;79 } else {80 Terminator terminator{sourceFile, sourceLine};81 terminator.Crash("host to host copy not supported");82 }83 // TODO: Use cudaMemcpyAsync when we have support for stream.84 CUDA_REPORT_IF_ERROR(cudaMemcpy(dst, src, bytes, kind));85}86 87void RTDEF(CUFDataTransferPtrDesc)(void *addr, Descriptor *desc,88 std::size_t bytes, unsigned mode, const char *sourceFile, int sourceLine) {89 Terminator terminator{sourceFile, sourceLine};90 terminator.Crash(91 "not yet implemented: CUDA data transfer from a descriptor to a pointer");92}93 94void RTDECL(CUFDataTransferDescDesc)(Descriptor *dstDesc, Descriptor *srcDesc,95 unsigned mode, const char *sourceFile, int sourceLine) {96 MemmoveFct memmoveFct;97 Terminator terminator{sourceFile, sourceLine};98 if (mode == kHostToDevice) {99 memmoveFct = &MemmoveHostToDevice;100 } else if (mode == kDeviceToHost) {101 memmoveFct = &MemmoveDeviceToHost;102 } else if (mode == kDeviceToDevice) {103 memmoveFct = &MemmoveDeviceToDevice;104 } else {105 terminator.Crash("host to host copy not supported");106 }107 // Allocate dst descriptor if not allocated.108 if (!dstDesc->IsAllocated()) {109 dstDesc->ApplyMold(*srcDesc, dstDesc->rank());110 dstDesc->Allocate(/*asyncObject=*/nullptr);111 }112 if ((srcDesc->rank() > 0) && (dstDesc->Elements() <= srcDesc->Elements()) &&113 srcDesc->IsContiguous() && dstDesc->IsContiguous()) {114 // Special case when rhs is bigger than lhs and both are contiguous arrays.115 // In this case we do a simple ptr to ptr transfer with the size of lhs.116 // This is be allowed in the reference compiler and it avoids error117 // triggered in the Assign runtime function used for the main case below.118 RTNAME(CUFDataTransferPtrPtr)(dstDesc->raw().base_addr,119 srcDesc->raw().base_addr, dstDesc->Elements() * dstDesc->ElementBytes(),120 mode, sourceFile, sourceLine);121 } else {122 Fortran::runtime::Assign(123 *dstDesc, *srcDesc, terminator, MaybeReallocate, memmoveFct);124 }125}126 127void RTDECL(CUFDataTransferCstDesc)(Descriptor *dstDesc, Descriptor *srcDesc,128 unsigned mode, const char *sourceFile, int sourceLine) {129 MemmoveFct memmoveFct;130 Terminator terminator{sourceFile, sourceLine};131 if (mode == kHostToDevice) {132 memmoveFct = &MemmoveHostToDevice;133 } else if (mode == kDeviceToHost) {134 memmoveFct = &MemmoveDeviceToHost;135 } else if (mode == kDeviceToDevice) {136 memmoveFct = &MemmoveDeviceToDevice;137 } else {138 terminator.Crash("host to host copy not supported");139 }140 141 Fortran::runtime::DoFromSourceAssign(142 *dstDesc, *srcDesc, terminator, memmoveFct);143}144 145void RTDECL(CUFDataTransferDescDescNoRealloc)(Descriptor *dstDesc,146 Descriptor *srcDesc, unsigned mode, const char *sourceFile,147 int sourceLine) {148 MemmoveFct memmoveFct;149 Terminator terminator{sourceFile, sourceLine};150 if (mode == kHostToDevice) {151 memmoveFct = &MemmoveHostToDevice;152 } else if (mode == kDeviceToHost) {153 memmoveFct = &MemmoveDeviceToHost;154 } else if (mode == kDeviceToDevice) {155 memmoveFct = &MemmoveDeviceToDevice;156 } else {157 terminator.Crash("host to host copy not supported");158 }159 Fortran::runtime::Assign(160 *dstDesc, *srcDesc, terminator, NoAssignFlags, memmoveFct);161}162 163void RTDECL(CUFDataTransferGlobalDescDesc)(Descriptor *dstDesc,164 Descriptor *srcDesc, unsigned mode, const char *sourceFile,165 int sourceLine) {166 RTNAME(CUFDataTransferDescDesc)167 (dstDesc, srcDesc, mode, sourceFile, sourceLine);168 if ((mode == kHostToDevice) || (mode == kDeviceToDevice)) {169 void *deviceAddr{170 RTNAME(CUFGetDeviceAddress)((void *)dstDesc, sourceFile, sourceLine)};171 RTNAME(CUFDescriptorSync)172 ((Descriptor *)deviceAddr, dstDesc, sourceFile, sourceLine);173 }174}175}176} // namespace Fortran::runtime::cuda177