wehub-resource-sync
128 lines
4.5 KiB
C++
128 lines
4.5 KiB
C++
#pragma once
|
|
|
|
#include <cuda.h>
|
|
#include <deep_ep/common/exception.cuh>
|
|
|
|
#include "lazy_driver.hpp"
|
|
|
|
namespace deep_ep::shared_memory {
|
|
|
|
union MemHandleInner {
|
|
cudaIpcMemHandle_t cuda_ipc_mem_handle;
|
|
CUmemFabricHandle cu_mem_fabric_handle;
|
|
};
|
|
|
|
struct MemHandle {
|
|
MemHandleInner inner;
|
|
size_t size;
|
|
};
|
|
|
|
static void cu_mem_set_access_all(void* ptr, size_t size) {
|
|
int device_count;
|
|
CUDA_RUNTIME_CHECK(cudaGetDeviceCount(&device_count));
|
|
|
|
constexpr int kMaxDeviceCount = 8;
|
|
EP_HOST_ASSERT(0 < device_count and device_count <= kMaxDeviceCount);
|
|
|
|
CUmemAccessDesc access_desc[kMaxDeviceCount];
|
|
for (int i = 0; i < device_count; ++ i) {
|
|
access_desc[i].location.type = CU_MEM_LOCATION_TYPE_DEVICE;
|
|
access_desc[i].location.id = i;
|
|
access_desc[i].flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
|
|
}
|
|
CUDA_DRIVER_CHECK(lazy_cuMemSetAccess(reinterpret_cast<CUdeviceptr>(ptr), size, access_desc, device_count));
|
|
}
|
|
|
|
static void cu_mem_free(void* ptr) {
|
|
CUmemGenericAllocationHandle handle;
|
|
CUDA_DRIVER_CHECK(lazy_cuMemRetainAllocationHandle(&handle, ptr));
|
|
|
|
size_t size = 0;
|
|
CUDA_DRIVER_CHECK(lazy_cuMemGetAddressRange_v2(nullptr, &size, reinterpret_cast<CUdeviceptr>(ptr)));
|
|
|
|
CUDA_DRIVER_CHECK(lazy_cuMemUnmap(reinterpret_cast<CUdeviceptr>(ptr), size));
|
|
CUDA_DRIVER_CHECK(lazy_cuMemAddressFree(reinterpret_cast<CUdeviceptr>(ptr), size));
|
|
CUDA_DRIVER_CHECK(lazy_cuMemRelease(handle));
|
|
}
|
|
|
|
class SharedMemoryAllocator {
|
|
public:
|
|
explicit SharedMemoryAllocator(const bool& use_fabric) : use_fabric(use_fabric) {}
|
|
|
|
void malloc(void** ptr, size_t size) const {
|
|
if (use_fabric) {
|
|
CUdevice device;
|
|
CUDA_DRIVER_CHECK(lazy_cuCtxGetDevice(&device));
|
|
|
|
CUmemAllocationProp prop = {};
|
|
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
|
|
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
|
|
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
|
|
prop.location.id = device;
|
|
|
|
size_t alignment = 0;
|
|
EP_HOST_ASSERT(size > 0);
|
|
CUDA_DRIVER_CHECK(lazy_cuMemGetAllocationGranularity(&alignment, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
|
|
size = ((size + alignment - 1) / alignment) * alignment;
|
|
|
|
CUmemGenericAllocationHandle handle;
|
|
CUDA_DRIVER_CHECK(lazy_cuMemCreate(&handle, size, &prop, 0));
|
|
CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(reinterpret_cast<CUdeviceptr*>(ptr), size, alignment, 0, 0));
|
|
CUDA_DRIVER_CHECK(lazy_cuMemMap(reinterpret_cast<CUdeviceptr>(*ptr), size, 0, handle, 0));
|
|
cu_mem_set_access_all(*ptr, size);
|
|
} else {
|
|
CUDA_RUNTIME_CHECK(cudaMalloc(ptr, size));
|
|
}
|
|
}
|
|
|
|
void free(void* ptr) const {
|
|
if (use_fabric) {
|
|
cu_mem_free(ptr);
|
|
} else {
|
|
CUDA_RUNTIME_CHECK(cudaFree(ptr));
|
|
}
|
|
}
|
|
|
|
void get_mem_handle(MemHandle* mem_handle, void* ptr) const {
|
|
size_t size = 0;
|
|
CUDA_DRIVER_CHECK(lazy_cuMemGetAddressRange_v2(nullptr, &size, reinterpret_cast<CUdeviceptr>(ptr)));
|
|
mem_handle->size = size;
|
|
|
|
if (use_fabric) {
|
|
CUmemGenericAllocationHandle handle;
|
|
CUDA_DRIVER_CHECK(lazy_cuMemRetainAllocationHandle(&handle, ptr));
|
|
CUDA_DRIVER_CHECK(lazy_cuMemExportToShareableHandle(&mem_handle->inner.cu_mem_fabric_handle, handle, CU_MEM_HANDLE_TYPE_FABRIC, 0));
|
|
} else {
|
|
CUDA_RUNTIME_CHECK(cudaIpcGetMemHandle(&mem_handle->inner.cuda_ipc_mem_handle, ptr));
|
|
}
|
|
}
|
|
|
|
void open_mem_handle(void** ptr, MemHandle* mem_handle) const {
|
|
if (use_fabric) {
|
|
size_t size = mem_handle->size;
|
|
|
|
CUmemGenericAllocationHandle handle;
|
|
CUDA_DRIVER_CHECK(lazy_cuMemImportFromShareableHandle(&handle, &mem_handle->inner.cu_mem_fabric_handle, CU_MEM_HANDLE_TYPE_FABRIC));
|
|
|
|
CUDA_DRIVER_CHECK(lazy_cuMemAddressReserve(reinterpret_cast<CUdeviceptr*>(ptr), size, 0, 0, 0));
|
|
CUDA_DRIVER_CHECK(lazy_cuMemMap(reinterpret_cast<CUdeviceptr>(*ptr), size, 0, handle, 0));
|
|
cu_mem_set_access_all(*ptr, size);
|
|
} else {
|
|
CUDA_RUNTIME_CHECK(cudaIpcOpenMemHandle(ptr, mem_handle->inner.cuda_ipc_mem_handle, cudaIpcMemLazyEnablePeerAccess));
|
|
}
|
|
}
|
|
|
|
void close_mem_handle(void* ptr) const {
|
|
if (use_fabric) {
|
|
cu_mem_free(ptr);
|
|
} else {
|
|
CUDA_RUNTIME_CHECK(cudaIpcCloseMemHandle(ptr));
|
|
}
|
|
}
|
|
|
|
private:
|
|
bool use_fabric;
|
|
};
|
|
|
|
} // namespace deep_ep::shared_memory
|