184 lines · cpp
1//===--- cuda/dynamic_cuda/cuda.pp ------------------------------- 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// Implement subset of cuda api by calling into cuda library via dlopen10// Does the dlopen/dlsym calls as part of the call to cuInit11//12//===----------------------------------------------------------------------===//13 14#include "llvm/Support/DynamicLibrary.h"15 16#include "Shared/Debug.h"17 18#include "DLWrap.h"19#include "cuda.h"20 21#include <memory>22#include <string>23#include <unordered_map>24 25DLWRAP_INITIALIZE()26 27DLWRAP_INTERNAL(cuInit, 1)28 29DLWRAP(cuCtxGetDevice, 1)30DLWRAP(cuDeviceGet, 2)31DLWRAP(cuDeviceGetAttribute, 3)32DLWRAP(cuDeviceGetCount, 1)33DLWRAP(cuFuncGetAttribute, 3)34DLWRAP(cuFuncSetAttribute, 3)35 36// Device info37DLWRAP(cuDeviceGetName, 3)38DLWRAP(cuDeviceGetUuid, 2)39DLWRAP(cuDeviceTotalMem, 2)40DLWRAP(cuDriverGetVersion, 1)41 42DLWRAP(cuGetErrorString, 2)43DLWRAP(cuLaunchKernel, 11)44DLWRAP(cuLaunchHostFunc, 3)45 46DLWRAP(cuMemAlloc, 2)47DLWRAP(cuMemAllocHost, 2)48DLWRAP(cuMemAllocManaged, 3)49DLWRAP(cuMemAllocAsync, 3)50 51DLWRAP(cuMemcpyDtoDAsync, 4)52DLWRAP(cuMemcpyDtoH, 3)53DLWRAP(cuMemcpyDtoHAsync, 4)54DLWRAP(cuMemcpyHtoD, 3)55DLWRAP(cuMemcpyHtoDAsync, 4)56 57DLWRAP(cuMemsetD8Async, 4)58DLWRAP(cuMemsetD16Async, 4)59DLWRAP(cuMemsetD32Async, 4)60DLWRAP(cuMemsetD2D8Async, 6)61DLWRAP(cuMemsetD2D16Async, 6)62DLWRAP(cuMemsetD2D32Async, 6)63 64DLWRAP(cuMemFree, 1)65DLWRAP(cuMemFreeHost, 1)66DLWRAP(cuMemFreeAsync, 2)67 68DLWRAP(cuModuleGetFunction, 3)69DLWRAP(cuModuleGetGlobal, 4)70 71DLWRAP(cuModuleUnload, 1)72DLWRAP(cuStreamCreate, 2)73DLWRAP(cuStreamDestroy, 1)74DLWRAP(cuStreamSynchronize, 1)75DLWRAP(cuStreamQuery, 1)76DLWRAP(cuStreamAddCallback, 4)77DLWRAP(cuCtxSetCurrent, 1)78DLWRAP(cuDevicePrimaryCtxRelease, 1)79DLWRAP(cuDevicePrimaryCtxGetState, 3)80DLWRAP(cuDevicePrimaryCtxSetFlags, 2)81DLWRAP(cuDevicePrimaryCtxRetain, 2)82DLWRAP(cuModuleLoadDataEx, 5)83DLWRAP(cuOccupancyMaxPotentialBlockSize, 6)84 85DLWRAP(cuDeviceCanAccessPeer, 3)86DLWRAP(cuCtxEnablePeerAccess, 2)87DLWRAP(cuMemcpyPeerAsync, 6)88 89DLWRAP(cuCtxGetLimit, 2)90DLWRAP(cuCtxSetLimit, 2)91 92DLWRAP(cuEventCreate, 2)93DLWRAP(cuEventRecord, 2)94DLWRAP(cuEventQuery, 1)95DLWRAP(cuStreamWaitEvent, 3)96DLWRAP(cuEventSynchronize, 1)97DLWRAP(cuEventDestroy, 1)98 99DLWRAP_FINALIZE()100 101DLWRAP(cuMemUnmap, 2)102DLWRAP(cuMemRelease, 1)103DLWRAP(cuMemAddressFree, 2)104DLWRAP(cuMemGetInfo, 2)105DLWRAP(cuMemAddressReserve, 5)106DLWRAP(cuMemMap, 5)107DLWRAP(cuMemCreate, 4)108DLWRAP(cuMemSetAccess, 4)109DLWRAP(cuMemGetAllocationGranularity, 3)110 111#ifndef DYNAMIC_CUDA_PATH112#define DYNAMIC_CUDA_PATH "libcuda.so"113#endif114 115#ifndef TARGET_NAME116#define TARGET_NAME CUDA117#endif118#ifndef DEBUG_PREFIX119#define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"120#endif121 122static bool checkForCUDA() {123 // return true if dlopen succeeded and all functions found124 125 // Prefer _v2 versions of functions if found in the library126 std::unordered_map<std::string, const char *> TryFirst = {127 {"cuMemAlloc", "cuMemAlloc_v2"},128 {"cuMemFree", "cuMemFree_v2"},129 {"cuMemcpyDtoH", "cuMemcpyDtoH_v2"},130 {"cuMemcpyHtoD", "cuMemcpyHtoD_v2"},131 {"cuStreamDestroy", "cuStreamDestroy_v2"},132 {"cuModuleGetGlobal", "cuModuleGetGlobal_v2"},133 {"cuMemcpyDtoHAsync", "cuMemcpyDtoHAsync_v2"},134 {"cuMemcpyDtoDAsync", "cuMemcpyDtoDAsync_v2"},135 {"cuMemcpyHtoDAsync", "cuMemcpyHtoDAsync_v2"},136 {"cuDevicePrimaryCtxRelease", "cuDevicePrimaryCtxRelease_v2"},137 {"cuDevicePrimaryCtxSetFlags", "cuDevicePrimaryCtxSetFlags_v2"},138 };139 140 const char *CudaLib = DYNAMIC_CUDA_PATH;141 std::string ErrMsg;142 auto DynlibHandle = std::make_unique<llvm::sys::DynamicLibrary>(143 llvm::sys::DynamicLibrary::getPermanentLibrary(CudaLib, &ErrMsg));144 if (!DynlibHandle->isValid()) {145 DP("Unable to load library '%s': %s!\n", CudaLib, ErrMsg.c_str());146 return false;147 }148 149 for (size_t I = 0; I < dlwrap::size(); I++) {150 const char *Sym = dlwrap::symbol(I);151 152 auto It = TryFirst.find(Sym);153 if (It != TryFirst.end()) {154 const char *First = It->second;155 void *P = DynlibHandle->getAddressOfSymbol(First);156 if (P) {157 DP("Implementing %s with dlsym(%s) -> %p\n", Sym, First, P);158 *dlwrap::pointer(I) = P;159 continue;160 }161 }162 163 void *P = DynlibHandle->getAddressOfSymbol(Sym);164 if (P == nullptr) {165 DP("Unable to find '%s' in '%s'!\n", Sym, CudaLib);166 return false;167 }168 DP("Implementing %s with dlsym(%s) -> %p\n", Sym, Sym, P);169 170 *dlwrap::pointer(I) = P;171 }172 173 return true;174}175 176CUresult cuInit(unsigned X) {177 // Note: Called exactly once from cuda rtl.cpp in a global constructor so178 // does not need to handle being called repeatedly or concurrently179 if (!checkForCUDA()) {180 return CUDA_ERROR_INVALID_HANDLE;181 }182 return dlwrap_cuInit(X);183}184