202 lines · cpp
1//===-- lib/cuda/allocator.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/allocator.h"10#include "flang-rt/runtime/allocator-registry.h"11#include "flang-rt/runtime/derived.h"12#include "flang-rt/runtime/descriptor.h"13#include "flang-rt/runtime/environment.h"14#include "flang-rt/runtime/lock.h"15#include "flang-rt/runtime/stat.h"16#include "flang-rt/runtime/terminator.h"17#include "flang-rt/runtime/type-info.h"18#include "flang/Common/ISO_Fortran_binding_wrapper.h"19#include "flang/Runtime/CUDA/common.h"20#include "flang/Support/Fortran.h"21 22namespace Fortran::runtime::cuda {23 24struct DeviceAllocation {25 void *ptr;26 std::size_t size;27 cudaStream_t stream;28};29 30// Compare address values. nullptr will be sorted at the end of the array.31int compareDeviceAlloc(const void *a, const void *b) {32 const DeviceAllocation *deva = (const DeviceAllocation *)a;33 const DeviceAllocation *devb = (const DeviceAllocation *)b;34 if (deva->ptr == nullptr && devb->ptr == nullptr)35 return 0;36 if (deva->ptr == nullptr)37 return 1;38 if (devb->ptr == nullptr)39 return -1;40 return deva->ptr < devb->ptr ? -1 : (deva->ptr > devb->ptr ? 1 : 0);41}42 43// Dynamic array for tracking asynchronous allocations.44static DeviceAllocation *deviceAllocations = nullptr;45Lock lock;46static int maxDeviceAllocations{512}; // Initial size47static int numDeviceAllocations{0};48static constexpr int allocNotFound{-1};49 50static void initAllocations() {51 if (!deviceAllocations) {52 deviceAllocations = static_cast<DeviceAllocation *>(53 malloc(maxDeviceAllocations * sizeof(DeviceAllocation)));54 if (!deviceAllocations) {55 Terminator terminator{__FILE__, __LINE__};56 terminator.Crash("Failed to allocate tracking array");57 }58 }59}60 61static void doubleAllocationArray() {62 unsigned newSize = maxDeviceAllocations * 2;63 DeviceAllocation *newArray = static_cast<DeviceAllocation *>(64 realloc(deviceAllocations, newSize * sizeof(DeviceAllocation)));65 if (!newArray) {66 Terminator terminator{__FILE__, __LINE__};67 terminator.Crash("Failed to reallocate tracking array");68 }69 deviceAllocations = newArray;70 maxDeviceAllocations = newSize;71}72 73static unsigned findAllocation(void *ptr) {74 if (numDeviceAllocations == 0) {75 return allocNotFound;76 }77 78 int left{0};79 int right{numDeviceAllocations - 1};80 81 if (left == right) {82 return left;83 }84 85 while (left <= right) {86 int mid = left + (right - left) / 2;87 if (deviceAllocations[mid].ptr == ptr) {88 return mid;89 }90 if (deviceAllocations[mid].ptr < ptr) {91 left = mid + 1;92 } else {93 right = mid - 1;94 }95 }96 return allocNotFound;97}98 99static void insertAllocation(void *ptr, std::size_t size, cudaStream_t stream) {100 CriticalSection critical{lock};101 initAllocations();102 if (numDeviceAllocations >= maxDeviceAllocations) {103 doubleAllocationArray();104 }105 deviceAllocations[numDeviceAllocations].ptr = ptr;106 deviceAllocations[numDeviceAllocations].size = size;107 deviceAllocations[numDeviceAllocations].stream = stream;108 ++numDeviceAllocations;109 qsort(deviceAllocations, numDeviceAllocations, sizeof(DeviceAllocation),110 compareDeviceAlloc);111}112 113static void eraseAllocation(int pos) {114 deviceAllocations[pos].ptr = nullptr;115 deviceAllocations[pos].size = 0;116 deviceAllocations[pos].stream = (cudaStream_t)0;117 qsort(deviceAllocations, numDeviceAllocations, sizeof(DeviceAllocation),118 compareDeviceAlloc);119 --numDeviceAllocations;120}121 122extern "C" {123 124void RTDEF(CUFRegisterAllocator)() {125 allocatorRegistry.Register(126 kPinnedAllocatorPos, {&CUFAllocPinned, CUFFreePinned});127 allocatorRegistry.Register(128 kDeviceAllocatorPos, {&CUFAllocDevice, CUFFreeDevice});129 allocatorRegistry.Register(130 kManagedAllocatorPos, {&CUFAllocManaged, CUFFreeManaged});131 allocatorRegistry.Register(132 kUnifiedAllocatorPos, {&CUFAllocUnified, CUFFreeUnified});133}134 135cudaStream_t RTDECL(CUFGetAssociatedStream)(void *p) {136 int pos = findAllocation(p);137 if (pos >= 0) {138 cudaStream_t stream = deviceAllocations[pos].stream;139 return stream;140 }141 return nullptr;142}143}144 145void *CUFAllocPinned(146 std::size_t sizeInBytes, [[maybe_unused]] std::int64_t *asyncObject) {147 void *p;148 cudaMallocHost((void **)&p, sizeInBytes);149 return p;150}151 152void CUFFreePinned(void *p) { cudaFreeHost(p); }153 154void *CUFAllocDevice(std::size_t sizeInBytes, std::int64_t *asyncObject) {155 void *p;156 if (Fortran::runtime::executionEnvironment.cudaDeviceIsManaged) {157 cudaMallocManaged((void **)&p, sizeInBytes, cudaMemAttachGlobal);158 } else {159 if (asyncObject == nullptr) {160 cudaMalloc(&p, sizeInBytes);161 } else {162 cudaMallocAsync(&p, sizeInBytes, (cudaStream_t)*asyncObject);163 insertAllocation(p, sizeInBytes, (cudaStream_t)*asyncObject);164 }165 }166 return p;167}168 169void CUFFreeDevice(void *p) {170 CriticalSection critical{lock};171 int pos = findAllocation(p);172 if (pos >= 0) {173 cudaStream_t stream = deviceAllocations[pos].stream;174 eraseAllocation(pos);175 cudaFreeAsync(p, stream);176 } else {177 cudaFree(p);178 }179}180 181void *CUFAllocManaged(182 std::size_t sizeInBytes, [[maybe_unused]] std::int64_t *asyncObject) {183 void *p;184 cudaMallocManaged((void **)&p, sizeInBytes, cudaMemAttachGlobal);185 return reinterpret_cast<void *>(p);186}187 188void CUFFreeManaged(void *p) { cudaFree(p); }189 190void *CUFAllocUnified(191 std::size_t sizeInBytes, [[maybe_unused]] std::int64_t *asyncObject) {192 // Call alloc managed for the time being.193 return CUFAllocManaged(sizeInBytes, asyncObject);194}195 196void CUFFreeUnified(void *p) {197 // Call free managed for the time being.198 CUFFreeManaged(p);199}200 201} // namespace Fortran::runtime::cuda202