Skip to content

[gpu] Add acquire/release context accessors to all GPU backends #9431

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.

Sign up for GitHub

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

Open
derek-gerstmann wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[gpu] Add acquire/release context accessors to all GPU backends #9431

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
213 changes: 213 additions & 0 deletions src/JITModule.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -746,6 +746,36 @@ void merge_handlers(JITHandlers &base, const JITHandlers &addins) {
if (addins.custom_cuda_get_stream) {
base.custom_cuda_get_stream = addins.custom_cuda_get_stream;
}
if (addins.custom_opencl_acquire_context) {
base.custom_opencl_acquire_context = addins.custom_opencl_acquire_context;
}
if (addins.custom_opencl_release_context) {
base.custom_opencl_release_context = addins.custom_opencl_release_context;
}
if (addins.custom_metal_acquire_context) {
base.custom_metal_acquire_context = addins.custom_metal_acquire_context;
}
if (addins.custom_metal_release_context) {
base.custom_metal_release_context = addins.custom_metal_release_context;
}
if (addins.custom_d3d12compute_acquire_context) {
base.custom_d3d12compute_acquire_context = addins.custom_d3d12compute_acquire_context;
}
if (addins.custom_d3d12compute_release_context) {
base.custom_d3d12compute_release_context = addins.custom_d3d12compute_release_context;
}
if (addins.custom_vulkan_acquire_context) {
base.custom_vulkan_acquire_context = addins.custom_vulkan_acquire_context;
}
if (addins.custom_vulkan_release_context) {
base.custom_vulkan_release_context = addins.custom_vulkan_release_context;
}
if (addins.custom_webgpu_acquire_context) {
base.custom_webgpu_acquire_context = addins.custom_webgpu_acquire_context;
}
if (addins.custom_webgpu_release_context) {
base.custom_webgpu_release_context = addins.custom_webgpu_release_context;
}
}

void print_handler(JITUserContext *context, const char *msg) {
Expand Down Expand Up @@ -842,6 +872,93 @@ int cuda_get_stream_handler(JITUserContext *context, void *cuda_context, void **
}
}

int opencl_acquire_context_handler(JITUserContext *context, void **cl_context_ptr, void **cl_command_queue_ptr, bool create) {
if (context && context->handlers.custom_opencl_acquire_context) {
return context->handlers.custom_opencl_acquire_context(context, cl_context_ptr, cl_command_queue_ptr, create);
} else {
return active_handlers.custom_opencl_acquire_context(context, cl_context_ptr, cl_command_queue_ptr, create);
}
}

int opencl_release_context_handler(JITUserContext *context) {
if (context && context->handlers.custom_opencl_release_context) {
return context->handlers.custom_opencl_release_context(context);
} else {
return active_handlers.custom_opencl_release_context(context);
}
}

int metal_acquire_context_handler(JITUserContext *context, void **metal_device_ptr, void **metal_command_queue_ptr, bool create) {
if (context && context->handlers.custom_metal_acquire_context) {
return context->handlers.custom_metal_acquire_context(context, metal_device_ptr, metal_command_queue_ptr, create);
} else {
return active_handlers.custom_metal_acquire_context(context, metal_device_ptr, metal_command_queue_ptr, create);
}
}

int metal_release_context_handler(JITUserContext *context) {
if (context && context->handlers.custom_metal_release_context) {
return context->handlers.custom_metal_release_context(context);
} else {
return active_handlers.custom_metal_release_context(context);
}
}

int d3d12compute_acquire_context_handler(JITUserContext *context, void **d3d12_device_ptr, void **d3d12_command_queue_ptr, bool create) {
if (context && context->handlers.custom_d3d12compute_acquire_context) {
return context->handlers.custom_d3d12compute_acquire_context(context, d3d12_device_ptr, d3d12_command_queue_ptr, create);
} else {
return active_handlers.custom_d3d12compute_acquire_context(context, d3d12_device_ptr, d3d12_command_queue_ptr, create);
}
}

int d3d12compute_release_context_handler(JITUserContext *context) {
if (context && context->handlers.custom_d3d12compute_release_context) {
return context->handlers.custom_d3d12compute_release_context(context);
} else {
return active_handlers.custom_d3d12compute_release_context(context);
}
}

