xref: /llvm-project/flang/runtime/CUDA/memory.cpp (revision 5802367ddb46bcdeb0befeffbc99a1d72a5d9082)
1 //===-- runtime/CUDA/memory.cpp -------------------------------------------===//
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-exception
6 //
7 //===----------------------------------------------------------------------===//
8 
9 #include "flang/Runtime/CUDA/memory.h"
10 #include "../assign-impl.h"
11 #include "../terminator.h"
12 #include "flang/Runtime/CUDA/common.h"
13 #include "flang/Runtime/CUDA/descriptor.h"
14 #include "flang/Runtime/CUDA/memmove-function.h"
15 #include "flang/Runtime/assign.h"
16 
17 #include "cuda_runtime.h"
18 
19 namespace Fortran::runtime::cuda {
20 
21 extern "C" {
22 
23 void *RTDEF(CUFMemAlloc)(
24     std::size_t bytes, unsigned type, const char *sourceFile, int sourceLine) {
25   void *ptr = nullptr;
26   if (bytes != 0) {
27     if (type == kMemTypeDevice) {
28       CUDA_REPORT_IF_ERROR(cudaMalloc((void **)&ptr, bytes));
29     } else if (type == kMemTypeManaged || type == kMemTypeUnified) {
30       CUDA_REPORT_IF_ERROR(
31           cudaMallocManaged((void **)&ptr, bytes, cudaMemAttachGlobal));
32     } else if (type == kMemTypePinned) {
33       CUDA_REPORT_IF_ERROR(cudaMallocHost((void **)&ptr, bytes));
34     } else {
35       Terminator terminator{sourceFile, sourceLine};
36       terminator.Crash("unsupported memory type");
37     }
38   }
39   return ptr;
40 }
41 
42 void RTDEF(CUFMemFree)(
43     void *ptr, unsigned type, const char *sourceFile, int sourceLine) {
44   if (!ptr)
45     return;
46   if (type == kMemTypeDevice || type == kMemTypeManaged ||
47       type == kMemTypeUnified) {
48     CUDA_REPORT_IF_ERROR(cudaFree(ptr));
49   } else if (type == kMemTypePinned) {
50     CUDA_REPORT_IF_ERROR(cudaFreeHost(ptr));
51   } else {
52     Terminator terminator{sourceFile, sourceLine};
53     terminator.Crash("unsupported memory type");
54   }
55 }
56 
57 void RTDEF(CUFMemsetDescriptor)(
58     Descriptor *desc, void *value, const char *sourceFile, int sourceLine) {
59   Terminator terminator{sourceFile, sourceLine};
60   terminator.Crash("not yet implemented: CUDA data transfer from a scalar "
61                    "value to a descriptor");
62 }
63 
64 void RTDEF(CUFDataTransferPtrPtr)(void *dst, void *src, std::size_t bytes,
65     unsigned mode, const char *sourceFile, int sourceLine) {
66   cudaMemcpyKind kind;
67   if (mode == kHostToDevice) {
68     kind = cudaMemcpyHostToDevice;
69   } else if (mode == kDeviceToHost) {
70     kind = cudaMemcpyDeviceToHost;
71   } else if (mode == kDeviceToDevice) {
72     kind = cudaMemcpyDeviceToDevice;
73   } else {
74     Terminator terminator{sourceFile, sourceLine};
75     terminator.Crash("host to host copy not supported");
76   }
77   // TODO: Use cudaMemcpyAsync when we have support for stream.
78   CUDA_REPORT_IF_ERROR(cudaMemcpy(dst, src, bytes, kind));
79 }
80 
81 void RTDEF(CUFDataTransferPtrDesc)(void *addr, Descriptor *desc,
82     std::size_t bytes, unsigned mode, const char *sourceFile, int sourceLine) {
83   Terminator terminator{sourceFile, sourceLine};
84   terminator.Crash(
85       "not yet implemented: CUDA data transfer from a descriptor to a pointer");
86 }
87 
88 void RTDECL(CUFDataTransferDescDesc)(Descriptor *dstDesc, Descriptor *srcDesc,
89     unsigned mode, const char *sourceFile, int sourceLine) {
90   MemmoveFct memmoveFct;
91   Terminator terminator{sourceFile, sourceLine};
92   if (mode == kHostToDevice) {
93     memmoveFct = &MemmoveHostToDevice;
94   } else if (mode == kDeviceToHost) {
95     memmoveFct = &MemmoveDeviceToHost;
96   } else if (mode == kDeviceToDevice) {
97     memmoveFct = &MemmoveDeviceToDevice;
98   } else {
99     terminator.Crash("host to host copy not supported");
100   }
101   Fortran::runtime::Assign(
102       *dstDesc, *srcDesc, terminator, MaybeReallocate, memmoveFct);
103 }
104 
105 void RTDECL(CUFDataTransferCstDesc)(Descriptor *dstDesc, Descriptor *srcDesc,
106     unsigned mode, const char *sourceFile, int sourceLine) {
107   MemmoveFct memmoveFct;
108   Terminator terminator{sourceFile, sourceLine};
109   if (mode == kHostToDevice) {
110     memmoveFct = &MemmoveHostToDevice;
111   } else if (mode == kDeviceToHost) {
112     memmoveFct = &MemmoveDeviceToHost;
113   } else if (mode == kDeviceToDevice) {
114     memmoveFct = &MemmoveDeviceToDevice;
115   } else {
116     terminator.Crash("host to host copy not supported");
117   }
118 
119   Fortran::runtime::DoFromSourceAssign(
120       *dstDesc, *srcDesc, terminator, memmoveFct);
121 }
122 
123 void RTDECL(CUFDataTransferDescDescNoRealloc)(Descriptor *dstDesc,
124     Descriptor *srcDesc, unsigned mode, const char *sourceFile,
125     int sourceLine) {
126   MemmoveFct memmoveFct;
127   Terminator terminator{sourceFile, sourceLine};
128   if (mode == kHostToDevice) {
129     memmoveFct = &MemmoveHostToDevice;
130   } else if (mode == kDeviceToHost) {
131     memmoveFct = &MemmoveDeviceToHost;
132   } else if (mode == kDeviceToDevice) {
133     memmoveFct = &MemmoveDeviceToDevice;
134   } else {
135     terminator.Crash("host to host copy not supported");
136   }
137   Fortran::runtime::Assign(
138       *dstDesc, *srcDesc, terminator, NoAssignFlags, memmoveFct);
139 }
140 
141 void RTDECL(CUFDataTransferGlobalDescDesc)(Descriptor *dstDesc,
142     Descriptor *srcDesc, unsigned mode, const char *sourceFile,
143     int sourceLine) {
144   RTNAME(CUFDataTransferDescDesc)
145   (dstDesc, srcDesc, mode, sourceFile, sourceLine);
146   if ((mode == kHostToDevice) || (mode == kDeviceToDevice)) {
147     void *deviceAddr{
148         RTNAME(CUFGetDeviceAddress)((void *)dstDesc, sourceFile, sourceLine)};
149     RTNAME(CUFDescriptorSync)
150     ((Descriptor *)deviceAddr, srcDesc, sourceFile, sourceLine);
151   }
152 }
153 }
154 } // namespace Fortran::runtime::cuda
155