-
Notifications
You must be signed in to change notification settings - Fork 55
refactor: decouple Context from Stream and Event objects
#579
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
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -1329,19 +1329,19 @@ def unload_module(self, module): | |
|
|
||
| def get_default_stream(self): | ||
| handle = drvapi.cu_stream(int(binding.CUstream(CU_STREAM_DEFAULT))) | ||
| return Stream(weakref.proxy(self), handle, None) | ||
| return Stream(handle) | ||
|
|
||
| def get_legacy_default_stream(self): | ||
| handle = drvapi.cu_stream( | ||
| int(binding.CUstream(binding.CU_STREAM_LEGACY)) | ||
| ) | ||
| return Stream(weakref.proxy(self), handle, None) | ||
| return Stream(handle) | ||
|
|
||
| def get_per_thread_default_stream(self): | ||
| handle = drvapi.cu_stream( | ||
| int(binding.CUstream(binding.CU_STREAM_PER_THREAD)) | ||
| ) | ||
| return Stream(weakref.proxy(self), handle, None) | ||
| return Stream(handle) | ||
|
|
||
| def create_stream(self): | ||
| # The default stream creation flag, specifying that the created | ||
|
|
@@ -1351,26 +1351,22 @@ def create_stream(self): | |
| flags = binding.CUstream_flags.CU_STREAM_DEFAULT.value | ||
| handle = drvapi.cu_stream(int(driver.cuStreamCreate(flags))) | ||
| return Stream( | ||
| weakref.proxy(self), | ||
| handle, | ||
| _stream_finalizer(self.deallocations, handle), | ||
| handle, finalizer=_stream_finalizer(self.deallocations, handle) | ||
| ) | ||
|
|
||
| def create_external_stream(self, ptr): | ||
| if not isinstance(ptr, int): | ||
| raise TypeError("ptr for external stream must be an int") | ||
| handle = drvapi.cu_stream(int(binding.CUstream(ptr))) | ||
| return Stream(weakref.proxy(self), handle, None, external=True) | ||
| return Stream(handle, external=True) | ||
|
|
||
| def create_event(self, timing=True): | ||
| flags = 0 | ||
| if not timing: | ||
| flags |= enums.CU_EVENT_DISABLE_TIMING | ||
| handle = drvapi.cu_event(int(driver.cuEventCreate(flags))) | ||
| return Event( | ||
| weakref.proxy(self), | ||
| handle, | ||
| finalizer=_event_finalizer(self.deallocations, handle), | ||
| handle, finalizer=_event_finalizer(self.deallocations, handle) | ||
| ) | ||
|
|
||
| def synchronize(self): | ||
|
|
@@ -1383,7 +1379,7 @@ def defer_cleanup(self): | |
| yield | ||
|
|
||
| def __repr__(self): | ||
| return "<CUDA context %s of device %d>" % (self.handle, self.device.id) | ||
| return f"<CUDA context {self.handle} of device {self.device.id:d}>" | ||
|
|
||
| def __eq__(self, other): | ||
| if isinstance(other, Context): | ||
|
|
@@ -2058,9 +2054,8 @@ class ManagedOwnedPointer(OwnedPointer, mviewbuf.MemAlloc): | |
| pass | ||
|
|
||
|
|
||
| class Stream(object): | ||
| def __init__(self, context, handle, finalizer, external=False): | ||
| self.context = context | ||
| class Stream: | ||
| def __init__(self, handle, finalizer=None, external=False): | ||
|
Comment on lines
+2057
to
+2058
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There is unfortunately public code using the @gmarkall what are your thoughts?
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We should probably lay out somewhere what constructors are public, since hand constructing these doesn't appear to be part of any explicit public API, just the one you get from the default Python conventions.
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you highlight the code that's using the constructor explicitly? The way it's presented in the documentation is maybe not ideal, but users are directed to use one of:
to construct a |
||
| self.handle = handle | ||
| self.external = external | ||
| if finalizer is not None: | ||
|
|
@@ -2077,18 +2072,18 @@ def __cuda_stream__(self): | |
|
|
||
| def __repr__(self): | ||
| default_streams = { | ||
| drvapi.CU_STREAM_DEFAULT: "<Default CUDA stream on %s>", | ||
| drvapi.CU_STREAM_LEGACY: "<Legacy default CUDA stream on %s>", | ||
| drvapi.CU_STREAM_PER_THREAD: "<Per-thread default CUDA stream on %s>", | ||
| drvapi.CU_STREAM_DEFAULT: "<Default CUDA stream>", | ||
| drvapi.CU_STREAM_LEGACY: "<Legacy default CUDA stream>", | ||
| drvapi.CU_STREAM_PER_THREAD: "<Per-thread default CUDA stream>", | ||
| } | ||
| ptr = self.handle.value or drvapi.CU_STREAM_DEFAULT | ||
|
|
||
| if ptr in default_streams: | ||
| return default_streams[ptr] % self.context | ||
| return default_streams[ptr] | ||
| elif self.external: | ||
| return "<External CUDA stream %d on %s>" % (ptr, self.context) | ||
| return f"<External CUDA stream {ptr:d}>" | ||
| else: | ||
| return "<CUDA stream %d on %s>" % (ptr, self.context) | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The context was only being used for repring. The complexity of weak references doesn't seem justifiable for pretty printing, hence this PR :) |
||
| return f"<CUDA stream {ptr:d}>" | ||
|
|
||
| def synchronize(self): | ||
| """ | ||
|
|
@@ -2190,9 +2185,8 @@ def callback(stream, status, future): | |
| return future | ||
|
|
||
|
|
||
| class Event(object): | ||
| def __init__(self, context, handle, finalizer=None): | ||
| self.context = context | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This was just never used in the class. |
||
| class Event: | ||
| def __init__(self, handle, finalizer=None): | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. +1 on this being something that has been used with the
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. It is currently doing nothing in the class and to my eye wasting space + adding the complexity of a
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I agree about it not serving a purpose currently, but there's something to be said about not breaking existing code using it...
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. No use of Same is true for If there's some usage of it that we can point to that would be useful information.
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Similarly to streams, users shoudn't be creating |
||
| self.handle = handle | ||
| if finalizer is not None: | ||
| weakref.finalize(self, finalizer) | ||
|
|
||
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.
This is not related to the other changes, but it was causing problems in CI, and isn't necessary anymore since we're getting all dependencies from
pixiin this script, so I removed it.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.
I think the reason this remained is that the CUDA pathfinder can mistakenly find system libraries ahead of ones in the environment, so we remove them to make sure we're testing with the right libraries. This was discussed in NVIDIA/cuda-python#839 - note that although the issue is closed, I believe the fix only makes it less likely to find the library from the system install of the toolkit, not impossible.
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.
The behavior now should be correct. pip wheels will be found first, then conda packages, then a default system search using
dlopenwhich has its own hierarchy.As far as I know the only issue that remains is that if you're trying to find system libraries (as opposed to pip wheels or conda packages) that nvvm doesn't end up in
ldconfig/LD_LIBRARY_PATHor any of the standard search paths. We hope to address this in the near future.