From ea6f54c380757efd66318f6e0aebecd93788b262 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 25 Nov 2021 10:14:09 -0500 Subject: [PATCH 01/34] Draft of callback_memory_resource --- .../mr/device/callback_memory_resource.hpp | 60 +++++++++++++++++++ .../mr/device/callback_resource_adaptor.hpp | 0 2 files changed, 60 insertions(+) create mode 100644 include/rmm/mr/device/callback_memory_resource.hpp create mode 100644 include/rmm/mr/device/callback_resource_adaptor.hpp diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp new file mode 100644 index 000000000..09cf9249b --- /dev/null +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -0,0 +1,60 @@ +/* + * Copyright (c) 2020-2021, 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. + */ +#pragma once + +#include + +#include +#include +#include + +namespace rmm::mr { + +using allocate_callback_t = std::function; +using deallocate_callback_t = std::function; + +template +class callback_memory_resource final : public device_memory_resource { + +public: + callback_memory_resource(allocate_callback_t allocate_callback, + deallocate_callback_t deallocate_callback, + void* allocate_callback_arg, + void* deallocate_callback_arg) + : allocate_callback_(allocate_callback), + deallocate_callback_(deallocate_callback), + allocate_callback_arg_(allocate_callback_arg), + deallocate_callback_arg_(deallocate_callback_arg) {} + +private: + + void* do_allocate(std::size_t bytes, cuda_stream_view stream) override + { + return allocate_callback_(bytes, stream, allocate_callback_arg_); + } + + void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override + { + deallocate_callback_(bytes, stream, deallocate_callback_arg_); + } + + allocate_callback_t allocate_callback_; + deallocate_callback_t deallocate_callback_; + void* allocate_callback_arg_; + void* deallocate_callback_arg_; +}; + +} // namespace rmm::mr diff --git a/include/rmm/mr/device/callback_resource_adaptor.hpp b/include/rmm/mr/device/callback_resource_adaptor.hpp new file mode 100644 index 000000000..e69de29bb From ba5a43d9be4cce9ca60298673924c7852ad15507 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 25 Nov 2021 11:36:09 -0500 Subject: [PATCH 02/34] building... --- .../mr/device/callback_memory_resource.hpp | 17 ++++--- python/rmm/_lib/memory_resource.pxd | 4 ++ python/rmm/_lib/memory_resource.pyx | 44 ++++++++++++++++++- 3 files changed, 59 insertions(+), 6 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 09cf9249b..6738ad418 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -23,10 +23,9 @@ namespace rmm::mr { -using allocate_callback_t = std::function; -using deallocate_callback_t = std::function; +using allocate_callback_t = std::function; // TODO: can we accept a stream here? +using deallocate_callback_t = std::function; -template class callback_memory_resource final : public device_memory_resource { public: @@ -43,14 +42,22 @@ class callback_memory_resource final : public device_memory_resource { void* do_allocate(std::size_t bytes, cuda_stream_view stream) override { - return allocate_callback_(bytes, stream, allocate_callback_arg_); + return allocate_callback_(bytes, allocate_callback_arg_); } void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override { - deallocate_callback_(bytes, stream, deallocate_callback_arg_); + deallocate_callback_(ptr, bytes, deallocate_callback_arg_); } + [[nodiscard]] std::pair do_get_mem_info(cuda_stream_view) const override + { + throw std::runtime_error("cannot get free / total memory"); + } + + [[nodiscard]] virtual bool supports_streams() const noexcept { return false; } + [[nodiscard]] virtual bool supports_get_mem_info() const noexcept { return false; } + allocate_callback_t allocate_callback_; deallocate_callback_t deallocate_callback_; void* allocate_callback_arg_; diff --git a/python/rmm/_lib/memory_resource.pxd b/python/rmm/_lib/memory_resource.pxd index 3a71fd500..135d28787 100644 --- a/python/rmm/_lib/memory_resource.pxd +++ b/python/rmm/_lib/memory_resource.pxd @@ -57,6 +57,10 @@ cdef class BinningMemoryResource(UpstreamResourceAdaptor): size_t allocation_size, DeviceMemoryResource bin_resource=*) +cdef class CallbackMemoryResource(DeviceMemoryResource): + cdef object _allocate_func + cdef object _deallocate_func + cdef class LoggingResourceAdaptor(UpstreamResourceAdaptor): cdef object _log_file_name cpdef get_file_name(self) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index d7711ed8a..1a53463de 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -17,7 +17,7 @@ import warnings from collections import defaultdict from cython.operator cimport dereference as deref -from libc.stdint cimport int8_t, int64_t +from libc.stdint cimport int8_t, int64_t, uintptr_t from libcpp cimport bool from libcpp.cast cimport dynamic_cast from libcpp.memory cimport make_shared, make_unique, shared_ptr, unique_ptr @@ -73,6 +73,19 @@ cdef extern from "rmm/mr/device/fixed_size_memory_resource.hpp" \ size_t block_size, size_t block_to_preallocate) except + +cdef extern from "rmm/mr/device/callback_memory_resource.hpp" \ + namespace "rmm::mr" nogil: + ctypedef void* (*allocate_callback_t)(size_t, void*) + ctypedef void (*deallocate_callback_t)(void*, size_t, void*) + + cdef cppclass callback_memory_resource(device_memory_resource): + callback_memory_resource( + allocate_callback_t allocate_callback, + deallocate_callback_t deallocate_callback, + void* allocate_callback_arg, + void* deallocate_callback_arg + ) except + + cdef extern from "rmm/mr/device/binning_memory_resource.hpp" \ namespace "rmm::mr" nogil: cdef cppclass binning_memory_resource[Upstream](device_memory_resource): @@ -436,6 +449,35 @@ cdef class BinningMemoryResource(UpstreamResourceAdaptor): bin_resource.get_mr()) +cdef void* _allocate_callback_wrapper(size_t nbytes, void *ctx) with gil: + return ((ctx)(nbytes)) + +cdef void _deallocate_callback_wrapper( + void* ptr, + size_t nbytes, + void *ctx +) with gil: + (ctx)((ptr), nbytes) + + +cdef class CallbackMemoryResource: + def __init__( + self, + allocate_func, + deallocate_func, + ): + self._allocate_func = allocate_func + self._deallocate_func = deallocate_func + self.c_obj.reset( + new callback_memory_resource( + (_allocate_callback_wrapper), + (_deallocate_callback_wrapper), + (allocate_func), + (deallocate_func) + ) + ) + + def _append_id(filename, id): """ Append ".dev" onto a filename before the extension From 46c6507f528d2010ae7117a204a84912b211ffa8 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 1 Dec 2021 10:10:12 -0500 Subject: [PATCH 03/34] Some fixes --- python/rmm/_lib/memory_resource.pyx | 2 +- python/rmm/mr.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index 1a53463de..eed66ac10 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -457,7 +457,7 @@ cdef void _deallocate_callback_wrapper( size_t nbytes, void *ctx ) with gil: - (ctx)((ptr), nbytes) + (ctx)((ptr), nbytes) cdef class CallbackMemoryResource: diff --git a/python/rmm/mr.py b/python/rmm/mr.py index eaa9705eb..ea0a0cf8c 100644 --- a/python/rmm/mr.py +++ b/python/rmm/mr.py @@ -13,6 +13,7 @@ # limitations under the License. from rmm._lib.memory_resource import ( BinningMemoryResource, + CallbackMemoryResource, CudaAsyncMemoryResource, CudaMemoryResource, DeviceMemoryResource, @@ -39,6 +40,7 @@ __all__ = [ "BinningMemoryResource", + "CallbackMemoryResource", "CudaAsyncMemoryResource", "CudaMemoryResource", "DeviceMemoryResource", From eb8f9e119ec2d9cc172744b80c118ac92afe6a03 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 25 Nov 2021 10:14:09 -0500 Subject: [PATCH 04/34] Draft of callback_memory_resource --- .../mr/device/callback_memory_resource.hpp | 60 +++++++++++++++++++ .../mr/device/callback_resource_adaptor.hpp | 0 2 files changed, 60 insertions(+) create mode 100644 include/rmm/mr/device/callback_memory_resource.hpp create mode 100644 include/rmm/mr/device/callback_resource_adaptor.hpp diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp new file mode 100644 index 000000000..09cf9249b --- /dev/null +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -0,0 +1,60 @@ +/* + * Copyright (c) 2020-2021, 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. + */ +#pragma once + +#include + +#include +#include +#include + +namespace rmm::mr { + +using allocate_callback_t = std::function; +using deallocate_callback_t = std::function; + +template +class callback_memory_resource final : public device_memory_resource { + +public: + callback_memory_resource(allocate_callback_t allocate_callback, + deallocate_callback_t deallocate_callback, + void* allocate_callback_arg, + void* deallocate_callback_arg) + : allocate_callback_(allocate_callback), + deallocate_callback_(deallocate_callback), + allocate_callback_arg_(allocate_callback_arg), + deallocate_callback_arg_(deallocate_callback_arg) {} + +private: + + void* do_allocate(std::size_t bytes, cuda_stream_view stream) override + { + return allocate_callback_(bytes, stream, allocate_callback_arg_); + } + + void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override + { + deallocate_callback_(bytes, stream, deallocate_callback_arg_); + } + + allocate_callback_t allocate_callback_; + deallocate_callback_t deallocate_callback_; + void* allocate_callback_arg_; + void* deallocate_callback_arg_; +}; + +} // namespace rmm::mr diff --git a/include/rmm/mr/device/callback_resource_adaptor.hpp b/include/rmm/mr/device/callback_resource_adaptor.hpp new file mode 100644 index 000000000..e69de29bb From 8f1a62d05916da821255a6e5f7014affa7127c47 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 25 Nov 2021 11:36:09 -0500 Subject: [PATCH 05/34] building... --- .../mr/device/callback_memory_resource.hpp | 17 ++++--- python/rmm/_lib/memory_resource.pxd | 4 ++ python/rmm/_lib/memory_resource.pyx | 44 ++++++++++++++++++- 3 files changed, 59 insertions(+), 6 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 09cf9249b..6738ad418 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -23,10 +23,9 @@ namespace rmm::mr { -using allocate_callback_t = std::function; -using deallocate_callback_t = std::function; +using allocate_callback_t = std::function; // TODO: can we accept a stream here? +using deallocate_callback_t = std::function; -template class callback_memory_resource final : public device_memory_resource { public: @@ -43,14 +42,22 @@ class callback_memory_resource final : public device_memory_resource { void* do_allocate(std::size_t bytes, cuda_stream_view stream) override { - return allocate_callback_(bytes, stream, allocate_callback_arg_); + return allocate_callback_(bytes, allocate_callback_arg_); } void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override { - deallocate_callback_(bytes, stream, deallocate_callback_arg_); + deallocate_callback_(ptr, bytes, deallocate_callback_arg_); } + [[nodiscard]] std::pair do_get_mem_info(cuda_stream_view) const override + { + throw std::runtime_error("cannot get free / total memory"); + } + + [[nodiscard]] virtual bool supports_streams() const noexcept { return false; } + [[nodiscard]] virtual bool supports_get_mem_info() const noexcept { return false; } + allocate_callback_t allocate_callback_; deallocate_callback_t deallocate_callback_; void* allocate_callback_arg_; diff --git a/python/rmm/_lib/memory_resource.pxd b/python/rmm/_lib/memory_resource.pxd index 3a71fd500..135d28787 100644 --- a/python/rmm/_lib/memory_resource.pxd +++ b/python/rmm/_lib/memory_resource.pxd @@ -57,6 +57,10 @@ cdef class BinningMemoryResource(UpstreamResourceAdaptor): size_t allocation_size, DeviceMemoryResource bin_resource=*) +cdef class CallbackMemoryResource(DeviceMemoryResource): + cdef object _allocate_func + cdef object _deallocate_func + cdef class LoggingResourceAdaptor(UpstreamResourceAdaptor): cdef object _log_file_name cpdef get_file_name(self) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index e02f1a3ae..0aaa5d93b 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -17,7 +17,7 @@ import warnings from collections import defaultdict from cython.operator cimport dereference as deref -from libc.stdint cimport int8_t, int64_t +from libc.stdint cimport int8_t, int64_t, uintptr_t from libcpp cimport bool from libcpp.cast cimport dynamic_cast from libcpp.memory cimport make_shared, make_unique, shared_ptr, unique_ptr @@ -76,6 +76,19 @@ cdef extern from "rmm/mr/device/fixed_size_memory_resource.hpp" \ size_t block_size, size_t block_to_preallocate) except + +cdef extern from "rmm/mr/device/callback_memory_resource.hpp" \ + namespace "rmm::mr" nogil: + ctypedef void* (*allocate_callback_t)(size_t, void*) + ctypedef void (*deallocate_callback_t)(void*, size_t, void*) + + cdef cppclass callback_memory_resource(device_memory_resource): + callback_memory_resource( + allocate_callback_t allocate_callback, + deallocate_callback_t deallocate_callback, + void* allocate_callback_arg, + void* deallocate_callback_arg + ) except + + cdef extern from "rmm/mr/device/binning_memory_resource.hpp" \ namespace "rmm::mr" nogil: cdef cppclass binning_memory_resource[Upstream](device_memory_resource): @@ -444,6 +457,35 @@ cdef class BinningMemoryResource(UpstreamResourceAdaptor): bin_resource.get_mr()) +cdef void* _allocate_callback_wrapper(size_t nbytes, void *ctx) with gil: + return ((ctx)(nbytes)) + +cdef void _deallocate_callback_wrapper( + void* ptr, + size_t nbytes, + void *ctx +) with gil: + (ctx)((ptr), nbytes) + + +cdef class CallbackMemoryResource: + def __init__( + self, + allocate_func, + deallocate_func, + ): + self._allocate_func = allocate_func + self._deallocate_func = deallocate_func + self.c_obj.reset( + new callback_memory_resource( + (_allocate_callback_wrapper), + (_deallocate_callback_wrapper), + (allocate_func), + (deallocate_func) + ) + ) + + def _append_id(filename, id): """ Append ".dev" onto a filename before the extension From e7c8c887030b7ce5a156f7640b293fe24371ce45 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 1 Dec 2021 10:10:12 -0500 Subject: [PATCH 06/34] Some fixes --- python/rmm/_lib/memory_resource.pyx | 2 +- python/rmm/mr.py | 2 ++ 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index 0aaa5d93b..cee4b775c 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -465,7 +465,7 @@ cdef void _deallocate_callback_wrapper( size_t nbytes, void *ctx ) with gil: - (ctx)((ptr), nbytes) + (ctx)((ptr), nbytes) cdef class CallbackMemoryResource: diff --git a/python/rmm/mr.py b/python/rmm/mr.py index eaa9705eb..ea0a0cf8c 100644 --- a/python/rmm/mr.py +++ b/python/rmm/mr.py @@ -13,6 +13,7 @@ # limitations under the License. from rmm._lib.memory_resource import ( BinningMemoryResource, + CallbackMemoryResource, CudaAsyncMemoryResource, CudaMemoryResource, DeviceMemoryResource, @@ -39,6 +40,7 @@ __all__ = [ "BinningMemoryResource", + "CallbackMemoryResource", "CudaAsyncMemoryResource", "CudaMemoryResource", "DeviceMemoryResource", From 1e2a10cd880e1d3719cc537a0fea1686c42f0e13 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 6 Apr 2022 07:08:02 -0400 Subject: [PATCH 07/34] Expose `allocate()` and `deallocate()` --- python/rmm/_lib/memory_resource.pxd | 4 ++-- python/rmm/_lib/memory_resource.pyx | 6 ++++++ 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/python/rmm/_lib/memory_resource.pxd b/python/rmm/_lib/memory_resource.pxd index 135d28787..c3d4084b5 100644 --- a/python/rmm/_lib/memory_resource.pxd +++ b/python/rmm/_lib/memory_resource.pxd @@ -21,11 +21,11 @@ from libcpp.vector cimport vector cdef extern from "rmm/mr/device/device_memory_resource.hpp" \ namespace "rmm::mr" nogil: cdef cppclass device_memory_resource: - pass + void* allocate(size_t bytes) except + + void deallocate(void* ptr, size_t bytes) except + cdef class DeviceMemoryResource: cdef shared_ptr[device_memory_resource] c_obj - cdef device_memory_resource* get_mr(self) cdef class UpstreamResourceAdaptor(DeviceMemoryResource): diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index cee4b775c..0bc02ef80 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -181,6 +181,12 @@ cdef class DeviceMemoryResource: cdef device_memory_resource* get_mr(self): return self.c_obj.get() + def allocate(self, size_t nbytes): + return self.c_obj.get().allocate(nbytes) + + def deallocate(self, uintptr_t ptr, size_t nbytes): + self.c_obj.get().deallocate((ptr), nbytes) + cdef class UpstreamResourceAdaptor(DeviceMemoryResource): From d2fca0d0766db8a9d6192b24f4034ba52499d00b Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 6 Apr 2022 11:21:51 -0400 Subject: [PATCH 08/34] Add first cpp test --- .../mr/device/callback_memory_resource.hpp | 39 ++++++++++----- .../mr/device/callback_resource_adaptor.hpp | 0 tests/CMakeLists.txt | 3 ++ tests/mr/device/callback_mr_tests.cpp | 48 +++++++++++++++++++ 4 files changed, 78 insertions(+), 12 deletions(-) delete mode 100644 include/rmm/mr/device/callback_resource_adaptor.hpp create mode 100644 tests/mr/device/callback_mr_tests.cpp diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 6738ad418..093c6a6f8 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -23,31 +23,46 @@ namespace rmm::mr { -using allocate_callback_t = std::function; // TODO: can we accept a stream here? -using deallocate_callback_t = std::function; +/** + * @brief Callback function type used by callback_memory_resource for allocation. + * + * The function signature must match that of `rmm::device_memory_resource::allocate()`. + */ +using allocate_callback_t = std::function; -class callback_memory_resource final : public device_memory_resource { +/** + * @brief Callback function type used by callback_memory_resource for deallocation. + * + * The function signature must match that of `rmm::device_memory_resource::deallocate()`. + */ +using deallocate_callback_t = std::function; -public: +/** + * @brief A device memory resource that uses the provided callbacks for memory allocation + * and deallocation. + */ +class callback_memory_resource final : public device_memory_resource { + public: callback_memory_resource(allocate_callback_t allocate_callback, - deallocate_callback_t deallocate_callback, - void* allocate_callback_arg, - void* deallocate_callback_arg) + deallocate_callback_t deallocate_callback, + void* allocate_callback_arg, + void* deallocate_callback_arg) : allocate_callback_(allocate_callback), deallocate_callback_(deallocate_callback), allocate_callback_arg_(allocate_callback_arg), - deallocate_callback_arg_(deallocate_callback_arg) {} - -private: + deallocate_callback_arg_(deallocate_callback_arg) + { + } + private: void* do_allocate(std::size_t bytes, cuda_stream_view stream) override { - return allocate_callback_(bytes, allocate_callback_arg_); + return allocate_callback_(bytes, allocate_callback_arg_, stream); } void do_deallocate(void* ptr, std::size_t bytes, cuda_stream_view stream) override { - deallocate_callback_(ptr, bytes, deallocate_callback_arg_); + deallocate_callback_(ptr, bytes, deallocate_callback_arg_, stream); } [[nodiscard]] std::pair do_get_mem_info(cuda_stream_view) const override diff --git a/include/rmm/mr/device/callback_resource_adaptor.hpp b/include/rmm/mr/device/callback_resource_adaptor.hpp deleted file mode 100644 index e69de29bb..000000000 diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 344f66db5..8f68141fd 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -150,3 +150,6 @@ ConfigureTest(ARENA_MR_TEST mr/device/arena_mr_tests.cpp) # binning MR tests ConfigureTest(BINNING_MR_TEST mr/device/binning_mr_tests.cpp) + +# callback memory resource tests +ConfigureTest(CALLBACK_MR_TEST mr/device/callback_mr_tests.cpp) diff --git a/tests/mr/device/callback_mr_tests.cpp b/tests/mr/device/callback_mr_tests.cpp new file mode 100644 index 000000000..cd1920664 --- /dev/null +++ b/tests/mr/device/callback_mr_tests.cpp @@ -0,0 +1,48 @@ +/* + * Copyright (c) 2021, 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 "../../byte_literals.hpp" +#include "rmm/cuda_stream_view.hpp" +#include "rmm/mr/device/device_memory_resource.hpp" + +#include +#include +#include +#include + +#include + +namespace rmm::test { +namespace { + +TEST(CallbackTest, CallbackTestSimple) +{ + auto base_mr = rmm::mr::get_current_device_resource(); + auto allocate_callback = [&base_mr](std::size_t size, void* arg, cuda_stream_view stream) { + return base_mr->allocate(size, stream); + }; + auto deallocate_callback = [&base_mr]( + void* ptr, std::size_t size, void* arg, cuda_stream_view stream) { + base_mr->deallocate(ptr, size, stream); + }; + auto mr = + rmm::mr::callback_memory_resource(allocate_callback, deallocate_callback, nullptr, nullptr); + auto ptr = mr.allocate(10_MiB); + mr.deallocate(ptr, 10_MiB); +} + +} // namespace +} // namespace rmm::test From 63b9690bd226403edafb685e73328aa0ce7f339c Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 6 Apr 2022 14:34:47 -0400 Subject: [PATCH 09/34] Add another test --- tests/mr/device/callback_mr_tests.cpp | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tests/mr/device/callback_mr_tests.cpp b/tests/mr/device/callback_mr_tests.cpp index cd1920664..d52d88492 100644 --- a/tests/mr/device/callback_mr_tests.cpp +++ b/tests/mr/device/callback_mr_tests.cpp @@ -28,7 +28,7 @@ namespace rmm::test { namespace { -TEST(CallbackTest, CallbackTestSimple) +TEST(CallbackTest, PassThroughTest) { auto base_mr = rmm::mr::get_current_device_resource(); auto allocate_callback = [&base_mr](std::size_t size, void* arg, cuda_stream_view stream) { @@ -44,5 +44,29 @@ TEST(CallbackTest, CallbackTestSimple) mr.deallocate(ptr, 10_MiB); } +TEST(CallbackTest, LoggingTest) +{ + testing::internal::CaptureStdout(); + + auto base_mr = rmm::mr::get_current_device_resource(); + auto allocate_callback = [&base_mr](std::size_t size, void* arg, cuda_stream_view stream) { + std::cout << "Allocating " << size << " bytes" << std::endl; + return base_mr->allocate(size, stream); + }; + auto deallocate_callback = [&base_mr]( + void* ptr, std::size_t size, void* arg, cuda_stream_view stream) { + std::cout << "Deallocating " << size << " bytes" << std::endl; + base_mr->deallocate(ptr, size, stream); + }; + auto mr = + rmm::mr::callback_memory_resource(allocate_callback, deallocate_callback, nullptr, nullptr); + auto ptr = mr.allocate(10_MiB); + mr.deallocate(ptr, 10_MiB); + + std::string output = testing::internal::GetCapturedStdout(); + std::string expect {"Allocating " + std::to_string(10_MiB) + " bytes" + "\n" + "Deallocating " + std::to_string(10_MiB) + " bytes\n"}; + ASSERT_EQ(expect, output); +} + } // namespace } // namespace rmm::test From cdc53e777cf2568174b415fa7ebbdcbd39662808 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Fri, 8 Apr 2022 11:00:53 -0400 Subject: [PATCH 10/34] Use fmt instead --- tests/mr/device/callback_mr_tests.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/mr/device/callback_mr_tests.cpp b/tests/mr/device/callback_mr_tests.cpp index d52d88492..4233d4490 100644 --- a/tests/mr/device/callback_mr_tests.cpp +++ b/tests/mr/device/callback_mr_tests.cpp @@ -23,6 +23,8 @@ #include #include +#include + #include namespace rmm::test { @@ -64,7 +66,7 @@ TEST(CallbackTest, LoggingTest) mr.deallocate(ptr, 10_MiB); std::string output = testing::internal::GetCapturedStdout(); - std::string expect {"Allocating " + std::to_string(10_MiB) + " bytes" + "\n" + "Deallocating " + std::to_string(10_MiB) + " bytes\n"}; + std::string expect = fmt::format("Allocating {} bytes\nDeallocating {} bytes\n", 10_MiB, 10_MiB); ASSERT_EQ(expect, output); } From e1014d7af167a677b97ebdbdb8169637a0553d89 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Fri, 8 Apr 2022 12:57:34 -0400 Subject: [PATCH 11/34] C++ docs --- .../mr/device/callback_memory_resource.hpp | 41 +++++++++++++++---- tests/mr/device/callback_mr_tests.cpp | 14 +++---- 2 files changed, 40 insertions(+), 15 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 093c6a6f8..45c639839 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -24,18 +24,28 @@ namespace rmm::mr { /** - * @brief Callback function type used by callback_memory_resource for allocation. + * @brief Callback function type used by callback memory resource for allocation. * - * The function signature must match that of `rmm::device_memory_resource::allocate()`. + * The signature of the callback function is: + * `void* allocate_callback_t(std::size_t bytes, cuda_stream_view stream, void* arg); + * + * `bytes` specifies the number of bytes to allocate. `stream` is the CUDA stream + * on which the memory is used. Additional information to the callback may be passed + * via `arg`. The function returns the pointer to the allocated memory. */ -using allocate_callback_t = std::function; +using allocate_callback_t = std::function; /** * @brief Callback function type used by callback_memory_resource for deallocation. * - * The function signature must match that of `rmm::device_memory_resource::deallocate()`. + * The signature of the callback function is: + * `void* deallocate_callback_t(void* ptr, std::size_t bytes, cuda_stream_view stream, void* arg); + * + * `ptr` specifies the pointer to the memory to be freed, while `bytes` specified + * the number of bytes to free. `stream` is the CUDA stream on which the memory is used. + * Additional information to the callback may be passed via `arg` */ -using deallocate_callback_t = std::function; +using deallocate_callback_t = std::function; /** * @brief A device memory resource that uses the provided callbacks for memory allocation @@ -43,10 +53,18 @@ using deallocate_callback_t = std::function do_get_mem_info(cuda_stream_view) const override diff --git a/tests/mr/device/callback_mr_tests.cpp b/tests/mr/device/callback_mr_tests.cpp index 4233d4490..c1d696bdc 100644 --- a/tests/mr/device/callback_mr_tests.cpp +++ b/tests/mr/device/callback_mr_tests.cpp @@ -33,11 +33,11 @@ namespace { TEST(CallbackTest, PassThroughTest) { auto base_mr = rmm::mr::get_current_device_resource(); - auto allocate_callback = [&base_mr](std::size_t size, void* arg, cuda_stream_view stream) { + auto allocate_callback = [&base_mr](std::size_t size, cuda_stream_view stream, void* arg) { return base_mr->allocate(size, stream); }; auto deallocate_callback = [&base_mr]( - void* ptr, std::size_t size, void* arg, cuda_stream_view stream) { + void* ptr, std::size_t size, cuda_stream_view stream, void* arg) { base_mr->deallocate(ptr, size, stream); }; auto mr = @@ -49,15 +49,15 @@ TEST(CallbackTest, PassThroughTest) TEST(CallbackTest, LoggingTest) { testing::internal::CaptureStdout(); - + auto base_mr = rmm::mr::get_current_device_resource(); - auto allocate_callback = [&base_mr](std::size_t size, void* arg, cuda_stream_view stream) { + auto allocate_callback = [&base_mr](std::size_t size, cuda_stream_view stream, void* arg) { std::cout << "Allocating " << size << " bytes" << std::endl; return base_mr->allocate(size, stream); }; auto deallocate_callback = [&base_mr]( - void* ptr, std::size_t size, void* arg, cuda_stream_view stream) { - std::cout << "Deallocating " << size << " bytes" << std::endl; + void* ptr, std::size_t size, cuda_stream_view stream, void* arg) { + std::cout << "Deallocating " << size << " bytes" << std::endl; base_mr->deallocate(ptr, size, stream); }; auto mr = @@ -69,6 +69,6 @@ TEST(CallbackTest, LoggingTest) std::string expect = fmt::format("Allocating {} bytes\nDeallocating {} bytes\n", 10_MiB, 10_MiB); ASSERT_EQ(expect, output); } - + } // namespace } // namespace rmm::test From f53ccf1292b666a4ab67dd21fccfd0236b534c82 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Fri, 8 Apr 2022 13:43:58 -0400 Subject: [PATCH 12/34] Add python test --- python/rmm/_lib/memory_resource.pyx | 25 ++++++++++++++++++++++++- python/rmm/tests/test_rmm.py | 22 ++++++++++++++++++++++ 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index 0bc02ef80..53554be57 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -26,6 +26,7 @@ from libcpp.string cimport string from cuda.cudart import cudaError_t from rmm._cuda.gpu import CUDARuntimeError, getDevice, setDevice +from rmm._lib.cuda_stream_view cimport cuda_stream_view # NOTE: Keep extern declarations in .pyx file as much as possible to avoid @@ -463,18 +464,40 @@ cdef class BinningMemoryResource(UpstreamResourceAdaptor): bin_resource.get_mr()) -cdef void* _allocate_callback_wrapper(size_t nbytes, void *ctx) with gil: +cdef void* _allocate_callback_wrapper( + size_t nbytes, + cuda_stream_view stream, + void *ctx +) with gil: return ((ctx)(nbytes)) cdef void _deallocate_callback_wrapper( void* ptr, size_t nbytes, + cuda_stream_view stream, void *ctx ) with gil: (ctx)((ptr), nbytes) cdef class CallbackMemoryResource: + """ + A memory resource that uses the user-provided callables to do + memory allocation and deallocation. + + The allocation function must accept a single integer argument, + representing the number of bytes to allocate, and return + an integer representing the pointer to the allocated memory. + + The deallocation function must accept two arguments, + an integer representing the pointer to the memory to free, + and a second integer representing the number of bytes to free. + + Parameters + ---------- + allocate_func: callable + deallocate_func: callable + """ def __init__( self, allocate_func, diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index 409ed6b72..329dffeba 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -719,3 +719,25 @@ def test_dev_buf_circle_ref_dealloc(): # deallocate `dbuf1` (which needs the MR alive), a segfault occurs. gc.collect() + + +def test_custom_mr(capsys): + base_mr = rmm.mr.CudaMemoryResource() + + def allocate_func(size): + print(f"Allocating {size} bytes") + return base_mr.allocate(size) + + def deallocate_func(ptr, size): + print(f"Deallocating {size} bytes") + return base_mr.deallocate(ptr, size) + + rmm.mr.set_current_device_resource( + rmm.mr.CallbackMemoryResource(allocate_func, deallocate_func) + ) + + dbuf = rmm.DeviceBuffer(size=256) + del dbuf + + captured = capsys.readouterr() + assert captured.out == "Allocating 256 bytes\nDeallocating 256 bytes\n" From e9221fadd9ba4ddc1ce7517e72d0a9b822ea6237 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Fri, 8 Apr 2022 14:01:41 -0400 Subject: [PATCH 13/34] Remove files --- include/rmm/mr/device/callback_resource_adaptor.hpp | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 include/rmm/mr/device/callback_resource_adaptor.hpp diff --git a/include/rmm/mr/device/callback_resource_adaptor.hpp b/include/rmm/mr/device/callback_resource_adaptor.hpp deleted file mode 100644 index e69de29bb..000000000 From f00084c54fea40f28b7c48f200de2699535e39fc Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 13 Apr 2022 11:02:35 -0400 Subject: [PATCH 14/34] Doc update --- .../mr/device/callback_memory_resource.hpp | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 45c639839..9d1a3b673 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -29,9 +29,12 @@ namespace rmm::mr { * The signature of the callback function is: * `void* allocate_callback_t(std::size_t bytes, cuda_stream_view stream, void* arg); * - * `bytes` specifies the number of bytes to allocate. `stream` is the CUDA stream - * on which the memory is used. Additional information to the callback may be passed - * via `arg`. The function returns the pointer to the allocated memory. + * * Returns a pointer to an allocation of at least `bytes` usable immediately on + * `stream`. The stream-ordered behavior requirements are identical to + * `device_memory_resource::allocate`. + * + * * Additional context to the callback may be passed via `arg` that is provided at + * construction of the `callback_memory_resource`. */ using allocate_callback_t = std::function; @@ -41,9 +44,13 @@ using allocate_callback_t = std::function; From ecad57f6e0b7d1044b07e7691c4f8212d7227405 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath <3190405+shwina@users.noreply.github.com> Date: Wed, 13 Apr 2022 11:04:52 -0400 Subject: [PATCH 15/34] Update tests/mr/device/callback_mr_tests.cpp Co-authored-by: Jake Hemstad --- tests/mr/device/callback_mr_tests.cpp | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/tests/mr/device/callback_mr_tests.cpp b/tests/mr/device/callback_mr_tests.cpp index c1d696bdc..a06b5c2b5 100644 --- a/tests/mr/device/callback_mr_tests.cpp +++ b/tests/mr/device/callback_mr_tests.cpp @@ -51,17 +51,22 @@ TEST(CallbackTest, LoggingTest) testing::internal::CaptureStdout(); auto base_mr = rmm::mr::get_current_device_resource(); - auto allocate_callback = [&base_mr](std::size_t size, cuda_stream_view stream, void* arg) { + auto allocate_callback = [expected_base_mr](std::size_t size, cuda_stream_view stream, void* arg) { std::cout << "Allocating " << size << " bytes" << std::endl; + auto base_mr = static_cast(arg); + ASSERT_EQ(base_mr, expected_base_mr); return base_mr->allocate(size, stream); }; - auto deallocate_callback = [&base_mr]( + + auto deallocate_callback = [expected_base_mr]( void* ptr, std::size_t size, cuda_stream_view stream, void* arg) { std::cout << "Deallocating " << size << " bytes" << std::endl; + auto base_mr = static_cast(arg); + ASSERT_EQ(base_mr, expected_base_mr); base_mr->deallocate(ptr, size, stream); }; auto mr = - rmm::mr::callback_memory_resource(allocate_callback, deallocate_callback, nullptr, nullptr); + rmm::mr::callback_memory_resource(allocate_callback, deallocate_callback, base_mr, base_mr); auto ptr = mr.allocate(10_MiB); mr.deallocate(ptr, 10_MiB); From 01e109f76c0588c1354ea121f97744acc8a7938e Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 13 Apr 2022 11:14:08 -0400 Subject: [PATCH 16/34] Add note about performance hit --- python/rmm/_lib/memory_resource.pyx | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index 53554be57..246e0739a 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -486,12 +486,17 @@ cdef class CallbackMemoryResource: memory allocation and deallocation. The allocation function must accept a single integer argument, - representing the number of bytes to allocate, and return - an integer representing the pointer to the allocated memory. + representing the number of bytes to allocate, and return an + integer representing the pointer to the allocated memory. - The deallocation function must accept two arguments, - an integer representing the pointer to the memory to free, - and a second integer representing the number of bytes to free. + The deallocation function must accept two arguments, an integer + representing the pointer to the memory to free, and a second + integer representing the number of bytes to free. + + ``CallbackMemoryResource`` should really only be used for + debugging memory issues, as there is a significant performance + penalty associated with using a Python function for each memory + allocation and deallocation. Parameters ---------- From ad46071fcda8b392e56ef8ba9d730f04f746a715 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 13 Apr 2022 11:21:09 -0400 Subject: [PATCH 17/34] More doc --- include/rmm/mr/device/callback_memory_resource.hpp | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 9d1a3b673..d99814826 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -63,10 +63,17 @@ class callback_memory_resource final : public device_memory_resource { /** * @brief Construct a new callback memory resource. * + * Constructs a callback memory resource that uses the user-provided callbacks + * `allocate_callback` for allocation and `deallocate_callback` for deallocation. + * * @throws Nothing. * * @param allocate_callback The callback function used for allocation * @param deallocate_callback The callback function used for deallocation + * @param allocate_callback_arg Additional context passed to `allocate_callback`. + * It is the caller's responsibility to maintain the lifetime of the pointed-to data. + * @param deallocate_callback_arg Additional context passed to `deallocate_callback`. + * It is the caller's responsibility to maintain the lifetime of the pointed-to data. */ callback_memory_resource(allocate_callback_t allocate_callback, deallocate_callback_t deallocate_callback, From 3cafa3a85488526dfd70dd2c8c69af924fc44e28 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 13 Apr 2022 12:01:08 -0400 Subject: [PATCH 18/34] Doc update --- include/rmm/mr/device/callback_memory_resource.hpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index d99814826..0ab36a3c2 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -71,9 +71,11 @@ class callback_memory_resource final : public device_memory_resource { * @param allocate_callback The callback function used for allocation * @param deallocate_callback The callback function used for deallocation * @param allocate_callback_arg Additional context passed to `allocate_callback`. - * It is the caller's responsibility to maintain the lifetime of the pointed-to data. + * It is the caller's responsibility to maintain the lifetime of the pointed-to data + * for the duration of the lifetime of the `callback_memory_resource`. * @param deallocate_callback_arg Additional context passed to `deallocate_callback`. - * It is the caller's responsibility to maintain the lifetime of the pointed-to data. + * It is the caller's responsibility to maintain the lifetime of the pointed-to data + * for the duration of the lifetime of the `callback_memory_resource`. */ callback_memory_resource(allocate_callback_t allocate_callback, deallocate_callback_t deallocate_callback, From 4008617a77ba6000eb95118c669bd9624b7332fb Mon Sep 17 00:00:00 2001 From: Ashwin Srinath <3190405+shwina@users.noreply.github.com> Date: Wed, 13 Apr 2022 12:18:58 -0400 Subject: [PATCH 19/34] Update include/rmm/mr/device/callback_memory_resource.hpp Co-authored-by: Bradley Dice --- include/rmm/mr/device/callback_memory_resource.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 0ab36a3c2..2f96b3b67 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -50,7 +50,7 @@ using allocate_callback_t = std::function; From 8dc9ecabe537df0e65c0a1cf324968f4950b4e30 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath <3190405+shwina@users.noreply.github.com> Date: Wed, 13 Apr 2022 12:19:12 -0400 Subject: [PATCH 20/34] Update include/rmm/mr/device/callback_memory_resource.hpp Co-authored-by: Bradley Dice --- include/rmm/mr/device/callback_memory_resource.hpp | 2 -- 1 file changed, 2 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 2f96b3b67..acb1106bd 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -66,8 +66,6 @@ class callback_memory_resource final : public device_memory_resource { * Constructs a callback memory resource that uses the user-provided callbacks * `allocate_callback` for allocation and `deallocate_callback` for deallocation. * - * @throws Nothing. - * * @param allocate_callback The callback function used for allocation * @param deallocate_callback The callback function used for deallocation * @param allocate_callback_arg Additional context passed to `allocate_callback`. From a776631c6b73bafa7f1af8bcdc787a193a4cefa5 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Wed, 13 Apr 2022 14:45:58 -0400 Subject: [PATCH 21/34] Correctly pass base_mr using ctx --- tests/mr/device/callback_mr_tests.cpp | 22 ++++++++++------------ 1 file changed, 10 insertions(+), 12 deletions(-) diff --git a/tests/mr/device/callback_mr_tests.cpp b/tests/mr/device/callback_mr_tests.cpp index b43d35830..9fc333fc5 100644 --- a/tests/mr/device/callback_mr_tests.cpp +++ b/tests/mr/device/callback_mr_tests.cpp @@ -33,16 +33,18 @@ namespace { TEST(CallbackTest, PassThroughTest) { auto base_mr = rmm::mr::get_current_device_resource(); - auto allocate_callback = [&base_mr](std::size_t size, cuda_stream_view stream, void* arg) { + auto allocate_callback = [](std::size_t size, cuda_stream_view stream, void* arg) { + auto base_mr = static_cast(arg); return base_mr->allocate(size, stream); }; - auto deallocate_callback = [&base_mr]( - void* ptr, std::size_t size, cuda_stream_view stream, void* arg) { + auto deallocate_callback = [](void* ptr, std::size_t size, cuda_stream_view stream, void* arg) { + auto base_mr = static_cast(arg); base_mr->deallocate(ptr, size, stream); }; auto mr = - rmm::mr::callback_memory_resource(allocate_callback, deallocate_callback, nullptr, nullptr); + rmm::mr::callback_memory_resource(allocate_callback, deallocate_callback, base_mr, base_mr); auto ptr = mr.allocate(10_MiB); + ASSERT_NE(nullptr, ptr); mr.deallocate(ptr, 10_MiB); } @@ -51,19 +53,15 @@ TEST(CallbackTest, LoggingTest) testing::internal::CaptureStdout(); auto base_mr = rmm::mr::get_current_device_resource(); - auto allocate_callback = [expected_base_mr]( - std::size_t size, cuda_stream_view stream, void* arg) { + auto allocate_callback = [](std::size_t size, cuda_stream_view stream, void* arg) { std::cout << "Allocating " << size << " bytes" << std::endl; - auto base_mr = static_cast(arg); - ASSERT_EQ(base_mr, expected_base_mr); + auto base_mr = static_cast(arg); return base_mr->allocate(size, stream); }; - auto deallocate_callback = [expected_base_mr]( - void* ptr, std::size_t size, cuda_stream_view stream, void* arg) { + auto deallocate_callback = [](void* ptr, std::size_t size, cuda_stream_view stream, void* arg) { std::cout << "Deallocating " << size << " bytes" << std::endl; - auto base_mr = static_cast(arg); - ASSERT_EQ(base_mr, expected_base_mr); + auto base_mr = static_cast(arg); base_mr->deallocate(ptr, size, stream); }; auto mr = From 169fa1fa41610f95b0979d81118a9b7d96ed235b Mon Sep 17 00:00:00 2001 From: Ashwin Srinath <3190405+shwina@users.noreply.github.com> Date: Wed, 13 Apr 2022 14:47:19 -0400 Subject: [PATCH 22/34] Update include/rmm/mr/device/callback_memory_resource.hpp Co-authored-by: Bradley Dice --- include/rmm/mr/device/callback_memory_resource.hpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index acb1106bd..66cdce955 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -42,7 +42,7 @@ using allocate_callback_t = std::function Date: Wed, 13 Apr 2022 14:51:19 -0400 Subject: [PATCH 23/34] Address various review comments --- include/rmm/mr/device/callback_memory_resource.hpp | 2 +- python/rmm/_lib/memory_resource.pyx | 4 ++-- tests/mr/device/callback_mr_tests.cpp | 10 +++++----- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 66cdce955..8ae8ce065 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2021, NVIDIA CORPORATION. + * Copyright (c) 2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index 246e0739a..70409d9ab 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -467,7 +467,7 @@ cdef class BinningMemoryResource(UpstreamResourceAdaptor): cdef void* _allocate_callback_wrapper( size_t nbytes, cuda_stream_view stream, - void *ctx + void* ctx ) with gil: return ((ctx)(nbytes)) @@ -475,7 +475,7 @@ cdef void _deallocate_callback_wrapper( void* ptr, size_t nbytes, cuda_stream_view stream, - void *ctx + void* ctx ) with gil: (ctx)((ptr), nbytes) diff --git a/tests/mr/device/callback_mr_tests.cpp b/tests/mr/device/callback_mr_tests.cpp index 9fc333fc5..60103b024 100644 --- a/tests/mr/device/callback_mr_tests.cpp +++ b/tests/mr/device/callback_mr_tests.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -15,13 +15,13 @@ */ #include "../../byte_literals.hpp" -#include "rmm/cuda_stream_view.hpp" -#include "rmm/mr/device/device_memory_resource.hpp" - -#include +#include #include #include #include +#include + +#include #include From 8d69937e411a4bb1aa2189a579807aa44498909e Mon Sep 17 00:00:00 2001 From: Ashwin Srinath <3190405+shwina@users.noreply.github.com> Date: Thu, 14 Apr 2022 08:53:17 -0400 Subject: [PATCH 24/34] Update include/rmm/mr/device/callback_memory_resource.hpp Co-authored-by: Vyas Ramasubramani --- include/rmm/mr/device/callback_memory_resource.hpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 8ae8ce065..aeefe0fd7 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -33,8 +33,10 @@ namespace rmm::mr { * `stream`. The stream-ordered behavior requirements are identical to * `device_memory_resource::allocate`. * - * * Additional context to the callback may be passed via `arg` that is provided at - * construction of the `callback_memory_resource`. + * * This signature is compatible with `do_allocate` but adds the extra function + * parameter `arg`. The `arg` is provided to the constructor of the + * `callback_memory_resource` and will be forwarded along to every invocation + * of the callback function. */ using allocate_callback_t = std::function; From ba3c1bde5d4e6e033a43be17bef8595dc63e7d82 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath <3190405+shwina@users.noreply.github.com> Date: Thu, 14 Apr 2022 08:53:25 -0400 Subject: [PATCH 25/34] Update include/rmm/mr/device/callback_memory_resource.hpp Co-authored-by: Vyas Ramasubramani --- include/rmm/mr/device/callback_memory_resource.hpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index aeefe0fd7..7f0a517b7 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -51,8 +51,10 @@ using allocate_callback_t = std::function; From c1ba1879e59885fb09776f4d4cebfa7e0f91d3d8 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath <3190405+shwina@users.noreply.github.com> Date: Thu, 14 Apr 2022 08:53:39 -0400 Subject: [PATCH 26/34] Update python/rmm/_lib/memory_resource.pyx Co-authored-by: Vyas Ramasubramani --- python/rmm/_lib/memory_resource.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index 70409d9ab..75d3580b7 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -480,7 +480,7 @@ cdef void _deallocate_callback_wrapper( (ctx)((ptr), nbytes) -cdef class CallbackMemoryResource: +cdef class CallbackMemoryResource(DeviceMemoryResource): """ A memory resource that uses the user-provided callables to do memory allocation and deallocation. From 9d9486123f7e576a64dde32f811e0c96c1c7f001 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 14 Apr 2022 08:54:11 -0400 Subject: [PATCH 27/34] Default to nullptr --- include/rmm/mr/device/callback_memory_resource.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 8ae8ce065..ea359b3fd 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -77,8 +77,8 @@ class callback_memory_resource final : public device_memory_resource { */ callback_memory_resource(allocate_callback_t allocate_callback, deallocate_callback_t deallocate_callback, - void* allocate_callback_arg, - void* deallocate_callback_arg) noexcept + void* allocate_callback_arg = nullptr, + void* deallocate_callback_arg = nullptr) noexcept : allocate_callback_(allocate_callback), deallocate_callback_(deallocate_callback), allocate_callback_arg_(allocate_callback_arg), From 347dc9cd73a6c1fafa34b559b4d2a3ec7b1ab779 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 14 Apr 2022 08:55:14 -0400 Subject: [PATCH 28/34] Move docs --- python/rmm/_lib/memory_resource.pyx | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index 75d3580b7..c70d49456 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -485,14 +485,6 @@ cdef class CallbackMemoryResource(DeviceMemoryResource): A memory resource that uses the user-provided callables to do memory allocation and deallocation. - The allocation function must accept a single integer argument, - representing the number of bytes to allocate, and return an - integer representing the pointer to the allocated memory. - - The deallocation function must accept two arguments, an integer - representing the pointer to the memory to free, and a second - integer representing the number of bytes to free. - ``CallbackMemoryResource`` should really only be used for debugging memory issues, as there is a significant performance penalty associated with using a Python function for each memory @@ -501,7 +493,13 @@ cdef class CallbackMemoryResource(DeviceMemoryResource): Parameters ---------- allocate_func: callable + The allocation function must accept a single integer argument, + representing the number of bytes to allocate, and return an + integer representing the pointer to the allocated memory. deallocate_func: callable + The deallocation function must accept two arguments, an integer + representing the pointer to the memory to free, and a second + integer representing the number of bytes to free. """ def __init__( self, From b32dfa2cd443d4d68e42e94a9a5df6556b8e63db Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Thu, 14 Apr 2022 09:38:54 -0400 Subject: [PATCH 29/34] Use mock class in test --- .../rmm/mr/device/callback_memory_resource.hpp | 2 +- tests/mock_resource.hpp | 17 +++++++++++++++++ tests/mr/device/aligned_mr_tests.cpp | 11 +---------- tests/mr/device/callback_mr_tests.cpp | 14 ++++++++++---- 4 files changed, 29 insertions(+), 15 deletions(-) create mode 100644 tests/mock_resource.hpp diff --git a/include/rmm/mr/device/callback_memory_resource.hpp b/include/rmm/mr/device/callback_memory_resource.hpp index 8d60582b5..e2deef46b 100644 --- a/include/rmm/mr/device/callback_memory_resource.hpp +++ b/include/rmm/mr/device/callback_memory_resource.hpp @@ -81,7 +81,7 @@ class callback_memory_resource final : public device_memory_resource { */ callback_memory_resource(allocate_callback_t allocate_callback, deallocate_callback_t deallocate_callback, - void* allocate_callback_arg = nullptr, + void* allocate_callback_arg = nullptr, void* deallocate_callback_arg = nullptr) noexcept : allocate_callback_(allocate_callback), deallocate_callback_(deallocate_callback), diff --git a/tests/mock_resource.hpp b/tests/mock_resource.hpp new file mode 100644 index 000000000..6501b82b0 --- /dev/null +++ b/tests/mock_resource.hpp @@ -0,0 +1,17 @@ +#include + +#include + +namespace rmm::test { + +class mock_resource : public rmm::mr::device_memory_resource { + public: + MOCK_METHOD(bool, supports_streams, (), (const, override, noexcept)); + MOCK_METHOD(bool, supports_get_mem_info, (), (const, override, noexcept)); + MOCK_METHOD(void*, do_allocate, (std::size_t, cuda_stream_view), (override)); + MOCK_METHOD(void, do_deallocate, (void*, std::size_t, cuda_stream_view), (override)); + using size_pair = std::pair; + MOCK_METHOD(size_pair, do_get_mem_info, (cuda_stream_view), (const, override)); +}; + +} // namespace rmm::test diff --git a/tests/mr/device/aligned_mr_tests.cpp b/tests/mr/device/aligned_mr_tests.cpp index f1ed561f4..395c9c02a 100644 --- a/tests/mr/device/aligned_mr_tests.cpp +++ b/tests/mr/device/aligned_mr_tests.cpp @@ -14,6 +14,7 @@ * limitations under the License. */ +#include "../../mock_resource.hpp" #include #include #include @@ -28,16 +29,6 @@ namespace { using ::testing::Return; -class mock_resource : public rmm::mr::device_memory_resource { - public: - MOCK_METHOD(bool, supports_streams, (), (const, override, noexcept)); - MOCK_METHOD(bool, supports_get_mem_info, (), (const, override, noexcept)); - MOCK_METHOD(void*, do_allocate, (std::size_t, cuda_stream_view), (override)); - MOCK_METHOD(void, do_deallocate, (void*, std::size_t, cuda_stream_view), (override)); - using size_pair = std::pair; - MOCK_METHOD(size_pair, do_get_mem_info, (cuda_stream_view), (const, override)); -}; - using aligned_mock = rmm::mr::aligned_resource_adaptor; using aligned_real = rmm::mr::aligned_resource_adaptor; diff --git a/tests/mr/device/callback_mr_tests.cpp b/tests/mr/device/callback_mr_tests.cpp index 60103b024..101a75fc8 100644 --- a/tests/mr/device/callback_mr_tests.cpp +++ b/tests/mr/device/callback_mr_tests.cpp @@ -15,6 +15,7 @@ */ #include "../../byte_literals.hpp" +#include "../../mock_resource.hpp" #include #include #include @@ -25,14 +26,20 @@ #include +#include #include namespace rmm::test { namespace { -TEST(CallbackTest, PassThroughTest) +using ::testing::_; + +TEST(CallbackTest, TestCallbacksAreInvoked) { - auto base_mr = rmm::mr::get_current_device_resource(); + auto base_mr = mock_resource(); + EXPECT_CALL(base_mr, do_allocate(10_MiB, cuda_stream_view{})).Times(1); + EXPECT_CALL(base_mr, do_deallocate(_, 10_MiB, cuda_stream_view{})).Times(1); + auto allocate_callback = [](std::size_t size, cuda_stream_view stream, void* arg) { auto base_mr = static_cast(arg); return base_mr->allocate(size, stream); @@ -42,9 +49,8 @@ TEST(CallbackTest, PassThroughTest) base_mr->deallocate(ptr, size, stream); }; auto mr = - rmm::mr::callback_memory_resource(allocate_callback, deallocate_callback, base_mr, base_mr); + rmm::mr::callback_memory_resource(allocate_callback, deallocate_callback, &base_mr, &base_mr); auto ptr = mr.allocate(10_MiB); - ASSERT_NE(nullptr, ptr); mr.deallocate(ptr, 10_MiB); } From 02a05d6e0b48fba8bbbd0c7b2083fa139702d64b Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Tue, 19 Apr 2022 11:59:22 -0400 Subject: [PATCH 30/34] Add a test for allocate/deallocate --- python/rmm/tests/test_rmm.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index 329dffeba..c232048e3 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -721,6 +721,15 @@ def test_dev_buf_circle_ref_dealloc(): gc.collect() +def test_mr_allocate_deallocate(): + mr = rmm.mr.TrackingResourceAdaptor(rmm.mr.get_current_device_resource()) + size = 1 << 23 # 8 MiB + ptr = mr.allocate(size) + assert mr.get_allocated_bytes() == 1 << 23 + mr.deallocate(ptr, size) + assert mr.get_allocated_bytes() == 0 + + def test_custom_mr(capsys): base_mr = rmm.mr.CudaMemoryResource() From f9e445bcd346dda34d5892881d010c5a6db52242 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Tue, 19 Apr 2022 12:00:12 -0400 Subject: [PATCH 31/34] Add header --- tests/mock_resource.hpp | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/tests/mock_resource.hpp b/tests/mock_resource.hpp index 6501b82b0..0436e2a2a 100644 --- a/tests/mock_resource.hpp +++ b/tests/mock_resource.hpp @@ -1,3 +1,20 @@ +/* + * Copyright (c) 2022, 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. + */ +#pragma once + #include #include From d080157b6e17c55425d128983f035f258208533d Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Tue, 19 Apr 2022 13:23:51 -0400 Subject: [PATCH 32/34] Add example --- python/rmm/_lib/memory_resource.pyx | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index c70d49456..d4a4442d2 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -500,6 +500,26 @@ cdef class CallbackMemoryResource(DeviceMemoryResource): The deallocation function must accept two arguments, an integer representing the pointer to the memory to free, and a second integer representing the number of bytes to free. + + Examples + ------- + >>> import rmm + >>> base_mr = rmm.mr.CudaMemoryResource() + >>> def allocate_func(size): + ... print(f"Allocating {size} bytes") + ... return base_mr.allocate(size) + ... + >>> def deallocate_func(ptr, size): + ... print(f"Deallocating {size} bytes") + ... return base_mr.deallocate(ptr, size) + ... + >>> rmm.mr.set_current_device_resource( + rmm.mr.CallbackMemoryResource(allocate_func, deallocate_func) + ) + >>> dbuf = rmm.DeviceBuffer(size=256) + Allocating 256 bytes + >>> del dbuf + Deallocating 256 bytes """ def __init__( self, From 5529ca16dbd85480b86cb80295423b63f2506d06 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Tue, 19 Apr 2022 16:06:18 -0400 Subject: [PATCH 33/34] Copyright --- python/rmm/_lib/memory_resource.pxd | 2 +- python/rmm/_lib/memory_resource.pyx | 2 +- python/rmm/tests/test_rmm.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/python/rmm/_lib/memory_resource.pxd b/python/rmm/_lib/memory_resource.pxd index c3d4084b5..92ce1ce54 100644 --- a/python/rmm/_lib/memory_resource.pxd +++ b/python/rmm/_lib/memory_resource.pxd @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2021, NVIDIA CORPORATION. +# Copyright (c) 2020-2022, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/python/rmm/_lib/memory_resource.pyx b/python/rmm/_lib/memory_resource.pyx index d4a4442d2..a110bff97 100644 --- a/python/rmm/_lib/memory_resource.pyx +++ b/python/rmm/_lib/memory_resource.pyx @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2021, NVIDIA CORPORATION. +# Copyright (c) 2020-2022, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index c232048e3..fa38ee9fe 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -1,4 +1,4 @@ -# Copyright (c) 2020-2021, NVIDIA CORPORATION. +# Copyright (c) 2020-2022, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. From 3c73ca160e1e3adac04fd239751abac99dd0f596 Mon Sep 17 00:00:00 2001 From: Ashwin Srinath Date: Tue, 19 Apr 2022 16:06:41 -0400 Subject: [PATCH 34/34] Copyright --- tests/mr/device/aligned_mr_tests.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/mr/device/aligned_mr_tests.cpp b/tests/mr/device/aligned_mr_tests.cpp index 395c9c02a..62b941d6f 100644 --- a/tests/mr/device/aligned_mr_tests.cpp +++ b/tests/mr/device/aligned_mr_tests.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2021, NVIDIA CORPORATION. + * Copyright (c) 2021-2022, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License.