int vulkan_acquire_context_handler(JITUserContext *context, void **allocator_ptr, void **instance_ptr, void **device_ptr,
void **physical_device_ptr, void **queue_ptr, uint32_t *queue_family_index_ptr,
void **messenger_ptr, bool create) {
if (context && context->handlers.custom_vulkan_acquire_context) {
return context->handlers.custom_vulkan_acquire_context(context, allocator_ptr, instance_ptr, device_ptr,
physical_device_ptr, queue_ptr, queue_family_index_ptr,
messenger_ptr, create);
} else {
return active_handlers.custom_vulkan_acquire_context(context, allocator_ptr, instance_ptr, device_ptr,
physical_device_ptr, queue_ptr, queue_family_index_ptr,
messenger_ptr, create);
}
}

int vulkan_release_context_handler(JITUserContext *context, void *instance, void *device, void *queue, uint64_t messenger) {
if (context && context->handlers.custom_vulkan_release_context) {
return context->handlers.custom_vulkan_release_context(context, instance, device, queue, messenger);
} else {
return active_handlers.custom_vulkan_release_context(context, instance, device, queue, messenger);
}
}

int webgpu_acquire_context_handler(JITUserContext *context, void **instance_ptr, void **adapter_ptr,
void **device_ptr, void **staging_buffer_ptr, bool create) {
if (context && context->handlers.custom_webgpu_acquire_context) {
return context->handlers.custom_webgpu_acquire_context(context, instance_ptr, adapter_ptr, device_ptr, staging_buffer_ptr, create);
} else {
return active_handlers.custom_webgpu_acquire_context(context, instance_ptr, adapter_ptr, device_ptr, staging_buffer_ptr, create);
}
}

int webgpu_release_context_handler(JITUserContext *context) {
if (context && context->handlers.custom_webgpu_release_context) {
return context->handlers.custom_webgpu_release_context(context);
} else {
return active_handlers.custom_webgpu_release_context(context);
}
}

template<typename function_t>
function_t hook_function(const std::map<std::string, JITModule::Symbol> &exports, const char *hook_name, function_t hook) {
auto iter = exports.find(hook_name);
Expand Down Expand Up @@ -1107,6 +1224,102 @@ JITModule &make_module(llvm::Module *for_module, Target target,
runtime.add_dependency(shared_runtimes(CUDA));
}
}

// The same reasoning as for CUDA above applies to each of
// the other GPU backends: the debug and non-debug modules
// share a single set of context management handlers, and
// whichever module is created second declares a dependency
// on the first so things are destroyed in the correct order.

if (runtime_kind == OpenCL || runtime_kind == OpenCLDebug) {
if (!runtime_internal_handlers.custom_opencl_acquire_context) {
// Neither module has been created.
runtime_internal_handlers.custom_opencl_acquire_context =
hook_function(runtime.exports(), "halide_set_acquire_cl_context", opencl_acquire_context_handler);

runtime_internal_handlers.custom_opencl_release_context =
hook_function(runtime.exports(), "halide_set_release_cl_context", opencl_release_context_handler);

active_handlers = runtime_internal_handlers;
merge_handlers(active_handlers, default_handlers);
} else if (runtime_kind == OpenCL) {
runtime.add_dependency(shared_runtimes(OpenCLDebug));
} else {
runtime.add_dependency(shared_runtimes(OpenCL));
}
}

if (runtime_kind == Metal || runtime_kind == MetalDebug) {
if (!runtime_internal_handlers.custom_metal_acquire_context) {
// Neither module has been created.
runtime_internal_handlers.custom_metal_acquire_context =
hook_function(runtime.exports(), "halide_set_metal_acquire_context", metal_acquire_context_handler);

runtime_internal_handlers.custom_metal_release_context =
hook_function(runtime.exports(), "halide_set_metal_release_context", metal_release_context_handler);

active_handlers = runtime_internal_handlers;
merge_handlers(active_handlers, default_handlers);
} else if (runtime_kind == Metal) {
runtime.add_dependency(shared_runtimes(MetalDebug));
} else {
runtime.add_dependency(shared_runtimes(Metal));
}
}

