-
Notifications
You must be signed in to change notification settings - Fork 3.6k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
ARROW-3451: [C++/Python] pyarrow and numba CUDA interop #2732
Conversation
Author: Kouhei Sutou <kou@clear-code.com> Closes apache#2701 from kou/add-workaround-missing-gemfile and squashes the following commits: fc80d6e <Kouhei Sutou> Add workaround to verify 0.11.0
It's helpful to setup test environment. Author: Kouhei Sutou <kou@clear-code.com> Closes apache#2702 from kou/c-glib-add-missing-gemfile-to-archive and squashes the following commits: 5c97dad <Kouhei Sutou> Include Gemfile to archive
Author: Kouhei Sutou <kou@clear-code.com> Closes apache#2703 from kou/expand-variables-in-commit-message and squashes the following commits: da4f3bc <Kouhei Sutou> Expand variables in commit message
Author: Pindikura Ravindra <ravindra@dremio.com> Closes apache#2695 from pravindra/re2 and squashes the following commits: 4408b85 <Pindikura Ravindra> ARROW-3331: Add re2 to toolchain
Author: Kouhei Sutou <kou@clear-code.com> Closes apache#2706 from kou/fix-too-much-escape-changelog and squashes the following commits: a635b86 <Kouhei Sutou> Fix too much Markdown escape in CHANGELOG
I'm going to have to squash this branch to clean things up. Please avoid using |
Change-Id: I6057fdc68720c5d1c215ed16eb6aaedf0b67818e
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM so far, let me know when you want me to review again
Sorry, I had force-pushed your branch. |
@wesm , please review. My git skills are not advanced, so, excuse my clumsiness with creating PRs :) |
@pearu thanks -- I'm quite backlogged this week since at a conference but I will endeavor to review in the next 48h. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you, this looks great. Here are a few minor comments.
Codecov Report
@@ Coverage Diff @@
## master #2732 +/- ##
==========================================
+ Coverage 87.57% 88.44% +0.86%
==========================================
Files 402 343 -59
Lines 61454 58554 -2900
==========================================
- Hits 53821 51787 -2034
+ Misses 7561 6767 -794
+ Partials 72 0 -72
Continue to review full report at Codecov.
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Thank you! A few more comments still.
By the way, I'm getting the following warning when building:
|
Also, all the tests fail here:
The following patch fixes the issue: diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx
index 1887d54a..84845889 100644
--- a/python/pyarrow/_cuda.pyx
+++ b/python/pyarrow/_cuda.pyx
@@ -26,7 +26,7 @@ cdef class Context:
""" CUDA driver context.
"""
- def __cinit__(self, int device_number=0, int handle=0):
+ def __cinit__(self, int device_number=0, intptr_t handle=0):
"""Construct the shared CUDA driver context for a particular device.
Parameters |
I'd favour approach 2. The following patch gets the Numba interop to work. But it also fails the other tests, since we need to expose pushing / popping in some way (and/or some kind of guard like ContextSaver). Some global APIs like Note the CUDA docs explicitly say about the primary context API:
diff --git a/cpp/src/arrow/gpu/cuda_context.cc b/cpp/src/arrow/gpu/cuda_context.cc
index d07b8b00..6952bf73 100644
--- a/cpp/src/arrow/gpu/cuda_context.cc
+++ b/cpp/src/arrow/gpu/cuda_context.cc
@@ -40,16 +40,13 @@ struct CudaDevice {
class ContextSaver {
public:
- explicit ContextSaver(CUcontext new_context) : context_(NULL) {
- cuCtxGetCurrent(&context_);
- cuCtxSetCurrent(new_context);
+ explicit ContextSaver(CUcontext new_context) {
+ cuCtxPushCurrent(new_context);
}
~ContextSaver() {
- if (context_ != NULL) cuCtxSetCurrent(context_);
+ CUcontext unused;
+ cuCtxPopCurrent(&unused);
}
-
- private:
- CUcontext context_;
};
class CudaContext::CudaContextImpl {
@@ -59,7 +56,7 @@ class CudaContext::CudaContextImpl {
Status Init(const CudaDevice& device) {
device_ = device;
own_context_ = true;
- CU_RETURN_NOT_OK(cuCtxCreate(&context_, 0, device_.handle));
+ CU_RETURN_NOT_OK(cuDevicePrimaryCtxRetain(&context_, device_.handle));
is_open_ = true;
return Status::OK();
}
@@ -74,7 +71,7 @@ class CudaContext::CudaContextImpl {
Status Close() {
if (is_open_ && own_context_) {
- CU_RETURN_NOT_OK(cuCtxDestroy(context_));
+ CU_RETURN_NOT_OK(cuDevicePrimaryCtxRelease(device_.handle));
}
is_open_ = false;
return Status::OK(); |
Defining
and using
OTH, memory created by AllocateHost should be available for all contexts. So, I think something else is needed here.. |
What is needed is to initialize a context for the device, with |
Yeah, I am now using |
Yes, the |
I got more tests to work using |
OK. I got all tests passing for push-pop context management pattern. Will clean up and commit. |
@pitrou , now we are using push-pop pattern for setting the context. I am curious, will the tests pass in your machine when using Now that pyarrow and numba use primary context that is unique to the device, it looks like there is no need to store context handles anymore - one can always retrieve it with |
Currently, AllocateHost uses the primary context of |
There is one failure remaining in
|
No, I am not testing in debug mode. |
You should :-) |
Re IPC test failure: I added a synchronize call that might fix it (although exit code -6 is worrying). I'll use debug mode. In debug mode, I can now reproduce the IPC test failure.. The likely solution is to provide context when calling |
@pitrou , I believe the last commit fixes the IPC test, please review. |
@pearu Yes, the tests pass now! I'll start a review. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks good to me except some very minor comments. Thanks for being persistent!
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Just one nit. Thank you!
cpp/src/arrow/gpu/cuda_context.h
Outdated
@@ -36,15 +36,24 @@ class ARROW_EXPORT CudaDeviceManager { | |||
public: | |||
static Status GetInstance(CudaDeviceManager** manager); | |||
|
|||
/// \brief Get the shared CUDA driver context for a particular device | |||
/// \brief Get the cached CUDA driver context for a particular device |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Need removing: "cached"
PR failed because of unrelated flake8 issues :-( |
I'm gonna merge anyway. Don't want to mess up this PR by trying to rebase and push to it. |
This PR implements the following features:
pyarrow.cuda.Context.handle
property returnsCUcontext
valuepyarrow.cuda.Context.from_numba
method to createContext
that shares numba context handlerpyarrow.cuda.Context.to_numba
method to create a numba context instance that sharesContext
handlerpyarrow.cuda.foreign_buffer
method to create CudaBuffer of device memory defined by address and size. Resolved ARROW-3451.pyarrow.cuda.CudaBuffer.to_numba
method to create a numbaMemoryPointer
instance so thatCudaBuffer
memory is used within the arguments to numba jitted functions.pyarrow.cuda.CudaBuffer.from_numba
method to create a CudaBuffer view of device data represented by numbaMemoryPointer
.CudaDeviceManager::CreateSharedContext
used in 2. Resolves ARROW-1423.CudaContext::View
used in 4.