From 7fdd70d741a05bc6cf2aa44f0231007291d366b1 Mon Sep 17 00:00:00 2001 From: Derek Gerstmann Date: Sat, 5 Sep 2026 09:18:31 -0700 Subject: [PATCH] Add acquire/release context accessors to all GPU backends The acquire_context()/release_context() functions in each GPU runtime cannot be overridden on Windows: the runtime-library symbols are strong and are only referenced internally within the runtime libs, so the weak-symbol replacement used on other platforms does not work there. The CUDA backend already worked around this with halide_set_cuda_acquire_context()/halide_set_cuda_release_context() accessors that swap a function pointer at runtime. Apply the same pattern to the OpenCL, Metal, D3D12Compute, Vulkan, and WebGPU backends so they are consistent with CUDA: - Split each backend's acquire/release context function into a halide_default_* implementation plus a dispatcher that calls through a settable function pointer, and add halide_set_*_acquire_context / halide_set_*_release_context accessors with matching typedefs in the public runtime headers. - Wire the handlers into the JIT: add JITHandlers fields, merge_handlers entries, per-context handler shims, and make_module() hookup, mirroring the CUDA handling. - Register the new setters (and the previously-missing CUDA ones) in halide_runtime_api_functions[]. - Add per-backend correctness tests modeled on custom_cuda_context.cpp. Co-authored-by: Claude Opus 4.8 --- src/JITModule.cpp | 213 ++++++++++++++++++ src/JITModule.h | 58 +++++ src/runtime/HalideRuntimeD3D12Compute.h | 12 + src/runtime/HalideRuntimeMetal.h | 12 + src/runtime/HalideRuntimeOpenCL.h | 14 ++ src/runtime/HalideRuntimeVulkan.h | 21 ++ src/runtime/HalideRuntimeWebGPU.h | 16 ++ src/runtime/d3d12compute.cpp | 48 +++- src/runtime/metal.cpp | 44 +++- src/runtime/opencl.cpp | 41 +++- src/runtime/runtime_api.cpp | 13 ++ src/runtime/vulkan.cpp | 66 +++++- src/runtime/webgpu.cpp | 57 ++++- test/correctness/CMakeLists.txt | 5 + .../custom_d3d12compute_context.cpp | 94 ++++++++ test/correctness/custom_metal_context.cpp | 94 ++++++++ test/correctness/custom_opencl_context.cpp | 94 ++++++++ test/correctness/custom_vulkan_context.cpp | 101 +++++++++ test/correctness/custom_webgpu_context.cpp | 96 ++++++++ 19 files changed, 1073 insertions(+), 26 deletions(-) create mode 100644 test/correctness/custom_d3d12compute_context.cpp create mode 100644 test/correctness/custom_metal_context.cpp create mode 100644 test/correctness/custom_opencl_context.cpp create mode 100644 test/correctness/custom_vulkan_context.cpp create mode 100644 test/correctness/custom_webgpu_context.cpp diff --git a/src/JITModule.cpp b/src/JITModule.cpp index a6dfc77486b6..6db659ff5b07 100644 --- a/src/JITModule.cpp +++ b/src/JITModule.cpp @@ -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) { @@ -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 function_t hook_function(const std::map &exports, const char *hook_name, function_t hook) { auto iter = exports.find(hook_name); @@ -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")) diff --git a/src/JITModule.h b/src/JITModule.h index ccfefd02b3f4..9429b8c2bc89 100644 --- a/src/JITModule.h +++ b/src/JITModule.h @@ -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 { diff --git a/src/runtime/HalideRuntimeD3D12Compute.h b/src/runtime/HalideRuntimeD3D12Compute.h index 9edb30c30cd1..feb22898b635 100644 --- a/src/runtime/HalideRuntimeD3D12Compute.h +++ b/src/runtime/HalideRuntimeD3D12Compute.h @@ -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 diff --git a/src/runtime/HalideRuntimeMetal.h b/src/runtime/HalideRuntimeMetal.h index 30762e07d8ae..46ef9bfaf51a 100644 --- a/src/runtime/HalideRuntimeMetal.h +++ b/src/runtime/HalideRuntimeMetal.h @@ -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. diff --git a/src/runtime/HalideRuntimeOpenCL.h b/src/runtime/HalideRuntimeOpenCL.h index 510dc6f1ba8e..fb70de075c16 100644 --- a/src/runtime/HalideRuntimeOpenCL.h +++ b/src/runtime/HalideRuntimeOpenCL.h @@ -112,6 +112,20 @@ extern uintptr_t halide_opencl_get_cl_mem(void *user_context, struct halide_buff /** Returns the offset associated with the OpenCL memory allocation via device_crop or device_slice. */ extern uint64_t halide_opencl_get_crop_offset(void *user_context, halide_buffer_t *buf); +// These typedefs treat both a cl_context and a cl_command_queue as a +// void *, to avoid dependencies on the OpenCL headers. +typedef int (*halide_acquire_cl_context_t)(void *, // user_context + void **, // cl_context out parameter + void **, // cl_command_queue out parameter + bool); // should create a context if none exist +typedef int (*halide_release_cl_context_t)(void * /* user_context */); + +/** Set custom methods to acquire and release OpenCL contexts and command queues */ +// @{ +extern halide_acquire_cl_context_t halide_set_acquire_cl_context(halide_acquire_cl_context_t handler); +extern halide_release_cl_context_t halide_set_release_cl_context(halide_release_cl_context_t handler); +// @} + #ifdef __cplusplus } // End extern "C" #endif diff --git a/src/runtime/HalideRuntimeVulkan.h b/src/runtime/HalideRuntimeVulkan.h index e150b7c6d00b..b6e5a5f901a5 100644 --- a/src/runtime/HalideRuntimeVulkan.h +++ b/src/runtime/HalideRuntimeVulkan.h @@ -105,6 +105,27 @@ extern int halide_vulkan_release_context(void *user_context, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger); + +typedef int (*halide_vulkan_acquire_context_t)(void *, // user_context + struct halide_vulkan_memory_allocator **, // allocator out parameter + VkInstance *, // instance out parameter + VkDevice *, // device out parameter + VkPhysicalDevice *, // physical device out parameter + VkQueue *, // queue out parameter + uint32_t *, // queue family index out parameter + VkDebugUtilsMessengerEXT *, // debug messenger out parameter + bool); // should create a context if none exist +typedef int (*halide_vulkan_release_context_t)(void *, // user_context + VkInstance, // instance + VkDevice, // device + VkQueue, // queue + VkDebugUtilsMessengerEXT); // debug messenger + +/** Set custom methods to acquire and release Vulkan contexts */ +// @{ +extern halide_vulkan_acquire_context_t halide_set_vulkan_acquire_context(halide_vulkan_acquire_context_t handler); +extern halide_vulkan_release_context_t halide_set_vulkan_release_context(halide_vulkan_release_context_t handler); +// @} // -- // Override the default allocation callbacks (default uses Vulkan runtime implementation) diff --git a/src/runtime/HalideRuntimeWebGPU.h b/src/runtime/HalideRuntimeWebGPU.h index f1ae55f9f4f3..83147259621e 100644 --- a/src/runtime/HalideRuntimeWebGPU.h +++ b/src/runtime/HalideRuntimeWebGPU.h @@ -37,6 +37,22 @@ extern int halide_webgpu_run(void *user_context, extern void halide_webgpu_finalize_kernels(void *user_context, void *state_ptr); // @} +// These typedefs treat the WGPUInstance, WGPUAdapter, WGPUDevice, and +// WGPUBuffer handles as a void *, to avoid dependencies on the WebGPU headers. +typedef int (*halide_webgpu_acquire_context_t)(void *, // user_context + void **, // WGPUInstance out parameter + void **, // WGPUAdapter out parameter + void **, // WGPUDevice out parameter + void **, // WGPUBuffer staging buffer out parameter + bool); // should create a context if none exist +typedef int (*halide_webgpu_release_context_t)(void * /* user_context */); + +/** Set custom methods to acquire and release WebGPU contexts */ +// @{ +extern halide_webgpu_acquire_context_t halide_set_webgpu_acquire_context(halide_webgpu_acquire_context_t handler); +extern halide_webgpu_release_context_t halide_set_webgpu_release_context(halide_webgpu_release_context_t handler); +// @} + #ifdef __cplusplus } // End extern "C" #endif diff --git a/src/runtime/d3d12compute.cpp b/src/runtime/d3d12compute.cpp index 0fcf58cf4e37..56dd774b54d2 100644 --- a/src/runtime/d3d12compute.cpp +++ b/src/runtime/d3d12compute.cpp @@ -3090,6 +3090,8 @@ WEAK halide_error_code_t d3d12_create_context(void *user_context) { } } +WEAK int halide_default_d3d12compute_release_context(void *user_context); + // The default implementation of halide_d3d12compute_acquire_context uses the global // pointers above, and serializes access with a mutex. // Overriding implementations of acquire/release must implement the following @@ -3099,8 +3101,8 @@ WEAK halide_error_code_t d3d12_create_context(void *user_context) { // - A call to halide_acquire_d3d12compute_context is followed by a matching call to // halide_release_d3d12compute_context. halide_acquire_d3d12compute_context should block while a // previous call (if any) has not yet been released via halide_release_d3d12compute_context. -WEAK int halide_d3d12compute_acquire_context(void *user_context, halide_d3d12compute_device **device_ret, - halide_d3d12compute_command_queue **queue_ret, bool create) { +WEAK int halide_default_d3d12compute_acquire_context(void *user_context, halide_d3d12compute_device **device_ret, + halide_d3d12compute_command_queue **queue_ret, bool create) { TRACELOG; #ifdef HALIDE_D3D12_TRACE_TIME @@ -3120,7 +3122,7 @@ WEAK int halide_d3d12compute_acquire_context(void *user_context, halide_d3d12com if (create && (device == nullptr)) { auto error = d3d12_create_context(user_context); if (error) { - (void)halide_d3d12compute_release_context(user_context); // ignore error + (void)halide_default_d3d12compute_release_context(user_context); // ignore error return error; } } @@ -3135,7 +3137,7 @@ WEAK int halide_d3d12compute_acquire_context(void *user_context, halide_d3d12com return halide_error_code_success; } -WEAK int halide_d3d12compute_release_context(void *user_context) { +WEAK int halide_default_d3d12compute_release_context(void *user_context) { TRACELOG; halide_mutex_unlock(&thread_lock); return halide_error_code_success; @@ -3143,6 +3145,44 @@ WEAK int halide_d3d12compute_release_context(void *user_context) { } // extern "C" +namespace Halide { +namespace Runtime { +namespace Internal { +namespace D3D12Compute { + +WEAK halide_d3d12compute_acquire_context_t acquire_context = halide_default_d3d12compute_acquire_context; +WEAK halide_d3d12compute_release_context_t release_context = halide_default_d3d12compute_release_context; + +} // namespace D3D12Compute +} // namespace Internal +} // namespace Runtime +} // namespace Halide + +extern "C" { + +WEAK int halide_d3d12compute_acquire_context(void *user_context, halide_d3d12compute_device **device_ret, + halide_d3d12compute_command_queue **queue_ret, bool create) { + return D3D12Compute::acquire_context(user_context, device_ret, queue_ret, create); +} + +WEAK halide_d3d12compute_acquire_context_t halide_set_d3d12compute_acquire_context(halide_d3d12compute_acquire_context_t handler) { + halide_d3d12compute_acquire_context_t result = D3D12Compute::acquire_context; + D3D12Compute::acquire_context = handler; + return result; +} + +WEAK int halide_d3d12compute_release_context(void *user_context) { + return D3D12Compute::release_context(user_context); +} + +WEAK halide_d3d12compute_release_context_t halide_set_d3d12compute_release_context(halide_d3d12compute_release_context_t handler) { + halide_d3d12compute_release_context_t result = D3D12Compute::release_context; + D3D12Compute::release_context = handler; + return result; +} + +} // extern "C" + WEAK void d3d12_debug_dump() { error err(nullptr); diff --git a/src/runtime/metal.cpp b/src/runtime/metal.cpp index 95bbf31a8f29..f3f73450eaa6 100644 --- a/src/runtime/metal.cpp +++ b/src/runtime/metal.cpp @@ -363,8 +363,8 @@ extern "C" { // - A call to halide_acquire_metal_context is followed by a matching call to // halide_release_metal_context. halide_acquire_metal_context should block while a // previous call (if any) has not yet been released via halide_release_metal_context. -WEAK int halide_metal_acquire_context(void *user_context, mtl_device **device_ret, - mtl_command_queue **queue_ret, bool create) { +WEAK int halide_default_metal_acquire_context(void *user_context, mtl_device **device_ret, + mtl_command_queue **queue_ret, bool create) { halide_debug_assert(user_context, &thread_lock != nullptr); halide_mutex_lock(&thread_lock); @@ -403,13 +403,51 @@ WEAK int halide_metal_acquire_context(void *user_context, mtl_device **device_re return halide_error_code_success; } -WEAK int halide_metal_release_context(void *user_context) { +WEAK int halide_default_metal_release_context(void *user_context) { halide_mutex_unlock(&thread_lock); return halide_error_code_success; } } // extern "C" +namespace Halide { +namespace Runtime { +namespace Internal { +namespace Metal { + +WEAK halide_metal_acquire_context_t acquire_context = halide_default_metal_acquire_context; +WEAK halide_metal_release_context_t release_context = halide_default_metal_release_context; + +} // namespace Metal +} // namespace Internal +} // namespace Runtime +} // namespace Halide + +extern "C" { + +WEAK int halide_metal_acquire_context(void *user_context, mtl_device **device_ret, + mtl_command_queue **queue_ret, bool create) { + return Metal::acquire_context(user_context, device_ret, queue_ret, create); +} + +WEAK halide_metal_acquire_context_t halide_set_metal_acquire_context(halide_metal_acquire_context_t handler) { + halide_metal_acquire_context_t result = Metal::acquire_context; + Metal::acquire_context = handler; + return result; +} + +WEAK int halide_metal_release_context(void *user_context) { + return Metal::release_context(user_context); +} + +WEAK halide_metal_release_context_t halide_set_metal_release_context(halide_metal_release_context_t handler) { + halide_metal_release_context_t result = Metal::release_context; + Metal::release_context = handler; + return result; +} + +} // extern "C" + extern "C" size_t strnlen(const char *s, size_t maxlen); namespace Halide { diff --git a/src/runtime/opencl.cpp b/src/runtime/opencl.cpp index 8ccb827152f2..55b30e95c485 100644 --- a/src/runtime/opencl.cpp +++ b/src/runtime/opencl.cpp @@ -222,7 +222,7 @@ WEAK const char *halide_opencl_get_build_options(void *user_context) { // - A call to halide_acquire_cl_context is followed by a matching call to // halide_release_cl_context. halide_acquire_cl_context should block while a // previous call (if any) has not yet been released via halide_release_cl_context. -WEAK int halide_acquire_cl_context(void *user_context, cl_context *ctx, cl_command_queue *q, bool create = true) { +WEAK int halide_default_acquire_cl_context(void *user_context, cl_context *ctx, cl_command_queue *q, bool create = true) { // TODO: Should we use a more "assertive" assert? These asserts do // not block execution on failure. halide_abort_if_false(user_context, ctx != nullptr); @@ -248,7 +248,7 @@ WEAK int halide_acquire_cl_context(void *user_context, cl_context *ctx, cl_comma return halide_error_code_success; } -WEAK int halide_release_cl_context(void *user_context) { +WEAK int halide_default_release_cl_context(void *user_context) { __atomic_clear(&thread_lock, __ATOMIC_RELEASE); return halide_error_code_success; } @@ -260,6 +260,43 @@ namespace Runtime { namespace Internal { namespace OpenCL { +WEAK halide_acquire_cl_context_t acquire_context = (halide_acquire_cl_context_t)halide_default_acquire_cl_context; +WEAK halide_release_cl_context_t release_context = (halide_release_cl_context_t)halide_default_release_cl_context; + +} // namespace OpenCL +} // namespace Internal +} // namespace Runtime +} // namespace Halide + +extern "C" { + +WEAK int halide_acquire_cl_context(void *user_context, cl_context *ctx, cl_command_queue *q, bool create = true) { + return OpenCL::acquire_context(user_context, (void **)ctx, (void **)q, create); +} + +WEAK halide_acquire_cl_context_t halide_set_acquire_cl_context(halide_acquire_cl_context_t handler) { + halide_acquire_cl_context_t result = OpenCL::acquire_context; + OpenCL::acquire_context = handler; + return result; +} + +WEAK int halide_release_cl_context(void *user_context) { + return OpenCL::release_context(user_context); +} + +WEAK halide_release_cl_context_t halide_set_release_cl_context(halide_release_cl_context_t handler) { + halide_release_cl_context_t result = OpenCL::release_context; + OpenCL::release_context = handler; + return result; +} + +} // extern "C" + +namespace Halide { +namespace Runtime { +namespace Internal { +namespace OpenCL { + // Helper object to acquire and release the OpenCL context. class ClContext { void *const user_context; diff --git a/src/runtime/runtime_api.cpp b/src/runtime/runtime_api.cpp index ed516958dfda..466d051c5be2 100644 --- a/src/runtime/runtime_api.cpp +++ b/src/runtime/runtime_api.cpp @@ -182,6 +182,10 @@ extern "C" __attribute__((used)) void *halide_runtime_api_functions[] = { (void *)&halide_semaphore_init, (void *)&halide_semaphore_release, (void *)&halide_semaphore_try_acquire, + (void *)&halide_set_acquire_cl_context, + (void *)&halide_set_cuda_acquire_context, + (void *)&halide_set_cuda_get_stream, + (void *)&halide_set_cuda_release_context, (void *)&halide_set_custom_can_use_target_features, (void *)&halide_set_custom_do_par_for, (void *)&halide_set_custom_do_loop_task, @@ -193,10 +197,19 @@ extern "C" __attribute__((used)) void *halide_runtime_api_functions[] = { (void *)&halide_set_custom_malloc, (void *)&halide_set_custom_print, (void *)&halide_set_custom_trace, + (void *)&halide_set_d3d12compute_acquire_context, + (void *)&halide_set_d3d12compute_release_context, (void *)&halide_set_error_handler, (void *)&halide_set_gpu_device, + (void *)&halide_set_metal_acquire_context, + (void *)&halide_set_metal_release_context, (void *)&halide_set_num_threads, + (void *)&halide_set_release_cl_context, (void *)&halide_set_trace_file, + (void *)&halide_set_vulkan_acquire_context, + (void *)&halide_set_vulkan_release_context, + (void *)&halide_set_webgpu_acquire_context, + (void *)&halide_set_webgpu_release_context, (void *)&halide_shutdown_thread_pool, (void *)&halide_shutdown_trace, (void *)&halide_sleep_us, diff --git a/src/runtime/vulkan.cpp b/src/runtime/vulkan.cpp index f44b60b90e94..1a310a775f5b 100644 --- a/src/runtime/vulkan.cpp +++ b/src/runtime/vulkan.cpp @@ -29,15 +29,15 @@ extern "C" { // call to halide_release_vulkan_context. halide_acquire_vulkan_context // should block while a previous call (if any) has not yet been // released via halide_release_vulkan_context. -WEAK int halide_vulkan_acquire_context(void *user_context, - halide_vulkan_memory_allocator **allocator, - VkInstance *instance, - VkDevice *device, - VkPhysicalDevice *physical_device, - VkQueue *queue, - uint32_t *queue_family_index, - VkDebugUtilsMessengerEXT *messenger, - bool create) { +WEAK int halide_default_vulkan_acquire_context(void *user_context, + halide_vulkan_memory_allocator **allocator, + VkInstance *instance, + VkDevice *device, + VkPhysicalDevice *physical_device, + VkQueue *queue, + uint32_t *queue_family_index, + VkDebugUtilsMessengerEXT *messenger, + bool create) { #ifdef DEBUG_RUNTIME halide_start_clock(user_context); #endif @@ -74,11 +74,57 @@ WEAK int halide_vulkan_acquire_context(void *user_context, return halide_error_code_success; } -WEAK int halide_vulkan_release_context(void *user_context, VkInstance instance, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger) { +WEAK int halide_default_vulkan_release_context(void *user_context, VkInstance instance, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger) { halide_mutex_unlock(&thread_lock); return halide_error_code_success; } +} // extern "C" + +namespace Halide { +namespace Runtime { +namespace Internal { +namespace Vulkan { + +WEAK halide_vulkan_acquire_context_t acquire_context = halide_default_vulkan_acquire_context; +WEAK halide_vulkan_release_context_t release_context = halide_default_vulkan_release_context; + +} // namespace Vulkan +} // namespace Internal +} // namespace Runtime +} // namespace Halide + +extern "C" { + +WEAK int halide_vulkan_acquire_context(void *user_context, + halide_vulkan_memory_allocator **allocator, + VkInstance *instance, + VkDevice *device, + VkPhysicalDevice *physical_device, + VkQueue *queue, + uint32_t *queue_family_index, + VkDebugUtilsMessengerEXT *messenger, + bool create) { + return Vulkan::acquire_context(user_context, allocator, instance, device, + physical_device, queue, queue_family_index, messenger, create); +} + +WEAK halide_vulkan_acquire_context_t halide_set_vulkan_acquire_context(halide_vulkan_acquire_context_t handler) { + halide_vulkan_acquire_context_t result = Vulkan::acquire_context; + Vulkan::acquire_context = handler; + return result; +} + +WEAK int halide_vulkan_release_context(void *user_context, VkInstance instance, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger) { + return Vulkan::release_context(user_context, instance, device, queue, messenger); +} + +WEAK halide_vulkan_release_context_t halide_set_vulkan_release_context(halide_vulkan_release_context_t handler) { + halide_vulkan_release_context_t result = Vulkan::release_context; + Vulkan::release_context = handler; + return result; +} + WEAK bool halide_vulkan_is_initialized() { halide_mutex_lock(&thread_lock); bool is_initialized = (cached_instance != nullptr) && (cached_device != nullptr); diff --git a/src/runtime/webgpu.cpp b/src/runtime/webgpu.cpp index 5b1f4d930fc9..48754e914c69 100644 --- a/src/runtime/webgpu.cpp +++ b/src/runtime/webgpu.cpp @@ -71,12 +71,12 @@ void wgpuDeviceTick(WGPUDevice) { // halide_webgpu_release_context. halide_webgpu_acquire_context should block // while a previous call (if any) has not yet been released via // halide_webgpu_release_context. -WEAK int halide_webgpu_acquire_context(void *user_context, - WGPUInstance *instance_ret, - WGPUAdapter *adapter_ret, - WGPUDevice *device_ret, - WGPUBuffer *staging_buffer_ret, - bool create = true) { +WEAK int halide_default_webgpu_acquire_context(void *user_context, + WGPUInstance *instance_ret, + WGPUAdapter *adapter_ret, + WGPUDevice *device_ret, + WGPUBuffer *staging_buffer_ret, + bool create = true) { debug(user_context) << "WGPU: halide_webgpu_acquire_context (user_context: " << user_context << ")\n"; @@ -104,7 +104,7 @@ WEAK int halide_webgpu_acquire_context(void *user_context, return halide_error_code_success; } -WEAK int halide_webgpu_release_context(void *user_context) { +WEAK int halide_default_webgpu_release_context(void *user_context) { __atomic_clear(&context_lock, __ATOMIC_RELEASE); return halide_error_code_success; } @@ -116,6 +116,49 @@ namespace Runtime { namespace Internal { namespace WebGPU { +WEAK halide_webgpu_acquire_context_t acquire_context = (halide_webgpu_acquire_context_t)halide_default_webgpu_acquire_context; +WEAK halide_webgpu_release_context_t release_context = (halide_webgpu_release_context_t)halide_default_webgpu_release_context; + +} // namespace WebGPU +} // namespace Internal +} // namespace Runtime +} // namespace Halide + +extern "C" { + +WEAK int halide_webgpu_acquire_context(void *user_context, + WGPUInstance *instance_ret, + WGPUAdapter *adapter_ret, + WGPUDevice *device_ret, + WGPUBuffer *staging_buffer_ret, + bool create = true) { + return WebGPU::acquire_context(user_context, (void **)instance_ret, (void **)adapter_ret, + (void **)device_ret, (void **)staging_buffer_ret, create); +} + +WEAK halide_webgpu_acquire_context_t halide_set_webgpu_acquire_context(halide_webgpu_acquire_context_t handler) { + halide_webgpu_acquire_context_t result = WebGPU::acquire_context; + WebGPU::acquire_context = handler; + return result; +} + +WEAK int halide_webgpu_release_context(void *user_context) { + return WebGPU::release_context(user_context); +} + +WEAK halide_webgpu_release_context_t halide_set_webgpu_release_context(halide_webgpu_release_context_t handler) { + halide_webgpu_release_context_t result = WebGPU::release_context; + WebGPU::release_context = handler; + return result; +} + +} // extern "C" linkage + +namespace Halide { +namespace Runtime { +namespace Internal { +namespace WebGPU { + // Helper object to acquire and release the WebGPU context. class WgpuContext { void *user_context; diff --git a/test/correctness/CMakeLists.txt b/test/correctness/CMakeLists.txt index fa463f242e7e..5d1a37cef990 100644 --- a/test/correctness/CMakeLists.txt +++ b/test/correctness/CMakeLists.txt @@ -78,9 +78,14 @@ tests( custom_allocator.cpp custom_auto_scheduler.cpp custom_cuda_context.cpp + custom_d3d12compute_context.cpp custom_error_reporter.cpp custom_jit_context.cpp custom_lowering_pass.cpp + custom_metal_context.cpp + custom_opencl_context.cpp + custom_vulkan_context.cpp + custom_webgpu_context.cpp d3d12compute_sm6x.cpp d3d12compute_strict_float.cpp dead_realization_in_specialization.cpp diff --git a/test/correctness/custom_d3d12compute_context.cpp b/test/correctness/custom_d3d12compute_context.cpp new file mode 100644 index 000000000000..ddca4245bedf --- /dev/null +++ b/test/correctness/custom_d3d12compute_context.cpp @@ -0,0 +1,94 @@ +#include "Halide.h" + +#include + +using namespace Halide; + +// The runtime's default context-management functions, looked up from the JIT'd +// Direct3D 12 Compute runtime module. Our custom handlers delegate to these so +// that we can verify the custom accessors are actually invoked without needing +// to link against the Direct3D SDK. +static int (*default_d3d12compute_acquire_context)(void *, void **, void **, bool) = nullptr; +static int (*default_d3d12compute_release_context)(void *) = nullptr; + +struct D3D12ComputeState : public Halide::JITUserContext { + std::atomic acquires = 0, releases = 0; + + static int my_acquire_context(JITUserContext *ctx, void **device, void **queue, bool create) { + D3D12ComputeState *state = (D3D12ComputeState *)ctx; + state->acquires++; + return default_d3d12compute_acquire_context(ctx, device, queue, create); + } + + static int my_release_context(JITUserContext *ctx) { + D3D12ComputeState *state = (D3D12ComputeState *)ctx; + state->releases++; + return default_d3d12compute_release_context(ctx); + } + + D3D12ComputeState() { + handlers.custom_d3d12compute_acquire_context = my_acquire_context; + handlers.custom_d3d12compute_release_context = my_release_context; + } +}; + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::D3D12Compute)) { + printf("[SKIP] D3D12Compute not enabled.\n"); + return 0; + } + + // Force-initialize the D3D12 runtime module by running something trivial, + // then extract the default context-management functions from it. + evaluate_may_gpu(Expr(0.f)); + + default_d3d12compute_acquire_context = (int (*)(void *, void **, void **, bool)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_d3d12compute_acquire_context"); + default_d3d12compute_release_context = (int (*)(void *)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_d3d12compute_release_context"); + + if (default_d3d12compute_acquire_context == nullptr || default_d3d12compute_release_context == nullptr) { + printf("Failed to extract default D3D12Compute context functions from runtime\n"); + return 1; + } + + // Run a kernel on multiple threads, using our custom context accessors on + // every acquire/release. This would likely crash or produce incorrect + // results if the accessors were not being invoked consistently. + const int width = 32, height = 256; + Buffer in(width, height); + in.fill(4.0f); + + D3D12ComputeState state; + + Func f, g; + Var x, xi, y; + f(x, y) = sqrt(in(x, y)); + g(x, y) = f(x, y); + f.gpu_tile(x, x, xi, 32).compute_at(g, y); + g.parallel(y); + + for (int i = 0; i < 10; i++) { + Buffer out = g.realize(&state, {width, height}); + out.copy_to_host(&state); + for (int y = 0; y < height; y++) { + for (int x = 0; x < width; x++) { + float correct = 2.0f; + if (out(x, y) != correct) { + printf("out(%d, %d) = %f instead of %f\n", x, y, out(x, y), correct); + return 1; + } + } + } + } + + if (state.acquires.load() != state.releases.load() || state.acquires.load() == 0) { + printf("Context acquires: %d releases: %d\n", state.acquires.load(), state.releases.load()); + printf("Expected these to match and be nonzero\n"); + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/custom_metal_context.cpp b/test/correctness/custom_metal_context.cpp new file mode 100644 index 000000000000..74b416e4689b --- /dev/null +++ b/test/correctness/custom_metal_context.cpp @@ -0,0 +1,94 @@ +#include "Halide.h" + +#include + +using namespace Halide; + +// The runtime's default context-management functions, looked up from the JIT'd +// Metal runtime module. Our custom handlers delegate to these so that we can +// verify the custom accessors are actually invoked without needing to link +// against the Metal SDK. +static int (*default_metal_acquire_context)(void *, void **, void **, bool) = nullptr; +static int (*default_metal_release_context)(void *) = nullptr; + +struct MetalState : public Halide::JITUserContext { + std::atomic acquires = 0, releases = 0; + + static int my_acquire_context(JITUserContext *ctx, void **device, void **queue, bool create) { + MetalState *state = (MetalState *)ctx; + state->acquires++; + return default_metal_acquire_context(ctx, device, queue, create); + } + + static int my_release_context(JITUserContext *ctx) { + MetalState *state = (MetalState *)ctx; + state->releases++; + return default_metal_release_context(ctx); + } + + MetalState() { + handlers.custom_metal_acquire_context = my_acquire_context; + handlers.custom_metal_release_context = my_release_context; + } +}; + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::Metal)) { + printf("[SKIP] Metal not enabled.\n"); + return 0; + } + + // Force-initialize the Metal runtime module by running something trivial, + // then extract the default context-management functions from it. + evaluate_may_gpu(Expr(0.f)); + + default_metal_acquire_context = (int (*)(void *, void **, void **, bool)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_metal_acquire_context"); + default_metal_release_context = (int (*)(void *)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_metal_release_context"); + + if (default_metal_acquire_context == nullptr || default_metal_release_context == nullptr) { + printf("Failed to extract default Metal context functions from runtime\n"); + return 1; + } + + // Run a kernel on multiple threads, using our custom context accessors on + // every acquire/release. This would likely crash or produce incorrect + // results if the accessors were not being invoked consistently. + const int width = 32, height = 256; + Buffer in(width, height); + in.fill(4.0f); + + MetalState state; + + Func f, g; + Var x, xi, y; + f(x, y) = sqrt(in(x, y)); + g(x, y) = f(x, y); + f.gpu_tile(x, x, xi, 32).compute_at(g, y); + g.parallel(y); + + for (int i = 0; i < 10; i++) { + Buffer out = g.realize(&state, {width, height}); + out.copy_to_host(&state); + for (int y = 0; y < height; y++) { + for (int x = 0; x < width; x++) { + float correct = 2.0f; + if (out(x, y) != correct) { + printf("out(%d, %d) = %f instead of %f\n", x, y, out(x, y), correct); + return 1; + } + } + } + } + + if (state.acquires.load() != state.releases.load() || state.acquires.load() == 0) { + printf("Context acquires: %d releases: %d\n", state.acquires.load(), state.releases.load()); + printf("Expected these to match and be nonzero\n"); + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/custom_opencl_context.cpp b/test/correctness/custom_opencl_context.cpp new file mode 100644 index 000000000000..93177d8d9870 --- /dev/null +++ b/test/correctness/custom_opencl_context.cpp @@ -0,0 +1,94 @@ +#include "Halide.h" + +#include + +using namespace Halide; + +// The runtime's default context-management functions, looked up from the JIT'd +// OpenCL runtime module. Our custom handlers delegate to these so that we can +// verify the custom accessors are actually invoked without needing to link +// against the OpenCL SDK. +static int (*default_acquire_cl_context)(void *, void **, void **, bool) = nullptr; +static int (*default_release_cl_context)(void *) = nullptr; + +struct OpenCLState : public Halide::JITUserContext { + std::atomic acquires = 0, releases = 0; + + static int my_acquire_context(JITUserContext *ctx, void **cl_ctx, void **cl_q, bool create) { + OpenCLState *state = (OpenCLState *)ctx; + state->acquires++; + return default_acquire_cl_context(ctx, cl_ctx, cl_q, create); + } + + static int my_release_context(JITUserContext *ctx) { + OpenCLState *state = (OpenCLState *)ctx; + state->releases++; + return default_release_cl_context(ctx); + } + + OpenCLState() { + handlers.custom_opencl_acquire_context = my_acquire_context; + handlers.custom_opencl_release_context = my_release_context; + } +}; + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::OpenCL)) { + printf("[SKIP] OpenCL not enabled.\n"); + return 0; + } + + // Force-initialize the OpenCL runtime module by running something trivial, + // then extract the default context-management functions from it. + evaluate_may_gpu(Expr(0.f)); + + default_acquire_cl_context = (int (*)(void *, void **, void **, bool)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_acquire_cl_context"); + default_release_cl_context = (int (*)(void *)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_release_cl_context"); + + if (default_acquire_cl_context == nullptr || default_release_cl_context == nullptr) { + printf("Failed to extract default OpenCL context functions from runtime\n"); + return 1; + } + + // Run a kernel on multiple threads, using our custom context accessors on + // every acquire/release. This would likely crash or produce incorrect + // results if the accessors were not being invoked consistently. + const int width = 32, height = 256; + Buffer in(width, height); + in.fill(4.0f); + + OpenCLState state; + + Func f, g; + Var x, xi, y; + f(x, y) = sqrt(in(x, y)); + g(x, y) = f(x, y); + f.gpu_tile(x, x, xi, 32).compute_at(g, y); + g.parallel(y); + + for (int i = 0; i < 10; i++) { + Buffer out = g.realize(&state, {width, height}); + out.copy_to_host(&state); + for (int y = 0; y < height; y++) { + for (int x = 0; x < width; x++) { + float correct = 2.0f; + if (out(x, y) != correct) { + printf("out(%d, %d) = %f instead of %f\n", x, y, out(x, y), correct); + return 1; + } + } + } + } + + if (state.acquires.load() != state.releases.load() || state.acquires.load() == 0) { + printf("Context acquires: %d releases: %d\n", state.acquires.load(), state.releases.load()); + printf("Expected these to match and be nonzero\n"); + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/custom_vulkan_context.cpp b/test/correctness/custom_vulkan_context.cpp new file mode 100644 index 000000000000..9fd86e2b8844 --- /dev/null +++ b/test/correctness/custom_vulkan_context.cpp @@ -0,0 +1,101 @@ +#include "Halide.h" + +#include +#include + +using namespace Halide; + +// The runtime's default context-management functions, looked up from the JIT'd +// Vulkan runtime module. Our custom handlers delegate to these so that we can +// verify the custom accessors are actually invoked without needing to link +// against the Vulkan SDK. All Vulkan handles are modelled as void * here to +// avoid a dependence on the Vulkan headers. +static int (*default_vulkan_acquire_context)(void *, void **, void **, void **, void **, + void **, uint32_t *, void **, bool) = nullptr; +static int (*default_vulkan_release_context)(void *, void *, void *, void *, uint64_t) = nullptr; + +struct VulkanState : public Halide::JITUserContext { + std::atomic acquires = 0, releases = 0; + + static int my_acquire_context(JITUserContext *ctx, void **allocator, void **instance, void **device, + void **physical_device, void **queue, uint32_t *queue_family_index, + void **messenger, bool create) { + VulkanState *state = (VulkanState *)ctx; + state->acquires++; + return default_vulkan_acquire_context(ctx, allocator, instance, device, physical_device, + queue, queue_family_index, messenger, create); + } + + static int my_release_context(JITUserContext *ctx, void *instance, void *device, void *queue, uint64_t messenger) { + VulkanState *state = (VulkanState *)ctx; + state->releases++; + return default_vulkan_release_context(ctx, instance, device, queue, messenger); + } + + VulkanState() { + handlers.custom_vulkan_acquire_context = my_acquire_context; + handlers.custom_vulkan_release_context = my_release_context; + } +}; + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::Vulkan)) { + printf("[SKIP] Vulkan not enabled.\n"); + return 0; + } + + // Force-initialize the Vulkan runtime module by running something trivial, + // then extract the default context-management functions from it. + evaluate_may_gpu(Expr(0.f)); + + default_vulkan_acquire_context = + (int (*)(void *, void **, void **, void **, void **, void **, uint32_t *, void **, bool)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_vulkan_acquire_context"); + default_vulkan_release_context = (int (*)(void *, void *, void *, void *, uint64_t)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_vulkan_release_context"); + + if (default_vulkan_acquire_context == nullptr || default_vulkan_release_context == nullptr) { + printf("Failed to extract default Vulkan context functions from runtime\n"); + return 1; + } + + // Run a kernel on multiple threads, using our custom context accessors on + // every acquire/release. This would likely crash or produce incorrect + // results if the accessors were not being invoked consistently. + const int width = 32, height = 256; + Buffer in(width, height); + in.fill(4.0f); + + VulkanState state; + + Func f, g; + Var x, xi, y; + f(x, y) = sqrt(in(x, y)); + g(x, y) = f(x, y); + f.gpu_tile(x, x, xi, 32).compute_at(g, y); + g.parallel(y); + + for (int i = 0; i < 10; i++) { + Buffer out = g.realize(&state, {width, height}); + out.copy_to_host(&state); + for (int y = 0; y < height; y++) { + for (int x = 0; x < width; x++) { + float correct = 2.0f; + if (out(x, y) != correct) { + printf("out(%d, %d) = %f instead of %f\n", x, y, out(x, y), correct); + return 1; + } + } + } + } + + if (state.acquires.load() != state.releases.load() || state.acquires.load() == 0) { + printf("Context acquires: %d releases: %d\n", state.acquires.load(), state.releases.load()); + printf("Expected these to match and be nonzero\n"); + return 1; + } + + printf("Success!\n"); + return 0; +} diff --git a/test/correctness/custom_webgpu_context.cpp b/test/correctness/custom_webgpu_context.cpp new file mode 100644 index 000000000000..2f79e494f64f --- /dev/null +++ b/test/correctness/custom_webgpu_context.cpp @@ -0,0 +1,96 @@ +#include "Halide.h" + +#include + +using namespace Halide; + +// The runtime's default context-management functions, looked up from the JIT'd +// WebGPU runtime module. Our custom handlers delegate to these so that we can +// verify the custom accessors are actually invoked without needing to link +// against the WebGPU SDK. All WebGPU handles are modelled as void * here to +// avoid a dependence on the WebGPU headers. +static int (*default_webgpu_acquire_context)(void *, void **, void **, void **, void **, bool) = nullptr; +static int (*default_webgpu_release_context)(void *) = nullptr; + +struct WebGPUState : public Halide::JITUserContext { + std::atomic acquires = 0, releases = 0; + + static int my_acquire_context(JITUserContext *ctx, void **instance, void **adapter, + void **device, void **staging_buffer, bool create) { + WebGPUState *state = (WebGPUState *)ctx; + state->acquires++; + return default_webgpu_acquire_context(ctx, instance, adapter, device, staging_buffer, create); + } + + static int my_release_context(JITUserContext *ctx) { + WebGPUState *state = (WebGPUState *)ctx; + state->releases++; + return default_webgpu_release_context(ctx); + } + + WebGPUState() { + handlers.custom_webgpu_acquire_context = my_acquire_context; + handlers.custom_webgpu_release_context = my_release_context; + } +}; + +int main(int argc, char **argv) { + Target target = get_jit_target_from_environment(); + if (!target.has_feature(Target::WebGPU)) { + printf("[SKIP] WebGPU not enabled.\n"); + return 0; + } + + // Force-initialize the WebGPU runtime module by running something trivial, + // then extract the default context-management functions from it. + evaluate_may_gpu(Expr(0.f)); + + default_webgpu_acquire_context = (int (*)(void *, void **, void **, void **, void **, bool)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_webgpu_acquire_context"); + default_webgpu_release_context = (int (*)(void *)) + Internal::JITSharedRuntime::find_symbol(target, "halide_default_webgpu_release_context"); + + if (default_webgpu_acquire_context == nullptr || default_webgpu_release_context == nullptr) { + printf("Failed to extract default WebGPU context functions from runtime\n"); + return 1; + } + + // Run a kernel on multiple threads, using our custom context accessors on + // every acquire/release. This would likely crash or produce incorrect + // results if the accessors were not being invoked consistently. + const int width = 32, height = 256; + Buffer in(width, height); + in.fill(4.0f); + + WebGPUState state; + + Func f, g; + Var x, xi, y; + f(x, y) = sqrt(in(x, y)); + g(x, y) = f(x, y); + f.gpu_tile(x, x, xi, 32).compute_at(g, y); + g.parallel(y); + + for (int i = 0; i < 10; i++) { + Buffer out = g.realize(&state, {width, height}); + out.copy_to_host(&state); + for (int y = 0; y < height; y++) { + for (int x = 0; x < width; x++) { + float correct = 2.0f; + if (out(x, y) != correct) { + printf("out(%d, %d) = %f instead of %f\n", x, y, out(x, y), correct); + return 1; + } + } + } + } + + if (state.acquires.load() != state.releases.load() || state.acquires.load() == 0) { + printf("Context acquires: %d releases: %d\n", state.acquires.load(), state.releases.load()); + printf("Expected these to match and be nonzero\n"); + return 1; + } + + printf("Success!\n"); + return 0; +}