if (runtime_kind == D3D12Compute || runtime_kind == D3D12ComputeDebug) {
if (!runtime_internal_handlers.custom_d3d12compute_acquire_context) {
// Neither module has been created.
runtime_internal_handlers.custom_d3d12compute_acquire_context =
hook_function(runtime.exports(), "halide_set_d3d12compute_acquire_context", d3d12compute_acquire_context_handler);

runtime_internal_handlers.custom_d3d12compute_release_context =
hook_function(runtime.exports(), "halide_set_d3d12compute_release_context", d3d12compute_release_context_handler);

active_handlers = runtime_internal_handlers;
merge_handlers(active_handlers, default_handlers);
} else if (runtime_kind == D3D12Compute) {
runtime.add_dependency(shared_runtimes(D3D12ComputeDebug));
} else {
runtime.add_dependency(shared_runtimes(D3D12Compute));
}
}

if (runtime_kind == Vulkan || runtime_kind == VulkanDebug) {
if (!runtime_internal_handlers.custom_vulkan_acquire_context) {
// Neither module has been created.
runtime_internal_handlers.custom_vulkan_acquire_context =
hook_function(runtime.exports(), "halide_set_vulkan_acquire_context", vulkan_acquire_context_handler);

runtime_internal_handlers.custom_vulkan_release_context =
hook_function(runtime.exports(), "halide_set_vulkan_release_context", vulkan_release_context_handler);

active_handlers = runtime_internal_handlers;
merge_handlers(active_handlers, default_handlers);
} else if (runtime_kind == Vulkan) {
runtime.add_dependency(shared_runtimes(VulkanDebug));
} else {
runtime.add_dependency(shared_runtimes(Vulkan));
}
}

if (runtime_kind == WebGPU || runtime_kind == WebGPUDebug) {
if (!runtime_internal_handlers.custom_webgpu_acquire_context) {
// Neither module has been created.
runtime_internal_handlers.custom_webgpu_acquire_context =
hook_function(runtime.exports(), "halide_set_webgpu_acquire_context", webgpu_acquire_context_handler);

runtime_internal_handlers.custom_webgpu_release_context =
hook_function(runtime.exports(), "halide_set_webgpu_release_context", webgpu_release_context_handler);

active_handlers = runtime_internal_handlers;
merge_handlers(active_handlers, default_handlers);
} else if (runtime_kind == WebGPU) {
runtime.add_dependency(shared_runtimes(WebGPUDebug));
} else {
runtime.add_dependency(shared_runtimes(WebGPU));
}
}
}

