From fe84d63a0b00b011100a55c848ed79737feba185 Mon Sep 17 00:00:00 2001 From: Lawrence Mitchell Date: Tue, 12 Dec 2023 14:23:24 +0000 Subject: [PATCH] WIP: Correct signatures for torch allocator plug in Since https://github.com/pytorch/pytorch/pull/91398, the signature of the pluggable allocate and deallocate functions must accept the device id. The current version only accepts a device id for allocate, which means that when using a stream ordered allocator with devices other than device zero, we pass an invalid stream into the deallocation function. To fix this, adapt the signature to match the one pytorch expects. Now, since we have the device available during allocation and deallocation, we would like to use that device to obtain the appropriate memory resource. Unfortunately, since RMM's cuda_device_id does not have a nullary constructor, we can't use it in Cython without some hacky workarounds. However, since we don't actually need to build a Python module, but rather just a single shared library that offers two extern "C" functions, let's just write our allocator hooks directly in C++. - Closes #1405 --- python/rmm/_lib/CMakeLists.txt | 5 +- python/rmm/_lib/torch_allocator.pyx | 24 ---------- python/rmm/allocators/_torch_allocator.cpp | 56 ++++++++++++++++++++++ python/rmm/allocators/torch.py | 6 +-- 4 files changed, 60 insertions(+), 31 deletions(-) delete mode 100644 python/rmm/_lib/torch_allocator.pyx create mode 100644 python/rmm/allocators/_torch_allocator.cpp diff --git a/python/rmm/_lib/CMakeLists.txt b/python/rmm/_lib/CMakeLists.txt index 852dd87c4..4ad0fa4a5 100644 --- a/python/rmm/_lib/CMakeLists.txt +++ b/python/rmm/_lib/CMakeLists.txt @@ -12,12 +12,9 @@ # the License. # ============================================================================= -set(cython_sources device_buffer.pyx lib.pyx logger.pyx memory_resource.pyx cuda_stream.pyx - torch_allocator.pyx) +set(cython_sources device_buffer.pyx lib.pyx logger.pyx memory_resource.pyx cuda_stream.pyx) set(linked_libraries rmm::rmm) # Build all of the Cython targets rapids_cython_create_modules(SOURCE_FILES "${cython_sources}" LINKED_LIBRARIES "${linked_libraries}" CXX) -# The cdef public functions in this file need to have a C ABI -target_compile_definitions(torch_allocator PRIVATE CYTHON_EXTERN_C=extern\ "C") diff --git a/python/rmm/_lib/torch_allocator.pyx b/python/rmm/_lib/torch_allocator.pyx deleted file mode 100644 index 12dc9fe11..000000000 --- a/python/rmm/_lib/torch_allocator.pyx +++ /dev/null @@ -1,24 +0,0 @@ -from cuda.ccudart cimport cudaStream_t - -from rmm._lib.cuda_stream_view cimport cuda_stream_view -from rmm._lib.memory_resource cimport device_memory_resource -from rmm._lib.per_device_resource cimport get_current_device_resource - - -cdef public void* allocate( - ssize_t size, int device, void* stream -) except * with gil: - cdef device_memory_resource* mr = get_current_device_resource() - cdef cuda_stream_view stream_view = cuda_stream_view( - (stream) - ) - return mr[0].allocate(size, stream_view) - -cdef public void deallocate( - void* ptr, ssize_t size, void* stream -) except * with gil: - cdef device_memory_resource* mr = get_current_device_resource() - cdef cuda_stream_view stream_view = cuda_stream_view( - (stream) - ) - mr[0].deallocate(ptr, size, stream_view) diff --git a/python/rmm/allocators/_torch_allocator.cpp b/python/rmm/allocators/_torch_allocator.cpp new file mode 100644 index 000000000..5b33a5135 --- /dev/null +++ b/python/rmm/allocators/_torch_allocator.cpp @@ -0,0 +1,56 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +#include +#include +#include + +// These signatures must match those required by CUDAPluggableAllocator in +// github.com/pytorch/pytorch/blob/main/torch/csrc/cuda/CUDAPluggableAllocator.h +// Since the loading is done at runtime via dlopen, no error checking +// can before performed. + +/** + * @brief Allocate memory of at least \p size bytes. + * + * @throws rmm::bad_alloc When the requested allocation cannot be satisfied. + * + * @param size The number of bytes to allocate + * @param device The device whose memory resource one should use + * @param stream CUDA stream to perform allocation on + * @return void* Pointer to the newly allocated memory + */ +extern "C" void* allocate(std::size_t size, int device, void* stream) +{ + auto mr = rmm::mr::get_per_device_resource(rmm::cuda_device_id{device}); + return mr->allocate(size, rmm::cuda_stream_view{static_cast(stream)}); +} + +/** + * @brief Deallocate memory pointed to by \p ptr. + * + * @param ptr Pointer to be deallocated + * @param size The number of bytes in the allocation + * @param device The device whose memory resource one should use + * @param stream CUDA stream to perform deallocation on + */ +extern "C" void deallocate(void* ptr, std::size_t size, int device, void* stream) +{ + auto mr = rmm::mr::get_per_device_resource(rmm::cuda_device_id{device}); + mr->deallocate(ptr, size, rmm::cuda_stream_view{static_cast(stream)}); +} diff --git a/python/rmm/allocators/torch.py b/python/rmm/allocators/torch.py index 65b310a89..d2761f2b7 100644 --- a/python/rmm/allocators/torch.py +++ b/python/rmm/allocators/torch.py @@ -16,11 +16,11 @@ except ImportError: rmm_torch_allocator = None else: - import rmm._lib.torch_allocator + import pathlib - _alloc_free_lib_path = rmm._lib.torch_allocator.__file__ rmm_torch_allocator = CUDAPluggableAllocator( - _alloc_free_lib_path, + pathlib.Path(__file__).absolute().parent / "_torch_allocator.so", alloc_fn_name="allocate", free_fn_name="deallocate", ) + del pathlib