From 65c4d92f8ea9e5679022abc9a8e8c3acb7a75fab Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Mon, 10 Oct 2022 09:06:00 -0700 Subject: [PATCH 1/7] add DeviceBuffer.copy --- python/rmm/_lib/device_buffer.pyx | 26 +++++++++++++++++++++++--- python/rmm/tests/test_rmm.py | 19 +++++++++++++++++++ 2 files changed, 42 insertions(+), 3 deletions(-) diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index 9ac034774..8ba36abff 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -133,6 +133,27 @@ cdef class DeviceBuffer: } return intf + def copy(self): + """Returns a deep copy of DeviceBuffer + + Returns + ------- + A deep copy of existing ``DeviceBuffer`` + + Examples + -------- + >>> import rmm + >>> db = rmm.DeviceBuffer.to_device(b"abc") + >>> db_copy = db.copy() + >>> db.copy_to_host() + array([97, 98, 99], dtype=uint8) + >>> db_copy.copy_to_host() + array([97, 98, 99], dtype=uint8) + """ + new_copy = DeviceBuffer(size=self.size) + copy_device_to_ptr(self.ptr, new_copy.ptr, self.size) + return new_copy + @staticmethod cdef DeviceBuffer c_from_unique_ptr(unique_ptr[device_buffer] ptr): cdef DeviceBuffer buf = DeviceBuffer.__new__(DeviceBuffer) @@ -475,13 +496,12 @@ cpdef void copy_device_to_ptr(uintptr_t d_src, Examples -------- >>> import rmm - >>> import numpy as np >>> db = rmm.DeviceBuffer(size=5) >>> db2 = rmm.DeviceBuffer.to_device(b"abc") >>> rmm._lib.device_buffer.copy_device_to_ptr(db2.ptr, db.ptr, db2.size) >>> hb = db.copy_to_host() - >>> print(hb) - array([10, 11, 12, 0, 0], dtype=uint8) + >>> hb + array([97, 98, 99, 0, 0], dtype=uint8) """ with nogil: diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index e4b503c08..72f30ee57 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -871,3 +871,22 @@ def func_with_arg(x): rmm.unregister_reinitialize_hook(func_with_arg) rmm.reinitialize() assert L == [2] + + +@pytest.mark.parametrize( + "cuda_ary", + [ + lambda: rmm.DeviceBuffer.to_device(b"abc"), + lambda: cuda.to_device(np.array([97, 98, 99, 0, 0], dtype="u1")), + ], +) +def test_rmm_device_buffer_copy(cuda_ary): + cuda_ary = cuda_ary() + db = rmm.DeviceBuffer.to_device(np.zeros(5, dtype="u1")) + db.copy_from_device(cuda_ary) + db_copy = db.copy() + + expected = np.array([97, 98, 99, 0, 0], dtype="u1") + result = db_copy.copy_to_host() + + np.testing.assert_equal(expected, result) From 02f0d520d8c5b7d682cfc5d182b97e8c2d33f828 Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Wed, 12 Oct 2022 10:27:59 -0700 Subject: [PATCH 2/7] address first pass of reviews --- python/rmm/_lib/device_buffer.pxd | 2 ++ python/rmm/_lib/device_buffer.pyx | 13 +++++++++---- python/rmm/tests/test_rmm.py | 4 ++++ 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/python/rmm/_lib/device_buffer.pxd b/python/rmm/_lib/device_buffer.pxd index fee6b126a..42c2b4fef 100644 --- a/python/rmm/_lib/device_buffer.pxd +++ b/python/rmm/_lib/device_buffer.pxd @@ -26,6 +26,8 @@ cdef extern from "rmm/device_buffer.hpp" namespace "rmm" nogil: device_buffer(size_t size, cuda_stream_view stream) except + device_buffer(const void* source_data, size_t size, cuda_stream_view stream) except + + device_buffer(const device_buffer buf, + cuda_stream_view stream) except + void reserve(size_t new_capacity, cuda_stream_view stream) except + void resize(size_t new_size, cuda_stream_view stream) except + void shrink_to_fit(cuda_stream_view stream) except + diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index 8ba36abff..d77c0616e 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -17,7 +17,7 @@ cimport cython from cpython.bytes cimport PyBytes_AS_STRING, PyBytes_FromStringAndSize from cython.operator cimport dereference from libc.stdint cimport uintptr_t -from libcpp.memory cimport unique_ptr +from libcpp.memory cimport make_unique, unique_ptr from libcpp.utility cimport move from rmm._cuda.stream cimport Stream @@ -149,10 +149,15 @@ cdef class DeviceBuffer: array([97, 98, 99], dtype=uint8) >>> db_copy.copy_to_host() array([97, 98, 99], dtype=uint8) + >>> assert db is not db_copy + >>> assert db.ptr != db_copy.ptr """ - new_copy = DeviceBuffer(size=self.size) - copy_device_to_ptr(self.ptr, new_copy.ptr, self.size) - return new_copy + ret = DeviceBuffer.c_from_unique_ptr( + make_unique[device_buffer](self.c_obj.get()[0], self.stream.view()) + ) + ret.mr = self.mr + ret.stream = self.stream + return ret @staticmethod cdef DeviceBuffer c_from_unique_ptr(unique_ptr[device_buffer] ptr): diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index 72f30ee57..7e98856cd 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -886,6 +886,10 @@ def test_rmm_device_buffer_copy(cuda_ary): db.copy_from_device(cuda_ary) db_copy = db.copy() + assert db is not db_copy + assert db.ptr != db_copy.ptr + assert len(db) == len(db_copy) + expected = np.array([97, 98, 99, 0, 0], dtype="u1") result = db_copy.copy_to_host() From a1185e39a108e6bcd3c7f66157c6cd04e36ec08a Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Wed, 12 Oct 2022 10:38:37 -0700 Subject: [PATCH 3/7] add __deep_copy__ --- python/rmm/_lib/device_buffer.pyx | 3 +++ python/rmm/tests/test_rmm.py | 8 ++++++-- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index d77c0616e..ff8b7d609 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -159,6 +159,9 @@ cdef class DeviceBuffer: ret.stream = self.stream return ret + def __deepcopy__(self, memo): + return self.copy() + @staticmethod cdef DeviceBuffer c_from_unique_ptr(unique_ptr[device_buffer] ptr): cdef DeviceBuffer buf = DeviceBuffer.__new__(DeviceBuffer) diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index 7e98856cd..02dcec946 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import gc import os import sys @@ -880,11 +881,14 @@ def func_with_arg(x): lambda: cuda.to_device(np.array([97, 98, 99, 0, 0], dtype="u1")), ], ) -def test_rmm_device_buffer_copy(cuda_ary): +@pytest.mark.parametrize( + "make_copy", [lambda db: db.copy(), lambda db: copy.deepcopy(db)] +) +def test_rmm_device_buffer_copy(cuda_ary, make_copy): cuda_ary = cuda_ary() db = rmm.DeviceBuffer.to_device(np.zeros(5, dtype="u1")) db.copy_from_device(cuda_ary) - db_copy = db.copy() + db_copy = make_copy(db) assert db is not db_copy assert db.ptr != db_copy.ptr From 33cd0e3811d9c2cf9269dbfca38ad362e6a9af82 Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Wed, 12 Oct 2022 12:40:12 -0500 Subject: [PATCH 4/7] Update python/rmm/_lib/device_buffer.pyx Co-authored-by: Bradley Dice --- python/rmm/_lib/device_buffer.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index ff8b7d609..e42111971 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -134,7 +134,7 @@ cdef class DeviceBuffer: return intf def copy(self): - """Returns a deep copy of DeviceBuffer + """Returns a deep copy of DeviceBuffer. Returns ------- From 35345d050d916bfcc88dc0d27bcc5baf1827b23e Mon Sep 17 00:00:00 2001 From: galipremsagar Date: Mon, 17 Oct 2022 11:34:48 -0700 Subject: [PATCH 5/7] address reviews --- python/rmm/_lib/device_buffer.pyx | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index ff8b7d609..3f744c4d3 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -152,14 +152,11 @@ cdef class DeviceBuffer: >>> assert db is not db_copy >>> assert db.ptr != db_copy.ptr """ - ret = DeviceBuffer.c_from_unique_ptr( - make_unique[device_buffer](self.c_obj.get()[0], self.stream.view()) - ) + ret = DeviceBuffer(ptr=self.ptr, size=self.size, stream=self.stream) ret.mr = self.mr - ret.stream = self.stream return ret - def __deepcopy__(self, memo): + def __copy__(self): return self.copy() @staticmethod From 470dc5ae7f373b77d25b11a84bece08ba36873de Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Mon, 17 Oct 2022 13:35:35 -0500 Subject: [PATCH 6/7] Update python/rmm/_lib/device_buffer.pyx Co-authored-by: Bradley Dice --- python/rmm/_lib/device_buffer.pyx | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/rmm/_lib/device_buffer.pyx b/python/rmm/_lib/device_buffer.pyx index e42111971..ce6c744af 100644 --- a/python/rmm/_lib/device_buffer.pyx +++ b/python/rmm/_lib/device_buffer.pyx @@ -134,7 +134,7 @@ cdef class DeviceBuffer: return intf def copy(self): - """Returns a deep copy of DeviceBuffer. + """Returns a copy of DeviceBuffer. Returns ------- From 816667c7927097c0f751a87dfb00f0f217096fc5 Mon Sep 17 00:00:00 2001 From: GALI PREM SAGAR Date: Mon, 17 Oct 2022 18:22:02 -0500 Subject: [PATCH 7/7] Update python/rmm/tests/test_rmm.py Co-authored-by: Bradley Dice --- python/rmm/tests/test_rmm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/rmm/tests/test_rmm.py b/python/rmm/tests/test_rmm.py index d2bc073e5..931ff5336 100644 --- a/python/rmm/tests/test_rmm.py +++ b/python/rmm/tests/test_rmm.py @@ -898,7 +898,7 @@ def func_with_arg(x): ], ) @pytest.mark.parametrize( - "make_copy", [lambda db: db.copy(), lambda db: copy.deepcopy(db)] + "make_copy", [lambda db: db.copy(), lambda db: copy.copy(db)] ) def test_rmm_device_buffer_copy(cuda_ary, make_copy): cuda_ary = cuda_ary()