pitrou commented on code in PR #50923:
URL: https://github.com/apache/arrow/pull/50923#discussion_r3948574372
##########
python/pyarrow/tests/test_cuda_numba_interop.py:
##########
@@ -50,8 +52,14 @@ def teardown_module(module):
ids=context_choice_ids)
def test_context(c):
ctx, nb_ctx = context_choices[c]
- assert ctx.handle == nb_ctx.handle.value
- assert ctx.handle == ctx.to_numba().handle.value
+ converted = ctx.to_numba()
+ assert type(converted.handle) is type(nb_ctx.handle)
+ if isinstance(converted.handle, ctypes.c_void_p):
Review Comment:
Here as well, can you comment on which path corresponds to which Numba
version(s)?
##########
python/pyarrow/_cuda.pyx:
##########
@@ -16,13 +16,43 @@
# under the License.
+import ctypes
+
+from pyarrow.vendored.version import Version
+
from pyarrow.lib cimport *
from pyarrow.includes.libarrow_cuda cimport *
from pyarrow.lib import allocate_buffer, as_buffer, ArrowTypeError
from pyarrow.util import get_contiguous_span
cimport cpython as cp
+_NUMBA_CUDA_NATIVE_CONTEXT_VERSION = Version("0.28")
+
+
+def _numba_context_handle_value(handle):
+ """Return the integer value of a legacy or native Numba context handle."""
+ if hasattr(handle, "value"):
+ return handle.value
+ return int(handle)
+
+
+def _make_numba_context_handle(uintptr_t handle):
+ """Create the context handle representation expected by Numba."""
+ import numba.cuda
+
+ # numba-cuda 0.28 replaced ctypes context handles with CUContext type.
+ if getattr(numba.cuda, "implementation", None) == "NVIDIA":
+ import numba_cuda
+
+ version = Version(numba_cuda.__version__)
+ if version >= _NUMBA_CUDA_NATIVE_CONTEXT_VERSION:
+ from cuda.bindings.driver import CUcontext
+ return CUcontext(handle)
+
+ return ctypes.c_void_p(handle)
Review Comment:
This is the legacy path, right?
##########
python/pyarrow/_cuda.pyx:
##########
@@ -16,13 +16,43 @@
# under the License.
+import ctypes
+
+from pyarrow.vendored.version import Version
+
from pyarrow.lib cimport *
from pyarrow.includes.libarrow_cuda cimport *
from pyarrow.lib import allocate_buffer, as_buffer, ArrowTypeError
from pyarrow.util import get_contiguous_span
cimport cpython as cp
+_NUMBA_CUDA_NATIVE_CONTEXT_VERSION = Version("0.28")
+
+
+def _numba_context_handle_value(handle):
+ """Return the integer value of a legacy or native Numba context handle."""
+ if hasattr(handle, "value"):
+ return handle.value
+ return int(handle)
Review Comment:
Can you add comments telling which path is legacy or native?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
