This is an automated email from the ASF dual-hosted git repository. wesm pushed a commit to branch master in repository https://gitbox.apache.org/repos/asf/arrow.git
The following commit(s) were added to refs/heads/master by this push: new 2b574f9 ARROW-4486: [Python][CUDA] Add base argument to foreign_buffer 2b574f9 is described below commit 2b574f913784e77cc79ec7d002c94b1ae51116f1 Author: Pearu Peterson <pearu.peter...@gmail.com> AuthorDate: Thu Mar 14 09:47:20 2019 -0500 ARROW-4486: [Python][CUDA] Add base argument to foreign_buffer Resolves [ARROW-4486](https://issues.apache.org/jira/browse/ARROW-4486) Author: Pearu Peterson <pearu.peter...@gmail.com> Closes #3850 from pearu/arrow-4486 and squashes the following commits: e96265df0 <Pearu Peterson> Add base argument to foreign_buffer, resolves ARROW-4486 --- python/pyarrow/_cuda.pxd | 5 ++++- python/pyarrow/_cuda.pyx | 26 +++++++++++++++++++------- python/pyarrow/tests/test_cuda.py | 27 +++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 8 deletions(-) diff --git a/python/pyarrow/_cuda.pxd b/python/pyarrow/_cuda.pxd index 3b8d966..1180601 100644 --- a/python/pyarrow/_cuda.pxd +++ b/python/pyarrow/_cuda.pxd @@ -41,8 +41,11 @@ cdef class IpcMemHandle: cdef class CudaBuffer(Buffer): cdef: shared_ptr[CCudaBuffer] cuda_buffer + object base - cdef void init_cuda(self, const shared_ptr[CCudaBuffer]& buffer) + cdef void init_cuda(self, + const shared_ptr[CCudaBuffer]& buffer, + object base) cdef class HostBuffer(Buffer): diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx index fa84fc6..87be0e6 100644 --- a/python/pyarrow/_cuda.pyx +++ b/python/pyarrow/_cuda.pyx @@ -190,7 +190,7 @@ cdef class Context: check_status(self.context.get().Allocate(nbytes, &cudabuf)) return pyarrow_wrap_cudabuffer(cudabuf) - def foreign_buffer(self, address, size): + def foreign_buffer(self, address, size, base=None): """Create device buffer from address and size as a view. The caller is responsible for allocating and freeing the @@ -206,6 +206,8 @@ cdef class Context: `get_device_address` method. size : int Specify the size of device buffer in bytes. + base : {None, object} + Specify object that owns the referenced memory. Returns ------- @@ -222,7 +224,7 @@ cdef class Context: check_status(self.context.get().View(<uint8_t*>c_addr, c_size, &cudabuf)) - return pyarrow_wrap_cudabuffer(cudabuf) + return pyarrow_wrap_cudabuffer_base(cudabuf, base) def open_ipc_buffer(self, ipc_handle): """ Open existing CUDA IPC memory handle @@ -309,7 +311,7 @@ cdef class Context: """ if isinstance(obj, HostBuffer): - return self.foreign_buffer(obj.address, obj.size) + return self.foreign_buffer(obj.address, obj.size, base=obj) elif isinstance(obj, Buffer): return CudaBuffer.from_buffer(obj) elif isinstance(obj, CudaBuffer): @@ -323,7 +325,7 @@ cdef class Context: start, end = get_contiguous_span( desc['shape'], desc.get('strides'), np.dtype(desc['typestr']).itemsize) - return self.foreign_buffer(addr + start, end - start) + return self.foreign_buffer(addr + start, end - start, base=obj) raise ArrowTypeError('cannot create device buffer view from' ' `%s` object' % (type(obj))) @@ -387,9 +389,12 @@ cdef class CudaBuffer(Buffer): "`<pyarrow.Context instance>.device_buffer`" " method instead.") - cdef void init_cuda(self, const shared_ptr[CCudaBuffer]& buffer): + cdef void init_cuda(self, + const shared_ptr[CCudaBuffer]& buffer, + object base): self.cuda_buffer = buffer self.init(<shared_ptr[CBuffer]> buffer) + self.base = base @staticmethod def from_buffer(buf): @@ -426,7 +431,7 @@ cdef class CudaBuffer(Buffer): ctx = Context.from_numba(mem.context) if mem.device_pointer.value is None and mem.size==0: return ctx.new_buffer(0) - return ctx.foreign_buffer(mem.device_pointer.value, mem.size) + return ctx.foreign_buffer(mem.device_pointer.value, mem.size, base=mem) def to_numba(self): """Return numba memory pointer of CudaBuffer instance. @@ -949,9 +954,16 @@ cdef public api bint pyarrow_is_cudabuffer(object buffer): cdef public api object \ + pyarrow_wrap_cudabuffer_base(const shared_ptr[CCudaBuffer]& buf, base): + cdef CudaBuffer result = CudaBuffer.__new__(CudaBuffer) + result.init_cuda(buf, base) + return result + + +cdef public api object \ pyarrow_wrap_cudabuffer(const shared_ptr[CCudaBuffer]& buf): cdef CudaBuffer result = CudaBuffer.__new__(CudaBuffer) - result.init_cuda(buf) + result.init_cuda(buf, None) return result diff --git a/python/pyarrow/tests/test_cuda.py b/python/pyarrow/tests/test_cuda.py index 7c56e33..c908e06 100644 --- a/python/pyarrow/tests/test_cuda.py +++ b/python/pyarrow/tests/test_cuda.py @@ -265,6 +265,33 @@ def test_context_from_object(size): ctx.buffer_from_object(np.array([1, 2, 3])) +def test_foreign_buffer(): + ctx = global_context + dtype = np.dtype(np.uint8) + size = 10 + hbuf = cuda.new_host_buffer(size * dtype.itemsize) + + # test host buffer memory reference counting + rc = sys.getrefcount(hbuf) + fbuf = ctx.foreign_buffer(hbuf.address, hbuf.size, hbuf) + assert sys.getrefcount(hbuf) == rc + 1 + del fbuf + assert sys.getrefcount(hbuf) == rc + + # test postponed dealloction of host buffer memory + fbuf = ctx.foreign_buffer(hbuf.address, hbuf.size, hbuf) + del hbuf + fbuf.copy_to_host() + + # test deallocating the host buffer memory making it inaccessible + hbuf = cuda.new_host_buffer(size * dtype.itemsize) + fbuf = ctx.foreign_buffer(hbuf.address, hbuf.size) + del hbuf + with pytest.raises(pa.ArrowIOError, + match=('Cuda Driver API call in')): + fbuf.copy_to_host() + + @pytest.mark.parametrize("size", [0, 1, 1000]) def test_CudaBuffer(size): arr, buf = make_random_buffer(size)