You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@arrow.apache.org by we...@apache.org on 2019/03/14 14:47:29 UTC
[arrow] branch master updated: ARROW-4486: [Python][CUDA] Add base
argument to foreign_buffer
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 <pe...@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 <pe...@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)