flang/runtime/CUDA/memory.cpp

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