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/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