Skip to content

Commit

Permalink
Accept memory resource (and stream) in copy
Browse files Browse the repository at this point in the history
Changes `copy` to allow a memory resource to be specified (using the
default device memory resource if unspecified). Also accepts a `stream`
argument (similar to other functions).
  • Loading branch information
jakirkham committed Oct 18, 2022
1 parent 295156f commit 1532bff
Showing 1 changed file with 6 additions and 4 deletions.
10 changes: 6 additions & 4 deletions python/rmm/_lib/device_buffer.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -146,7 +146,9 @@ cdef class DeviceBuffer:
}
return intf

def copy(self):
def copy(self, *,
Stream stream=DEFAULT_STREAM,
DeviceMemoryResource mr=None):
"""Returns a copy of DeviceBuffer.
Returns
Expand All @@ -165,9 +167,9 @@ cdef class DeviceBuffer:
>>> assert db is not db_copy
>>> assert db.ptr != db_copy.ptr
"""
ret = DeviceBuffer(ptr=self.ptr, size=self.size, stream=self.stream)
ret.mr = self.mr
return ret
return DeviceBuffer(
ptr=self.ptr, size=self.size, stream=stream, mr=mr
)

def __copy__(self):
return self.copy()
Expand Down

0 comments on commit 1532bff

Please sign in to comment.