Switch to use CUDA driver APIs in Device constructor
#460
Conversation
leofang
added
enhancement
leofang
changed the title
Device constructorDevice constructor
|
/ok to test |
|
| if err == 0: | ||
| device_id = int(dev) | ||
| else: | ||
| ctx = handle_return(driver.cuCtxGetCurrent()) |
There was a problem hiding this comment.
what's going on here? Is there some requirement from cudart which requires CtxGetCurrent() to be called before the device can be queried?
There was a problem hiding this comment.
Could it be helpful to raise a more specific error here?
It might be helpful to add a comment right after the else:
else:
# Emulate cudart behavior
err, ctx = driver.cuCtxGetCurrent()
if err != 0:
raise <Informative Error, what we really want is the current device (not primarily current context)>There was a problem hiding this comment.
This handles the case that no context is set to current. The logic is as follows:
- if a context is set to current, we can easily get the device ID associated with the current context
- if no context is set to current (which can happen right after
cuInit(0)and before anything else is called), we confirm it is the case by checkingctxpointer is zero (errwill always succeed), and then pick device 0
| ctx = handle_return(driver.cuCtxGetCurrent()) | ||
| assert int(ctx) == 0 | ||
| device_id = 0 # cudart behavior | ||
| assert isinstance(device_id, int), f"{device_id=}" |
There was a problem hiding this comment.
Is this intentionally not using the
assert_type(device_id, int)
helper?
There was a problem hiding this comment.
I think this PR predates the helper, I'll add it
| if err == 0: | ||
| device_id = int(dev) | ||
| else: | ||
| ctx = handle_return(driver.cuCtxGetCurrent()) |
There was a problem hiding this comment.
Could it be helpful to raise a more specific error here?
It might be helpful to add a comment right after the else:
else:
# Emulate cudart behavior
err, ctx = driver.cuCtxGetCurrent()
if err != 0:
raise <Informative Error, what we really want is the current device (not primarily current context)>
Blocked by #459 & #439 (comment).Before this PR:
With this PR:
(Bindings are built from the main branch.)