uint64_t arg_addr = llvm::cantFail(runtime.jit_module->JIT->lookup("halide_jit_module_argument"))
Expand Down
58 changes: 58 additions & 0 deletions src/JITModule.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -124,6 +124,64 @@ struct JITHandlers {
* stream to use. The cuda context and stream are both modelled
* as a void *, to avoid a dependence on the cuda headers. */
int32_t (*custom_cuda_get_stream)(JITUserContext *user_context, void *cuda_context, void **stream_ptr){nullptr};

/** A custom method for the Halide runtime to acquire an OpenCL
* context and command queue. Both are modelled as a void *, to
* avoid a dependence on the OpenCL headers. If the create argument
* is set to true, a context should be created if one does not
* already exist. */
int32_t (*custom_opencl_acquire_context)(JITUserContext *user_context, void **cl_context_ptr, void **cl_command_queue_ptr, bool create){nullptr};

/** The Halide runtime calls this when it is done with an OpenCL
* context. The default implementation does nothing. */
int32_t (*custom_opencl_release_context)(JITUserContext *user_context){nullptr};

/** A custom method for the Halide runtime to acquire a Metal device
* and command queue. Both are modelled as a void *, to avoid a
* dependence on the Metal headers. If the create argument is set to
* true, a context should be created if one does not already exist. */
int32_t (*custom_metal_acquire_context)(JITUserContext *user_context, void **metal_device_ptr, void **metal_command_queue_ptr, bool create){nullptr};

/** The Halide runtime calls this when it is done with a Metal
* context. The default implementation does nothing. */
int32_t (*custom_metal_release_context)(JITUserContext *user_context){nullptr};

/** A custom method for the Halide runtime to acquire a Direct3D 12
* device and command queue. Both are modelled as a void *, to
* avoid a dependence on the Direct3D headers. If the create
* argument is set to true, a context should be created if one does
* not already exist. */
int32_t (*custom_d3d12compute_acquire_context)(JITUserContext *user_context, void **d3d12_device_ptr, void **d3d12_command_queue_ptr, bool create){nullptr};

/** The Halide runtime calls this when it is done with a Direct3D 12
* context. The default implementation does nothing. */
int32_t (*custom_d3d12compute_release_context)(JITUserContext *user_context){nullptr};

/** A custom method for the Halide runtime to acquire a Vulkan
* context. The allocator and the various Vulkan handles are all
* modelled as a void *, to avoid a dependence on the Vulkan
* headers. If the create argument is set to true, a context should
* be created if one does not already exist. */
int32_t (*custom_vulkan_acquire_context)(JITUserContext *user_context, void **allocator_ptr, void **instance_ptr, void **device_ptr, void **physical_device_ptr, void **queue_ptr, uint32_t *queue_family_index_ptr, void **messenger_ptr, bool create){nullptr};

/** The Halide runtime calls this when it is done with a Vulkan
* context. The dispatchable Vulkan handles are modelled as a void *,
* to avoid a dependence on the Vulkan headers. The debug messenger is
* a non-dispatchable handle, which is a 64-bit integer rather than a
* pointer on 32-bit targets, so it is modelled as a uint64_t. The
* default implementation does nothing. */
int32_t (*custom_vulkan_release_context)(JITUserContext *user_context, void *instance, void *device, void *queue, uint64_t messenger){nullptr};

/** A custom method for the Halide runtime to acquire a WebGPU
* context. The instance, adapter, device, and staging buffer are
* all modelled as a void *, to avoid a dependence on the WebGPU
* headers. If the create argument is set to true, a context should
* be created if one does not already exist. */
int32_t (*custom_webgpu_acquire_context)(JITUserContext *user_context, void **instance_ptr, void **adapter_ptr, void **device_ptr, void **staging_buffer_ptr, bool create){nullptr};

/** The Halide runtime calls this when it is done with a WebGPU
* context. The default implementation does nothing. */
int32_t (*custom_webgpu_release_context)(JITUserContext *user_context){nullptr};
};

namespace Internal {
Expand Down
12 changes: 12 additions & 0 deletions src/runtime/HalideRuntimeD3D12Compute.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,18 @@ extern int halide_d3d12compute_acquire_context(void *user_context, struct halide
*/
extern int halide_d3d12compute_release_context(void *user_context);

typedef int (*halide_d3d12compute_acquire_context_t)(void *, // user_context
struct halide_d3d12compute_device **, // device out parameter
struct halide_d3d12compute_command_queue **, // command queue out parameter
bool); // should create a context if none exist
typedef int (*halide_d3d12compute_release_context_t)(void * /* user_context */);

/** Set custom methods to acquire and release D3D12 Compute contexts and command queues */
// @{
extern halide_d3d12compute_acquire_context_t halide_set_d3d12compute_acquire_context(halide_d3d12compute_acquire_context_t handler);
extern halide_d3d12compute_release_context_t halide_set_d3d12compute_release_context(halide_d3d12compute_release_context_t handler);
// @}

#ifdef __cplusplus
} // End extern "C"
#endif
Expand Down
12 changes: 12 additions & 0 deletions src/runtime/HalideRuntimeMetal.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,18 @@ extern int halide_metal_acquire_context(void *user_context, struct halide_metal_
*/
extern int halide_metal_release_context(void *user_context);

typedef int (*halide_metal_acquire_context_t)(void *, // user_context
struct halide_metal_device **, // device out parameter
struct halide_metal_command_queue **, // command queue out parameter
bool); // should create a context if none exist
typedef int (*halide_metal_release_context_t)(void * /* user_context */);

/** Set custom methods to acquire and release Metal contexts and command queues */
// @{
extern halide_metal_acquire_context_t halide_set_metal_acquire_context(halide_metal_acquire_context_t handler);
extern halide_metal_release_context_t halide_set_metal_release_context(halide_metal_release_context_t handler);
// @}

/** This function is called as part of the callback when a Metal command buffer completes.
* The return value, if not halide_error_code_success, will be stashed in Metal runtime and returned
* to the next call into the runtime, and the error string will be saved as well.
Expand Down
Loading
Loading