Allow users to use their own cuda contexts and streams in JIT mode #6345
Conversation
src/JITModule.h
Outdated
| @@ -28,16 +28,130 @@ struct JITUserContext; | |||
|
|
|||
| /** A set of custom overrides of runtime functions */ | |||
| struct JITHandlers { | |||
| /** Set the function called to print messages from the runtime. | |||
| * If you are compiling statically, you can also just define your | |||
There was a problem hiding this comment.
The comments about "if you are compiling statically" are arguably confusing and/or misplaced here. I'd suggest replacing everything after the first sentence with something like "(Note that this applies only when jitting code; if you are doing ahead-of-time compilation, see [README_foo.md or wherever we document this, rather than trying to replicate that documentation here])".
There was a problem hiding this comment.
Done (in underlying branch)
There was a problem hiding this comment.
Should we have something like
jit_handlers().reset()to set all the handlers back to their defaults?
It's just a struct, so you can say my_func.jit_handlers() = JITHandlers{} to reset it. Not sure if it's worth adding a method.
| @@ -65,6 +65,23 @@ extern uintptr_t halide_cuda_get_device_ptr(void *user_context, struct halide_bu | |||
| * driver. See halide_reuse_device_allocations. */ | |||
| extern int halide_cuda_release_unused_device_allocations(void *user_context); | |||
|
|
|||
| // These typedefs treat both a CUcontext and a CUstream as a void *, | |||
| // to avoid dependencies on cuda headers. | |||
| typedef int (*halide_cuda_acquire_context_t)(void *, // user_context | |||
There was a problem hiding this comment.
(Completely orthogonal to this PR, but IMHO we should consider migrating typedef bar foo; to using foo = bar; as I think it reads easier and is easier to search for)
There was a problem hiding this comment.
I believe this header is supposed to compile in C mode.
There was a problem hiding this comment.
and/or with janky legacy toolchains
There was a problem hiding this comment.
(Do we actually compile/run any tests in plain-C mode? If not, we should add one)
These can come up if a JITUserContext is passed to something like copy_to_device before getting fully populated by passing it to a call to realize.
and reuse the runtime's name resolution mechanism instead
This change means we'll only ever create one built-in cuda context in this circumstance.
|
Turns out the gpu_object_lifetime tests were creating both a cuda runtime module and a cuda-debug runtime module, and were thus creating two cuda contexts. The latest changes dedup this and add some comments explaining things. |
By exposing the acquire_cuda_context and friends methods in the same way as other JIT runtime overrides, we can let people replace them by swapping in their own handlers for those methods (see the new test)
This PR is a bit of a can-kick, because it's a piecemeal function-pointer-based exposure of three new runtime methods, rather than a single coherent way to replace any parts of the runtime you want. It should solve the problem in the short term, however.
Builds on #6344