From 48797431dae5dc0d1a029690c25ea0c64026dd62 Mon Sep 17 00:00:00 2001 From: T-Karu-smaecs Date: Sat, 12 Sep 2026 22:03:47 +0200 Subject: [PATCH 1/4] Add Context LaunchEngine for GPU dispatch CB/fence checkout. Process-lifetime launch command pool with per-job CB and fence checkout/return under mutex; stop creating a command pool per launch. --- docs/internals/gpu/context.md | 10 +- docs/internals/gpu/module.md | 2 +- src/cthreads/cpp/gpu/headers/context.hpp | 55 ++++++- src/cthreads/cpp/gpu/headers/module.hpp | 6 +- src/cthreads/cpp/gpu/impl/context.cpp | 183 ++++++++++++++++++++++- src/cthreads/cpp/gpu/impl/module.cpp | 98 +++--------- 6 files changed, 266 insertions(+), 88 deletions(-) diff --git a/docs/internals/gpu/context.md b/docs/internals/gpu/context.md index c72b065..26eeaa2 100644 --- a/docs/internals/gpu/context.md +++ b/docs/internals/gpu/context.md @@ -23,7 +23,7 @@ The `Context` struct holds all of that state for the whole process. There is one ### `void init()` -Opens the loader, creates instance and device, resolves entry points, marks `ready = true`, and creates the TransferEngine (command pool and fence). Throws typed-style error strings (for example `cthreads.gpu.VulkanLoaderNotFound`) that Python maps into exceptions in `cthreads.gpu`. +Opens the loader, creates instance and device, resolves entry points, marks `ready = true`, and creates the TransferEngine and LaunchEngine. Throws typed-style error strings (for example `cthreads.gpu.VulkanLoaderNotFound`) that Python maps into exceptions in `cthreads.gpu`. Call this when you need the GPU. Python `cthreads.gpu.init()` ends up here. @@ -31,7 +31,8 @@ Call this when you need the GPU. Python `cthreads.gpu.init()` ends up here. Destroys children first, then parents: -1. TransferEngine (staging buffer, fence, command pool) +1. LaunchEngine (free fences, then command pool) +2. TransferEngine (staging buffer, fence, command pool) 2. Shader cache entries (pipelines and layouts) 3. Logical device 4. Instance @@ -99,6 +100,10 @@ Used by [module / launch](./module.md): `vkCmdBindPipeline`, `vkCmdBindDescripto See the next section. `transfer_engine_mutex` serializes use of the single shared engine. +### LaunchEngine and mutex + +`LaunchEngine` owns the process-lifetime command pool used by `launch_gpu_kernel`. Jobs checkout a command buffer + fence, submit under `launch_engine_mutex`, wait their own fence on join, then return the CB/fence to free lists. Overlapping jobs are supported; one shared fence is not. + ## Technical terms - Vulkan loader: system library that discovers Installable Client Drivers (ICDs), which are the vendor GPU drivers. @@ -144,6 +149,7 @@ A future pool of engines is discussed for CPU threads launching GPU work. See [g 6. `vkCreateDevice` opens the logical device; `vkGetDeviceQueue` gets the queue. 7. Resolve device-level functions (buffers, commands, shaders, descriptors). 8. Create TransferEngine pool and fence. +9. Create LaunchEngine command pool (CB/fence free lists grow on demand). 9. Set `ready = true`. ## Key Vulkan calls during shutdown (story order) diff --git a/docs/internals/gpu/module.md b/docs/internals/gpu/module.md index e28edcd..2fee4f7 100644 --- a/docs/internals/gpu/module.md +++ b/docs/internals/gpu/module.md @@ -19,7 +19,7 @@ Public `gpu()` / `@Gpu` (later) will call these same types. Tests exercise them - Fence: CPU waits until the submitted dispatch has finished. - Writeback: copy device list SSBOs back into the kept Python `list` objects (`pass_as` ref). -- Per-job command pool: short-lived pool that owns the launch command buffer (not the TransferEngine pool). +- Per-job command buffer + fence: checked out from Context `LaunchEngine` for the job lifetime; returned on join (supports overlapping launches). The command **pool** is process-lifetime on Context. ## Struct `SpawnedGpuKernel` diff --git a/src/cthreads/cpp/gpu/headers/context.hpp b/src/cthreads/cpp/gpu/headers/context.hpp index 36c778e..a3e29e0 100644 --- a/src/cthreads/cpp/gpu/headers/context.hpp +++ b/src/cthreads/cpp/gpu/headers/context.hpp @@ -1,8 +1,9 @@ #pragma once #include -#include #include #include +#include +#include #include "memory.hpp" @@ -23,6 +24,30 @@ struct TransferEngine { memory::GpuBuffer staging; }; +/** + * Process-lifetime launch command pool with per-job CB + fence checkout. + * + * Overlapping jobs each hold one command_buffer and one fence until join. + * The pool itself is never created/destroyed per launch. Free lists grow on + * demand; checkout resets a returned CB/fence before reuse. + * + * Guard with Context::launch_engine_mutex for checkout, return, and queue submit. + */ +struct LaunchEngine { + VkCommandPool command_pool = VK_NULL_HANDLE; + std::vector free_command_buffers; + std::vector free_fences; +}; + +/** + * Per-job handles checked out from LaunchEngine (not owned by the job forever). + * Return via return_launch_resources after the fence has been waited. + */ +struct LaunchResources { + VkCommandBuffer command_buffer = VK_NULL_HANDLE; + VkFence fence = VK_NULL_HANDLE; +}; + struct Context { // OS handle to the loader shared library (HMODULE on Windows, void* on Linux). // Purpose: keep the DLL mapped and FreeLibrary/dlclose on shutdown. @@ -105,9 +130,11 @@ struct Context { // True only after init() fully succeeded. bool ready = false; - // These are both temporary until the cpu side @gpu calls are implemented (this however is a future poject and not on the current timeline) TransferEngine transfer_engine; std::mutex transfer_engine_mutex; + + LaunchEngine launch_engine; + std::mutex launch_engine_mutex; }; // Process-wide singleton accessor. Context& context(); @@ -120,4 +147,26 @@ bool available(); // Requires ready context; returns device_name. const std::string& device_name(); -} // namespace cthreads::gpu \ No newline at end of file +/** + * Checkout a command buffer + fence from Context::launch_engine. + * Caller records the CB, then submit_launch, then join waits the fence, then + * return_launch_resources. Holds launch_engine_mutex only for the checkout. + */ +LaunchResources checkout_launch_resources(Context& context); + +/** + * Return CB + fence to the free lists after the fence has been waited. + * Clears the handles in resources. Holds launch_engine_mutex. + */ +void return_launch_resources(Context& context, LaunchResources& resources); + +/** + * vkQueueSubmit for a launch CB under launch_engine_mutex (same lock domain as checkout). + */ +void submit_launch( + Context& context, + VkCommandBuffer command_buffer, + VkFence fence +); + +} // namespace cthreads::gpu diff --git a/src/cthreads/cpp/gpu/headers/module.hpp b/src/cthreads/cpp/gpu/headers/module.hpp index ca15fc7..904699e 100644 --- a/src/cthreads/cpp/gpu/headers/module.hpp +++ b/src/cthreads/cpp/gpu/headers/module.hpp @@ -37,9 +37,9 @@ struct Context; * - pack: GpuPack = device-local scalar + list SSBOs for this launch. * - descriptor_pool: DescriptorPool = pool that allocated descriptor_set (for free_set). * - descriptor_set: VkDescriptorSet = bindings wired to pack buffers. - * - command_buffer: VkCommandBuffer = recorded dispatch (optional until submit path). - * - command_pool: VkCommandPool = pool that owns command_buffer (for free). - * - fence: VkFence = signals when the submitted dispatch has finished. + * - command_buffer: VkCommandBuffer = checked out from Context LaunchEngine. + * - command_pool: VkCommandPool = Context launch pool (borrowed; not destroyed on join). + * - fence: VkFence = checked out per job; returned to LaunchEngine after wait. * - symbol: string = shader cache key for this kernel. * - group_count_x/y/z: uint32_t = vkCmdDispatch workgroup counts. * - values_keep: shared_ptr to py::list = Python args kept alive for list writeback. diff --git a/src/cthreads/cpp/gpu/impl/context.cpp b/src/cthreads/cpp/gpu/impl/context.cpp index 60fe4ad..944765d 100644 --- a/src/cthreads/cpp/gpu/impl/context.cpp +++ b/src/cthreads/cpp/gpu/impl/context.cpp @@ -115,6 +115,71 @@ namespace { } } + // ------ Hidden LaunchEngine helpers ------ + + void shutdown_launch_engine_unlocked(Context& c) { + LaunchEngine& le = c.launch_engine; + if (le.command_pool == VK_NULL_HANDLE && + le.free_command_buffers.empty() && + le.free_fences.empty()) { + return; + } + + if (c.device != VK_NULL_HANDLE && c.vkDestroyFence) { + for (VkFence fence : le.free_fences) { + if (fence != VK_NULL_HANDLE) { + c.vkDestroyFence(c.device, fence, nullptr); + } + } + } + le.free_fences.clear(); + // Destroying the pool frees every CB allocated from it (idle and any + // still checked out if shutdown races a live job — process teardown). + le.free_command_buffers.clear(); + if (le.command_pool != VK_NULL_HANDLE && c.device != VK_NULL_HANDLE && + c.vkDestroyCommandPool) { + c.vkDestroyCommandPool(c.device, le.command_pool, nullptr); + } + le.command_pool = VK_NULL_HANDLE; + } + + void shutdown_launch_engine(Context& c) { + std::lock_guard lock(c.launch_engine_mutex); + shutdown_launch_engine_unlocked(c); + } + + void init_launch_engine(Context& c) { + std::lock_guard lock(c.launch_engine_mutex); + if (c.launch_engine.command_pool != VK_NULL_HANDLE) { + return; + } + + if (c.device == VK_NULL_HANDLE || !c.vkCreateCommandPool || + !c.vkDestroyCommandPool) { + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: init_launch_engine missing " + "device or command pool entry points"); + } + + if (!c.launch_engine.free_command_buffers.empty() || + !c.launch_engine.free_fences.empty()) { + shutdown_launch_engine_unlocked(c); + } + + VkCommandPoolCreateInfo pool_info{}; + pool_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO; + pool_info.queueFamilyIndex = c.queue_family; + // RESET: checkout path calls vkResetCommandBuffer before re-record. + pool_info.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT; + if (c.vkCreateCommandPool( + c.device, &pool_info, nullptr, &c.launch_engine.command_pool) != + VK_SUCCESS) { + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: vkCreateCommandPool failed for " + "LaunchEngine"); + } + } + // ------ Hidden Context Helpers ------ // Look up one export inside the already-loaded loader module. @@ -357,12 +422,14 @@ namespace { c, c.instance, "vkCmdPipelineBarrier"); c.ready = true; - // After device + PFNs + ready: reusable copy pool/fence (staging grows later). + // After device + PFNs + ready: reusable copy + launch pools. init_transfer_engine(c); + init_launch_engine(c); } void shutdown_unlocked(Context& c) { - // Children before parents: transfer engine, shader cache, then device. + // Children before parents: launch engine, transfer engine, shader cache, then device. + shutdown_launch_engine(c); shutdown_transfer_engine(c); shader::ShaderCache::getInstance().clear(c); @@ -504,4 +571,114 @@ namespace { shutdown_unlocked(context()); } -} // namespace cthreads::gpu \ No newline at end of file + +LaunchResources checkout_launch_resources(Context& context) { + std::lock_guard lock(context.launch_engine_mutex); + LaunchEngine& le = context.launch_engine; + if (!context.ready || le.command_pool == VK_NULL_HANDLE) { + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: checkout_launch_resources needs a " + "ready LaunchEngine"); + } + if (!context.vkAllocateCommandBuffers || !context.vkResetCommandBuffer || + !context.vkCreateFence || !context.vkResetFences) { + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: checkout_launch_resources missing " + "command/fence entry points"); + } + + LaunchResources out{}; + + if (!le.free_command_buffers.empty()) { + out.command_buffer = le.free_command_buffers.back(); + le.free_command_buffers.pop_back(); + if (context.vkResetCommandBuffer(out.command_buffer, 0) != VK_SUCCESS) { + le.free_command_buffers.push_back(out.command_buffer); + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: vkResetCommandBuffer failed in " + "checkout_launch_resources"); + } + } else { + VkCommandBufferAllocateInfo alloc_info{}; + alloc_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; + alloc_info.commandPool = le.command_pool; + alloc_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; + alloc_info.commandBufferCount = 1; + if (context.vkAllocateCommandBuffers( + context.device, &alloc_info, &out.command_buffer) != + VK_SUCCESS) { + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: vkAllocateCommandBuffers failed " + "in checkout_launch_resources"); + } + } + + if (!le.free_fences.empty()) { + out.fence = le.free_fences.back(); + le.free_fences.pop_back(); + if (context.vkResetFences(context.device, 1, &out.fence) != VK_SUCCESS) { + le.free_fences.push_back(out.fence); + le.free_command_buffers.push_back(out.command_buffer); + out = LaunchResources{}; + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: vkResetFences failed in " + "checkout_launch_resources"); + } + } else { + VkFenceCreateInfo fence_info{}; + fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; + if (context.vkCreateFence( + context.device, &fence_info, nullptr, &out.fence) != + VK_SUCCESS) { + le.free_command_buffers.push_back(out.command_buffer); + out.command_buffer = VK_NULL_HANDLE; + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: vkCreateFence failed in " + "checkout_launch_resources"); + } + } + + return out; +} + +void return_launch_resources(Context& context, LaunchResources& resources) { + std::lock_guard lock(context.launch_engine_mutex); + LaunchEngine& le = context.launch_engine; + if (resources.command_buffer != VK_NULL_HANDLE) { + le.free_command_buffers.push_back(resources.command_buffer); + resources.command_buffer = VK_NULL_HANDLE; + } + if (resources.fence != VK_NULL_HANDLE) { + le.free_fences.push_back(resources.fence); + resources.fence = VK_NULL_HANDLE; + } +} + +void submit_launch( + Context& context, + VkCommandBuffer command_buffer, + VkFence fence +) { + std::lock_guard lock(context.launch_engine_mutex); + if (!context.ready || !context.queue || !context.vkQueueSubmit) { + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: submit_launch needs a ready queue"); + } + if (command_buffer == VK_NULL_HANDLE || fence == VK_NULL_HANDLE) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: submit_launch null command buffer " + "or fence"); + } + + VkSubmitInfo submit{}; + submit.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; + submit.commandBufferCount = 1; + submit.pCommandBuffers = &command_buffer; + if (context.vkQueueSubmit(context.queue, 1, &submit, fence) != VK_SUCCESS) { + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: vkQueueSubmit failed in " + "submit_launch"); + } +} + +} // namespace cthreads::gpu diff --git a/src/cthreads/cpp/gpu/impl/module.cpp b/src/cthreads/cpp/gpu/impl/module.cpp index 53d257e..93eca00 100644 --- a/src/cthreads/cpp/gpu/impl/module.cpp +++ b/src/cthreads/cpp/gpu/impl/module.cpp @@ -39,34 +39,23 @@ size_t std430_align_of(const std::string& kind) { } void release_inflight(Context& context, SpawnedGpuKernel& job) { - // Destroying the pool frees any CBs allocated from it; free first when we can. - if (job.command_buffer != VK_NULL_HANDLE && - job.command_pool != VK_NULL_HANDLE && - context.device != VK_NULL_HANDLE && - context.vkFreeCommandBuffers) { // free the cmd buffer when all relevant ressources are valid - context.vkFreeCommandBuffers( - context.device, job.command_pool, 1, &job.command_buffer); - } - job.command_buffer = VK_NULL_HANDLE; - - // Per-launch command pool (not the TransferEngine pool). - if (job.command_pool != VK_NULL_HANDLE && - context.device != VK_NULL_HANDLE && - context.vkDestroyCommandPool) { - context.vkDestroyCommandPool(context.device, job.command_pool, nullptr); + // Return checked-out CB + fence to LaunchEngine (do not destroy the pool). + if (job.command_buffer != VK_NULL_HANDLE || job.fence != VK_NULL_HANDLE) { + LaunchResources resources{}; + resources.command_buffer = job.command_buffer; + resources.fence = job.fence; + job.command_buffer = VK_NULL_HANDLE; + job.fence = VK_NULL_HANDLE; + job.command_pool = VK_NULL_HANDLE; + return_launch_resources(context, resources); + } else { + job.command_pool = VK_NULL_HANDLE; } - job.command_pool = VK_NULL_HANDLE; if (job.descriptor_set != VK_NULL_HANDLE) { // free the descriptors (if not freed yet) pack::free_set(context, job.descriptor_pool, job.descriptor_set); } pack::destroy_pool(context, job.descriptor_pool); - // destroy the fence if not already done - if (job.fence != VK_NULL_HANDLE && context.device != VK_NULL_HANDLE && - context.vkDestroyFence) { - context.vkDestroyFence(context.device, job.fence, nullptr); - } - job.fence = VK_NULL_HANDLE; // destroy the pack pack::destroy_gpu_pack(context, job.pack); job.symbol.clear(); // clear the symbol @@ -640,44 +629,20 @@ std::shared_ptr launch_gpu_kernel( pack::update_descriptors( context, job->descriptor_set, entry, job->pack); - // Need bind/dispatch/barrier + the usual CB/submit PFNs. - if (!context.vkCreateCommandPool || !context.vkDestroyCommandPool || - !context.vkAllocateCommandBuffers || !context.vkFreeCommandBuffers || - !context.vkBeginCommandBuffer || !context.vkEndCommandBuffer || + // Need bind/dispatch/barrier PFNs (CB/fence come from LaunchEngine). + if (!context.vkBeginCommandBuffer || !context.vkEndCommandBuffer || !context.vkCmdPipelineBarrier || !context.vkCmdBindPipeline || - !context.vkCmdBindDescriptorSets || !context.vkCmdDispatch || - !context.vkCreateFence || !context.vkDestroyFence || - !context.vkQueueSubmit || !context.queue) { + !context.vkCmdBindDescriptorSets || !context.vkCmdDispatch) { throw std::runtime_error( "cthreads.gpu.VulkanInitFailed: launch_gpu_kernel missing " - "dispatch/command/fence entry points or queue"); + "dispatch/command entry points"); } - // Per-job command pool: own lifetime, no TransferEngine mutex needed. - VkCommandPoolCreateInfo pool_info{}; - pool_info.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO; - pool_info.queueFamilyIndex = context.queue_family; - pool_info.flags = VK_COMMAND_POOL_CREATE_TRANSIENT_BIT; - if (context.vkCreateCommandPool( - context.device, &pool_info, nullptr, &job->command_pool) != - VK_SUCCESS) { - throw std::runtime_error( - "cthreads.gpu.VulkanInitFailed: vkCreateCommandPool failed in " - "launch_gpu_kernel"); - } - - VkCommandBufferAllocateInfo alloc_info{}; - alloc_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO; - alloc_info.commandPool = job->command_pool; - alloc_info.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY; - alloc_info.commandBufferCount = 1; - if (context.vkAllocateCommandBuffers( - context.device, &alloc_info, &job->command_buffer) != - VK_SUCCESS) { - throw std::runtime_error( - "cthreads.gpu.VulkanInitFailed: vkAllocateCommandBuffers failed " - "in launch_gpu_kernel"); - } + // Checkout CB + fence from Context LaunchEngine (pool is process-lifetime). + LaunchResources launch = checkout_launch_resources(context); + job->command_buffer = launch.command_buffer; + job->fence = launch.fence; + job->command_pool = context.launch_engine.command_pool; VkCommandBufferBeginInfo begin_info{}; begin_info.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO; @@ -734,27 +699,8 @@ std::shared_ptr launch_gpu_kernel( "launch_gpu_kernel"); } - // Per-job fence (not TransferEngine.fence). Unsignaled until submit done. - VkFenceCreateInfo fence_info{}; - fence_info.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO; - if (context.vkCreateFence( - context.device, &fence_info, nullptr, &job->fence) != - VK_SUCCESS) { - throw std::runtime_error( - "cthreads.gpu.VulkanInitFailed: vkCreateFence failed in " - "launch_gpu_kernel"); - } - - VkSubmitInfo submit{}; - submit.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO; - submit.commandBufferCount = 1; - submit.pCommandBuffers = &job->command_buffer; - if (context.vkQueueSubmit( - context.queue, 1, &submit, job->fence) != VK_SUCCESS) { - throw std::runtime_error( - "cthreads.gpu.VulkanInitFailed: vkQueueSubmit failed in " - "launch_gpu_kernel"); - } + // Fence was checked out unsignaled; submit under launch_engine_mutex. + submit_launch(context, job->command_buffer, job->fence); // Do not wait here — join() waits on job->fence. } catch (...) { // Tear down any handles already stashed; then rethrow to Python. From 8abfc53f044973157d9c4a716ac392de1f6b49f2 Mon Sep 17 00:00:00 2001 From: T-Karu-smaecs Date: Sun, 13 Sep 2026 21:31:52 +0200 Subject: [PATCH 2/4] Ship @Gpu public path: glslang emit, gpu()/prepare, and SPH demo. Complete the Vulkan GPU user path on top of LaunchEngine: vendored glslang SPIR-V, Signature/codegen, ShaderCache register, resident-aware prepare after shutdown, lists-only bindings, list[bool] marshal, and broad unit/pipeline coverage. --- .gitignore | 1 + LICENSE | 127 +++ docs/STYLE.md | 4 +- docs/internals/gpu/module.md | 8 +- docs/internals/gpu/pack.md | 5 +- docs/internals/gpu/shader.md | 3 +- docs/vk_guide/05-dynamic-loading.md | 2 +- docs/vk_guide/13-map-to-our-code.md | 4 +- pyproject.toml | 3 + src/cthreads/cpp/CMakeLists.txt | 86 +- src/cthreads/cpp/bindings/gpu_module.cpp | 100 ++ src/cthreads/cpp/bindings/gpu_module.hpp | 7 +- .../cpp/bindings/gpu_testing_module.cpp | 12 +- src/cthreads/cpp/gpu/headers/compile_glsl.hpp | 26 + src/cthreads/cpp/gpu/headers/descriptors.hpp | 12 +- src/cthreads/cpp/gpu/headers/module.hpp | 3 +- src/cthreads/cpp/gpu/headers/shader.hpp | 6 +- src/cthreads/cpp/gpu/headers/shader_cache.hpp | 47 +- src/cthreads/cpp/gpu/impl/compile_glsl.cpp | 91 ++ src/cthreads/cpp/gpu/impl/descriptors.cpp | 13 +- src/cthreads/cpp/gpu/impl/module.cpp | 34 +- src/cthreads/cpp/gpu/impl/shader_cache.cpp | 12 + .../cpp/gpu/third_party_notices/README.md | 12 + src/cthreads/python/.gitignore | 11 + src/cthreads/python/cthreads/cache.py | 19 +- .../python/cthreads/frontend/Gpu/Wrapper.py | 4 +- .../cthreads/frontend/Registry/registry.py | 14 +- src/cthreads/python/cthreads/gpu/__init__.py | 103 +- .../python/cthreads/gpu/_ext_gpu_api.py | 173 +++ .../gpu/compiler/orchestrator/__init__.py | 7 + .../orchestrator/gpu_compile_session.py | 93 ++ .../gpu/compiler/orchestrator/gpu_unit.py | 119 ++ .../cthreads/gpu/compiler/translation/Glsl.py | 86 ++ .../gpu/compiler/translation/Signature.py | 101 ++ .../gpu/compiler/translation/Typeof.py | 0 .../gpu/compiler/translation/assemble.py | 40 + .../gpu/compiler/translation/context.py | 27 + .../compiler/translation/plugins/__init__.py | 107 ++ .../gpu/compiler/translation/plugins/base.py | 64 ++ .../compiler/translation/plugins/indexes.py | 61 + .../translation/plugins/math_calls.py | 42 + .../gpu/compiler/translation/result.py | 42 + .../gpu/compiler/translation/spirv.py | 133 +++ .../gpu/compiler/translation/syntax/Assign.py | 144 +++ .../gpu/compiler/translation/syntax/Flow.py | 180 +++ .../gpu/compiler/translation/syntax/Index.py | 51 + .../gpu/compiler/translation/syntax/Name.py | 67 ++ .../gpu/compiler/translation/syntax/Op.py | 102 ++ .../gpu/compiler/translation/syntax/Syntax.py | 106 ++ .../gpu/compiler/translation/translate.py | 68 ++ .../python/cthreads/gpu/frontend/__init__.py | 49 + .../cthreads/gpu/{ => frontend}/errors.py | 20 +- .../python/cthreads/gpu/frontend/indexes.py | 110 ++ .../python/cthreads/gpu/frontend/lib.py | 106 ++ .../python/cthreads/gpu/frontend/wrapper.py | 78 ++ .../python/cthreads/gpu/gpu_kernel_meta.py | 379 ++++++ .../python/cthreads/gpu/gpu_marshal.py | 75 ++ src/cthreads/python/cthreads/gpu/runtime.py | 160 +++ .../gpu/third_party_notices/README.md | 12 + .../third_party_notices/glslang-LICENSE.txt | 1016 +++++++++++++++++ tests/helpers_gpu.py | 16 + tests/unit/.gitignore | 11 + tests/unit/test_cache.py | 32 +- tests/unit/test_gpu_assemble.py | 50 + tests/unit/test_gpu_context.py | 22 +- tests/unit/test_gpu_decorator.py | 101 ++ tests/unit/test_gpu_errors.py | 51 + tests/unit/test_gpu_glsl.py | 99 ++ tests/unit/test_gpu_indexes.py | 47 + tests/unit/test_gpu_kernel_meta.py | 190 +++ tests/unit/test_gpu_marshal.py | 153 +++ tests/unit/test_gpu_pack.py | 6 +- tests/unit/test_gpu_pipeline.py | 460 ++++++++ tests/unit/test_gpu_prepare.py | 270 +++++ tests/unit/test_gpu_reserved_names.py | 73 ++ tests/unit/test_gpu_shader.py | 59 +- tests/unit/test_gpu_signature.py | 214 ++++ tests/unit/test_gpu_spirv.py | 117 ++ tests/unit/test_gpu_syntax.py | 701 ++++++++++++ tests/unit/test_gpu_translate.py | 135 +++ 80 files changed, 7243 insertions(+), 151 deletions(-) create mode 100644 LICENSE create mode 100644 src/cthreads/cpp/gpu/headers/compile_glsl.hpp create mode 100644 src/cthreads/cpp/gpu/impl/compile_glsl.cpp create mode 100644 src/cthreads/cpp/gpu/third_party_notices/README.md create mode 100644 src/cthreads/python/.gitignore create mode 100644 src/cthreads/python/cthreads/gpu/_ext_gpu_api.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/orchestrator/__init__.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/orchestrator/gpu_compile_session.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/orchestrator/gpu_unit.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/Glsl.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/Signature.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/Typeof.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/assemble.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/context.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/plugins/__init__.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/plugins/base.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/plugins/indexes.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/plugins/math_calls.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/result.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/spirv.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Assign.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Flow.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Index.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Name.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Op.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Syntax.py create mode 100644 src/cthreads/python/cthreads/gpu/compiler/translation/translate.py create mode 100644 src/cthreads/python/cthreads/gpu/frontend/__init__.py rename src/cthreads/python/cthreads/gpu/{ => frontend}/errors.py (77%) create mode 100644 src/cthreads/python/cthreads/gpu/frontend/indexes.py create mode 100644 src/cthreads/python/cthreads/gpu/frontend/lib.py create mode 100644 src/cthreads/python/cthreads/gpu/frontend/wrapper.py create mode 100644 src/cthreads/python/cthreads/gpu/gpu_kernel_meta.py create mode 100644 src/cthreads/python/cthreads/gpu/gpu_marshal.py create mode 100644 src/cthreads/python/cthreads/gpu/runtime.py create mode 100644 src/cthreads/python/cthreads/gpu/third_party_notices/README.md create mode 100644 src/cthreads/python/cthreads/gpu/third_party_notices/glslang-LICENSE.txt create mode 100644 tests/helpers_gpu.py create mode 100644 tests/unit/.gitignore create mode 100644 tests/unit/test_gpu_assemble.py create mode 100644 tests/unit/test_gpu_decorator.py create mode 100644 tests/unit/test_gpu_errors.py create mode 100644 tests/unit/test_gpu_glsl.py create mode 100644 tests/unit/test_gpu_indexes.py create mode 100644 tests/unit/test_gpu_kernel_meta.py create mode 100644 tests/unit/test_gpu_marshal.py create mode 100644 tests/unit/test_gpu_pipeline.py create mode 100644 tests/unit/test_gpu_prepare.py create mode 100644 tests/unit/test_gpu_reserved_names.py create mode 100644 tests/unit/test_gpu_signature.py create mode 100644 tests/unit/test_gpu_spirv.py create mode 100644 tests/unit/test_gpu_syntax.py create mode 100644 tests/unit/test_gpu_translate.py diff --git a/.gitignore b/.gitignore index bfabda3..06432db 100644 --- a/.gitignore +++ b/.gitignore @@ -52,6 +52,7 @@ api/ # >>> cthreads (auto) __Thread__/ __Threadable__/ +__Gpu__/ .cthreads_cache.json cthreads_kernels.dll cthreads_kernels.so diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..f98bfb4 --- /dev/null +++ b/LICENSE @@ -0,0 +1,127 @@ +MIT License + +Copyright (c) 2026 Tobias Karusseit + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. + +================================================================================ +Third-party notices +================================================================================ + +cthreads depends on (or may optionally link / redistribute) the following +third-party components. Their licenses apply to those components only. +Full license texts are available from the upstream projects. + +-------------------------------------------------------------------------------- +pybind11 +-------------------------------------------------------------------------------- +Used to bind the native `_ext` module to Python (FetchContent or system +package via CMake). + +License: BSD-3-Clause +Copyright (c) 2016 Wenzel Jakob , All rights reserved. +Homepage: https://github.com/pybind/pybind11 + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +3. Neither the name of the copyright holder nor the names of its contributors + may be used to endorse or promote products derived from this software + without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +-------------------------------------------------------------------------------- +Vulkan (Khronos Group) — optional, CTHREADS_GPU=ON +-------------------------------------------------------------------------------- +GPU builds use Vulkan headers from the Vulkan SDK / CMake Vulkan package. +At runtime, cthreads loads the system Vulkan loader (for example vulkan-1) +dynamically; the loader and ICD are provided by the platform / GPU vendor +and are not redistributed as part of cthreads by default. + +Vulkan-Headers / related Khronos materials are typically licensed under the +Apache License, Version 2.0. +Homepage: https://github.com/KhronosGroup/Vulkan-Headers +License reference: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- +Khronos glslang — GPU GLSL -> SPIR-V (CTHREADS_GPU=ON) +-------------------------------------------------------------------------------- +When built with CTHREADS_GPU=ON, cthreads FetchContent-vendors and statically +links Khronos glslang into `_ext` to implement `_ext.gpu.compile_glsl`. +This is the same compiler engine Google shaderc wraps. End users of a GPU +wheel do not need glslc or the Vulkan SDK shader tools. + +License: Apache License, Version 2.0 (with BSD-style components in the tree) +Homepage: https://github.com/KhronosGroup/glslang +License reference: https://www.apache.org/licenses/LICENSE-2.0 + +Upstream license text is copied at build time into: + + cthreads/gpu/third_party_notices/ + +(see README.md there). Redistribute that directory with any binary that +includes the GPU extension. + +-------------------------------------------------------------------------------- +scikit-build-core +-------------------------------------------------------------------------------- +Used as the Python build backend to drive CMake (build-time dependency; not +part of the runtime import of cthreads). + +License: Apache License, Version 2.0 +Homepage: https://github.com/scikit-build/scikit-build-core +License reference: https://www.apache.org/licenses/LICENSE-2.0 + +-------------------------------------------------------------------------------- +Apache License, Version 2.0 (summary for Apache-licensed deps above) +-------------------------------------------------------------------------------- +You may reproduce and distribute copies of Apache-2.0 works under the terms +of that license. A copy of the full license text is available at: + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the Apache License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the Apache License for the specific language governing permissions and +limitations under the License. + +-------------------------------------------------------------------------------- +Python +-------------------------------------------------------------------------------- +cthreads is designed to run on CPython. The Python interpreter and standard +library are governed by the Python Software Foundation License. +Homepage: https://www.python.org/ diff --git a/docs/STYLE.md b/docs/STYLE.md index 3b6f514..647a2ec 100644 --- a/docs/STYLE.md +++ b/docs/STYLE.md @@ -77,7 +77,9 @@ def example_fn() -> None: 1. the use of special utf characters is not permitted ```latex -Exmaple: +Example: — should be - or depending on ctx ,. etc. → should be -> ``` + +2. `from __future__ import annotations` is only permitted iff its used to avoid import errors, improve import performance, with typechecking, or to avoid any other error. Otherwise explicit imports are preffered to ensure easy dependency maintenace \ No newline at end of file diff --git a/docs/internals/gpu/module.md b/docs/internals/gpu/module.md index 2fee4f7..2af5a4f 100644 --- a/docs/internals/gpu/module.md +++ b/docs/internals/gpu/module.md @@ -13,7 +13,9 @@ Depends on: [Context](./context.md), [Pack](./pack.md), [Descriptors](./descript `launch_gpu_kernel` is the GPU analogue of CPU `spawn_from_meta`: build a `GpuPack`, wire descriptors, record bind+dispatch, submit with a fence, and return a job handle. `SpawnedGpuKernel::join` waits on that fence, downloads ref lists into the same Python objects, then releases Vulkan state. There is no OS worker thread and no mid-run `sync_state`. -Public `gpu()` / `@Gpu` (later) will call these same types. Tests exercise them today via `_ext.gpu.testing.smoke_launch_saxpy`. +Public `gpu()` / `@Gpu` (later) will call these same types. Product pybind: +`_ext.gpu.launch_gpu_kernel` + `_ext.gpu.GpuJob`. Tests register smoke SPIR-V +via `_ext.gpu.testing.register_smoke_saxpy`, then launch on the product path. ## Technical terms @@ -49,4 +51,6 @@ Value scalars are not written back. Threadable/schema marshal is later. ## Testing -`smoke_launch_saxpy` registers committed saxpy SPIR-V, launches, joins, and asserts `y` matches CPU saxpy. Pytest: `tests/unit/test_gpu_shader.py::test_live_smoke_launch_saxpy`. +`register_smoke_saxpy` puts committed saxpy SPIR-V in ShaderCache (test-only). +Pytest drives product `launch_gpu_kernel` + `GpuJob.join` and asserts `y`: +`tests/unit/test_gpu_shader.py::test_live_launch_saxpy_product_path`. diff --git a/docs/internals/gpu/pack.md b/docs/internals/gpu/pack.md index d57e8ae..d021ca0 100644 --- a/docs/internals/gpu/pack.md +++ b/docs/internals/gpu/pack.md @@ -97,4 +97,7 @@ destroy_gpu_pack ## Testing today -`_ext.gpu.testing` exposes pack round-trip helpers (float and int packs) and `smoke_launch_saxpy` (full launch + join writeback). Pytest: `tests/unit/test_gpu_pack.py`, `tests/unit/test_gpu_shader.py`. +`_ext.gpu.testing` exposes pack round-trip helpers (float and int packs) and +`register_smoke_saxpy` (SPIR-V cache only). Launch/join use product +`_ext.gpu.launch_gpu_kernel`. Pytest: `tests/unit/test_gpu_pack.py`, +`tests/unit/test_gpu_shader.py`. diff --git a/docs/internals/gpu/shader.md b/docs/internals/gpu/shader.md index 7ead72c..4628132 100644 --- a/docs/internals/gpu/shader.md +++ b/docs/internals/gpu/shader.md @@ -25,7 +25,8 @@ Per-launch objects (GpuPack buffers, descriptor sets, fences) are not stored her ## Access rights -- Writers (registry / testing, via `add` once friended): insert new entries. +- Writers: `ShaderRegistry::register_spirv` only (friend of `ShaderCache`; also + exposed as `_ext.gpu.register_shader`). Tests use the same writer. - Everyone else: `get` returns a const reference. - Context shutdown: `clear` destroys all Vulkan objects, then empties the map. diff --git a/docs/vk_guide/05-dynamic-loading.md b/docs/vk_guide/05-dynamic-loading.md index 1cd449a..593c524 100644 --- a/docs/vk_guide/05-dynamic-loading.md +++ b/docs/vk_guide/05-dynamic-loading.md @@ -81,4 +81,4 @@ When adding a new Vulkan call in `memory.cpp`, first check: is its `PFN_` on `Co | `VulkanNoDevice` | No compute-capable GPU | | `VulkanInitFailed: vkCreate… failed` | Driver rejected create | -Python maps these string prefixes to exception types in `cthreads.gpu.errors`. +Python maps these string prefixes to exception types in `cthreads.gpu.frontend.errors`. diff --git a/docs/vk_guide/13-map-to-our-code.md b/docs/vk_guide/13-map-to-our-code.md index 16fcc38..ffe83dd 100644 --- a/docs/vk_guide/13-map-to-our-code.md +++ b/docs/vk_guide/13-map-to-our-code.md @@ -46,8 +46,8 @@ Newcomer-oriented C++ module docs: [docs/internals/gpu/README.md](../internals/g 1. Read guide 09-12. 2. Open `pack.hpp` / `descriptors.hpp` / `shader_cache.hpp` / `module.hpp`. -3. Trace `smoke_launch_saxpy` in `gpu/testing/shader_smoke.cpp` end-to-end. -4. Run `tests/unit/test_gpu_shader.py::test_live_smoke_launch_saxpy`. +3. Trace product `launch_gpu_kernel` after `testing.register_smoke_saxpy`. +4. Run `tests/unit/test_gpu_shader.py::test_live_launch_saxpy_product_path`. ## Build flag reminder diff --git a/pyproject.toml b/pyproject.toml index 2c5c5a2..5b071bb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,10 +11,12 @@ name = "cthreads" version = "0.1.1" description = "Compile @Threadable / @Thread Python into native C++ kernels and run them off the GIL." requires-python = ">=3.10" +license = { file = "LICENSE" } keywords = ["threading", "multithreading", "codegen", "python compiler", "cpp", "pybind11"] classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", + "License :: OSI Approved :: MIT License", "Programming Language :: C++", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3.10", @@ -53,6 +55,7 @@ wheel.exclude = [ "**/*.pyc", ] sdist.include = [ + "LICENSE", "src/cthreads/cpp/**", "src/cthreads/python/cthreads/**", "src/cthreads/python/api/**", diff --git a/src/cthreads/cpp/CMakeLists.txt b/src/cthreads/cpp/CMakeLists.txt index df89632..3b06479 100644 --- a/src/cthreads/cpp/CMakeLists.txt +++ b/src/cthreads/cpp/CMakeLists.txt @@ -145,19 +145,41 @@ endif() if(CTHREADS_GPU) find_package(Vulkan REQUIRED) # vulkan sdk for headers/includes message(STATUS "cthreads GPU: ON (Vulkan)") + + # Vendor Khronos glslang (GLSL -> SPIR-V). Same compiler engine shaderc uses. + # Linked statically into _ext so end users need no glslc / extra SDK tools. + include(FetchContent) + set(ENABLE_GLSLANG_BINARIES OFF CACHE BOOL "" FORCE) + set(ENABLE_HLSL OFF CACHE BOOL "" FORCE) + set(ENABLE_OPT OFF CACHE BOOL "" FORCE) + set(ENABLE_SPVREMAPPER OFF CACHE BOOL "" FORCE) + set(ENABLE_CTEST OFF CACHE BOOL "" FORCE) + set(SKIP_GLSLANG_INSTALL ON CACHE BOOL "" FORCE) + set(BUILD_EXTERNAL OFF CACHE BOOL "" FORCE) + set(GLSLANG_TESTS OFF CACHE BOOL "" FORCE) + set(ENABLE_GLSLANG_JS OFF CACHE BOOL "" FORCE) + FetchContent_Declare( + glslang + GIT_REPOSITORY https://github.com/KhronosGroup/glslang.git + GIT_TAG 15.1.0 + GIT_SHALLOW TRUE + ) + FetchContent_MakeAvailable(glslang) + message(STATUS "cthreads GPU: vendored glslang for compile_glsl") + target_sources(_ext PRIVATE - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/context.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/memory.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/pack.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/shader_cache.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/shader.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/descriptors.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/module.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/testing/pack_roundtrip.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/gpu/testing/shader_smoke.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/bindings/gpu_module.cpp" - "${CMAKE_CURRENT_SOURCE_DIR}/bindings/gpu_testing_module.cpp" - # GPU: context, memory, GpuPack, shader cache/entry (+ test-only roundtrip) + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/context.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/memory.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/pack.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/shader_cache.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/shader.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/descriptors.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/module.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/compile_glsl.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/testing/pack_roundtrip.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/testing/shader_smoke.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/bindings/gpu_module.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/bindings/gpu_testing_module.cpp" ) target_include_directories(_ext PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/gpu/headers @@ -165,4 +187,42 @@ if(CTHREADS_GPU) ${Vulkan_INCLUDE_DIRS} ) target_compile_definitions(_ext PRIVATE CTHREADS_WITH_GPU=1) -endif() \ No newline at end of file + target_link_libraries(_ext PRIVATE + glslang + SPIRV + glslang-default-resource-limits + ) + + # Ship upstream license texts next to the Python package (Apache/BSD notices). + set(_cthreads_gpu_notices_out + "${_cthreads_py_out}/gpu/third_party_notices") + set(_cthreads_gpu_notices_src + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/third_party_notices") + file(MAKE_DIRECTORY "${_cthreads_gpu_notices_out}") + if(EXISTS "${_cthreads_gpu_notices_src}/README.md") + configure_file( + "${_cthreads_gpu_notices_src}/README.md" + "${_cthreads_gpu_notices_out}/README.md" + COPYONLY + ) + endif() + if(DEFINED glslang_SOURCE_DIR) + foreach(_lic IN ITEMS LICENSE.txt LICENSE.TXT LICENSE) + if(EXISTS "${glslang_SOURCE_DIR}/${_lic}") + configure_file( + "${glslang_SOURCE_DIR}/${_lic}" + "${_cthreads_gpu_notices_out}/glslang-${_lic}" + COPYONLY + ) + break() + endif() + endforeach() + endif() + + if(DEFINED SKBUILD AND NOT (DEFINED SKBUILD_STATE AND SKBUILD_STATE STREQUAL "editable")) + install(DIRECTORY "${_cthreads_gpu_notices_out}/" + DESTINATION cthreads/gpu/third_party_notices + OPTIONAL + ) + endif() +endif() diff --git a/src/cthreads/cpp/bindings/gpu_module.cpp b/src/cthreads/cpp/bindings/gpu_module.cpp index c63e42e..bc9e577 100644 --- a/src/cthreads/cpp/bindings/gpu_module.cpp +++ b/src/cthreads/cpp/bindings/gpu_module.cpp @@ -6,9 +6,19 @@ #include "gpu_testing_module.hpp" #include "../gpu/headers/context.hpp" +#include "../gpu/headers/compile_glsl.hpp" +#include "../gpu/headers/module.hpp" +#include "../gpu/headers/shader_cache.hpp" #include +#include +#include +#include +#include +#include +#include + namespace py = pybind11; void bind_gpu(py::module_& parent) { @@ -41,6 +51,96 @@ void bind_gpu(py::module_& parent) { "Destroy device/instance and unload the Vulkan loader." ); + // Product launch handle (mirror of CPU SpawnedKernel / Job surface). + // join() always uses the process Context so Python never holds a Context&. + py::class_< + cthreads::gpu::SpawnedGpuKernel, + std::shared_ptr>(g, "GpuJob") + .def( + "start", + &cthreads::gpu::SpawnedGpuKernel::start, + "No-op: work is submitted at launch_gpu_kernel time." + ) + .def( + "join", + [](cthreads::gpu::SpawnedGpuKernel& self) { + self.join(cthreads::gpu::context()); + }, + "Wait for the GPU fence, download ref lists, release inflight state." + ) + .def( + "done", + &cthreads::gpu::SpawnedGpuKernel::done, + "True after join (or failure) has marked the job complete." + ) + .def( + "wait", + &cthreads::gpu::SpawnedGpuKernel::wait, + py::call_guard(), + "Block until done_flag; does not download. Prefer join()." + ); + + g.def( + "compile_glsl", + [](const std::string& source) { + // glslang can take noticeable time; release the GIL while compiling. + std::vector words; + { + py::gil_scoped_release release; + words = cthreads::gpu::compile_glsl_to_spirv(source); + } + const char* raw = reinterpret_cast(words.data()); + const py::ssize_t nbytes = + static_cast(words.size() * sizeof(std::uint32_t)); + return py::bytes(raw, nbytes); + }, + py::arg("source"), + "Compile a GLSL compute shader string to SPIR-V bytes (vendored glslang). " + "No external glslc / Vulkan SDK tools required." + ); + + g.def( + "register_shader", + [](const std::string& symbol, const py::bytes& spirv, uint32_t binding_count) { + cthreads::gpu::init(); + const std::string raw = spirv; + if (raw.empty() || (raw.size() % 4) != 0) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: register_shader spirv must " + "be a non-empty multiple of 4 bytes"); + } + if (binding_count == 0) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: register_shader " + "binding_count must be >= 1"); + } + // SPIR-V is a stream of little-endian uint32 words. + std::vector words(raw.size() / 4); + std::memcpy(words.data(), raw.data(), raw.size()); + (void)cthreads::gpu::shader::ShaderRegistry::register_spirv( + cthreads::gpu::context(), + symbol, + words.data(), + words.size(), + binding_count); + }, + py::arg("symbol"), + py::arg("spirv"), + py::arg("binding_count"), + "Create a compute pipeline from SPIR-V and insert it into ShaderCache. " + "Sole product writer path (ShaderRegistry). Duplicate symbol throws." + ); + + g.def( + "launch_gpu_kernel", + &cthreads::gpu::launch_gpu_kernel, + py::arg("meta"), + py::arg("ordered_values"), + "Marshal args from meta + ordered_values, submit compute, return GpuJob. " + "Does not wait; call job.join() for fence wait and list writeback. " + "Requires the kernel symbol to already be in ShaderCache." + ); + // Test-only pack round-trips live in a separate submodule / translation unit // so product bindings stay small. Not re-exported by cthreads.gpu. bind_gpu_testing(g); diff --git a/src/cthreads/cpp/bindings/gpu_module.hpp b/src/cthreads/cpp/bindings/gpu_module.hpp index 60e4d1c..d66e9d5 100644 --- a/src/cthreads/cpp/bindings/gpu_module.hpp +++ b/src/cthreads/cpp/bindings/gpu_module.hpp @@ -4,5 +4,8 @@ namespace py = pybind11; -/** Register ``cthreads._ext.gpu`` (probe API + test-only ``testing`` submodule). */ -void bind_gpu(py::module_& parent); \ No newline at end of file +/** + * Register ``cthreads._ext.gpu`` (probe API, launch_gpu_kernel / GpuJob, + * and test-only ``testing`` submodule). + */ +void bind_gpu(py::module_& parent); diff --git a/src/cthreads/cpp/bindings/gpu_testing_module.cpp b/src/cthreads/cpp/bindings/gpu_testing_module.cpp index 81984f0..e71af90 100644 --- a/src/cthreads/cpp/bindings/gpu_testing_module.cpp +++ b/src/cthreads/cpp/bindings/gpu_testing_module.cpp @@ -106,8 +106,14 @@ void bind_gpu_testing(py::module_& gpu_parent) { "Raises GpuInvalidArgument for binding_count 0. Test-only." ); t.def( - "smoke_launch_saxpy", - &cthreads::gpu::testing::smoke_launch_saxpy, - "Register smoke SPIR-V, launch_gpu_kernel saxpy, join writeback, check y. Test-only." + "clear_shader_cache", + &cthreads::gpu::testing::clear_shader_cache, + "Clear process ShaderCache. Test-only teardown." + ); + t.def( + "register_smoke_saxpy", + &cthreads::gpu::testing::register_smoke_saxpy, + "Register smoke saxpy SPIR-V in ShaderCache; returns symbol key. " + "Does not launch - use _ext.gpu.launch_gpu_kernel. Test-only." ); } diff --git a/src/cthreads/cpp/gpu/headers/compile_glsl.hpp b/src/cthreads/cpp/gpu/headers/compile_glsl.hpp new file mode 100644 index 0000000..60961c9 --- /dev/null +++ b/src/cthreads/cpp/gpu/headers/compile_glsl.hpp @@ -0,0 +1,26 @@ +#pragma once + +#include +#include +#include + +namespace cthreads::gpu { + +/** + * Compile a GLSL compute shader string to SPIR-V words. + * + * Uses the vendored Khronos glslang library (same compiler engine shaderc + * wraps). No external glslc / Vulkan SDK tools required at runtime. + * + * #### Args: + * - source: std::string = full compute GLSL (`#version` + buffers + main) + * + * #### Returns + * - std::vector = SPIR-V words + * + * #### Raises + * - std::runtime_error = parse/link failure (message includes glslang log) + */ +std::vector compile_glsl_to_spirv(const std::string& source); + +} // namespace cthreads::gpu diff --git a/src/cthreads/cpp/gpu/headers/descriptors.hpp b/src/cthreads/cpp/gpu/headers/descriptors.hpp index 3a2ebbd..f64dd37 100644 --- a/src/cthreads/cpp/gpu/headers/descriptors.hpp +++ b/src/cthreads/cpp/gpu/headers/descriptors.hpp @@ -123,12 +123,12 @@ void free_set( ); /** - * Writes binding-convention buffer bindings from a GpuPack into a descriptor set. + * Writes buffer bindings from a GpuPack into a descriptor set. * - * Binding 0 -> pack.scalar_buffer. Bindings 1..N -> pack.container_slots[0..N-1]. - * binding_count must equal 1 + pack.container_slots.size() and match the set - * layout. Every binding must have a non-null VkBuffer (empty list slots are - * not supported here yet; use a non-empty buffer or a future dummy SSBO). + * If pack has a scalar buffer: binding 0 = scalars, 1..N = lists. + * If not: binding 0..N-1 = lists (no scalar descriptor). + * binding_count must equal (has_scalars ? 1 : 0) + container_slots.size(). + * Every written binding must have a non-null VkBuffer. * * Called by the launch path after allocate_set and before recording bind/dispatch. * @@ -136,7 +136,7 @@ void free_set( * - context: Context& = initialized GPU context with vkUpdateDescriptorSets. * - set: VkDescriptorSet = destination set from allocate_set. * - binding_count: uint32_t = number of STORAGE_BUFFER bindings to write. - * - pack: const GpuPack& = source buffers in binding order (scalars then lists). + * - pack: const GpuPack& = source buffers (optional scalars, then lists). * * #### Throws: * - runtime_error if set is null, binding_count mismatches the pack, any diff --git a/src/cthreads/cpp/gpu/headers/module.hpp b/src/cthreads/cpp/gpu/headers/module.hpp index 904699e..e5dfa20 100644 --- a/src/cthreads/cpp/gpu/headers/module.hpp +++ b/src/cthreads/cpp/gpu/headers/module.hpp @@ -61,7 +61,8 @@ struct SpawnedGpuKernel { size_t value_index = 0; // index into values_keep size_t container_index = 0; // index into pack.container_slots size_t numel = 0; - std::string elem_kind; // "float" / "int" / "double" + std::string elem_kind; // "float" / "int" / "double" / "bool" + }; pack::GpuPack pack{}; diff --git a/src/cthreads/cpp/gpu/headers/shader.hpp b/src/cthreads/cpp/gpu/headers/shader.hpp index 960fae9..8f53466 100644 --- a/src/cthreads/cpp/gpu/headers/shader.hpp +++ b/src/cthreads/cpp/gpu/headers/shader.hpp @@ -22,7 +22,7 @@ struct Context; * - SPIR-V: binary compute shader (uint32 words). Built for tests as committed * bytes or via shaderc; @Gpu emit feeds the same helper later. * - binding_count: STORAGE_BUFFER bindings for the binding convention - * (1 = scalars only, or 1 + number of list SSBOs). + * ((scalars ? 1 : 0) + number of list SSBOs; at least 1 total). * - ShaderCacheEntry: module + set layout + pipeline layout + compute pipeline. */ namespace cthreads::gpu::shader { @@ -34,8 +34,8 @@ namespace cthreads::gpu::shader { * binding_count-1 as STORAGE_BUFFER), pipeline layout, compute pipeline. * On failure, destroys any objects already created and throws. * - * Does not insert into ShaderCache; the registry (or test harness) calls - * ShaderCache::add with the returned entry. + * Does not insert into ShaderCache; ShaderRegistry::register_spirv calls + * create_entry then ShaderCache::add. * * #### Parameters: * - context: Context& = initialized GPU context with device and create entry points. diff --git a/src/cthreads/cpp/gpu/headers/shader_cache.hpp b/src/cthreads/cpp/gpu/headers/shader_cache.hpp index a120ea6..7e99eeb 100644 --- a/src/cthreads/cpp/gpu/headers/shader_cache.hpp +++ b/src/cthreads/cpp/gpu/headers/shader_cache.hpp @@ -10,10 +10,6 @@ namespace cthreads::gpu { struct Context; } -namespace cthreads::gpu::testing { -struct ShaderCacheTestAccess; -} - /** * Shader cache: reusable per-symbol Vulkan pipeline objects. * @@ -35,10 +31,10 @@ namespace cthreads::gpu::shader { * * #### Fields: * - shader_module: VkShaderModule = SPIR-V module (may be null after pipeline create). - * - set_layout: VkDescriptorSetLayout = bindings 0 scalars, 1..N lists. + * - set_layout: VkDescriptorSetLayout = bindings 0..N-1 (scalars at 0 if present, then lists). * - pipeline_layout: VkPipelineLayout = layout used to create the compute pipeline. * - pipeline: VkPipeline = compute pipeline ready to bind. - * - binding_count: uint32_t = number of STORAGE_BUFFER bindings (1 + list count). + * - binding_count: uint32_t = STORAGE_BUFFER bindings ((scalars?1:0) + list count). */ struct ShaderCacheEntry { VkShaderModule shader_module = VK_NULL_HANDLE; @@ -54,10 +50,42 @@ struct ShaderCacheEntry { ShaderCacheEntry& operator=(ShaderCacheEntry&& other) noexcept; }; +/** + * Sole writer into ShaderCache (create_entry + add). + * + * Python reaches this via `_ext.gpu.register_shader`. Tests and future native + * callers use the same type - do not friend other writers or call add() directly. + */ +struct ShaderRegistry { + /** + * Build a pipeline entry from SPIR-V and insert it under `symbol`. + * + * #### Parameters: + * - context: Context& = ready GPU context + * - symbol: string = ShaderCache key (kernel name) + * - spirv: const uint32_t* = SPIR-V words + * - spirv_word_count: size_t = word count + * - binding_count: uint32_t = SSBO binding count (>= 1) + * + * #### Returns: + * - const ShaderCacheEntry& = entry stored in the cache + * + * #### Throws: + * - runtime_error = create_entry failure or duplicate symbol + */ + static const ShaderCacheEntry& register_spirv( + cthreads::gpu::Context& context, + const std::string& symbol, + const uint32_t* spirv, + size_t spirv_word_count, + uint32_t binding_count + ); +}; + /** * Process-wide map of kernel symbol -> reusable pipeline objects. * - * Writers (registry, later) call add. Everyone else only get / clear. + * Writers: ShaderRegistry only. Everyone else: get / clear. * Entries are immutable after insert; clear runs on Context shutdown. */ class ShaderCache { @@ -68,7 +96,7 @@ class ShaderCache { ShaderCache() = default; ~ShaderCache(); - // Registry-only once a friend exists. Duplicate key throws. + // ShaderRegistry-only. Duplicate key throws. const ShaderCacheEntry& add(const std::string& key, ShaderCacheEntry&& entry); public: @@ -87,8 +115,7 @@ class ShaderCache { void clear(cthreads::gpu::Context& context); friend struct cthreads::gpu::Context; // shutdown / future access - // Test-only access to private add (see gpu/testing/shader_smoke.cpp). - friend struct cthreads::gpu::testing::ShaderCacheTestAccess; + friend struct ShaderRegistry; }; } // namespace cthreads::gpu::shader diff --git a/src/cthreads/cpp/gpu/impl/compile_glsl.cpp b/src/cthreads/cpp/gpu/impl/compile_glsl.cpp new file mode 100644 index 0000000..4101abd --- /dev/null +++ b/src/cthreads/cpp/gpu/impl/compile_glsl.cpp @@ -0,0 +1,91 @@ +// Copyright (c) 2026 Tobias Karusseit +// This source code is licensed under the MIT license found in the +// LICENSE file in the root directory of this source tree. + +#include "../headers/compile_glsl.hpp" + +#include +#include +#include + +#include +#include +#include + +namespace cthreads::gpu { +namespace { + +std::once_flag g_glslang_once; + +void ensure_glslang_initialized() { + std::call_once(g_glslang_once, []() { + if (!glslang::InitializeProcess()) { + throw std::runtime_error( + "cthreads.gpu.VulkanInitFailed: glslang InitializeProcess failed"); + } + }); +} + +} // namespace + +std::vector compile_glsl_to_spirv(const std::string& source) { + if (source.empty()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: compile_glsl source is empty"); + } + + ensure_glslang_initialized(); + + const EShLanguage stage = EShLangCompute; + glslang::TShader shader(stage); + const char* strings[] = {source.c_str()}; + shader.setStrings(strings, 1); + + // Vulkan 1.0 / SPIR-V 1.0 is enough for our compute SSBOs + builtins. + shader.setEnvInput( + glslang::EShSourceGlsl, stage, glslang::EShClientVulkan, 100); + shader.setEnvClient(glslang::EShClientVulkan, glslang::EShTargetVulkan_1_0); + shader.setEnvTarget(glslang::EShTargetSpv, glslang::EShTargetSpv_1_0); + + const EShMessages messages = + static_cast(EShMsgSpvRules | EShMsgVulkanRules); + const TBuiltInResource* resources = GetDefaultResources(); + + std::string log; + if (!shader.parse(resources, 100, false, messages)) { + log = shader.getInfoLog(); + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GLSL parse failed:\n" + log); + } + + glslang::TProgram program; + program.addShader(&shader); + if (!program.link(messages)) { + log = program.getInfoLog(); + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GLSL link failed:\n" + log); + } + + const glslang::TIntermediate* intermediate = program.getIntermediate(stage); + if (intermediate == nullptr) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GLSL link produced no intermediate"); + } + + std::vector spirv; + spv::SpvBuildLogger logger; + glslang::SpvOptions options; + options.generateDebugInfo = false; + options.disableOptimizer = true; + options.optimizeSize = false; + glslang::GlslangToSpv(*intermediate, spirv, &logger, &options); + + if (spirv.empty()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GlslangToSpv produced empty SPIR-V: " + + logger.getAllMessages()); + } + return spirv; +} + +} // namespace cthreads::gpu diff --git a/src/cthreads/cpp/gpu/impl/descriptors.cpp b/src/cthreads/cpp/gpu/impl/descriptors.cpp index db43163..39f50cb 100644 --- a/src/cthreads/cpp/gpu/impl/descriptors.cpp +++ b/src/cthreads/cpp/gpu/impl/descriptors.cpp @@ -142,10 +142,14 @@ void update_descriptors( "cthreads.gpu.GpuInvalidArgument: update_descriptors binding_count " "must be >= 1"); } - if (binding_count != 1u + static_cast(pack.container_slots.size())) { + const bool has_scalars = (pack.scalar_buffer.buffer != VK_NULL_HANDLE); + const uint32_t expected = + (has_scalars ? 1u : 0u) + + static_cast(pack.container_slots.size()); + if (binding_count != expected) { throw std::runtime_error( "cthreads.gpu.GpuInvalidArgument: update_descriptors binding_count " - "must equal 1 + container_slots.size()"); + "must equal (scalars?1:0) + container_slots.size()"); } if (!context.vkUpdateDescriptorSets) { throw std::runtime_error( @@ -160,11 +164,12 @@ void update_descriptors( for (uint32_t i = 0; i < binding_count; ++i) { VkBuffer buffer = VK_NULL_HANDLE; VkDeviceSize size = 0; - if (i == 0) { + if (has_scalars && i == 0) { buffer = pack.scalar_buffer.buffer; size = pack.scalar_buffer.size; } else { - const ContainerSlot& slot = pack.container_slots[i - 1]; + const uint32_t list_i = has_scalars ? (i - 1u) : i; + const ContainerSlot& slot = pack.container_slots[list_i]; buffer = slot.buffer.buffer; size = slot.buffer.size; } diff --git a/src/cthreads/cpp/gpu/impl/module.cpp b/src/cthreads/cpp/gpu/impl/module.cpp index 93eca00..0927d93 100644 --- a/src/cthreads/cpp/gpu/impl/module.cpp +++ b/src/cthreads/cpp/gpu/impl/module.cpp @@ -219,6 +219,18 @@ void writeback_ref_lists(Context& context, SpawnedGpuKernel& job) { for (size_t j = 0; j < slot.numel; ++j) { list_val[j] = host[j]; } + } else if (slot.elem_kind == "bool") { + // GLSL bool is std430 32-bit 0/1 (same packing as scalar bool). + std::vector host(slot.numel); + pack::download_container( + context, + job.pack, + slot.container_index, + host.data(), + host.size() * sizeof(std::int32_t)); + for (size_t j = 0; j < slot.numel; ++j) { + list_val[j] = host[j] != 0; + } } else if (slot.elem_kind == "double") { std::vector host(slot.numel); pack::download_container( @@ -551,7 +563,7 @@ std::shared_ptr launch_gpu_kernel( context, job->pack, scalar_host.data(), scalar_bytes); } - // Upload each list container (binding 1..N) from ordered_values. + // Upload each list container from ordered_values (pack slot order). for (size_t c = 0; c < container_plans.size(); ++c) { const ContainerSlotPlan& plan = container_plans[c]; if (plan.numel == 0) { @@ -580,6 +592,18 @@ std::shared_ptr launch_gpu_kernel( c, host.data(), host.size() * sizeof(std::int32_t)); + } else if (plan.elem_kind == "bool") { + // std430 bool = 4 bytes; coerce Python bool to 0/1 int32. + std::vector host(plan.numel); + for (size_t j = 0; j < plan.numel; ++j) { + host[j] = list_val[j].cast() ? 1 : 0; + } + pack::upload_container( + context, + job->pack, + c, + host.data(), + host.size() * sizeof(std::int32_t)); } else if (plan.elem_kind == "double") { std::vector host(plan.numel); for (size_t j = 0; j < plan.numel; ++j) { @@ -602,13 +626,15 @@ std::shared_ptr launch_gpu_kernel( const shader::ShaderCacheEntry& entry = shader::ShaderCache::getInstance().get(symbol); - // binding_count on the entry must match 1 + number of list slots + // binding_count = (scalars ? 1 : 0) + list count (matches Python Signature) const uint32_t expected_bindings = - 1u + static_cast(container_specs.size()); + (scalar_bytes > 0 ? 1u : 0u) + + static_cast(container_specs.size()); if (entry.binding_count != expected_bindings) { throw std::runtime_error( "cthreads.gpu.GpuInvalidArgument: ShaderCacheEntry binding_count (" + - std::to_string(entry.binding_count) + ") != 1 + list count (" + + std::to_string(entry.binding_count) + + ") != (scalars?1:0) + list count (" + std::to_string(expected_bindings) + ")"); } if (entry.pipeline == VK_NULL_HANDLE || diff --git a/src/cthreads/cpp/gpu/impl/shader_cache.cpp b/src/cthreads/cpp/gpu/impl/shader_cache.cpp index 288a078..a953b61 100644 --- a/src/cthreads/cpp/gpu/impl/shader_cache.cpp +++ b/src/cthreads/cpp/gpu/impl/shader_cache.cpp @@ -1,5 +1,6 @@ #include "../headers/shader_cache.hpp" #include "../headers/context.hpp" +#include "../headers/shader.hpp" #include #include @@ -120,4 +121,15 @@ void ShaderCache::clear(Context& context) { _cache.clear(); } +const ShaderCacheEntry& ShaderRegistry::register_spirv( + Context& context, + const std::string& symbol, + const uint32_t* spirv, + size_t spirv_word_count, + uint32_t binding_count +) { + ShaderCacheEntry entry = create_entry(context, spirv, spirv_word_count, binding_count); + return ShaderCache::getInstance().add(symbol, std::move(entry)); +} + } // namespace cthreads::gpu::shader diff --git a/src/cthreads/cpp/gpu/third_party_notices/README.md b/src/cthreads/cpp/gpu/third_party_notices/README.md new file mode 100644 index 0000000..8dd0642 --- /dev/null +++ b/src/cthreads/cpp/gpu/third_party_notices/README.md @@ -0,0 +1,12 @@ +# Third-party notices for native GLSL -> SPIR-V + +When `CTHREADS_GPU=ON`, cthreads links **Khronos glslang** into `_ext` so +`compile_glsl` works without installing `glslc` or other Vulkan SDK tools. + +glslang is the same compiler engine Google **shaderc** wraps. License texts +copied here at build time (see also the root `LICENSE` third-party section): + +- `glslang-LICENSE*` — Khronos glslang (Apache-2.0 / BSD-style components) + +End-user wheels that include the GPU extension must redistribute these notices +alongside the binary. diff --git a/src/cthreads/python/.gitignore b/src/cthreads/python/.gitignore new file mode 100644 index 0000000..0bf994e --- /dev/null +++ b/src/cthreads/python/.gitignore @@ -0,0 +1,11 @@ +# >>> cthreads (auto) +__Thread__/ +__Threadable__/ +__Gpu__/ +.cthreads_cache.json +cthreads_kernels.dll +cthreads_kernels.so +cthreads_kernels.lib +libcthreads_kernels.so +libcthreads_kernels.dylib +# <<< cthreads (auto) diff --git a/src/cthreads/python/cthreads/cache.py b/src/cthreads/python/cthreads/cache.py index fce73d4..55d6324 100644 --- a/src/cthreads/python/cthreads/cache.py +++ b/src/cthreads/python/cthreads/cache.py @@ -40,6 +40,7 @@ def source_fingerprint(*objs: Any) -> str: _GITIGNORE_PATTERNS = ( "__Thread__/", "__Threadable__/", + "__Gpu__/", ".cthreads_cache.json", "cthreads_kernels.dll", "cthreads_kernels.so", @@ -66,18 +67,30 @@ def cache_path_for_root(root: Path) -> Path: return root / CACHE_FILENAME +def _empty_cache(version: str) -> dict[str, Any]: + """Fresh cache document: CPU units + GPU units share one file, separate maps.""" + return { + "version": version, + "units": {}, + "gpu_units": {}, + "link_hash": None, + "binary": None, + } + + def load_cache(root: Path) -> dict[str, Any]: path = cache_path_for_root(root) version = REGISTRY.VERSION if not path.is_file(): - return {"version": version, "units": {}, "link_hash": None, "binary": None} + return _empty_cache(version) try: data = json.loads(path.read_text(encoding="utf-8")) except (OSError, json.JSONDecodeError): - return {"version": version, "units": {}, "link_hash": None, "binary": None} + return _empty_cache(version) if data.get("version") != version: - return {"version": version, "units": {}, "link_hash": None, "binary": None} + return _empty_cache(version) data.setdefault("units", {}) + data.setdefault("gpu_units", {}) return data diff --git a/src/cthreads/python/cthreads/frontend/Gpu/Wrapper.py b/src/cthreads/python/cthreads/frontend/Gpu/Wrapper.py index 10e9542..1d0d064 100644 --- a/src/cthreads/python/cthreads/frontend/Gpu/Wrapper.py +++ b/src/cthreads/python/cthreads/frontend/Gpu/Wrapper.py @@ -1,3 +1,3 @@ +from ...gpu.frontend.wrapper import Gpu - -def Gpu(fn, device=None): pass \ No newline at end of file +__all__ = ["Gpu"] \ No newline at end of file diff --git a/src/cthreads/python/cthreads/frontend/Registry/registry.py b/src/cthreads/python/cthreads/frontend/Registry/registry.py index 62f26c3..37c5183 100644 --- a/src/cthreads/python/cthreads/frontend/Registry/registry.py +++ b/src/cthreads/python/cthreads/frontend/Registry/registry.py @@ -10,7 +10,7 @@ if TYPE_CHECKING: from ...compiler.orchestrator.units import ThreadableUnit, ThreadUnit - + from ...gpu.compiler.orchestrator.gpu_unit import GpuUnit class Registry: """ @@ -29,6 +29,10 @@ def __init__(self) -> None: self.threadable_units: dict[str, ThreadableUnit] = {} self.thread_units: dict[str, ThreadUnit] = {} + # gpu stuff + self.gpu_functions: dict[str, object] = {} # function qualname -> function (@Gpu) + self.gpu_function_units: dict[str, GpuUnit] = {} # function qualname -> GpuUnit + def register_threadable(self, cls: type) -> None: """Registers a threadable class for code generation""" self.threadables[cls.__name__] = cls @@ -37,15 +41,23 @@ def register_thread(self, fn: object) -> None: """Registers a thread function for code generation""" self.threads[fn.__qualname__] = fn + def register_gpu_function(self, fn: object) -> None: + """Registers a gpu function for code generation""" + self.gpu_functions[fn.__qualname__] = fn + def clear(self) -> None: """Clears the registry""" self.threadables.clear() self.threads.clear() self.threadable_units.clear() self.thread_units.clear() + self.gpu_functions.clear() + self.gpu_function_units.clear() from ...kernel_meta import KERNELS + from ...gpu.gpu_kernel_meta import GPU_KERNELS KERNELS.clear() + GPU_KERNELS.clear() REGISTRY = Registry() diff --git a/src/cthreads/python/cthreads/gpu/__init__.py b/src/cthreads/python/cthreads/gpu/__init__.py index 108b975..2455106 100644 --- a/src/cthreads/python/cthreads/gpu/__init__.py +++ b/src/cthreads/python/cthreads/gpu/__init__.py @@ -1,86 +1,53 @@ -"""Vulkan GPU runtime probe API (Issue 1). - -Soft-imports ``cthreads._ext.gpu`` so CPU-only builds still import cleanly. """ +Vulkan GPU runtime probe API. + +Native access goes through `_ext_gpu_api` (`cthreads._ext.gpu`). +Public helpers and error types are re-exported from `frontend`. -from __future__ import annotations +Launch helpers (`prepare` / `gpu` / `compile`) live in `runtime` so the +callable name `prepare` does not shadow a submodule. +""" -from .errors import ( +from . import _ext_gpu_api +from .frontend import ( + BlockDim, + BlockIdx, CThreadsGPUError, GPUNotAvailable, + GlobalIdx, + GridDim, + Gpu, GpuInvalidArgument, GpuUseAfterDestroy, + ThreadIdx, VulkanInitFailed, VulkanLoaderNotFound, VulkanNoDevice, VulkanNotBuiltError, VulkanOutOfMemory, + _map_error, + available, + device_name, + init, + shutdown, ) - -try: - from cthreads._ext import gpu as _gpu -except ImportError: - _gpu = None - - -def _map_error(exc: BaseException) -> CThreadsGPUError: - msg = str(exc) - if "VulkanLoaderNotFound" in msg: - return VulkanLoaderNotFound(msg) - if "VulkanNoDevice" in msg: - return VulkanNoDevice(msg) - if "VulkanOutOfMemory" in msg: - return VulkanOutOfMemory(msg) - if "GpuUseAfterDestroy" in msg: - return GpuUseAfterDestroy(msg) - if "GpuInvalidArgument" in msg: - return GpuInvalidArgument(msg) - if "VulkanNotBuilt" in msg: - return VulkanNotBuiltError(msg) - if "VulkanInitFailed" in msg: - return VulkanInitFailed(msg) - return VulkanInitFailed(msg) - - -def available() -> bool: - """Return True if Vulkan loader + compute device can be initialized.""" - if _gpu is None: - return False - return bool(_gpu.available()) - - -def device_name() -> str: - """Return the active GPU name (calls init). Raises on failure / not built.""" - if _gpu is None: - raise VulkanNotBuiltError( - "cthreads built without CTHREADS_GPU; rebuild with -DCTHREADS_GPU=ON" - ) - try: - return str(_gpu.device_name()) - except Exception as exc: - raise _map_error(exc) from exc - - -def init() -> None: - """Explicitly initialize the Vulkan context.""" - if _gpu is None: - raise VulkanNotBuiltError( - "cthreads built without CTHREADS_GPU; rebuild with -DCTHREADS_GPU=ON" - ) - try: - _gpu.init() - except Exception as exc: - raise _map_error(exc) from exc +from .runtime import GpuJob, compile, gpu, prepare -def shutdown() -> None: - """Destroy device/instance and unload the Vulkan loader (no-op if not built).""" - if _gpu is None: - return - _gpu.shutdown() +def __getattr__(name: str): + if name == "_gpu": + return _ext_gpu_api._gpu + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") __all__ = [ + "Gpu", + "GpuJob", + "BlockDim", + "BlockIdx", + "GlobalIdx", + "GridDim", + "ThreadIdx", "CThreadsGPUError", "GPUNotAvailable", "GpuInvalidArgument", @@ -90,8 +57,14 @@ def shutdown() -> None: "VulkanNoDevice", "VulkanNotBuiltError", "VulkanOutOfMemory", + "_map_error", + "_gpu", + "_ext_gpu_api", "available", + "compile", "device_name", + "gpu", "init", + "prepare", "shutdown", ] diff --git a/src/cthreads/python/cthreads/gpu/_ext_gpu_api.py b/src/cthreads/python/cthreads/gpu/_ext_gpu_api.py new file mode 100644 index 0000000..eebc9f0 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/_ext_gpu_api.py @@ -0,0 +1,173 @@ +""" +Lazy binding to `cthreads._ext.gpu` (native Vulkan submodule). + +Central entry for all Python code that talks to the C++ GPU package. +Soft-imports so CPU-only builds still import `cthreads.gpu` cleanly. +""" + +from __future__ import annotations + +from typing import Any + +try: + from cthreads._ext import gpu as _gpu +except ImportError: + _gpu = None # type: ignore[assignment] + + +def _require_ext_gpu(): + """ + Return the native `_ext.gpu` module. + + #### Returns + - module = pybind `cthreads._ext.gpu` submodule + + #### Raises + - RuntimeError = extension was built without CTHREADS_GPU + """ + if _gpu is None: + raise RuntimeError( + "cthreads._ext.gpu is not available - rebuild with -DCTHREADS_GPU=ON" + ) + return _gpu + + +def available() -> bool: + """ + Report whether the Vulkan loader and a compute device can initialize. + + Never raises. Returns False when the GPU extension is missing or init + would fail. + + #### Returns + - bool = True when a compute-capable GPU context can be created + """ + if _gpu is None: + return False + return bool(_gpu.available()) + + +def device_name() -> str: + """ + Return the active GPU device name (may call init first). + + #### Returns + - str = Vulkan device name string + + #### Raises + - RuntimeError = `_ext.gpu` is not built + - Exception = native init / device query failures (unmapped) + """ + return str(_require_ext_gpu().device_name()) + + +def init() -> None: + """ + Explicitly initialize the process-wide Vulkan context. + + #### Returns + - None + + #### Raises + - RuntimeError = `_ext.gpu` is not built + - Exception = native init failures (unmapped) + """ + _require_ext_gpu().init() + + +def shutdown() -> None: + """ + Destroy the Vulkan device/instance and unload the loader. + + No-op when the GPU extension is not built. + + #### Returns + - None + """ + if _gpu is None: + return + _gpu.shutdown() + + +def testing() -> Any | None: + """ + Return the test-only `_ext.gpu.testing` submodule when present. + + #### Returns + - module | None = testing helpers, or None if missing / not built + """ + if _gpu is None: + return None + return getattr(_gpu, "testing", None) + + +def launch_gpu_kernel(meta: dict[str, Any], ordered_values: list[Any]) -> Any: + """ + Submit one GPU kernel from metadata and ordered Python arguments. + + Returns a native GpuJob handle. Does not wait; the caller joins that + handle for fence wait and list writeback. The kernel `symbol` must already + be registered in the ShaderCache. + + #### Args: + - meta: dict[str, Any] = kernel metadata (`symbol`, `params`, layout fields) + - ordered_values: list[Any] = arguments in parameter order matching `meta` + + #### Returns + - Any = native `_ext.gpu.GpuJob` (SpawnedGpuKernel) + + #### Raises + - RuntimeError = `_ext.gpu` is not built + - Exception = native launch failures (unmapped) + + #### Technical terms: + - GpuJob: per-launch GPU job handle (fence, pack, writeback plan) + - writeback: download of ref list buffers into the same Python list objects + """ + return _require_ext_gpu().launch_gpu_kernel(meta, ordered_values) + + +def compile_glsl(source: str) -> bytes: + """ + Compile GLSL compute source to SPIR-V via vendored glslang in `_ext`. + + #### Args: + - source: str = full compute shader text + + #### Returns + - bytes = SPIR-V binary + + #### Raises + - RuntimeError = `_ext.gpu` is not built or lacks compile_glsl + - Exception = native compile failures (unmapped) + """ + gpu = _require_ext_gpu() + compile_native = getattr(gpu, "compile_glsl", None) + if not callable(compile_native): + raise RuntimeError( + "cthreads._ext.gpu.compile_glsl is missing - rebuild with " + "CTHREADS_GPU=ON (glslang vendored into _ext)" + ) + return bytes(compile_native(source)) + + +def register_shader(symbol: str, spirv: bytes, binding_count: int) -> None: + """ + Insert a compute pipeline into the process ShaderCache from SPIR-V bytes. + + Calls native `ShaderRegistry::register_spirv` (sole writer). Duplicate + symbols raise. + + #### Args: + - symbol: str = ShaderCache key / kernel name + - spirv: bytes = SPIR-V binary (multiple of 4 bytes) + - binding_count: int = number of STORAGE_BUFFER bindings (>= 1) + + #### Returns + - None + + #### Raises + - RuntimeError = `_ext.gpu` is not built + - Exception = native create/insert failures (unmapped) + """ + _require_ext_gpu().register_shader(symbol, spirv, binding_count) diff --git a/src/cthreads/python/cthreads/gpu/compiler/orchestrator/__init__.py b/src/cthreads/python/cthreads/gpu/compiler/orchestrator/__init__.py new file mode 100644 index 0000000..53a1657 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/orchestrator/__init__.py @@ -0,0 +1,7 @@ +from .gpu_unit import GpuUnit +from .gpu_compile_session import GpuCompileSession + +__all__ = [ + "GpuUnit", + "GpuCompileSession", +] diff --git a/src/cthreads/python/cthreads/gpu/compiler/orchestrator/gpu_compile_session.py b/src/cthreads/python/cthreads/gpu/compiler/orchestrator/gpu_compile_session.py new file mode 100644 index 0000000..e0ba5f0 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/orchestrator/gpu_compile_session.py @@ -0,0 +1,93 @@ +""" +Drain `REGISTRY.gpu_functions` into `GpuUnit`s and run emit. + +Does not keep its own unit maps; units live on REGISTRY. +""" + +from __future__ import annotations + +import inspect +from pathlib import Path +from typing import Any, get_type_hints + +from ....cache import ensure_gitignore, load_cache, save_cache +from ....compiler.orchestrator.units.handle import Handle +from ....frontend.Registry import REGISTRY +from ....gpu.gpu_kernel_meta import GPU_KERNELS, build_gpu_kernel_meta +from ....types import PyType, hint_to_pytype +from .gpu_unit import GpuUnit + + +class GpuCompileSession: + """ + Stateless GPU compile pass: registry functions -> units -> emit -> cache. + """ + + @staticmethod + def compile(force: bool = False) -> dict[str, Any]: + """ + Build GpuUnits for every registered `@Gpu` function and emit. + + #### Args: + - force: bool = passed through to `GpuUnit.emit` (default False) + + #### Returns + - dict[str, Any] = `root`, `cache`, `rewritten` (unit names that rewrote) + + #### Raises + - RuntimeError = nothing registered + - TypeError = bad annotations / non-None return + """ + REGISTRY.gpu_function_units.clear() + GPU_KERNELS.clear() + + if not REGISTRY.gpu_functions: + raise RuntimeError("Nothing registered to compile") + + sample: Any = next(iter(REGISTRY.gpu_functions.values())) + root = Path(inspect.getfile(sample)).resolve().parent + ensure_gitignore(root) + cache = load_cache(root) + rewritten: list[str] = [] + + for qualname, fn in list(REGISTRY.gpu_functions.items()): + src_file = Path(inspect.getfile(fn)).resolve() + hints = get_type_hints(fn) + params: list[tuple[str, PyType]] = [] + for pname in inspect.signature(fn).parameters: + if pname not in hints: + raise TypeError( + f"GPU function {qualname}: " + f"parameter {pname!r} needs a type annotation" + ) + params.append((pname, hint_to_pytype(hints[pname]))) + + ret_hint = hints.get("return") + if ret_hint not in (None, type(None)): + raise TypeError( + f"GPU function {qualname}: " + f"return type {ret_hint!r} is not allowed; use -> None " + f"with in-place list updates" + ) + return_type = None + + # Rebuild meta after GPU_KERNELS.clear (decorator may have built it earlier). + build_gpu_kernel_meta(fn) + + gpu_unit = GpuUnit( + handle=Handle( + name=qualname, + path=str(src_file), + target=fn, + ), + params=params, + return_type=return_type, + ) + REGISTRY.gpu_function_units[qualname] = gpu_unit + + for unit in REGISTRY.gpu_function_units.values(): + if unit.emit(force=force, cache=cache): + rewritten.append(unit.handle.name) + + save_cache(root, cache) + return {"root": root, "cache": cache, "rewritten": rewritten} diff --git a/src/cthreads/python/cthreads/gpu/compiler/orchestrator/gpu_unit.py b/src/cthreads/python/cthreads/gpu/compiler/orchestrator/gpu_unit.py new file mode 100644 index 0000000..198f96f --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/orchestrator/gpu_unit.py @@ -0,0 +1,119 @@ +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from ....cache import source_fingerprint, write_if_changed +from ....compiler.orchestrator.units.baseUnit import BaseUnit +from ....types.pyType import PyType +from ... import _ext_gpu_api +from ...gpu_kernel_meta import build_gpu_kernel_meta +from ..translation.translate import translate_function_for_gpu + + +@dataclass +class GpuUnit(BaseUnit): + """ + One `@Gpu` free function: translate GLSL, compile SPIR-V, register shader. + """ + + params: list[tuple[str, PyType]] + return_type: PyType | None + + def validate(self) -> None: + """ + Ensure the handle target is a `@Gpu` function. + + #### Raises + - TypeError = target is not marked `@Gpu` + """ + fn = self.handle.target + if not getattr(fn, "__gpu__", False): + raise TypeError(f"{self.handle.name} is not a @Gpu function") + + def emit(self, *, force: bool = False, cache: dict[str, Any] | None = None) -> bool: + """ + Translate, compile with shaderc, write artifacts, register SPIR-V. + + Always registers into the process ShaderCache (Vulkan state is not on + disk). The source fingerprint only skips rewriting `__Gpu__` files. + + #### Args: + - force: bool = rewrite `__Gpu__` artifacts even if the hash matches + - cache: dict[str, Any] | None = shared `.cthreads_cache.json` document + + #### Returns + - bool = True if `__Gpu__` files were rewritten + + #### Raises + - RuntimeError = GLSL/SPIR-V compile failed or GPU ext missing + """ + self.validate() + fn = self.handle.target + meta = getattr(fn, "__gpu_kernel_meta__", None) + if not isinstance(meta, dict): + meta = build_gpu_kernel_meta(fn).to_dict() + + src_hash: str = source_fingerprint(fn) + local_size_x: int = int(meta.get("local_size_x", 64)) + result = translate_function_for_gpu( + fn, local_size_x=local_size_x, compile_spirv=True + ) + if result.spirv is None: + raise RuntimeError( + f"GPU unit {self.handle.name}: SPIR-V compile produced no bytes" + ) + + src_file: Path = Path(self.handle.path).resolve() + out_dir: Path = src_file.parent / "__Gpu__" + out_dir.mkdir(parents=True, exist_ok=True) + comp_path: Path = out_dir / f"{result.func_name}.comp" + spv_path: Path = out_dir / f"{result.func_name}.spv" + + rewritten: bool = False + gpu_units: dict[str, Any] = {} + if cache is not None: + gpu_units = cache.setdefault("gpu_units", {}) + prev = gpu_units.get(self.handle.name) + hash_ok = ( + isinstance(prev, dict) + and prev.get("hash") == src_hash + and not force + ) + else: + hash_ok = False + + if not hash_ok: + rewritten = write_if_changed(comp_path, result.source) + prev_spv: bytes | None = None + if spv_path.is_file(): + try: + prev_spv = spv_path.read_bytes() + except OSError: + prev_spv = None + if prev_spv != result.spirv: + spv_path.write_bytes(result.spirv) + rewritten = True + + symbol: str = str(meta.get("symbol", result.func_name)) + binding_count: int = int(meta.get("binding_count", result.binding_count)) + # Always populate process ShaderCache (also after shutdown released it). + # Disk fingerprint only gates `__Gpu__` file rewrites above. + try: + _ext_gpu_api.register_shader(symbol, result.spirv, binding_count) + except Exception as exc: + msg: str = str(exc) + if "already exists" not in msg: + raise + + if cache is not None: + gpu_units[self.handle.name] = { + "hash": src_hash, + "symbol": symbol, + "binding_count": binding_count, + "registered": True, + "comp": str(comp_path), + "spv": str(spv_path), + } + return rewritten diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/Glsl.py b/src/cthreads/python/cthreads/gpu/compiler/translation/Glsl.py new file mode 100644 index 0000000..3e3d8cd --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/Glsl.py @@ -0,0 +1,86 @@ +""" +Map cthreads PyTypes to GLSL type names and std430 sizes for GPU Signature. + +Python float -> GLSL float (f32). CPU @Thread uses double; GPU v1 does not. +Lists are not a GLSL type string — use type_name on the element type. +""" + +from __future__ import annotations + +from ....types import ( + PyBool, + PyDict, + PyFloat, + PyInt, + PyList, + PyString, + PyThreadable, + PyType, + is_shared_pytype, + is_sync_pytype, + is_tbuffer_pytype, +) + +# GLSL / host pack sizes (must match gpu_kernel_meta._GPU_SCALAR_BYTES). +_SCALAR: dict[type, tuple[str, int]] = { + PyInt: ("int", 4), + PyFloat: ("float", 4), + PyBool: ("bool", 4), +} + + +def _scalar_entry(py_type: PyType) -> tuple[str, int]: + for cls, entry in _SCALAR.items(): + if isinstance(py_type, cls): + return entry + raise TypeError( + f"GPU GLSL: unsupported scalar type {py_type.name!r} " + f"(allowed: int, float, bool)" + ) + + +def type_name(py_type: PyType) -> str: + """ + GLSL type name for a scalar PyType (`int` / `float` / `bool`). + + For lists, pass `py_type.inner_type` (or use `elem_type_name`). + """ + if isinstance(py_type, PyList): + raise TypeError( + "GPU GLSL: list is not a scalar type name; " + "use elem_type_name(py_type) or type_name(py_type.inner_type)" + ) + if ( + is_sync_pytype(py_type) + or is_shared_pytype(py_type) + or is_tbuffer_pytype(py_type) + or isinstance(py_type, (PyDict, PyString, PyThreadable)) + ): + raise TypeError( + f"GPU GLSL: {py_type.name!r} is not supported on the GPU path" + ) + return _scalar_entry(py_type)[0] + + +def elem_type_name(py_type: PyList) -> str: + """GLSL element type for a `list[...]` PyType (e.g. `float` for list[float]).""" + if not isinstance(py_type, PyList): + raise TypeError( + f"GPU GLSL: elem_type_name expects PyList, got {type(py_type)!r}" + ) + return type_name(py_type.inner_type) + + +def size_bytes(py_type: PyType) -> int: + """std430 size in bytes for a scalar (lists: use size_bytes on the element).""" + if isinstance(py_type, PyList): + raise TypeError( + "GPU GLSL: list has no single scalar size; " + "use size_bytes(py_type.inner_type) for the element" + ) + return _scalar_entry(py_type)[1] + + +def align_bytes(py_type: PyType) -> int: + """std430 alignment for a v1 scalar (same as size for int/float/bool).""" + return size_bytes(py_type) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/Signature.py b/src/cthreads/python/cthreads/gpu/compiler/translation/Signature.py new file mode 100644 index 0000000..5ce7e65 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/Signature.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import ast +from typing import get_type_hints + +from ....types import PyList, hint_to_pytype +from .Glsl import align_bytes, elem_type_name, size_bytes, type_name +from .context import GpuTranslationContext +from .result import GpuSignatureResult, ListField, ScalarField + + +class GpuSignature: + """Emit GLSL buffer preamble; fill ctx.symbols and binding_names.""" + + @staticmethod + def translate( + func_def: ast.FunctionDef, ctx: GpuTranslationContext + ) -> GpuSignatureResult: + # No Threadable owners yet — no localns. + hints = get_type_hints(ctx.fn) + args = list(func_def.args.args) + + if func_def.args.vararg or func_def.args.kwarg or func_def.args.kwonlyargs: + raise TypeError( + f"GPU function {ctx.func_name}: " + "*args/**kwargs/kw-only args are not supported" + ) + + ret_hint = hints.get("return") + if ret_hint is not None and ret_hint is not type(None): + raise TypeError( + f"GPU function {ctx.func_name}: return must be None " + f"(in-place list writeback only), got {ret_hint!r}" + ) + + scalar_fields: list[ScalarField] = [] + # (name, glsl_elem) — binding numbers assigned after we know if scalars exist. + pending_lists: list[tuple[str, str]] = [] + scalar_bytes = 0 + + for arg in args: + if arg.arg not in hints: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"parameter {arg.arg!r} needs a type annotation" + ) + py_type = hint_to_pytype(hints[arg.arg]) + ctx.symbols[arg.arg] = py_type + + if isinstance(py_type, PyList): + glsl_elem = elem_type_name(py_type) + pending_lists.append((arg.arg, glsl_elem)) + ctx.list_params.add(arg.arg) + else: + glsl_ty = type_name(py_type) + align = align_bytes(py_type) + scalar_bytes = (scalar_bytes + align - 1) & ~(align - 1) + scalar_bytes += size_bytes(py_type) + scalar_fields.append((arg.arg, glsl_ty)) + ctx.scalar_params.add(arg.arg) + + # Scalars at binding 0 when present; lists start at 1 or 0 accordingly. + list_base = 1 if scalar_fields else 0 + list_fields: list[ListField] = [] + for i, (name, glsl_elem) in enumerate(pending_lists): + binding = list_base + i + list_fields.append((binding, name, glsl_elem)) + ctx.binding_names.append((binding, name, "list")) + + if scalar_fields: + ctx.binding_names.insert(0, (0, "scalars", "scalar")) + + binding_count = (1 if scalar_fields else 0) + len(list_fields) + + lines: list[str] = [f"layout(local_size_x = {ctx.local_size_x}) in;", ""] + if scalar_fields: + lines.append("layout(set = 0, binding = 0, std430) buffer Scalars {") + for name, glsl_ty in scalar_fields: + lines.append(f" {glsl_ty} {name};") + lines.append("} scalars;") + lines.append("") + for binding, name, glsl_elem in list_fields: + block = name[:1].upper() + name[1:] + lines.append( + f"layout(set = 0, binding = {binding}, std430) buffer {block} {{" + ) + lines.append(f" {glsl_elem} data[];") + lines.append(f"}} {name};") + lines.append("") + + preamble = "\n".join(lines).rstrip() + "\n" + + return GpuSignatureResult( + func_name=func_def.name, + preamble=preamble, + binding_count=binding_count, + scalar_bytes=scalar_bytes, + local_size_x=ctx.local_size_x, + scalar_fields=scalar_fields, + list_fields=list_fields, + ) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/Typeof.py b/src/cthreads/python/cthreads/gpu/compiler/translation/Typeof.py new file mode 100644 index 0000000..e69de29 diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/assemble.py b/src/cthreads/python/cthreads/gpu/compiler/translation/assemble.py new file mode 100644 index 0000000..6e6dc63 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/assemble.py @@ -0,0 +1,40 @@ +""" +Assemble a full GLSL compute shader (`.comp`) from preamble + body. +""" + +from __future__ import annotations + +_DEFAULT_GLSL_VERSION: int = 450 + + +def assemble_comp( + preamble: str, + body: str, + *, + version: int = _DEFAULT_GLSL_VERSION, +) -> str: + """ + Build a complete compute shader string: `#version`, buffers, `void main()`. + + Body lines from GpuSyntax are already indented; they become the main body. + + #### Args: + - preamble: str = Signature output (local_size + buffer layouts) + - body: str = lowered statement lines (already indented) + - version: int = GLSL version directive (default 450) + + #### Returns + - str = full `.comp` source text + """ + pre: str = preamble.rstrip() + bod: str = body.rstrip() + lines: list[str] = [f"#version {version}", ""] + if pre: + lines.append(pre) + lines.append("") + lines.append("void main() {") + if bod: + lines.append(bod) + lines.append("}") + lines.append("") + return "\n".join(lines) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/context.py b/src/cthreads/python/cthreads/gpu/compiler/translation/context.py new file mode 100644 index 0000000..94a9663 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/context.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable + +from ....types import PyType + +# (binding index, param name, kind) — kind is "scalar" or "list" +BindingName = tuple[int, str, str] + + +@dataclass +class GpuTranslationContext: + """Mutable codegen state for one @Gpu function.""" + + fn: Callable + local_size_x: int = 64 + symbols: dict[str, PyType] = field(default_factory=dict) + binding_names: list[BindingName] = field(default_factory=list) + # Param names that live in the binding-0 Scalars SSBO (not bare GLSL ids). + scalar_params: set[str] = field(default_factory=set) + # Param names that are list SSBO instances (x -> x.data[i] via Index). + list_params: set[str] = field(default_factory=set) + + @property + def func_name(self) -> str: + return self.fn.__name__ diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/__init__.py b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/__init__.py new file mode 100644 index 0000000..f8b218e --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/__init__.py @@ -0,0 +1,107 @@ +""" +Ordered GPU plugin lists. GpuSyntax.expr tries these for Call / Attribute. +""" + +from __future__ import annotations + +import ast + +from ..context import GpuTranslationContext +from .base import AttrPlugin, CallPlugin, TranslateExpr + +CALL_PLUGINS: list[CallPlugin] = [] +ATTR_PLUGINS: list[AttrPlugin] = [] + + +def register_call(plugin: CallPlugin) -> CallPlugin: + """ + Append a CallPlugin to the GPU call registry. + + #### Args: + - plugin: CallPlugin = plugin instance to register + + #### Returns + - CallPlugin = the same plugin (for chaining) + """ + CALL_PLUGINS.append(plugin) + return plugin + + +def register_attr(plugin: AttrPlugin) -> AttrPlugin: + """ + Append an AttrPlugin to the GPU attribute registry. + + #### Args: + - plugin: AttrPlugin = plugin instance to register + + #### Returns + - AttrPlugin = the same plugin (for chaining) + """ + ATTR_PLUGINS.append(plugin) + return plugin + + +def lower_call( + node: ast.Call, + ctx: GpuTranslationContext, + translate_expr: TranslateExpr, +) -> str | None: + """ + Try each CallPlugin until one returns GLSL text. + + #### Args: + - node: ast.Call = call expression + - ctx: GpuTranslationContext = current GPU translation state + - translate_expr: TranslateExpr = nested expression lowerer + + #### Returns + - str | None = GLSL text, or None if no plugin matched + """ + for plugin in CALL_PLUGINS: + out = plugin.try_lower(node, ctx, translate_expr) + if out is not None: + return out + return None + + +def lower_attr( + node: ast.Attribute, + ctx: GpuTranslationContext, + translate_expr: TranslateExpr, +) -> str | None: + """ + Try each AttrPlugin until one returns GLSL text. + + #### Args: + - node: ast.Attribute = attribute expression + - ctx: GpuTranslationContext = current GPU translation state + - translate_expr: TranslateExpr = nested expression lowerer + + #### Returns + - str | None = GLSL text, or None if no plugin matched + """ + for plugin in ATTR_PLUGINS: + out = plugin.try_lower(node, ctx, translate_expr) + if out is not None: + return out + return None + + +__all__ = [ + "CALL_PLUGINS", + "ATTR_PLUGINS", + "AttrPlugin", + "CallPlugin", + "TranslateExpr", + "register_call", + "register_attr", + "lower_call", + "lower_attr", +] + +# Side-effect: register concrete plugins. +from .indexes import IndexAttrPlugin # noqa: E402 +from .math_calls import MathCallPlugin # noqa: E402 + +register_attr(IndexAttrPlugin()) +register_call(MathCallPlugin()) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/base.py b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/base.py new file mode 100644 index 0000000..745bc9b --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/base.py @@ -0,0 +1,64 @@ +""" +Plugin bases for GPU Call / Attribute lowering. +""" + +from __future__ import annotations + +import ast +from abc import ABC, abstractmethod +from collections.abc import Callable + +from ..context import GpuTranslationContext + +# Already-translated subexpressions; typically GpuSyntax.expr +TranslateExpr = Callable[[ast.expr, GpuTranslationContext], str] + + +class AttrPlugin(ABC): + """ + Handles `ast.Attribute` (index builtins, props - not calls). + """ + + @abstractmethod + def try_lower( + self, + node: ast.Attribute, + ctx: GpuTranslationContext, + translate_expr: TranslateExpr, + ) -> str | None: + """ + Return GLSL text if this plugin handles `node`, else None. + + #### Args: + - node: ast.Attribute = attribute expression + - ctx: GpuTranslationContext = current GPU translation state + - translate_expr: TranslateExpr = nested expression lowerer + + #### Returns + - str | None = GLSL expression, or None to try the next plugin + """ + + +class CallPlugin(ABC): + """ + Handles `ast.Call` (math builtins later). + """ + + @abstractmethod + def try_lower( + self, + node: ast.Call, + ctx: GpuTranslationContext, + translate_expr: TranslateExpr, + ) -> str | None: + """ + Return GLSL text if this plugin handles `node`, else None. + + #### Args: + - node: ast.Call = call expression + - ctx: GpuTranslationContext = current GPU translation state + - translate_expr: TranslateExpr = nested expression lowerer + + #### Returns + - str | None = GLSL expression, or None to try the next plugin + """ diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/indexes.py b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/indexes.py new file mode 100644 index 0000000..9c3c87b --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/indexes.py @@ -0,0 +1,61 @@ +""" +AttrPlugin: ThreadIdx / BlockIdx / GlobalIdx / ... .x|.y|.z -> GLSL builtins. +""" +import ast +from typing import Any + +from ....frontend.indexes import GpuIndexBuiltin +from ..context import GpuTranslationContext +from .base import AttrPlugin, TranslateExpr + +_AXES: frozenset[str] = frozenset({"x", "y", "z"}) + + +def _globals(ctx: GpuTranslationContext) -> dict: + g = getattr(ctx.fn, "__globals__", None) + return g if isinstance(g, dict) else {} + + +def _resolve_value_obj(node: ast.expr, globals_ns: dict) -> Any: + """ + Resolve a simple Name or module.Name expression to a Python object. + """ + if isinstance(node, ast.Name): + return globals_ns.get(node.id) + if isinstance(node, ast.Attribute) and isinstance(node.value, ast.Name): + mod = globals_ns.get(node.value.id) + if mod is None: + return None + return getattr(mod, node.attr, None) + return None + + +def _glsl_base_for(obj: Any) -> str | None: + if obj is None: + return None + cls = obj if isinstance(obj, type) else type(obj) + if not isinstance(cls, type) or not issubclass(cls, GpuIndexBuiltin): + return None + base: str = getattr(cls, "_glsl_base", "") + return base or None + + +class IndexAttrPlugin(AttrPlugin): + """ + Lower `GlobalIdx.x` / `gpu.ThreadIdx.y` to `gl_*Invocation*.*`. + """ + + def try_lower( + self, + node: ast.Attribute, + ctx: GpuTranslationContext, + translate_expr: TranslateExpr, + ) -> str | None: + if node.attr not in _AXES: + return None + obj = _resolve_value_obj(node.value, _globals(ctx)) + base = _glsl_base_for(obj) + if base is None: + return None + # GLSL invocation IDs are uint; cast to int so `i: int = GlobalIdx.x` typechecks. + return f"int({base}.{node.attr})" diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/math_calls.py b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/math_calls.py new file mode 100644 index 0000000..979cdd8 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/math_calls.py @@ -0,0 +1,42 @@ +""" +Minimal GLSL math CallPlugins for @Gpu (sqrt first — needed for SPH forces). +""" + +from __future__ import annotations + +import ast + +from ..context import GpuTranslationContext +from .base import CallPlugin, TranslateExpr + + +class MathCallPlugin(CallPlugin): + """ + Lower `sqrt(x)` / `math.sqrt(x)` to GLSL `sqrt(...)`. + """ + + def try_lower( + self, + node: ast.Call, + ctx: GpuTranslationContext, + translate_expr: TranslateExpr, + ) -> str | None: + if node.keywords: + return None + if len(node.args) != 1: + return None + fn = node.func + is_sqrt = False + if isinstance(fn, ast.Name) and fn.id == "sqrt": + is_sqrt = True + elif ( + isinstance(fn, ast.Attribute) + and fn.attr == "sqrt" + and isinstance(fn.value, ast.Name) + and fn.value.id == "math" + ): + is_sqrt = True + if not is_sqrt: + return None + arg = translate_expr(node.args[0], ctx) + return f"sqrt({arg})" diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/result.py b/src/cthreads/python/cthreads/gpu/compiler/translation/result.py new file mode 100644 index 0000000..5a2dda1 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/result.py @@ -0,0 +1,42 @@ +from __future__ import annotations + +from dataclasses import dataclass + +# (param_name, glsl_type) e.g. ("n", "int") +ScalarField = tuple[str, str] +# (binding, param_name, glsl_elem_type) e.g. (1, "x", "float") +ListField = tuple[int, str, str] + + +@dataclass +class GpuSignatureResult: + """GLSL buffer preamble + layout facts for one @Gpu Signature.""" + + func_name: str + preamble: str + binding_count: int + scalar_bytes: int + local_size_x: int + scalar_fields: list[ScalarField] + list_fields: list[ListField] + + +@dataclass +class GpuTranslationResult: + """ + Emit contract for one translated `@Gpu` function. + + `source` is the full `.comp` text (`#version` + preamble + `main`). + `spirv` is set when compile_spirv was requested and compilation succeeded. + """ + + func_name: str + preamble: str + body: str + source: str + binding_count: int + scalar_bytes: int + local_size_x: int + scalar_fields: list[ScalarField] + list_fields: list[ListField] + spirv: bytes | None = None diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/spirv.py b/src/cthreads/python/cthreads/gpu/compiler/translation/spirv.py new file mode 100644 index 0000000..9f9d32d --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/spirv.py @@ -0,0 +1,133 @@ +""" +Compile GLSL compute source to SPIR-V. + +Prefers in-process `_ext.gpu.compile_glsl` (vendored Khronos glslang linked +into the GPU extension — no extra user tools). Falls back to the Vulkan SDK +`glslc` CLI when the native binding is unavailable (CPU-only builds / dev). +""" + +from __future__ import annotations + +import os +import shutil +import subprocess +import tempfile +from pathlib import Path + + +def _find_glslc() -> str | None: + """ + Locate the shaderc `glslc` executable (dev fallback only). + + #### Returns + - str | None = path to glslc, or None if not found + """ + found: str | None = shutil.which("glslc") + if found: + return found + sdk: str | None = os.environ.get("VULKAN_SDK") + if not sdk: + return None + for rel in ("Bin/glslc.exe", "Bin/glslc", "bin/glslc"): + candidate: Path = Path(sdk) / rel + if candidate.is_file(): + return str(candidate) + return None + + +def compile_glsl_to_spirv(source: str) -> bytes: + """ + Compile a GLSL compute shader string to SPIR-V bytes. + + Uses native `cthreads._ext.gpu.compile_glsl` when the GPU extension was + built with glslang. Otherwise invokes `glslc` if present. + + #### Args: + - source: str = full `.comp` text (`#version` + buffers + `main`) + + #### Returns + - bytes = SPIR-V binary (multiple of 4 bytes, magic 0x07230203) + + #### Raises + - RuntimeError = no compiler available, or compile failed + + #### Technical terms: + - SPIR-V: intermediate binary Vulkan drivers consume + - glslang: Khronos GLSL compiler linked into `_ext` for GPU builds + """ + try: + from cthreads._ext import gpu as _gpu # type: ignore + + compile_native = getattr(_gpu, "compile_glsl", None) + if callable(compile_native): + try: + out = compile_native(source) + except Exception as exc: + raise RuntimeError(_format_glsl_compile_error(str(exc))) from exc + if not isinstance(out, (bytes, bytearray)): + raise RuntimeError( + "cthreads.gpu: _ext.gpu.compile_glsl did not return bytes" + ) + data: bytes = bytes(out) + if not data or (len(data) % 4) != 0: + raise RuntimeError( + "cthreads.gpu: native compile_glsl returned invalid SPIR-V" + ) + if data[:4] != b"\x03\x02\x23\x07": + raise RuntimeError( + "cthreads.gpu: native compile_glsl missing SPIR-V magic" + ) + return data + except ImportError: + pass + + glslc: str | None = _find_glslc() + if glslc is None: + raise RuntimeError( + "cthreads.gpu: no GLSL compiler found. Rebuild with " + "-DCTHREADS_GPU=ON (vendors glslang into _ext), or install the " + "Vulkan SDK glslc for a temporary CLI fallback" + ) + + with tempfile.TemporaryDirectory(prefix="cthreads_glsl_") as tmp: + tmp_path: Path = Path(tmp) + comp_path: Path = tmp_path / "kernel.comp" + spv_path: Path = tmp_path / "kernel.spv" + comp_path.write_text(source, encoding="utf-8") + cmd: list[str] = [ + glslc, + "-fshader-stage=compute", + str(comp_path), + "-o", + str(spv_path), + ] + proc = subprocess.run(cmd, capture_output=True, text=True) + if proc.returncode != 0: + err: str = (proc.stderr or proc.stdout or "").strip() + raise RuntimeError( + _format_glsl_compile_error( + "glslc (shaderc CLI fallback) failed:\n" + err + ) + ) + data = spv_path.read_bytes() + + if not data or (len(data) % 4) != 0: + raise RuntimeError("cthreads.gpu: glslc produced invalid SPIR-V") + if data[:4] != b"\x03\x02\x23\x07": + raise RuntimeError("cthreads.gpu: glslc output missing SPIR-V magic") + return data + + +def _format_glsl_compile_error(detail: str) -> str: + """ + Wrap a glslang/glslc failure with a short identifier hint. + + Reserved GLSL words used as buffer instance names (e.g. `out`, `in`) + fail at compile time; we surface that instead of maintaining a keyword list. + """ + return ( + "cthreads.gpu: GLSL compile failed:\n" + f"{detail.strip()}\n" + "Hint: kernel / parameter names become GLSL identifiers. Avoid " + "reserved words (e.g. out, in, buffer, shared, uniform, flat)." + ) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Assign.py b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Assign.py new file mode 100644 index 0000000..4d2b3c7 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Assign.py @@ -0,0 +1,144 @@ +""" +Assignment lowering for @Gpu bodies. + +Same shape as CPU Assign, without C++ includes / to_cpp / std::pow. +Annotated locals use Glsl.type_name. Power (`**`) is deferred to a future +math builtin mirror (like CPU stdlib -> C++). +""" +import ast + +from .....compiler.translation.Source import Source +from .....types import PyList, hint_to_pytype +from ..Glsl import type_name +from ..context import GpuTranslationContext + + +class GpuAssign: + """ + Lower ast.AnnAssign / Assign / AugAssign for GPU shader bodies. + """ + + @staticmethod + def ann_assign(node: ast.AnnAssign, ctx: GpuTranslationContext) -> list[str]: + """ + Declare a typed local and optionally initialize it. + + #### Args: + - node: ast.AnnAssign = annotated assignment statement + - ctx: GpuTranslationContext = symbols table for later Name lowering + + #### Returns + - list[str] = one indented GLSL declaration line + """ + from .Syntax import GpuSyntax + + if not isinstance(node.target, ast.Name): + raise TypeError( + f"GPU function {ctx.func_name}: " + "AnnAssign target must be a plain name" + ) + var_name: str = node.target.id + if var_name in ctx.symbols: + raise TypeError( + f"GPU function {ctx.func_name}: redeclaration of {var_name!r}" + ) + + hint = Source.resolve_annotation(node.annotation, ctx.fn.__globals__) + py_type = hint_to_pytype(hint) + if isinstance(py_type, PyList): + raise TypeError( + f"GPU function {ctx.func_name}: " + "local list declarations are not supported" + ) + + glsl_ty: str = type_name(py_type) + # Locals are bare ids; do not add to scalar_params (those are SSBO fields). + ctx.symbols[var_name] = py_type + + if node.value is None: + return [f" {glsl_ty} {var_name};"] + rhs: str = GpuSyntax.expr(node.value, ctx) + return [f" {glsl_ty} {var_name} = {rhs};"] + + @staticmethod + def assign(node: ast.Assign, ctx: GpuTranslationContext) -> list[str]: + """ + Lower a single-target assignment (`lhs = rhs`). + + Name and Index already rewrite scalar/list SSBO accessors. + + #### Args: + - node: ast.Assign = assignment statement + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - list[str] = one indented GLSL assignment line + """ + from .Syntax import GpuSyntax + + if len(node.targets) != 1: + raise TypeError( + f"GPU function {ctx.func_name}: " + "only single-target assignment is supported" + ) + target = node.targets[0] + if isinstance(target, ast.Name): + if target.id not in ctx.symbols: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"assign to unknown name {target.id!r} " + "(declare it with an annotated assignment first)" + ) + if target.id in ctx.list_params: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"cannot assign to list parameter {target.id!r} " + "(assign elements via indexing)" + ) + elif isinstance(target, ast.Subscript): + if isinstance(target.slice, ast.Slice): + raise TypeError( + f"GPU function {ctx.func_name}: " + "slice assignment is not supported" + ) + else: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"unsupported assign target {type(target).__name__}" + ) + + lhs: str = GpuSyntax.expr(target, ctx) + rhs: str = GpuSyntax.expr(node.value, ctx) + return [f" {lhs} = {rhs};"] + + @staticmethod + def aug_assign(node: ast.AugAssign, ctx: GpuTranslationContext) -> list[str]: + """ + Lower augmented assignment (`lhs op= rhs`). Power is not supported yet. + + #### Args: + - node: ast.AugAssign = augmented assignment statement + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - list[str] = one indented GLSL aug-assign line + """ + from .Op import GpuOp + from .Syntax import GpuSyntax + + if isinstance(node.op, ast.Pow): + raise TypeError( + f"GPU function {ctx.func_name}: " + "** / pow is not supported yet " + "(planned via a math builtin mirror)" + ) + + target: str = GpuSyntax.expr(node.target, ctx) + value: str = GpuSyntax.expr(node.value, ctx) + op = GpuOp.BINOPS.get(type(node.op)) + if not op: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"unsupported aug-assign operator {type(node.op).__name__}" + ) + return [f" {target} {op}= {value};"] diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Flow.py b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Flow.py new file mode 100644 index 0000000..5e2d56e --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Flow.py @@ -0,0 +1,180 @@ +""" +Control-flow lowering for @Gpu bodies. + +Subclasses CPU Flow to reuse nest / pass / break / continue. +Overrides handlers that import Syntax so they use GpuSyntax. +for-in over lists (C++ auto&) is rejected; range-for is kept. +""" +import ast + +from .....compiler.translation.syntax.Flow import Flow +from .....types import PyInt +from ..context import GpuTranslationContext +from .Op import GpuOp + + +class GpuFlow(Flow): + """ + Lower if / while / for / return for GPU shader bodies. Reuses CPU Flow for shared helpers. + """ + + @staticmethod + def if_stmt(node: ast.If, ctx: GpuTranslationContext) -> list[str]: + """ + Lower an if/else statement to GLSL. + + #### Args: + - node: ast.If = if statement + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - list[str] = indented GLSL if/else lines + """ + from .Syntax import GpuSyntax + + test: str = GpuSyntax.expr(node.test, ctx) + lines: list[str] = [f" if ({test}) {{"] + for stmt in node.body: + lines.extend(GpuFlow.nest(GpuSyntax.stmt(stmt, ctx))) + lines.append(" }") + if node.orelse: + lines.append(" else {") + for stmt in node.orelse: + lines.extend(GpuFlow.nest(GpuSyntax.stmt(stmt, ctx))) + lines.append(" }") + return lines + + @staticmethod + def while_stmt(node: ast.While, ctx: GpuTranslationContext) -> list[str]: + """ + Lower a while loop to GLSL (no while/else). + + #### Args: + - node: ast.While = while statement + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - list[str] = indented GLSL while lines + """ + from .Syntax import GpuSyntax + + if node.orelse: + raise TypeError( + f"GPU function {ctx.func_name}: while/else is not supported" + ) + test: str = GpuSyntax.expr(node.test, ctx) + lines: list[str] = [f" while ({test}) {{"] + for stmt in node.body: + lines.extend(GpuFlow.nest(GpuSyntax.stmt(stmt, ctx))) + lines.append(" }") + return lines + + @staticmethod + def for_stmt(node: ast.For, ctx: GpuTranslationContext) -> list[str]: + """ + Lower `for i in range(...)` to a C-style GLSL for loop. + + Iteration over list parameters is not supported (no C++ range-for). + + #### Args: + - node: ast.For = for statement + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - list[str] = indented GLSL for-loop lines + """ + from .Syntax import GpuSyntax + + if node.orelse: + raise TypeError( + f"GPU function {ctx.func_name}: for/else is not supported" + ) + if not isinstance(node.target, ast.Name): + raise TypeError( + f"GPU function {ctx.func_name}: " + "for-loop target must be a plain name" + ) + loop_var: str = node.target.id + if loop_var in ctx.symbols: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"for-loop rebinds existing name {loop_var!r}" + ) + + it = node.iter + if not GpuOp.is_builtin_call(it, "range"): + raise TypeError( + f"GPU function {ctx.func_name}: " + "for-iter must be range(...) " + "(list iteration is not supported on the GPU path)" + ) + assert isinstance(it, ast.Call) + if it.keywords: + raise TypeError( + f"GPU function {ctx.func_name}: " + "range() keyword args are not supported" + ) + n: int = len(it.args) + if n == 1: + start, stop, step = "0", GpuSyntax.expr(it.args[0], ctx), "1" + elif n == 2: + start = GpuSyntax.expr(it.args[0], ctx) + stop = GpuSyntax.expr(it.args[1], ctx) + step = "1" + elif n == 3: + start = GpuSyntax.expr(it.args[0], ctx) + stop = GpuSyntax.expr(it.args[1], ctx) + step = GpuSyntax.expr(it.args[2], ctx) + else: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"range() expects 1..3 args, got {n}" + ) + + ctx.symbols[loop_var] = PyInt() + lines: list[str] = [ + f" for (int {loop_var} = {start}; " + f"{loop_var} < {stop}; " + f"{loop_var} += {step}) {{" + ] + for stmt in node.body: + lines.extend(GpuFlow.nest(GpuSyntax.stmt(stmt, ctx))) + lines.append(" }") + del ctx.symbols[loop_var] + return lines + + @staticmethod + def return_stmt(node: ast.Return, ctx: GpuTranslationContext) -> list[str]: + """ + Lower return to a void early exit. + + Valued returns are not lowered yet; always emit bare `return;`. + + #### Args: + - node: ast.Return = return statement + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - list[str] = one indented `return;` line + """ + return [" return;"] + + @staticmethod + def expr_stmt(node: ast.Expr, ctx: GpuTranslationContext) -> list[str]: + """ + Lower expression statements; string doc-exprs are ignored. + + Calls are not supported until GPU builtins / math mirrors exist. + + #### Args: + - node: ast.Expr = expression statement + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - list[str] = empty, or a comment for unsupported forms + """ + if isinstance(node.value, ast.Constant) and isinstance(node.value.value, str): + return [] + return [ + f" // unsupported statement: Expr ({type(node.value).__name__})" + ] diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Index.py b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Index.py new file mode 100644 index 0000000..fc2e9ca --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Index.py @@ -0,0 +1,51 @@ +""" +ast.Subscript lowering for @Gpu bodies. + +List indexing becomes SSBO unsized-array access: x[i] -> x.data[i]. +""" + +import ast + +from ..context import GpuTranslationContext + + +class GpuIndex: + """ + Lower ast.Subscript for GPU list SSBOs (not slices). + """ + + @staticmethod + def subscript(node: ast.Subscript, ctx: GpuTranslationContext) -> str: + """ + Lower list indexing to GLSL `base.data[index]`. + + #### Args: + - node: ast.Subscript = Python subscript expression + - ctx: GpuTranslationContext = symbols and list param set + + #### Returns + - str = GLSL text such as `(x.data[i])` + + #### Raises + - TypeError = slice syntax, or subscript of a non-list param + + #### Technical terms: + - SSBO: shader storage buffer object; list params use `T data[]` + """ + from .Syntax import GpuSyntax + + if isinstance(node.slice, ast.Slice): + raise TypeError( + f"GPU function {ctx.func_name}: slice syntax is not supported" + ) + + # Only list kernel params expose a `data[]` member in the preamble. + if not isinstance(node.value, ast.Name) or node.value.id not in ctx.list_params: + raise TypeError( + f"GPU function {ctx.func_name}: " + "subscript is only supported on list parameters" + ) + + base: str = GpuSyntax.expr(node.value, ctx) + index: str = GpuSyntax.expr(node.slice, ctx) + return f"({base}.data[{index}])" diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Name.py b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Name.py new file mode 100644 index 0000000..2ba81fa --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Name.py @@ -0,0 +1,67 @@ +""" +ast.Name lowering for @Gpu bodies. + +Scalar params become fields on the binding-0 Scalars SSBO. +List params stay bare buffer instance names (Index adds .data[i]). +Locals stay bare identifiers. +self is stubbed until method kernels exist. +""" +import ast + +from ..context import GpuTranslationContext + + +class GpuName: + """ + Lower ast.Name nodes for GPU shader bodies with SSBO-aware rewriting. + """ + + @staticmethod + def name(node: ast.Name, ctx: GpuTranslationContext) -> str: + """ + Lower a name to a GLSL identifier or Scalars SSBO field access. + + #### Args: + - node: ast.Name = Python name expression + - ctx: GpuTranslationContext = symbols and scalar/list param sets + + #### Returns + - str = GLSL text (`scalars.n`, bare local/list id, or self hook) + + #### Technical terms: + - SSBO: shader storage buffer object holding kernel params on device + """ + if node.id == "self": + # Separate hook so method kernels can land here without rewriting name(). + return GpuName.lower_self(ctx) + + if node.id not in ctx.symbols: + raise TypeError( + f"GPU function {ctx.func_name}: unknown name {node.id!r}" + ) + + # Scalar kernel params live in the binding-0 block instance `scalars`. + if node.id in ctx.scalar_params: + return f"scalars.{node.id}" + + # List params are SSBO instance names; locals are ordinary GLSL ids. + return node.id + + @staticmethod + def lower_self(ctx: GpuTranslationContext) -> str: + """ + Lower `self` for method kernels (not implemented yet). + + #### Args: + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - str = GLSL receiver expression (when supported) + + #### Raises + - TypeError = self / method kernels are not supported yet + """ + raise TypeError( + f"GPU function {ctx.func_name}: " + "self / method kernels are not supported yet" + ) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Op.py b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Op.py new file mode 100644 index 0000000..1a29634 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Op.py @@ -0,0 +1,102 @@ +""" +GLSL operator lowering. + +Reuses CppOp operator tables (BINOPS / UNARYOPS / CMPOPS / BOOLOPS). +Handlers are overridden so they call GPU Syntax and emit GLSL (pow, no includes). +""" +import ast + +from .....compiler.translation.syntax.Op import Op as CppOp +from ..context import GpuTranslationContext + + +class GpuOp(CppOp): + """ast.BinOp / UnaryOp / Compare / BoolOp for @Gpu bodies.""" + + # No CPU sync builtin on the GPU path. + BUILTINS: frozenset[str] = frozenset({"range", "len"}) + + @staticmethod + def is_builtin_call(node: ast.AST, name: str) -> bool: + return ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == name + and name in GpuOp.BUILTINS + ) + + @staticmethod + def bin_op(node: ast.BinOp, ctx: GpuTranslationContext) -> str: + from .Syntax import GpuSyntax + + left = GpuSyntax.expr(node.left, ctx) + right = GpuSyntax.expr(node.right, ctx) + # Power stays out until math builtins mirror Python -> GLSL (like CPU). + if isinstance(node.op, ast.Pow): + raise TypeError( + f"GPU function {ctx.func_name}: " + "** / pow is not supported yet " + "(planned via a math builtin mirror)" + ) + op = GpuOp.BINOPS.get(type(node.op)) + if not op: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"unsupported binary operator {type(node.op).__name__}" + ) + return f"({left} {op} {right})" + + @staticmethod + def unary_op(node: ast.UnaryOp, ctx: GpuTranslationContext) -> str: + from .Syntax import GpuSyntax + + op = GpuOp.UNARYOPS.get(type(node.op)) + if not op: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"unsupported unary operator {type(node.op).__name__}" + ) + operand = GpuSyntax.expr(node.operand, ctx) + return f"({op}{operand})" + + @staticmethod + def compare(node: ast.Compare, ctx: GpuTranslationContext) -> str: + from .Syntax import GpuSyntax + + if len(node.ops) != len(node.comparators): + raise TypeError( + f"GPU function {ctx.func_name}: malformed Compare node" + ) + left = GpuSyntax.expr(node.left, ctx) + parts: list[str] = [] + prev = left + for op_node, comparator in zip(node.ops, node.comparators): + op = GpuOp.CMPOPS.get(type(op_node)) + if not op: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"unsupported compare operator {type(op_node).__name__}" + ) + right = GpuSyntax.expr(comparator, ctx) + parts.append(f"({prev} {op} {right})") + prev = right + if len(parts) == 1: + return parts[0] + return "(" + " && ".join(parts) + ")" + + @staticmethod + def bool_op(node: ast.BoolOp, ctx: GpuTranslationContext) -> str: + from .Syntax import GpuSyntax + + op = GpuOp.BOOLOPS.get(type(node.op)) + if not op: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"unsupported bool operator {type(node.op).__name__}" + ) + if len(node.values) < 2: + raise TypeError( + f"GPU function {ctx.func_name}: BoolOp needs at least two values" + ) + parts = [GpuSyntax.expr(v, ctx) for v in node.values] + return "(" + f" {op} ".join(parts) + ")" diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Syntax.py b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Syntax.py new file mode 100644 index 0000000..032e67b --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/syntax/Syntax.py @@ -0,0 +1,106 @@ +""" +Syntax dispatcher for @Gpu bodies. + +Wires GpuOp / GpuName / GpuIndex / GpuAssign / GpuFlow / Literal.constant. +Attribute and Call go through the GPU plugin registry. +""" + +from __future__ import annotations + +import ast +from typing import Callable + +from .....compiler.translation.syntax.Literal import Literal +from ..context import GpuTranslationContext +from ..plugins import lower_attr, lower_call +from .Assign import GpuAssign +from .Flow import GpuFlow +from .Index import GpuIndex +from .Name import GpuName +from .Op import GpuOp + +ExprHandler = Callable[[ast.AST, GpuTranslationContext], str] +StmtHandler = Callable[[ast.AST, GpuTranslationContext], list[str]] + + +class GpuSyntax: + """ + Dispatcher: GpuSyntax.expr / GpuSyntax.stmt to area static methods. + """ + + _EXPR: dict[type, ExprHandler] = { + ast.Constant: Literal.constant, + ast.Name: GpuName.name, + ast.Subscript: GpuIndex.subscript, + ast.BinOp: GpuOp.bin_op, + ast.UnaryOp: GpuOp.unary_op, + ast.Compare: GpuOp.compare, + ast.BoolOp: GpuOp.bool_op, + } + _STMT: dict[type, StmtHandler] = { + ast.AnnAssign: GpuAssign.ann_assign, + ast.Assign: GpuAssign.assign, + ast.AugAssign: GpuAssign.aug_assign, + ast.Pass: GpuFlow.pass_stmt, + ast.Break: GpuFlow.break_stmt, + ast.Continue: GpuFlow.continue_stmt, + ast.Return: GpuFlow.return_stmt, + ast.Expr: GpuFlow.expr_stmt, + ast.If: GpuFlow.if_stmt, + ast.For: GpuFlow.for_stmt, + ast.While: GpuFlow.while_stmt, + } + + @staticmethod + def expr(node: ast.expr, ctx: GpuTranslationContext) -> str: + """ + Lower one expression AST node to GLSL text. + + #### Args: + - node: ast.expr = expression node + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - str = GLSL expression text + """ + if isinstance(node, ast.Call): + out = lower_call(node, ctx, GpuSyntax.expr) + if out is not None: + return out + raise TypeError( + f"GPU function {ctx.func_name}: " + "unsupported call (no CallPlugin matched)" + ) + if isinstance(node, ast.Attribute): + out = lower_attr(node, ctx, GpuSyntax.expr) + if out is not None: + return out + raise TypeError( + f"GPU function {ctx.func_name}: " + "unsupported attribute (no AttrPlugin matched)" + ) + + handler = GpuSyntax._EXPR.get(type(node)) + if handler is None: + raise TypeError( + f"GPU function {ctx.func_name}: " + f"unsupported expression {type(node).__name__}" + ) + return handler(node, ctx) + + @staticmethod + def stmt(node: ast.stmt, ctx: GpuTranslationContext) -> list[str]: + """ + Lower one statement AST node to indented GLSL lines. + + #### Args: + - node: ast.stmt = statement node + - ctx: GpuTranslationContext = current GPU translation state + + #### Returns + - list[str] = GLSL lines (comment stub if unsupported) + """ + handler = GpuSyntax._STMT.get(type(node)) + if handler is None: + return [f" // unsupported statement: {type(node).__name__}"] + return handler(node, ctx) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/translate.py b/src/cthreads/python/cthreads/gpu/compiler/translation/translate.py new file mode 100644 index 0000000..2ec8068 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/translate.py @@ -0,0 +1,68 @@ +""" +One-shot GPU translate: parse -> Signature -> Syntax -> assemble -> optional SPIR-V. +""" + +from __future__ import annotations + +import ast +from typing import Callable + +from ....compiler.translation.Source import Source +from .Signature import GpuSignature +from .assemble import assemble_comp +from .context import GpuTranslationContext +from .result import GpuTranslationResult +from .spirv import compile_glsl_to_spirv +from .syntax.Syntax import GpuSyntax + + +def translate_function_for_gpu( + fn: Callable, + *, + local_size_x: int = 64, + compile_spirv: bool = False, +) -> GpuTranslationResult: + """ + Translate one `@Gpu` function to GLSL, optionally compile to SPIR-V. + + #### Args: + - fn: Callable = annotated GPU kernel function + - local_size_x: int = workgroup size x (default 64) + - compile_spirv: bool = if True, run shaderc (`glslc` / native) on `source` + + #### Returns + - GpuTranslationResult = preamble, body, `source`, optional `spirv`, layout + + #### Raises + - RuntimeError = SPIR-V compile requested but compiler missing / failed + """ + ctx: GpuTranslationContext = GpuTranslationContext( + fn=fn, local_size_x=local_size_x + ) + func_def: ast.FunctionDef = Source.parse_function(fn) + sig = GpuSignature.translate(func_def, ctx) + + body_lines: list[str] = [] + for stmt in func_def.body: + body_lines.extend(GpuSyntax.stmt(stmt, ctx)) + body: str = "\n".join(body_lines) + if body and not body.endswith("\n"): + body += "\n" + + source: str = assemble_comp(sig.preamble, body) + spirv: bytes | None = None + if compile_spirv: + spirv = compile_glsl_to_spirv(source) + + return GpuTranslationResult( + func_name=sig.func_name, + preamble=sig.preamble, + body=body, + source=source, + binding_count=sig.binding_count, + scalar_bytes=sig.scalar_bytes, + local_size_x=sig.local_size_x, + scalar_fields=sig.scalar_fields, + list_fields=sig.list_fields, + spirv=spirv, + ) diff --git a/src/cthreads/python/cthreads/gpu/frontend/__init__.py b/src/cthreads/python/cthreads/gpu/frontend/__init__.py new file mode 100644 index 0000000..81c3f87 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/frontend/__init__.py @@ -0,0 +1,49 @@ +from .wrapper import Gpu +from .errors import ( + CThreadsGPUError, + GPUNotAvailable, + GpuInvalidArgument, + GpuUseAfterDestroy, + VulkanInitFailed, + VulkanLoaderNotFound, + VulkanNoDevice, + VulkanNotBuiltError, + VulkanOutOfMemory, + _map_error, +) +from .indexes import ( + BlockDim, + BlockIdx, + GlobalIdx, + GridDim, + ThreadIdx, +) +from .lib import ( + available, + device_name, + init, + shutdown, +) + +__all__ = [ + "Gpu", + "BlockDim", + "BlockIdx", + "GlobalIdx", + "GridDim", + "ThreadIdx", + "CThreadsGPUError", + "GPUNotAvailable", + "GpuInvalidArgument", + "GpuUseAfterDestroy", + "VulkanInitFailed", + "VulkanLoaderNotFound", + "VulkanNoDevice", + "VulkanNotBuiltError", + "VulkanOutOfMemory", + "_map_error", + "available", + "device_name", + "init", + "shutdown", +] diff --git a/src/cthreads/python/cthreads/gpu/errors.py b/src/cthreads/python/cthreads/gpu/frontend/errors.py similarity index 77% rename from src/cthreads/python/cthreads/gpu/errors.py rename to src/cthreads/python/cthreads/gpu/frontend/errors.py index 5808db9..f016a79 100644 --- a/src/cthreads/python/cthreads/gpu/errors.py +++ b/src/cthreads/python/cthreads/gpu/frontend/errors.py @@ -1,6 +1,5 @@ """ctypes-style error types for cthreads.gpu (mapped from C++ message prefixes).""" - class CThreadsGPUError(Exception): def __init__(self, detail: str = "Unknown Error") -> None: self.detail = detail @@ -65,3 +64,22 @@ class GPUNotAvailable(CThreadsGPUError): def __init__(self, detail: str = "GPU not available") -> None: super().__init__(detail) + + +def _map_error(exc: BaseException) -> CThreadsGPUError: + msg = str(exc) + if "VulkanLoaderNotFound" in msg: + return VulkanLoaderNotFound(msg) + if "VulkanNoDevice" in msg: + return VulkanNoDevice(msg) + if "VulkanOutOfMemory" in msg: + return VulkanOutOfMemory(msg) + if "GpuUseAfterDestroy" in msg: + return GpuUseAfterDestroy(msg) + if "GpuInvalidArgument" in msg: + return GpuInvalidArgument(msg) + if "VulkanNotBuilt" in msg: + return VulkanNotBuiltError(msg) + if "VulkanInitFailed" in msg: + return VulkanInitFailed(msg) + return VulkanInitFailed(msg) \ No newline at end of file diff --git a/src/cthreads/python/cthreads/gpu/frontend/indexes.py b/src/cthreads/python/cthreads/gpu/frontend/indexes.py new file mode 100644 index 0000000..9f93a12 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/frontend/indexes.py @@ -0,0 +1,110 @@ +""" +CUDA-style index builtins for `@Gpu` kernels. + +Frontend markers only. Codegen lowers `GlobalIdx.x` (and friends) via the +GPU AttrPlugin registry to GLSL invocation IDs. + +#### Mapping: +- ThreadIdx -> gl_LocalInvocationID +- BlockIdx -> gl_WorkGroupID +- BlockDim -> gl_WorkGroupSize +- GridDim -> gl_NumWorkGroups +- GlobalIdx -> gl_GlobalInvocationID +""" + +from __future__ import annotations + + +class _Axis: + """ + Placeholder for one axis (`.x` / `.y` / `.z`) of an index builtin. + """ + + __slots__ = ("_label",) + + def __init__(self, label: str) -> None: + self._label: str = label + + def __repr__(self) -> str: + return f"" + + +class GpuIndexBuiltin: + """ + Base for ThreadIdx / BlockIdx / GlobalIdx marker classes. + + #### Technical terms: + - GLSL: shading language used for Vulkan compute shaders + """ + + # Subclasses set the GLSL built-in vector name (without .x/.y/.z). + _glsl_base: str = "" + + +class ThreadIdx(GpuIndexBuiltin): + """ + Local invocation index within a workgroup (`gl_LocalInvocationID`). + """ + + _glsl_base: str = "gl_LocalInvocationID" + x: _Axis = _Axis("ThreadIdx.x") + y: _Axis = _Axis("ThreadIdx.y") + z: _Axis = _Axis("ThreadIdx.z") + + +class BlockIdx(GpuIndexBuiltin): + """ + Workgroup index (`gl_WorkGroupID`). + """ + + _glsl_base: str = "gl_WorkGroupID" + x: _Axis = _Axis("BlockIdx.x") + y: _Axis = _Axis("BlockIdx.y") + z: _Axis = _Axis("BlockIdx.z") + + +class BlockDim(GpuIndexBuiltin): + """ + Workgroup size (`gl_WorkGroupSize`). + """ + + _glsl_base: str = "gl_WorkGroupSize" + x: _Axis = _Axis("BlockDim.x") + y: _Axis = _Axis("BlockDim.y") + z: _Axis = _Axis("BlockDim.z") + + +class GridDim(GpuIndexBuiltin): + """ + Number of workgroups (`gl_NumWorkGroups`). + """ + + _glsl_base: str = "gl_NumWorkGroups" + x: _Axis = _Axis("GridDim.x") + y: _Axis = _Axis("GridDim.y") + z: _Axis = _Axis("GridDim.z") + + +class GlobalIdx(GpuIndexBuiltin): + """ + Global invocation index (`gl_GlobalInvocationID`). + + Prefer this for 1D element-wise kernels (for example saxpy). + + #### Example: + ``py + from cthreads.gpu import GlobalIdx, Gpu + + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + `` + """ + + _glsl_base: str = "gl_GlobalInvocationID" + x: _Axis = _Axis("GlobalIdx.x") + y: _Axis = _Axis("GlobalIdx.y") + z: _Axis = _Axis("GlobalIdx.z") diff --git a/src/cthreads/python/cthreads/gpu/frontend/lib.py b/src/cthreads/python/cthreads/gpu/frontend/lib.py new file mode 100644 index 0000000..2dd1df0 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/frontend/lib.py @@ -0,0 +1,106 @@ +""" +Public GPU probe wrappers. + +Maps native C++ error prefixes to `frontend.errors` types and calls through +`_ext_gpu_api`. +""" + +from .. import _ext_gpu_api +from .errors import ( + VulkanNotBuiltError, + _map_error, +) + + +def available() -> bool: + """ + Return True if the Vulkan loader and a compute device can be initialized. + + #### Returns + - bool = True when GPU init would succeed + + #### Example: + ``py + from cthreads import gpu + if gpu.available(): + print(gpu.device_name()) + `` + """ + return _ext_gpu_api.available() + + +def device_name() -> str: + """ + Return the active GPU name (calls init). Raises on failure or not built. + + #### Returns + - str = Vulkan device name + + #### Raises + - VulkanNotBuiltError = extension compiled without CTHREADS_GPU + - CThreadsGPUError = mapped native init / device errors + + #### Example: + ``py + from cthreads import gpu + name = gpu.device_name() + `` + """ + if _ext_gpu_api._gpu is None: + raise VulkanNotBuiltError( + "cthreads built without CTHREADS_GPU; rebuild with -DCTHREADS_GPU=ON" + ) + try: + return _ext_gpu_api.device_name() + except Exception as exc: + raise _map_error(exc) from exc + + +def init() -> None: + """ + Explicitly initialize the Vulkan context. + + #### Returns + - None + + #### Raises + - VulkanNotBuiltError = extension compiled without CTHREADS_GPU + - CThreadsGPUError = mapped native init errors + + #### Example: + ``py + from cthreads import gpu + gpu.init() + `` + """ + if _ext_gpu_api._gpu is None: + raise VulkanNotBuiltError( + "cthreads built without CTHREADS_GPU; rebuild with -DCTHREADS_GPU=ON" + ) + try: + _ext_gpu_api.init() + except Exception as exc: + raise _map_error(exc) from exc + + +def shutdown() -> None: + """ + Destroy the device/instance and unload the Vulkan loader. + + Native ShaderCache is released with the device (explicit memory management). + Marks `prepare` so the next `prepare()` / `gpu()` rewalks the registry and + re-registers SPIR-V. No-op when the GPU extension is not built. + + #### Returns + - None + + #### Example: + ``py + from cthreads import gpu + gpu.shutdown() + `` + """ + _ext_gpu_api.shutdown() + from .. import runtime as runtime_mod + + runtime_mod._gpu_prepared = False diff --git a/src/cthreads/python/cthreads/gpu/frontend/wrapper.py b/src/cthreads/python/cthreads/gpu/frontend/wrapper.py new file mode 100644 index 0000000..8955af7 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/frontend/wrapper.py @@ -0,0 +1,78 @@ +""" +`@Gpu` decorator: mark, validate, and register a GPU kernel function. + +Does not emit SPIR-V or launch. That happens later via GpuCompileSession / `gpu()`. +""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import Any, TypeVar + +from ...frontend.Registry import REGISTRY +from ..gpu_kernel_meta import build_gpu_kernel_meta +from .errors import GPUNotAvailable +from .lib import available, device_name + +F = TypeVar("F", bound=Callable[..., Any]) + + +def Gpu(fn: F | None = None, *, log: bool = False): + """ + Mark a function as a `@Gpu` kernel and register it for later compile/launch. + + Supports `@Gpu` and `@Gpu(log=True)`. Validates annotations against the GPU + allowlist (`-> None`, scalars and `list` of scalars), attaches + `fn.__gpu_kernel_meta__`, and returns the same function (still runs as + normal Python when called directly). + + #### Args: + - fn: Callable | None = function when used as `@Gpu`; omit for `@Gpu(log=...)` + - log: bool = if True, print the assigned device name (default False) + + #### Returns + - Callable = the marked function, or a decorator when `fn` is omitted + + #### Raises + - GPUNotAvailable = Vulkan GPU path is not usable in this process + - TypeError = missing or unsupported annotations + + #### Example: + ``py + from cthreads.gpu import Gpu + + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + + @Gpu(log=True) + def saxpy_logged(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + `` + """ + + def apply(f: F) -> F: + if not available(): + raise GPUNotAvailable( + "GPU is not available (build with CTHREADS_GPU=ON and a Vulkan " + "device)" + ) + + f.__gpu__ = True # type: ignore[attr-defined] + f.__gpu_version__ = REGISTRY.VERSION # type: ignore[attr-defined] + + # Validate annotations and attach launch meta (no SPIR-V yet). + build_gpu_kernel_meta(f) + + REGISTRY.register_gpu_function(f) + + if log: + print( + f"\033[92mGPU LOG:\033[0m function {f.__name__} is assigned to " + f"GPU {device_name()}" + ) + return f + + if fn is not None: + return apply(fn) + return apply diff --git a/src/cthreads/python/cthreads/gpu/gpu_kernel_meta.py b/src/cthreads/python/cthreads/gpu/gpu_kernel_meta.py new file mode 100644 index 0000000..e7f975f --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/gpu_kernel_meta.py @@ -0,0 +1,379 @@ +""" +Compile-time metadata for one `@Gpu` kernel. + +Shape matches `launch_gpu_kernel` in `module.hpp`: `symbol`, binding layout, +`params` with flat `kind` / `pass_as` / `elem_kind` / `elem_bytes`, and optional +dispatch overrides. + +v1 is in-place list writeback only (`pass_as="ref"` on lists). Kernels are +`-> None`; scalars are not written back on join. +""" + +from __future__ import annotations + +import inspect +from dataclasses import dataclass, field +from typing import Any, Callable, get_type_hints + +from ..types import ( + PyBool, + PyDict, + PyFloat, + PyInt, + PyList, + PyString, + PyThreadable, + PyType, + hint_to_pytype, + is_shared_pytype, + is_sync_pytype, + is_tbuffer_pytype, +) + +# Populated by build_gpu_kernel_meta(); keyed by symbol. +GPU_KERNELS: dict[str, "GpuKernelMeta"] = {} + +# Host / std430 sizes used by module.cpp py_size_of (GLSL bool is 32-bit). +_GPU_SCALAR_BYTES: dict[str, int] = { + "bool": 4, + "int": 4, + "float": 4, + "double": 8, +} + +_DEFAULT_LOCAL_SIZE_X = 64 + + +def _align_up(value: int, alignment: int) -> int: + """ + Round `value` up to the next multiple of `alignment` (std430 field packing). + + #### Args: + - value: int = current byte offset + - alignment: int = required alignment in bytes + + #### Returns + - int = aligned offset + """ + return (value + alignment - 1) & ~(alignment - 1) + + +def _gpu_kind_and_bytes(py_type: PyType) -> tuple[str, int]: + """ + Map a scalar PyType to the launch `kind` string and std430 byte size. + + #### Args: + - py_type: PyType = scalar type from hint_to_pytype + + #### Returns + - tuple[str, int] = (kind, elem_bytes) for int / float / bool + + #### Raises + - TypeError = type is not a supported GPU scalar + """ + if isinstance(py_type, PyInt): + return "int", _GPU_SCALAR_BYTES["int"] + if isinstance(py_type, PyFloat): + # Shader float (f32). CPU @Thread uses double; GPU v1 follows GLSL float. + return "float", _GPU_SCALAR_BYTES["float"] + if isinstance(py_type, PyBool): + return "bool", _GPU_SCALAR_BYTES["bool"] + raise TypeError( + f"GPU kernel: unsupported scalar type {py_type.name!r} " + f"(allowed: int, float, bool)" + ) + + +def pytype_to_gpu_schema(py_type: PyType) -> GpuTypeSchema: + """ + Convert a PyType into a GPU schema for marshal and codegen. + + Lists must be `list[int|float|bool]`. Dict, str, Threadable, sync, and + shared types are rejected. + + #### Args: + - py_type: PyType = type from hint_to_pytype + + #### Returns + - GpuTypeSchema = scalar or list schema with spirv_type / elem_bytes + + #### Raises + - TypeError = unsupported GPU type + """ + if is_sync_pytype(py_type) or is_shared_pytype(py_type) or is_tbuffer_pytype( + py_type + ): + raise TypeError( + f"GPU kernel: {py_type.name!r} is not supported on the GPU path" + ) + if isinstance(py_type, (PyDict, PyString, PyThreadable)): + raise TypeError( + f"GPU kernel: {py_type.name!r} is not supported on the GPU path" + ) + if isinstance(py_type, PyList): + kind, elem_bytes = _gpu_kind_and_bytes(py_type.inner_type) + return GpuTypeSchema( + kind="list", + spirv_type=f"{kind}[]", + inner=GpuTypeSchema( + kind=kind, + spirv_type=kind, + elem_bytes=elem_bytes, + ), + elem_bytes=elem_bytes, + ) + kind, elem_bytes = _gpu_kind_and_bytes(py_type) + return GpuTypeSchema(kind=kind, spirv_type=kind, elem_bytes=elem_bytes) + + +@dataclass +class GpuTypeSchema: + """ + Layout for one marshal or codegen slot (scalar or list of scalars). + + Serialized under each param's `schema` key. Launch also reads flat + `kind` / `elem_kind` / `elem_bytes` from GpuParamMeta.to_dict. + + #### Technical terms: + - std430: GLSL/SPIR-V storage buffer packing rules for scalar fields + """ + + kind: str # int, float, bool, list + spirv_type: str + elem_bytes: int = 0 + inner: GpuTypeSchema | None = None + + def to_dict(self) -> dict[str, Any]: + """ + Serialize this schema for `fn.__gpu_kernel_meta__`. + + #### Returns + - dict[str, Any] = JSON-friendly schema node for marshal and tests + + #### Raises + - ValueError = list schema is missing an inner element schema + """ + d: dict[str, Any] = { + "kind": self.kind, + "spirv_type": self.spirv_type, + "elem_bytes": self.elem_bytes, + } + if self.kind == "list": + if self.inner is None: + raise ValueError("list schema requires inner element schema") + d["inner"] = self.inner.to_dict() + return d + + +@dataclass +class GpuParamMeta: + """ + One kernel parameter: Python name, pass mode, and GpuTypeSchema. + + Scalars default to `pass_as="value"` (packed into binding 0). Lists default + to `pass_as="ref"` (downloaded into the same Python list on join). + + #### Technical terms: + - pass_as: value packs into the scalar SSBO; ref lists are written back on join + - binding 0: scalar storage buffer; list buffers use bindings 1..N + """ + + name: str + pass_as: str # value | ref + schema: GpuTypeSchema + + def __post_init__(self) -> None: + if self.pass_as not in ("value", "ref"): + raise TypeError( + f"GPU param {self.name!r}: pass_as must be 'value' or 'ref', " + f"got {self.pass_as!r}" + ) + if self.schema.kind == "list" and self.pass_as not in ("value", "ref"): + raise TypeError( + f"GPU list param {self.name!r}: pass_as must be 'value' or 'ref'" + ) + + @property + def kind(self) -> str: + """ + Return the top-level schema kind for this parameter. + + #### Returns + - str = `int`, `float`, `bool`, or `list` + """ + return self.schema.kind + + @property + def elem_kind(self) -> str | None: + """ + Return the list element kind, or None for scalars. + + #### Returns + - str | None = inner kind when `kind == "list"`, else None + """ + if self.schema.kind == "list" and self.schema.inner is not None: + return self.schema.inner.kind + return None + + @property + def elem_bytes(self) -> int | None: + """ + Return the list element size in bytes, or None for scalars. + + #### Returns + - int | None = element byte width when `kind == "list"`, else None + """ + if self.schema.kind == "list": + return self.schema.elem_bytes + return None + + def to_dict(self) -> dict[str, Any]: + """ + Flatten this parameter for `launch_gpu_kernel`. + + Emits top-level `kind` / `pass_as` and, for lists, `elem_kind` / + `elem_bytes` as read by module.cpp. + + #### Returns + - dict[str, Any] = parameter metadata consumed by marshal and launch + """ + d: dict[str, Any] = { + "name": self.name, + "pass_as": self.pass_as, + "kind": self.schema.kind, + "schema": self.schema.to_dict(), + } + if self.schema.kind == "list": + d["elem_kind"] = self.elem_kind + d["elem_bytes"] = self.elem_bytes + return d + + +@dataclass +class GpuKernelMeta: + """ + Compile-time record for one `@Gpu` kernel (ShaderCache key + launch meta). + + No trampolines or DLL symbols. `symbol` is the ShaderCache key. Return is + always void (`-> None`); results are ref-list writeback on join. + + #### Technical terms: + - ShaderCache: process map of symbol to reusable pipeline and set layout + - writeback: download of ref list buffers into caller-owned Python lists + """ + + symbol: str + binding_count: int + scalar_bytes: int + params: list[GpuParamMeta] + local_size_x: int = _DEFAULT_LOCAL_SIZE_X + # Optional dispatch overrides; None => launch may ceil(n / local_size_x). + group_count_x: int | None = None + group_count_y: int | None = 1 + group_count_z: int | None = 1 + types: dict[str, Any] = field(default_factory=dict) + schemas: dict[str, Any] = field(default_factory=dict) + + def to_dict(self) -> dict[str, Any]: + """ + Serialize kernel metadata for `fn.__gpu_kernel_meta__` and launch. + + #### Returns + - dict[str, Any] = dict accepted by `launch_gpu_kernel` + """ + return { + "symbol": self.symbol, + "binding_count": self.binding_count, + "scalar_bytes": self.scalar_bytes, + "local_size_x": self.local_size_x, + "group_count_x": self.group_count_x, + "group_count_y": self.group_count_y, + "group_count_z": self.group_count_z, + "params": [p.to_dict() for p in self.params], + "types": dict(self.types), + "schemas": dict(self.schemas), + } + + +def build_gpu_kernel_meta( + fn: Callable, + *, + symbol: str | None = None, + local_size_x: int = _DEFAULT_LOCAL_SIZE_X, +) -> GpuKernelMeta: + """ + Build and attach metadata for one `@Gpu` function from its annotations. + + Stores the record in `GPU_KERNELS[symbol]` and sets + `fn.__gpu_kernel_meta__` to `meta.to_dict()`. + + #### Args: + - fn: Callable = `@Gpu` function with type hints + - symbol: str | None = ShaderCache key (default: `fn.__name__`) + - local_size_x: int = compute workgroup size (default: 64) + + #### Returns + - GpuKernelMeta = full metadata record for this kernel + + #### Raises + - TypeError = missing annotations, non-None return, or unsupported types + """ + hints = get_type_hints(fn) + ret = hints.get("return", None) + if ret not in (None, type(None)): + raise TypeError( + f"GPU kernel {fn.__qualname__}: return must be None " + f"(in-place list writeback only), got {ret!r}" + ) + + sig = inspect.signature(fn) + if any( + p.kind + in ( + inspect.Parameter.VAR_POSITIONAL, + inspect.Parameter.VAR_KEYWORD, + inspect.Parameter.KEYWORD_ONLY, + ) + for p in sig.parameters.values() + ): + raise TypeError( + f"GPU kernel {fn.__qualname__}: *args / **kwargs / keyword-only " + "args are not supported" + ) + + params: list[GpuParamMeta] = [] + scalar_bytes = 0 + list_count = 0 + + for pname in sig.parameters: + if pname not in hints: + raise TypeError( + f"GPU kernel {fn.__qualname__}: parameter {pname!r} needs a " + "type annotation" + ) + py_type = hint_to_pytype(hints[pname]) + schema = pytype_to_gpu_schema(py_type) + if schema.kind == "list": + list_count += 1 + pass_as = "ref" + else: + align = schema.elem_bytes + scalar_bytes = _align_up(scalar_bytes, align) + scalar_bytes += schema.elem_bytes + pass_as = "value" + params.append(GpuParamMeta(name=pname, pass_as=pass_as, schema=schema)) + + # Binding 0 = scalar SSBO (reserved when any scalars); lists at 1..N. + binding_count = (1 if scalar_bytes > 0 else 0) + list_count + sym = symbol if symbol is not None else fn.__name__ + + meta = GpuKernelMeta( + symbol=sym, + binding_count=binding_count, + scalar_bytes=scalar_bytes, + params=params, + local_size_x=local_size_x, + ) + GPU_KERNELS[sym] = meta + fn.__gpu_kernel_meta__ = meta.to_dict() # type: ignore[attr-defined] + return meta diff --git a/src/cthreads/python/cthreads/gpu/gpu_marshal.py b/src/cthreads/python/cthreads/gpu/gpu_marshal.py new file mode 100644 index 0000000..48c2734 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/gpu_marshal.py @@ -0,0 +1,75 @@ +""" +Marshal helpers for `gpu()` launch (ordered args + dispatch sizing). +""" + +from typing import Any + + +def ordered_values_for_meta(meta: dict[str, Any], args: tuple[Any, ...]) -> list[Any]: + """ + Build the ordered argument list for `launch_gpu_kernel`. + + #### Args: + - meta: dict[str, Any] = kernel metadata (`params`, …) + - args: tuple[Any, ...] = positional args from `gpu(fn, *args)` + + #### Returns + - list[Any] = args in parameter order (same Python list objects for refs) + + #### Raises + - TypeError = arity mismatch + """ + params = meta.get("params") + if not isinstance(params, list): + raise TypeError("GPU meta is missing a params list") + if len(args) != len(params): + raise TypeError( + f"GPU kernel {meta.get('symbol', '?')!r}: expected {len(params)} " + f"args, got {len(args)}" + ) + return list(args) + + +def infer_group_count_x(meta: dict[str, Any], ordered_values: list[Any]) -> int: + """ + Choose `group_count_x` as ceil(n / local_size_x). + + Prefers a scalar parameter named `n`, else the longest list argument length. + + #### Args: + - meta: dict[str, Any] = kernel metadata + - ordered_values: list[Any] = launch args in param order + + #### Returns + - int = workgroup count in X (>= 1) + """ + local_size_x: int = int(meta.get("local_size_x") or 64) + if local_size_x < 1: + local_size_x = 64 + + params = meta.get("params") + if not isinstance(params, list): + return 1 + + n: int | None = None + for i, param in enumerate(params): + if not isinstance(param, dict): + continue + name = param.get("name") + kind = param.get("kind") + if name == "n" and kind == "int": + n = int(ordered_values[i]) + break + + if n is None: + best = 0 + for i, param in enumerate(params): + if isinstance(param, dict) and param.get("kind") == "list": + val = ordered_values[i] + if isinstance(val, list): + best = max(best, len(val)) + n = best + + if n is None or n < 1: + return 1 + return (int(n) + local_size_x - 1) // local_size_x diff --git a/src/cthreads/python/cthreads/gpu/runtime.py b/src/cthreads/python/cthreads/gpu/runtime.py new file mode 100644 index 0000000..0c98c92 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/runtime.py @@ -0,0 +1,160 @@ +""" +High-level GPU prepare + `gpu()` launch entry. + +Mirrors CPU `cthreads.prepare` / `thread`: compile registered `@Gpu` kernels, +then launch via `_ext.gpu.launch_gpu_kernel`. +""" + +from typing import Any, Callable + +from ..job import Job +from . import _ext_gpu_api +from .compiler.orchestrator import GpuCompileSession +from .frontend.errors import GPUNotAvailable, _map_error +from .gpu_kernel_meta import build_gpu_kernel_meta +from .gpu_marshal import infer_group_count_x, ordered_values_for_meta + +# True after a successful GpuCompileSession.compile in this process. +# Cleared by frontend shutdown() when native ShaderCache is released. +_gpu_prepared: bool = False + + +class GpuJob(Job): + """ + Job wrapper for native GpuJob handles (void kernels; `result()` is None). + """ + + def result(self) -> None: + """ + GPU kernels are writeback-only; there is no scalar return value. + + #### Returns + - None + """ + return None + + +def compile(force: bool = False) -> dict[str, Any]: + """ + Drain registered `@Gpu` functions through GpuCompileSession (SPIR-V emit). + + #### Args: + - force: bool = rewrite `__Gpu__` artifacts even when fingerprints match + + #### Returns + - dict[str, Any] = session result (`root`, `cache`, `rewritten`) + """ + global _gpu_prepared + info = GpuCompileSession.compile(force=force) + _gpu_prepared = True + return info + + +def prepare(force: bool = False) -> dict[str, Any]: + """ + Compile all registered `@Gpu` kernels and register SPIR-V in ShaderCache. + + #### Args: + - force: bool = if True, re-init Vulkan and force-rebuild GPU units + + #### Returns + - dict[str, Any] = compile session info + + #### Raises + - GPUNotAvailable = no usable Vulkan device / GPU extension + - RuntimeError = nothing registered or compile failed + """ + global _gpu_prepared + if not _ext_gpu_api.available(): + raise GPUNotAvailable( + "GPU is not available (build with CTHREADS_GPU=ON and a Vulkan device)" + ) + if force: + _ext_gpu_api.shutdown() + _ext_gpu_api.init() + _gpu_prepared = False + return compile(force=force) + + +def gpu( + fn: Callable[..., Any], + *args: Any, + force: bool = False, + **kwargs: Any, +) -> GpuJob: + """ + Launch a `@Gpu` kernel and return a joinable job handle. + + Ensures GPU compile/emit has run, then submits via `launch_gpu_kernel`. + List arguments are written back in place on `join()`. + + #### Args: + - fn: Callable = `@Gpu`-decorated kernel + - *args: Any = positional kernel arguments (param order) + - force: bool = force recompile + Vulkan re-init before launch + - **kwargs: Any = not supported yet + + #### Returns + - GpuJob = awaitable / joinable handle (result is always None) + + #### Raises + - TypeError = missing `@Gpu`, bad arity, or unexpected kwargs + - GPUNotAvailable = Vulkan path not usable + - Exception = mapped native launch failures + + #### Example: + ``py + from cthreads.gpu import Gpu, GlobalIdx, gpu + + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + x = [1.0, 2.0, 3.0, 4.0] + y = [10.0, 20.0, 30.0, 40.0] + gpu(saxpy, len(x), 2.0, x, y).join() + `` + """ + global _gpu_prepared + + if kwargs: + raise TypeError("gpu(): keyword arguments are not supported yet") + if not callable(fn): + raise TypeError("gpu(): fn must be callable") + if not getattr(fn, "__gpu__", False): + raise TypeError( + f"gpu(): {getattr(fn, '__qualname__', fn)!r} is not a @Gpu function" + ) + + if not _ext_gpu_api.available(): + raise GPUNotAvailable( + "GPU is not available (build with CTHREADS_GPU=ON and a Vulkan device)" + ) + + if force or not _gpu_prepared: + prepare(force=force) + + meta_obj = getattr(fn, "__gpu_kernel_meta__", None) + if not isinstance(meta_obj, dict): + meta_obj = build_gpu_kernel_meta(fn).to_dict() + meta: dict[str, Any] = dict(meta_obj) + + ordered = ordered_values_for_meta(meta, args) + if meta.get("group_count_x") is None: + meta["group_count_x"] = infer_group_count_x(meta, ordered) + if meta.get("group_count_y") is None: + meta["group_count_y"] = 1 + if meta.get("group_count_z") is None: + meta["group_count_z"] = 1 + + try: + raw = _ext_gpu_api.launch_gpu_kernel(meta, ordered) + except Exception as exc: + raise _map_error(exc) from exc + + job = GpuJob(raw) + job.start() + return job diff --git a/src/cthreads/python/cthreads/gpu/third_party_notices/README.md b/src/cthreads/python/cthreads/gpu/third_party_notices/README.md new file mode 100644 index 0000000..8dd0642 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/third_party_notices/README.md @@ -0,0 +1,12 @@ +# Third-party notices for native GLSL -> SPIR-V + +When `CTHREADS_GPU=ON`, cthreads links **Khronos glslang** into `_ext` so +`compile_glsl` works without installing `glslc` or other Vulkan SDK tools. + +glslang is the same compiler engine Google **shaderc** wraps. License texts +copied here at build time (see also the root `LICENSE` third-party section): + +- `glslang-LICENSE*` — Khronos glslang (Apache-2.0 / BSD-style components) + +End-user wheels that include the GPU extension must redistribute these notices +alongside the binary. diff --git a/src/cthreads/python/cthreads/gpu/third_party_notices/glslang-LICENSE.txt b/src/cthreads/python/cthreads/gpu/third_party_notices/glslang-LICENSE.txt new file mode 100644 index 0000000..054e68a --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/third_party_notices/glslang-LICENSE.txt @@ -0,0 +1,1016 @@ +Here, glslang proper means core GLSL parsing, HLSL parsing, and SPIR-V code +generation. Glslang proper requires use of a number of licenses, one that covers +preprocessing and others that covers non-preprocessing. + +Bison was removed long ago. You can build glslang from the source grammar, +using tools of your choice, without using bison or any bison files. + +Other parts, outside of glslang proper, include: + +- gl_types.h, only needed for OpenGL-like reflection, and can be left out of + a parse and codegen project. See it for its license. + +- update_glslang_sources.py, which is not part of the project proper and does + not need to be used. + +- the SPIR-V "remapper", which is optional, but has the same license as + glslang proper + +- Google tests and SPIR-V tools, and anything in the external subdirectory + are external and optional; see them for their respective licenses. + +-------------------------------------------------------------------------------- + +The core of glslang-proper, minus the preprocessor is licenced as follows: + +-------------------------------------------------------------------------------- +3-Clause BSD License +-------------------------------------------------------------------------------- + +// +// Copyright (C) 2015-2018 Google, Inc. +// Copyright (C) +// +// All rights reserved. +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions +// are met: +// +// Redistributions of source code must retain the above copyright +// notice, this list of conditions and the following disclaimer. +// +// Redistributions in binary form must reproduce the above +// copyright notice, this list of conditions and the following +// disclaimer in the documentation and/or other materials provided +// with the distribution. +// +// Neither the name of 3Dlabs Inc. Ltd. nor the names of its +// contributors may be used to endorse or promote products derived +// from this software without specific prior written permission. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS +// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT +// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS +// FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE +// COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, +// INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, +// BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +// LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN +// ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +// POSSIBILITY OF SUCH DAMAGE. +// + + +-------------------------------------------------------------------------------- +2-Clause BSD License +-------------------------------------------------------------------------------- + +Copyright 2020 The Khronos Group Inc + +Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +-------------------------------------------------------------------------------- +The MIT License +-------------------------------------------------------------------------------- + +Copyright 2020 The Khronos Group Inc + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. + + +-------------------------------------------------------------------------------- +APACHE LICENSE, VERSION 2.0 +-------------------------------------------------------------------------------- + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + +-------------------------------------------------------------------------------- +GPL 3 with special bison exception +-------------------------------------------------------------------------------- + + GNU GENERAL PUBLIC LICENSE + Version 3, 29 June 2007 + + Copyright (C) 2007 Free Software Foundation, Inc. + Everyone is permitted to copy and distribute verbatim copies + of this license document, but changing it is not allowed. + + Preamble + + The GNU General Public License is a free, copyleft license for +software and other kinds of works. + + The licenses for most software and other practical works are designed +to take away your freedom to share and change the works. By contrast, +the GNU General Public License is intended to guarantee your freedom to +share and change all versions of a program--to make sure it remains free +software for all its users. We, the Free Software Foundation, use the +GNU General Public License for most of our software; it applies also to +any other work released this way by its authors. You can apply it to +your programs, too. + + When we speak of free software, we are referring to freedom, not +price. Our General Public Licenses are designed to make sure that you +have the freedom to distribute copies of free software (and charge for +them if you wish), that you receive source code or can get it if you +want it, that you can change the software or use pieces of it in new +free programs, and that you know you can do these things. + + To protect your rights, we need to prevent others from denying you +these rights or asking you to surrender the rights. Therefore, you have +certain responsibilities if you distribute copies of the software, or if +you modify it: responsibilities to respect the freedom of others. + + For example, if you distribute copies of such a program, whether +gratis or for a fee, you must pass on to the recipients the same +freedoms that you received. You must make sure that they, too, receive +or can get the source code. And you must show them these terms so they +know their rights. + + Developers that use the GNU GPL protect your rights with two steps: +(1) assert copyright on the software, and (2) offer you this License +giving you legal permission to copy, distribute and/or modify it. + + For the developers' and authors' protection, the GPL clearly explains +that there is no warranty for this free software. For both users' and +authors' sake, the GPL requires that modified versions be marked as +changed, so that their problems will not be attributed erroneously to +authors of previous versions. + + Some devices are designed to deny users access to install or run +modified versions of the software inside them, although the manufacturer +can do so. This is fundamentally incompatible with the aim of +protecting users' freedom to change the software. The systematic +pattern of such abuse occurs in the area of products for individuals to +use, which is precisely where it is most unacceptable. Therefore, we +have designed this version of the GPL to prohibit the practice for those +products. If such problems arise substantially in other domains, we +stand ready to extend this provision to those domains in future versions +of the GPL, as needed to protect the freedom of users. + + Finally, every program is threatened constantly by software patents. +States should not allow patents to restrict development and use of +software on general-purpose computers, but in those that do, we wish to +avoid the special danger that patents applied to a free program could +make it effectively proprietary. To prevent this, the GPL assures that +patents cannot be used to render the program non-free. + + The precise terms and conditions for copying, distribution and +modification follow. + + TERMS AND CONDITIONS + + 0. Definitions. + + "This License" refers to version 3 of the GNU General Public License. + + "Copyright" also means copyright-like laws that apply to other kinds of +works, such as semiconductor masks. + + "The Program" refers to any copyrightable work licensed under this +License. Each licensee is addressed as "you". "Licensees" and +"recipients" may be individuals or organizations. + + To "modify" a work means to copy from or adapt all or part of the work +in a fashion requiring copyright permission, other than the making of an +exact copy. The resulting work is called a "modified version" of the +earlier work or a work "based on" the earlier work. + + A "covered work" means either the unmodified Program or a work based +on the Program. + + To "propagate" a work means to do anything with it that, without +permission, would make you directly or secondarily liable for +infringement under applicable copyright law, except executing it on a +computer or modifying a private copy. Propagation includes copying, +distribution (with or without modification), making available to the +public, and in some countries other activities as well. + + To "convey" a work means any kind of propagation that enables other +parties to make or receive copies. Mere interaction with a user through +a computer network, with no transfer of a copy, is not conveying. + + An interactive user interface displays "Appropriate Legal Notices" +to the extent that it includes a convenient and prominently visible +feature that (1) displays an appropriate copyright notice, and (2) +tells the user that there is no warranty for the work (except to the +extent that warranties are provided), that licensees may convey the +work under this License, and how to view a copy of this License. If +the interface presents a list of user commands or options, such as a +menu, a prominent item in the list meets this criterion. + + 1. Source Code. + + The "source code" for a work means the preferred form of the work +for making modifications to it. "Object code" means any non-source +form of a work. + + A "Standard Interface" means an interface that either is an official +standard defined by a recognized standards body, or, in the case of +interfaces specified for a particular programming language, one that +is widely used among developers working in that language. + + The "System Libraries" of an executable work include anything, other +than the work as a whole, that (a) is included in the normal form of +packaging a Major Component, but which is not part of that Major +Component, and (b) serves only to enable use of the work with that +Major Component, or to implement a Standard Interface for which an +implementation is available to the public in source code form. A +"Major Component", in this context, means a major essential component +(kernel, window system, and so on) of the specific operating system +(if any) on which the executable work runs, or a compiler used to +produce the work, or an object code interpreter used to run it. + + The "Corresponding Source" for a work in object code form means all +the source code needed to generate, install, and (for an executable +work) run the object code and to modify the work, including scripts to +control those activities. However, it does not include the work's +System Libraries, or general-purpose tools or generally available free +programs which are used unmodified in performing those activities but +which are not part of the work. For example, Corresponding Source +includes interface definition files associated with source files for +the work, and the source code for shared libraries and dynamically +linked subprograms that the work is specifically designed to require, +such as by intimate data communication or control flow between those +subprograms and other parts of the work. + + The Corresponding Source need not include anything that users +can regenerate automatically from other parts of the Corresponding +Source. + + The Corresponding Source for a work in source code form is that +same work. + + 2. Basic Permissions. + + All rights granted under this License are granted for the term of +copyright on the Program, and are irrevocable provided the stated +conditions are met. This License explicitly affirms your unlimited +permission to run the unmodified Program. The output from running a +covered work is covered by this License only if the output, given its +content, constitutes a covered work. This License acknowledges your +rights of fair use or other equivalent, as provided by copyright law. + + You may make, run and propagate covered works that you do not +convey, without conditions so long as your license otherwise remains +in force. You may convey covered works to others for the sole purpose +of having them make modifications exclusively for you, or provide you +with facilities for running those works, provided that you comply with +the terms of this License in conveying all material for which you do +not control copyright. Those thus making or running the covered works +for you must do so exclusively on your behalf, under your direction +and control, on terms that prohibit them from making any copies of +your copyrighted material outside their relationship with you. + + Conveying under any other circumstances is permitted solely under +the conditions stated below. Sublicensing is not allowed; section 10 +makes it unnecessary. + + 3. Protecting Users' Legal Rights From Anti-Circumvention Law. + + No covered work shall be deemed part of an effective technological +measure under any applicable law fulfilling obligations under article +11 of the WIPO copyright treaty adopted on 20 December 1996, or +similar laws prohibiting or restricting circumvention of such +measures. + + When you convey a covered work, you waive any legal power to forbid +circumvention of technological measures to the extent such circumvention +is effected by exercising rights under this License with respect to +the covered work, and you disclaim any intention to limit operation or +modification of the work as a means of enforcing, against the work's +users, your or third parties' legal rights to forbid circumvention of +technological measures. + + 4. Conveying Verbatim Copies. + + You may convey verbatim copies of the Program's source code as you +receive it, in any medium, provided that you conspicuously and +appropriately publish on each copy an appropriate copyright notice; +keep intact all notices stating that this License and any +non-permissive terms added in accord with section 7 apply to the code; +keep intact all notices of the absence of any warranty; and give all +recipients a copy of this License along with the Program. + + You may charge any price or no price for each copy that you convey, +and you may offer support or warranty protection for a fee. + + 5. Conveying Modified Source Versions. + + You may convey a work based on the Program, or the modifications to +produce it from the Program, in the form of source code under the +terms of section 4, provided that you also meet all of these conditions: + + a) The work must carry prominent notices stating that you modified + it, and giving a relevant date. + + b) The work must carry prominent notices stating that it is + released under this License and any conditions added under section + 7. This requirement modifies the requirement in section 4 to + "keep intact all notices". + + c) You must license the entire work, as a whole, under this + License to anyone who comes into possession of a copy. This + License will therefore apply, along with any applicable section 7 + additional terms, to the whole of the work, and all its parts, + regardless of how they are packaged. This License gives no + permission to license the work in any other way, but it does not + invalidate such permission if you have separately received it. + + d) If the work has interactive user interfaces, each must display + Appropriate Legal Notices; however, if the Program has interactive + interfaces that do not display Appropriate Legal Notices, your + work need not make them do so. + + A compilation of a covered work with other separate and independent +works, which are not by their nature extensions of the covered work, +and which are not combined with it such as to form a larger program, +in or on a volume of a storage or distribution medium, is called an +"aggregate" if the compilation and its resulting copyright are not +used to limit the access or legal rights of the compilation's users +beyond what the individual works permit. Inclusion of a covered work +in an aggregate does not cause this License to apply to the other +parts of the aggregate. + + 6. Conveying Non-Source Forms. + + You may convey a covered work in object code form under the terms +of sections 4 and 5, provided that you also convey the +machine-readable Corresponding Source under the terms of this License, +in one of these ways: + + a) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by the + Corresponding Source fixed on a durable physical medium + customarily used for software interchange. + + b) Convey the object code in, or embodied in, a physical product + (including a physical distribution medium), accompanied by a + written offer, valid for at least three years and valid for as + long as you offer spare parts or customer support for that product + model, to give anyone who possesses the object code either (1) a + copy of the Corresponding Source for all the software in the + product that is covered by this License, on a durable physical + medium customarily used for software interchange, for a price no + more than your reasonable cost of physically performing this + conveying of source, or (2) access to copy the + Corresponding Source from a network server at no charge. + + c) Convey individual copies of the object code with a copy of the + written offer to provide the Corresponding Source. This + alternative is allowed only occasionally and noncommercially, and + only if you received the object code with such an offer, in accord + with subsection 6b. + + d) Convey the object code by offering access from a designated + place (gratis or for a charge), and offer equivalent access to the + Corresponding Source in the same way through the same place at no + further charge. You need not require recipients to copy the + Corresponding Source along with the object code. If the place to + copy the object code is a network server, the Corresponding Source + may be on a different server (operated by you or a third party) + that supports equivalent copying facilities, provided you maintain + clear directions next to the object code saying where to find the + Corresponding Source. Regardless of what server hosts the + Corresponding Source, you remain obligated to ensure that it is + available for as long as needed to satisfy these requirements. + + e) Convey the object code using peer-to-peer transmission, provided + you inform other peers where the object code and Corresponding + Source of the work are being offered to the general public at no + charge under subsection 6d. + + A separable portion of the object code, whose source code is excluded +from the Corresponding Source as a System Library, need not be +included in conveying the object code work. + + A "User Product" is either (1) a "consumer product", which means any +tangible personal property which is normally used for personal, family, +or household purposes, or (2) anything designed or sold for incorporation +into a dwelling. In determining whether a product is a consumer product, +doubtful cases shall be resolved in favor of coverage. For a particular +product received by a particular user, "normally used" refers to a +typical or common use of that class of product, regardless of the status +of the particular user or of the way in which the particular user +actually uses, or expects or is expected to use, the product. A product +is a consumer product regardless of whether the product has substantial +commercial, industrial or non-consumer uses, unless such uses represent +the only significant mode of use of the product. + + "Installation Information" for a User Product means any methods, +procedures, authorization keys, or other information required to install +and execute modified versions of a covered work in that User Product from +a modified version of its Corresponding Source. The information must +suffice to ensure that the continued functioning of the modified object +code is in no case prevented or interfered with solely because +modification has been made. + + If you convey an object code work under this section in, or with, or +specifically for use in, a User Product, and the conveying occurs as +part of a transaction in which the right of possession and use of the +User Product is transferred to the recipient in perpetuity or for a +fixed term (regardless of how the transaction is characterized), the +Corresponding Source conveyed under this section must be accompanied +by the Installation Information. But this requirement does not apply +if neither you nor any third party retains the ability to install +modified object code on the User Product (for example, the work has +been installed in ROM). + + The requirement to provide Installation Information does not include a +requirement to continue to provide support service, warranty, or updates +for a work that has been modified or installed by the recipient, or for +the User Product in which it has been modified or installed. Access to a +network may be denied when the modification itself materially and +adversely affects the operation of the network or violates the rules and +protocols for communication across the network. + + Corresponding Source conveyed, and Installation Information provided, +in accord with this section must be in a format that is publicly +documented (and with an implementation available to the public in +source code form), and must require no special password or key for +unpacking, reading or copying. + + 7. Additional Terms. + + "Additional permissions" are terms that supplement the terms of this +License by making exceptions from one or more of its conditions. +Additional permissions that are applicable to the entire Program shall +be treated as though they were included in this License, to the extent +that they are valid under applicable law. If additional permissions +apply only to part of the Program, that part may be used separately +under those permissions, but the entire Program remains governed by +this License without regard to the additional permissions. + + When you convey a copy of a covered work, you may at your option +remove any additional permissions from that copy, or from any part of +it. (Additional permissions may be written to require their own +removal in certain cases when you modify the work.) You may place +additional permissions on material, added by you to a covered work, +for which you have or can give appropriate copyright permission. + + Notwithstanding any other provision of this License, for material you +add to a covered work, you may (if authorized by the copyright holders of +that material) supplement the terms of this License with terms: + + a) Disclaiming warranty or limiting liability differently from the + terms of sections 15 and 16 of this License; or + + b) Requiring preservation of specified reasonable legal notices or + author attributions in that material or in the Appropriate Legal + Notices displayed by works containing it; or + + c) Prohibiting misrepresentation of the origin of that material, or + requiring that modified versions of such material be marked in + reasonable ways as different from the original version; or + + d) Limiting the use for publicity purposes of names of licensors or + authors of the material; or + + e) Declining to grant rights under trademark law for use of some + trade names, trademarks, or service marks; or + + f) Requiring indemnification of licensors and authors of that + material by anyone who conveys the material (or modified versions of + it) with contractual assumptions of liability to the recipient, for + any liability that these contractual assumptions directly impose on + those licensors and authors. + + All other non-permissive additional terms are considered "further +restrictions" within the meaning of section 10. If the Program as you +received it, or any part of it, contains a notice stating that it is +governed by this License along with a term that is a further +restriction, you may remove that term. If a license document contains +a further restriction but permits relicensing or conveying under this +License, you may add to a covered work material governed by the terms +of that license document, provided that the further restriction does +not survive such relicensing or conveying. + + If you add terms to a covered work in accord with this section, you +must place, in the relevant source files, a statement of the +additional terms that apply to those files, or a notice indicating +where to find the applicable terms. + + Additional terms, permissive or non-permissive, may be stated in the +form of a separately written license, or stated as exceptions; +the above requirements apply either way. + + 8. Termination. + + You may not propagate or modify a covered work except as expressly +provided under this License. Any attempt otherwise to propagate or +modify it is void, and will automatically terminate your rights under +this License (including any patent licenses granted under the third +paragraph of section 11). + + However, if you cease all violation of this License, then your +license from a particular copyright holder is reinstated (a) +provisionally, unless and until the copyright holder explicitly and +finally terminates your license, and (b) permanently, if the copyright +holder fails to notify you of the violation by some reasonable means +prior to 60 days after the cessation. + + Moreover, your license from a particular copyright holder is +reinstated permanently if the copyright holder notifies you of the +violation by some reasonable means, this is the first time you have +received notice of violation of this License (for any work) from that +copyright holder, and you cure the violation prior to 30 days after +your receipt of the notice. + + Termination of your rights under this section does not terminate the +licenses of parties who have received copies or rights from you under +this License. If your rights have been terminated and not permanently +reinstated, you do not qualify to receive new licenses for the same +material under section 10. + + 9. Acceptance Not Required for Having Copies. + + You are not required to accept this License in order to receive or +run a copy of the Program. Ancillary propagation of a covered work +occurring solely as a consequence of using peer-to-peer transmission +to receive a copy likewise does not require acceptance. However, +nothing other than this License grants you permission to propagate or +modify any covered work. These actions infringe copyright if you do +not accept this License. Therefore, by modifying or propagating a +covered work, you indicate your acceptance of this License to do so. + + 10. Automatic Licensing of Downstream Recipients. + + Each time you convey a covered work, the recipient automatically +receives a license from the original licensors, to run, modify and +propagate that work, subject to this License. You are not responsible +for enforcing compliance by third parties with this License. + + An "entity transaction" is a transaction transferring control of an +organization, or substantially all assets of one, or subdividing an +organization, or merging organizations. If propagation of a covered +work results from an entity transaction, each party to that +transaction who receives a copy of the work also receives whatever +licenses to the work the party's predecessor in interest had or could +give under the previous paragraph, plus a right to possession of the +Corresponding Source of the work from the predecessor in interest, if +the predecessor has it or can get it with reasonable efforts. + + You may not impose any further restrictions on the exercise of the +rights granted or affirmed under this License. For example, you may +not impose a license fee, royalty, or other charge for exercise of +rights granted under this License, and you may not initiate litigation +(including a cross-claim or counterclaim in a lawsuit) alleging that +any patent claim is infringed by making, using, selling, offering for +sale, or importing the Program or any portion of it. + + 11. Patents. + + A "contributor" is a copyright holder who authorizes use under this +License of the Program or a work on which the Program is based. The +work thus licensed is called the contributor's "contributor version". + + A contributor's "essential patent claims" are all patent claims +owned or controlled by the contributor, whether already acquired or +hereafter acquired, that would be infringed by some manner, permitted +by this License, of making, using, or selling its contributor version, +but do not include claims that would be infringed only as a +consequence of further modification of the contributor version. For +purposes of this definition, "control" includes the right to grant +patent sublicenses in a manner consistent with the requirements of +this License. + + Each contributor grants you a non-exclusive, worldwide, royalty-free +patent license under the contributor's essential patent claims, to +make, use, sell, offer for sale, import and otherwise run, modify and +propagate the contents of its contributor version. + + In the following three paragraphs, a "patent license" is any express +agreement or commitment, however denominated, not to enforce a patent +(such as an express permission to practice a patent or covenant not to +sue for patent infringement). To "grant" such a patent license to a +party means to make such an agreement or commitment not to enforce a +patent against the party. + + If you convey a covered work, knowingly relying on a patent license, +and the Corresponding Source of the work is not available for anyone +to copy, free of charge and under the terms of this License, through a +publicly available network server or other readily accessible means, +then you must either (1) cause the Corresponding Source to be so +available, or (2) arrange to deprive yourself of the benefit of the +patent license for this particular work, or (3) arrange, in a manner +consistent with the requirements of this License, to extend the patent +license to downstream recipients. "Knowingly relying" means you have +actual knowledge that, but for the patent license, your conveying the +covered work in a country, or your recipient's use of the covered work +in a country, would infringe one or more identifiable patents in that +country that you have reason to believe are valid. + + If, pursuant to or in connection with a single transaction or +arrangement, you convey, or propagate by procuring conveyance of, a +covered work, and grant a patent license to some of the parties +receiving the covered work authorizing them to use, propagate, modify +or convey a specific copy of the covered work, then the patent license +you grant is automatically extended to all recipients of the covered +work and works based on it. + + A patent license is "discriminatory" if it does not include within +the scope of its coverage, prohibits the exercise of, or is +conditioned on the non-exercise of one or more of the rights that are +specifically granted under this License. You may not convey a covered +work if you are a party to an arrangement with a third party that is +in the business of distributing software, under which you make payment +to the third party based on the extent of your activity of conveying +the work, and under which the third party grants, to any of the +parties who would receive the covered work from you, a discriminatory +patent license (a) in connection with copies of the covered work +conveyed by you (or copies made from those copies), or (b) primarily +for and in connection with specific products or compilations that +contain the covered work, unless you entered into that arrangement, +or that patent license was granted, prior to 28 March 2007. + + Nothing in this License shall be construed as excluding or limiting +any implied license or other defenses to infringement that may +otherwise be available to you under applicable patent law. + + 12. No Surrender of Others' Freedom. + + If conditions are imposed on you (whether by court order, agreement or +otherwise) that contradict the conditions of this License, they do not +excuse you from the conditions of this License. If you cannot convey a +covered work so as to satisfy simultaneously your obligations under this +License and any other pertinent obligations, then as a consequence you may +not convey it at all. For example, if you agree to terms that obligate you +to collect a royalty for further conveying from those to whom you convey +the Program, the only way you could satisfy both those terms and this +License would be to refrain entirely from conveying the Program. + + 13. Use with the GNU Affero General Public License. + + Notwithstanding any other provision of this License, you have +permission to link or combine any covered work with a work licensed +under version 3 of the GNU Affero General Public License into a single +combined work, and to convey the resulting work. The terms of this +License will continue to apply to the part which is the covered work, +but the special requirements of the GNU Affero General Public License, +section 13, concerning interaction through a network will apply to the +combination as such. + + 14. Revised Versions of this License. + + The Free Software Foundation may publish revised and/or new versions of +the GNU General Public License from time to time. Such new versions will +be similar in spirit to the present version, but may differ in detail to +address new problems or concerns. + + Each version is given a distinguishing version number. If the +Program specifies that a certain numbered version of the GNU General +Public License "or any later version" applies to it, you have the +option of following the terms and conditions either of that numbered +version or of any later version published by the Free Software +Foundation. If the Program does not specify a version number of the +GNU General Public License, you may choose any version ever published +by the Free Software Foundation. + + If the Program specifies that a proxy can decide which future +versions of the GNU General Public License can be used, that proxy's +public statement of acceptance of a version permanently authorizes you +to choose that version for the Program. + + Later license versions may give you additional or different +permissions. However, no additional obligations are imposed on any +author or copyright holder as a result of your choosing to follow a +later version. + + 15. Disclaimer of Warranty. + + THERE IS NO WARRANTY FOR THE PROGRAM, TO THE EXTENT PERMITTED BY +APPLICABLE LAW. EXCEPT WHEN OTHERWISE STATED IN WRITING THE COPYRIGHT +HOLDERS AND/OR OTHER PARTIES PROVIDE THE PROGRAM "AS IS" WITHOUT WARRANTY +OF ANY KIND, EITHER EXPRESSED OR IMPLIED, INCLUDING, BUT NOT LIMITED TO, +THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR +PURPOSE. THE ENTIRE RISK AS TO THE QUALITY AND PERFORMANCE OF THE PROGRAM +IS WITH YOU. SHOULD THE PROGRAM PROVE DEFECTIVE, YOU ASSUME THE COST OF +ALL NECESSARY SERVICING, REPAIR OR CORRECTION. + + 16. Limitation of Liability. + + IN NO EVENT UNLESS REQUIRED BY APPLICABLE LAW OR AGREED TO IN WRITING +WILL ANY COPYRIGHT HOLDER, OR ANY OTHER PARTY WHO MODIFIES AND/OR CONVEYS +THE PROGRAM AS PERMITTED ABOVE, BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY +GENERAL, SPECIAL, INCIDENTAL OR CONSEQUENTIAL DAMAGES ARISING OUT OF THE +USE OR INABILITY TO USE THE PROGRAM (INCLUDING BUT NOT LIMITED TO LOSS OF +DATA OR DATA BEING RENDERED INACCURATE OR LOSSES SUSTAINED BY YOU OR THIRD +PARTIES OR A FAILURE OF THE PROGRAM TO OPERATE WITH ANY OTHER PROGRAMS), +EVEN IF SUCH HOLDER OR OTHER PARTY HAS BEEN ADVISED OF THE POSSIBILITY OF +SUCH DAMAGES. + + 17. Interpretation of Sections 15 and 16. + + If the disclaimer of warranty and limitation of liability provided +above cannot be given local legal effect according to their terms, +reviewing courts shall apply local law that most closely approximates +an absolute waiver of all civil liability in connection with the +Program, unless a warranty or assumption of liability accompanies a +copy of the Program in return for a fee. + +Bison Exception + +As a special exception, you may create a larger work that contains part or all +of the Bison parser skeleton and distribute that work under terms of your +choice, so long as that work isn't itself a parser generator using the skeleton +or a modified version thereof as a parser skeleton. Alternatively, if you +modify or redistribute the parser skeleton itself, you may (at your option) +remove this special exception, which will cause the skeleton and the resulting +Bison output files to be licensed under the GNU General Public License without +this special exception. + +This special exception was added by the Free Software Foundation in version +2.2 of Bison. + + END OF TERMS AND CONDITIONS + +-------------------------------------------------------------------------------- +================================================================================ +-------------------------------------------------------------------------------- + +The preprocessor has the core licenses stated above, plus additional licences: + +/****************************************************************************\ +Copyright (c) 2002, NVIDIA Corporation. + +NVIDIA Corporation("NVIDIA") supplies this software to you in +consideration of your agreement to the following terms, and your use, +installation, modification or redistribution of this NVIDIA software +constitutes acceptance of these terms. If you do not agree with these +terms, please do not use, install, modify or redistribute this NVIDIA +software. + +In consideration of your agreement to abide by the following terms, and +subject to these terms, NVIDIA grants you a personal, non-exclusive +license, under NVIDIA's copyrights in this original NVIDIA software (the +"NVIDIA Software"), to use, reproduce, modify and redistribute the +NVIDIA Software, with or without modifications, in source and/or binary +forms; provided that if you redistribute the NVIDIA Software, you must +retain the copyright notice of NVIDIA, this notice and the following +text and disclaimers in all such redistributions of the NVIDIA Software. +Neither the name, trademarks, service marks nor logos of NVIDIA +Corporation may be used to endorse or promote products derived from the +NVIDIA Software without specific prior written permission from NVIDIA. +Except as expressly stated in this notice, no other rights or licenses +express or implied, are granted by NVIDIA herein, including but not +limited to any patent rights that may be infringed by your derivative +works or by other works in which the NVIDIA Software may be +incorporated. No hardware is licensed hereunder. + +THE NVIDIA SOFTWARE IS BEING PROVIDED ON AN "AS IS" BASIS, WITHOUT +WARRANTIES OR CONDITIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED, +INCLUDING WITHOUT LIMITATION, WARRANTIES OR CONDITIONS OF TITLE, +NON-INFRINGEMENT, MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, OR +ITS USE AND OPERATION EITHER ALONE OR IN COMBINATION WITH OTHER +PRODUCTS. + +IN NO EVENT SHALL NVIDIA BE LIABLE FOR ANY SPECIAL, INDIRECT, +INCIDENTAL, EXEMPLARY, CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED +TO, LOST PROFITS; PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF +USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) OR ARISING IN ANY WAY +OUT OF THE USE, REPRODUCTION, MODIFICATION AND/OR DISTRIBUTION OF THE +NVIDIA SOFTWARE, HOWEVER CAUSED AND WHETHER UNDER THEORY OF CONTRACT, +TORT (INCLUDING NEGLIGENCE), STRICT LIABILITY OR OTHERWISE, EVEN IF +NVIDIA HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +\****************************************************************************/ + +/* +** Copyright (c) 2014-2016 The Khronos Group Inc. +** +** Permission is hereby granted, free of charge, to any person obtaining a copy +** of this software and/or associated documentation files (the "Materials"), +** to deal in the Materials without restriction, including without limitation +** the rights to use, copy, modify, merge, publish, distribute, sublicense, +** and/or sell copies of the Materials, and to permit persons to whom the +** Materials are furnished to do so, subject to the following conditions: +** +** The above copyright notice and this permission notice shall be included in +** all copies or substantial portions of the Materials. +** +** MODIFICATIONS TO THIS FILE MAY MEAN IT NO LONGER ACCURATELY REFLECTS KHRONOS +** STANDARDS. THE UNMODIFIED, NORMATIVE VERSIONS OF KHRONOS SPECIFICATIONS AND +** HEADER INFORMATION ARE LOCATED AT https://www.khronos.org/registry/ +** +** THE MATERIALS ARE PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +** OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +** FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +** THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +** LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +** FROM,OUT OF OR IN CONNECTION WITH THE MATERIALS OR THE USE OR OTHER DEALINGS +** IN THE MATERIALS. +*/ diff --git a/tests/helpers_gpu.py b/tests/helpers_gpu.py new file mode 100644 index 0000000..ee5d5ca --- /dev/null +++ b/tests/helpers_gpu.py @@ -0,0 +1,16 @@ +""" +Shared helpers for GPU unit / pipeline tests. +""" + +from __future__ import annotations + +from types import ModuleType + +import cthreads.gpu.runtime as runtime_mod + + +def prepare_module() -> ModuleType: + """ + Return `cthreads.gpu.runtime` (holds `_gpu_prepared` / prepare / gpu). + """ + return runtime_mod diff --git a/tests/unit/.gitignore b/tests/unit/.gitignore new file mode 100644 index 0000000..0bf994e --- /dev/null +++ b/tests/unit/.gitignore @@ -0,0 +1,11 @@ +# >>> cthreads (auto) +__Thread__/ +__Threadable__/ +__Gpu__/ +.cthreads_cache.json +cthreads_kernels.dll +cthreads_kernels.so +cthreads_kernels.lib +libcthreads_kernels.so +libcthreads_kernels.dylib +# <<< cthreads (auto) diff --git a/tests/unit/test_cache.py b/tests/unit/test_cache.py index 1ea9a22..303c16f 100644 --- a/tests/unit/test_cache.py +++ b/tests/unit/test_cache.py @@ -47,12 +47,14 @@ def test_write_if_changed_writes_then_skips(tmp_path: Path): def test_load_cache_missing_and_corrupt(tmp_path: Path): empty = load_cache(tmp_path) assert empty["units"] == {} + assert empty["gpu_units"] == {} assert empty["link_hash"] is None bad = tmp_path / CACHE_FILENAME bad.write_text("{not json", encoding="utf-8") recovered = load_cache(tmp_path) assert recovered["units"] == {} + assert recovered["gpu_units"] == {} def test_load_cache_wrong_version(tmp_path: Path): @@ -63,12 +65,39 @@ def test_load_cache_wrong_version(tmp_path: Path): ) data = load_cache(tmp_path) assert data["units"] == {} + assert data["gpu_units"] == {} + + +def test_load_cache_fills_missing_gpu_units(tmp_path: Path): + path = cache_path_for_root(tmp_path) + path.write_text( + json.dumps( + { + "version": REGISTRY.VERSION, + "units": {"move": {"src_hash": "abc"}}, + "link_hash": None, + "binary": None, + } + ), + encoding="utf-8", + ) + data = load_cache(tmp_path) + assert data["units"]["move"]["src_hash"] == "abc" + assert data["gpu_units"] == {} def test_save_and_load_cache_roundtrip(tmp_path: Path): - save_cache(tmp_path, {"units": {"move": {"src_hash": "abc"}}, "link_hash": "L"}) + save_cache( + tmp_path, + { + "units": {"move": {"src_hash": "abc"}}, + "gpu_units": {"saxpy": {"src_hash": "def"}}, + "link_hash": "L", + }, + ) data = load_cache(tmp_path) assert data["units"]["move"]["src_hash"] == "abc" + assert data["gpu_units"]["saxpy"]["src_hash"] == "def" assert data["link_hash"] == "L" assert data["version"] == REGISTRY.VERSION @@ -89,6 +118,7 @@ def test_ensure_gitignore_creates_and_is_idempotent(tmp_path: Path): assert ensure_gitignore(tmp_path) is True text = (tmp_path / ".gitignore").read_text(encoding="utf-8") assert "__Thread__/" in text + assert "__Gpu__/" in text assert ".cthreads_cache.json" in text assert "cthreads_kernels.dll" in text assert ensure_gitignore(tmp_path) is False diff --git a/tests/unit/test_gpu_assemble.py b/tests/unit/test_gpu_assemble.py new file mode 100644 index 0000000..50ab1b2 --- /dev/null +++ b/tests/unit/test_gpu_assemble.py @@ -0,0 +1,50 @@ +"""Unit tests for assemble_comp.""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu.compiler.translation.assemble import assemble_comp + + +def test_assemble_wraps_version_and_main(): + pre = "layout(local_size_x = 64) in;\n" + body = " int i = 0;\n" + src = assemble_comp(pre, body) + assert src.startswith("#version 450\n") + assert "layout(local_size_x = 64) in;" in src + assert "void main() {" in src + assert " int i = 0;" in src + assert src.rstrip().endswith("}") + + +@pytest.mark.parametrize("ver", [450, 460, 430]) +def test_assemble_custom_version(ver): + src = assemble_comp("", " return;\n", version=ver) + assert src.startswith(f"#version {ver}\n") + + +def test_assemble_empty_body(): + src = assemble_comp("layout(local_size_x = 1) in;\n", "") + assert "void main() {" in src + assert src.rstrip().endswith("}") + + +def test_assemble_strips_trailing_whitespace_on_parts(): + src = assemble_comp("preamble\n\n", " body;\n\n") + assert "\n\n\nvoid main" not in src + assert "void main() {" in src + + +def test_assemble_multiline_body_preserved(): + body = " int i = 0;\n i += 1;\n return;\n" + src = assemble_comp("layout(local_size_x = 1) in;", body) + assert "int i = 0;" in src + assert "i += 1;" in src + assert "return;" in src + + +def test_assemble_order_version_preamble_main(): + src = assemble_comp("AAA\n", " BBB;\n") + assert src.index("#version") < src.index("AAA") < src.index("void main") + assert src.index("void main") < src.index("BBB") diff --git a/tests/unit/test_gpu_context.py b/tests/unit/test_gpu_context.py index f29371c..4a7fe53 100644 --- a/tests/unit/test_gpu_context.py +++ b/tests/unit/test_gpu_context.py @@ -11,7 +11,7 @@ import pytest from cthreads import gpu -from cthreads.gpu.errors import ( +from cthreads.gpu.frontend.errors import ( CThreadsGPUError, GPUNotAvailable, GpuInvalidArgument, @@ -82,24 +82,24 @@ def test_gpu_module_exports(): def test_not_built_available_false(monkeypatch): - monkeypatch.setattr(gpu, "_gpu", None) + monkeypatch.setattr(gpu._ext_gpu_api, "_gpu", None) assert gpu.available() is False def test_not_built_device_name_raises(monkeypatch): - monkeypatch.setattr(gpu, "_gpu", None) + monkeypatch.setattr(gpu._ext_gpu_api, "_gpu", None) with pytest.raises(VulkanNotBuiltError, match="CTHREADS_GPU"): gpu.device_name() def test_not_built_init_raises(monkeypatch): - monkeypatch.setattr(gpu, "_gpu", None) + monkeypatch.setattr(gpu._ext_gpu_api, "_gpu", None) with pytest.raises(VulkanNotBuiltError, match="CTHREADS_GPU"): gpu.init() def test_not_built_shutdown_noop(monkeypatch): - monkeypatch.setattr(gpu, "_gpu", None) + monkeypatch.setattr(gpu._ext_gpu_api, "_gpu", None) gpu.shutdown() # must not raise @@ -148,7 +148,7 @@ def shutdown(self) -> None: ], ) def test_map_error_via_device_name(monkeypatch, msg, exc_type): - monkeypatch.setattr(gpu, "_gpu", _FakeGpu(RuntimeError(msg))) + monkeypatch.setattr(gpu._ext_gpu_api, "_gpu", _FakeGpu(RuntimeError(msg))) with pytest.raises(exc_type) as ei: gpu.device_name() assert msg in str(ei.value) @@ -156,7 +156,7 @@ def test_map_error_via_device_name(monkeypatch, msg, exc_type): def test_map_error_via_init(monkeypatch): monkeypatch.setattr( - gpu, + gpu._ext_gpu_api, "_gpu", _FakeGpu(RuntimeError("cthreads.gpu.VulkanLoaderNotFound: missing")), ) @@ -166,17 +166,17 @@ def test_map_error_via_init(monkeypatch): def test_fake_available_false_without_raising(monkeypatch): """Mirrors C++ available(): False when init would fail, no exception.""" - monkeypatch.setattr(gpu, "_gpu", _FakeGpu(ready=False)) + monkeypatch.setattr(gpu._ext_gpu_api, "_gpu", _FakeGpu(ready=False)) assert gpu.available() is False def test_fake_ready_device_name(monkeypatch): - monkeypatch.setattr(gpu, "_gpu", _FakeGpu(ready=True)) + monkeypatch.setattr(gpu._ext_gpu_api, "_gpu", _FakeGpu(ready=True)) assert gpu.available() is True assert gpu.device_name() == "FakeGPU" gpu.init() gpu.shutdown() - assert gpu._gpu.shutdown_calls == 1 + assert gpu._ext_gpu_api._gpu.shutdown_calls == 1 # --------------------------------------------------------------------------- @@ -185,7 +185,7 @@ def test_fake_ready_device_name(monkeypatch): def _ext_gpu_built() -> bool: - return gpu._gpu is not None + return gpu._ext_gpu_api._gpu is not None def _require_gpu(): diff --git a/tests/unit/test_gpu_decorator.py b/tests/unit/test_gpu_decorator.py new file mode 100644 index 0000000..6232105 --- /dev/null +++ b/tests/unit/test_gpu_decorator.py @@ -0,0 +1,101 @@ +"""@Gpu decorator registration (no launch / emit).""" + +from __future__ import annotations + +import pytest + +from cthreads.frontend.Registry import REGISTRY +from cthreads.gpu.frontend import Gpu +from cthreads.gpu.frontend.errors import GPUNotAvailable +from cthreads.gpu import frontend as gpu_frontend + + +@pytest.fixture(autouse=True) +def _clear_registry(): + REGISTRY.clear() + yield + REGISTRY.clear() + + +def test_gpu_decorator_requires_available(monkeypatch): + monkeypatch.setattr(gpu_frontend.wrapper, "available", lambda: False) + + with pytest.raises(GPUNotAvailable): + + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + + +def test_gpu_decorator_registers_and_attaches_meta(monkeypatch): + monkeypatch.setattr(gpu_frontend.wrapper, "available", lambda: True) + + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + + assert saxpy.__gpu__ is True + assert saxpy.__gpu_version__ == REGISTRY.VERSION + assert saxpy.__qualname__ in REGISTRY.gpu_functions + meta = saxpy.__gpu_kernel_meta__ + assert meta["symbol"] == "saxpy" + assert meta["binding_count"] == 3 + assert meta["scalar_bytes"] == 8 + # Still a normal Python function. + assert saxpy(4, 2.0, [1.0], [0.0]) is None + + +def test_gpu_decorator_log_factory(monkeypatch, capsys): + monkeypatch.setattr(gpu_frontend.wrapper, "available", lambda: True) + monkeypatch.setattr(gpu_frontend.wrapper, "device_name", lambda: "FakeGPU") + + @Gpu(log=True) + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + + assert saxpy.__gpu__ is True + assert "FakeGPU" in capsys.readouterr().out + + +def test_gpu_decorator_rejects_non_none_return(monkeypatch): + monkeypatch.setattr(gpu_frontend.wrapper, "available", lambda: True) + + with pytest.raises(TypeError, match="return must be None"): + + @Gpu + def bad(x: int) -> int: + return x + + +def test_gpu_decorator_rejects_dict(monkeypatch): + monkeypatch.setattr(gpu_frontend.wrapper, "available", lambda: True) + + with pytest.raises(TypeError): + + @Gpu + def bad(d: dict[str, int]) -> None: + pass + + +def test_gpu_decorator_rejects_varargs(monkeypatch): + monkeypatch.setattr(gpu_frontend.wrapper, "available", lambda: True) + + with pytest.raises(TypeError): + + @Gpu + def bad(*args: int) -> None: + pass + + +def test_gpu_decorator_bool_and_int_list(monkeypatch): + monkeypatch.setattr(gpu_frontend.wrapper, "available", lambda: True) + + @Gpu + def k(n: int, flag: bool, xs: list[int], ys: list[bool]) -> None: + pass + + meta = k.__gpu_kernel_meta__ + assert meta["binding_count"] == 3 + assert meta["scalar_bytes"] == 8 + kinds = [p["kind"] for p in meta["params"]] + assert kinds == ["int", "bool", "list", "list"] diff --git a/tests/unit/test_gpu_errors.py b/tests/unit/test_gpu_errors.py new file mode 100644 index 0000000..2e5c5c8 --- /dev/null +++ b/tests/unit/test_gpu_errors.py @@ -0,0 +1,51 @@ +"""Unit tests for GPU error mapping and types.""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu.frontend.errors import ( + CThreadsGPUError, + GPUNotAvailable, + GpuInvalidArgument, + GpuUseAfterDestroy, + VulkanInitFailed, + VulkanLoaderNotFound, + VulkanNoDevice, + VulkanNotBuiltError, + VulkanOutOfMemory, + _map_error, +) + + +@pytest.mark.parametrize( + "msg, cls", + [ + ("VulkanLoaderNotFound: missing", VulkanLoaderNotFound), + ("VulkanNoDevice: none", VulkanNoDevice), + ("VulkanOutOfMemory: oom", VulkanOutOfMemory), + ("GpuUseAfterDestroy: gone", GpuUseAfterDestroy), + ("GpuInvalidArgument: bad", GpuInvalidArgument), + ("VulkanNotBuilt: off", VulkanNotBuiltError), + ("VulkanInitFailed: boom", VulkanInitFailed), + ("something else entirely", VulkanInitFailed), + ], +) +def test_map_error_prefixes(msg, cls): + out = _map_error(RuntimeError(msg)) + assert isinstance(out, cls) + assert isinstance(out, CThreadsGPUError) + assert msg in str(out) + + +def test_error_str_contains_prefix(): + err = GPUNotAvailable("no gpu") + assert "cthreads gpu error" in str(err) + assert "no gpu" in str(err) + assert err.detail == "no gpu" + + +def test_default_details(): + assert "Vulkan loader not found" in VulkanLoaderNotFound().detail + assert "No Vulkan compute device" in VulkanNoDevice().detail + assert "out of memory" in VulkanOutOfMemory().detail.lower() diff --git a/tests/unit/test_gpu_glsl.py b/tests/unit/test_gpu_glsl.py new file mode 100644 index 0000000..3d75ab1 --- /dev/null +++ b/tests/unit/test_gpu_glsl.py @@ -0,0 +1,99 @@ +"""Unit tests for GPU Glsl type helpers.""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu.compiler.translation.Glsl import ( + align_bytes, + elem_type_name, + size_bytes, + type_name, +) +from cthreads.types import PyBool, PyDict, PyFloat, PyInt, PyList, PyString, PyThreadable + + +@pytest.mark.parametrize( + "py, name, nbytes", + [ + (PyInt(), "int", 4), + (PyFloat(), "float", 4), + (PyBool(), "bool", 4), + ], +) +def test_scalar_type_name_and_size(py, name, nbytes): + assert type_name(py) == name + assert size_bytes(py) == nbytes + assert align_bytes(py) == nbytes + + +def test_float_is_glsl_float_not_double(): + assert type_name(PyFloat()) == "float" + assert size_bytes(PyFloat()) == 4 + + +@pytest.mark.parametrize( + "inner, elem", + [ + (PyInt(), "int"), + (PyFloat(), "float"), + (PyBool(), "bool"), + ], +) +def test_elem_type_name(inner, elem): + assert elem_type_name(PyList(inner)) == elem + + +def test_type_name_rejects_list(): + with pytest.raises(TypeError, match="list is not a scalar"): + type_name(PyList(PyFloat())) + + +def test_size_bytes_rejects_list(): + with pytest.raises(TypeError, match="list has no single scalar size"): + size_bytes(PyList(PyInt())) + + +def test_align_bytes_rejects_list(): + with pytest.raises(TypeError, match="list has no single scalar size"): + align_bytes(PyList(PyBool())) + + +def test_elem_type_name_rejects_non_list(): + with pytest.raises(TypeError, match="expects PyList"): + elem_type_name(PyInt()) # type: ignore[arg-type] + + +@pytest.mark.parametrize( + "bad", + [ + PyString(), + PyDict(PyString(), PyInt()), + PyThreadable("T"), + ], +) +def test_unsupported_scalar_types(bad): + with pytest.raises(TypeError, match="not supported"): + type_name(bad) + + +@pytest.mark.parametrize( + "bad", + [ + PyString(), + PyDict(PyString(), PyInt()), + PyThreadable("T"), + ], +) +def test_unsupported_size_bytes(bad): + with pytest.raises(TypeError, match="unsupported|not supported"): + size_bytes(bad) + + +def test_nested_list_elem_rejected(): + with pytest.raises(TypeError): + elem_type_name(PyList(PyList(PyInt()))) + + +def test_type_name_inner_of_list_ok(): + assert type_name(PyList(PyFloat()).inner_type) == "float" diff --git a/tests/unit/test_gpu_indexes.py b/tests/unit/test_gpu_indexes.py new file mode 100644 index 0000000..3d683af --- /dev/null +++ b/tests/unit/test_gpu_indexes.py @@ -0,0 +1,47 @@ +"""Unit tests for GPU index frontend markers.""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu import BlockDim, BlockIdx, GlobalIdx, GridDim, ThreadIdx +from cthreads.gpu.frontend.indexes import GpuIndexBuiltin, _Axis + + +@pytest.mark.parametrize( + "cls, base", + [ + (ThreadIdx, "gl_LocalInvocationID"), + (BlockIdx, "gl_WorkGroupID"), + (BlockDim, "gl_WorkGroupSize"), + (GridDim, "gl_NumWorkGroups"), + (GlobalIdx, "gl_GlobalInvocationID"), + ], +) +def test_index_classes_glsl_base(cls, base): + assert issubclass(cls, GpuIndexBuiltin) + assert cls._glsl_base == base + assert isinstance(cls.x, _Axis) + assert isinstance(cls.y, _Axis) + assert isinstance(cls.z, _Axis) + + +@pytest.mark.parametrize( + "cls", + [ThreadIdx, BlockIdx, BlockDim, GridDim, GlobalIdx], +) +@pytest.mark.parametrize("axis", ["x", "y", "z"]) +def test_axis_repr(cls, axis): + obj = getattr(cls, axis) + assert isinstance(obj, _Axis) + assert f"{cls.__name__}.{axis}" in repr(obj) + + +def test_indexes_reexported_from_package(): + import cthreads.gpu as g + + assert g.GlobalIdx is GlobalIdx + assert g.ThreadIdx is ThreadIdx + assert g.BlockIdx is BlockIdx + assert g.BlockDim is BlockDim + assert g.GridDim is GridDim diff --git a/tests/unit/test_gpu_kernel_meta.py b/tests/unit/test_gpu_kernel_meta.py new file mode 100644 index 0000000..dbf37bd --- /dev/null +++ b/tests/unit/test_gpu_kernel_meta.py @@ -0,0 +1,190 @@ +"""Unit tests for build_gpu_kernel_meta and schemas.""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu.gpu_kernel_meta import ( + GPU_KERNELS, + GpuParamMeta, + GpuTypeSchema, + build_gpu_kernel_meta, + pytype_to_gpu_schema, +) +from cthreads.types import PyBool, PyDict, PyFloat, PyInt, PyList, PyString + + +@pytest.fixture(autouse=True) +def _clear_kernels(): + GPU_KERNELS.clear() + yield + GPU_KERNELS.clear() + + +def test_saxpy_meta_shape(): + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + + meta = build_gpu_kernel_meta(saxpy) + assert meta.symbol == "saxpy" + assert meta.binding_count == 3 + assert meta.scalar_bytes == 8 + assert meta.local_size_x == 64 + assert [p.name for p in meta.params] == ["n", "a", "x", "y"] + assert meta.params[0].pass_as == "value" + assert meta.params[2].pass_as == "ref" + assert meta.params[2].elem_kind == "float" + assert meta.params[2].elem_bytes == 4 + d = meta.to_dict() + assert d["symbol"] == "saxpy" + assert saxpy.__gpu_kernel_meta__["binding_count"] == 3 # type: ignore[attr-defined] + assert GPU_KERNELS["saxpy"] is meta + + +def test_custom_symbol_and_local_size(): + def k(n: int) -> None: + pass + + meta = build_gpu_kernel_meta(k, symbol="my_k", local_size_x=32) + assert meta.symbol == "my_k" + assert meta.local_size_x == 32 + assert "my_k" in GPU_KERNELS + + +@pytest.mark.parametrize( + "fn_src, binding, scalar", + [ + ("def k(a: int) -> None: pass", 1, 4), + ("def k(a: float, b: bool) -> None: pass", 1, 8), + ("def k(x: list[int]) -> None: pass", 1, 0), + ("def k(x: list[float], y: list[bool]) -> None: pass", 2, 0), + ("def k(n: int, x: list[int]) -> None: pass", 2, 4), + ], +) +def test_binding_and_scalar_bytes(fn_src, binding, scalar): + ns: dict = {} + exec(fn_src, ns) + meta = build_gpu_kernel_meta(ns["k"]) + assert meta.binding_count == binding + assert meta.scalar_bytes == scalar + + +def test_rejects_return_int(): + def k(n: int) -> int: + return n + + with pytest.raises(TypeError, match="return must be None"): + build_gpu_kernel_meta(k) + + +def test_rejects_missing_annotation(): + def k(n: int, a) -> None: # type: ignore[no-untyped-def] + pass + + with pytest.raises(TypeError, match="type annotation"): + build_gpu_kernel_meta(k) + + +def test_rejects_varargs(): + def k(*args: int) -> None: + pass + + with pytest.raises(TypeError, match=r"\*args"): + build_gpu_kernel_meta(k) + + +def test_rejects_kwargs(): + def k(**kwargs: int) -> None: + pass + + with pytest.raises(TypeError, match=r"\*\*kwargs"): + build_gpu_kernel_meta(k) + + +def test_rejects_kwonly(): + def k(*, n: int) -> None: + pass + + with pytest.raises(TypeError, match="keyword-only"): + build_gpu_kernel_meta(k) + + +def test_rejects_str_param(): + def k(x: str) -> None: + pass + + with pytest.raises(TypeError): + build_gpu_kernel_meta(k) + + +def test_rejects_dict_param(): + def k(x: dict[str, int]) -> None: + pass + + with pytest.raises(TypeError): + build_gpu_kernel_meta(k) + + +def test_rejects_list_str(): + def k(x: list[str]) -> None: + pass + + with pytest.raises(TypeError): + build_gpu_kernel_meta(k) + + +def test_rejects_nested_list(): + def k(x: list[list[int]]) -> None: + pass + + with pytest.raises(TypeError): + build_gpu_kernel_meta(k) + + +def test_pytype_to_gpu_schema_scalars(): + assert pytype_to_gpu_schema(PyInt()).kind == "int" + assert pytype_to_gpu_schema(PyFloat()).spirv_type == "float" + assert pytype_to_gpu_schema(PyBool()).elem_bytes == 4 + + +def test_pytype_to_gpu_schema_list(): + s = pytype_to_gpu_schema(PyList(PyFloat())) + assert s.kind == "list" + assert s.spirv_type == "float[]" + assert s.inner is not None + assert s.inner.kind == "float" + + +def test_pytype_rejects_dict_string(): + with pytest.raises(TypeError): + pytype_to_gpu_schema(PyDict(PyString(), PyInt())) + with pytest.raises(TypeError): + pytype_to_gpu_schema(PyString()) + + +def test_param_meta_pass_as_invalid(): + with pytest.raises(TypeError, match="pass_as"): + GpuParamMeta( + name="x", + pass_as="inout", # type: ignore[arg-type] + schema=GpuTypeSchema(kind="int", spirv_type="int", elem_bytes=4), + ) + + +def test_list_schema_to_dict_requires_inner(): + bad = GpuTypeSchema(kind="list", spirv_type="int[]", elem_bytes=4, inner=None) + with pytest.raises(ValueError, match="inner"): + bad.to_dict() + + +def test_param_to_dict_list_flattens_elem(): + p = GpuParamMeta( + name="xs", + pass_as="ref", + schema=pytype_to_gpu_schema(PyList(PyInt())), + ) + d = p.to_dict() + assert d["kind"] == "list" + assert d["elem_kind"] == "int" + assert d["elem_bytes"] == 4 + assert "schema" in d diff --git a/tests/unit/test_gpu_marshal.py b/tests/unit/test_gpu_marshal.py new file mode 100644 index 0000000..89311f6 --- /dev/null +++ b/tests/unit/test_gpu_marshal.py @@ -0,0 +1,153 @@ +"""Unit tests for gpu_marshal helpers.""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu.gpu_kernel_meta import build_gpu_kernel_meta +from cthreads.gpu.gpu_marshal import infer_group_count_x, ordered_values_for_meta + + +def _saxpy_meta(): + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + + return build_gpu_kernel_meta(saxpy).to_dict() + + +def test_ordered_values_ok(): + meta = _saxpy_meta() + args = (4, 2.0, [1.0], [0.0]) + out = ordered_values_for_meta(meta, args) + assert out == [4, 2.0, [1.0], [0.0]] + assert out[2] is args[2] + + +def test_ordered_values_arity(): + meta = _saxpy_meta() + with pytest.raises(TypeError, match="expected 4"): + ordered_values_for_meta(meta, (1, 2.0)) + + +@pytest.mark.parametrize("extra", [0, 1, 5, 10]) +def test_ordered_values_too_many(extra): + meta = _saxpy_meta() + args = (1, 2.0, [0.0], [0.0]) + (0,) * extra + if extra == 0: + assert ordered_values_for_meta(meta, args) == list(args) + else: + with pytest.raises(TypeError, match="expected 4"): + ordered_values_for_meta(meta, args) + + +def test_ordered_values_missing_params(): + with pytest.raises(TypeError, match="params"): + ordered_values_for_meta({"symbol": "k"}, (1,)) + + +def test_ordered_values_params_not_list(): + with pytest.raises(TypeError, match="params"): + ordered_values_for_meta({"params": "bad"}, (1,)) + + +@pytest.mark.parametrize( + "n, local, expect", + [ + (1, 64, 1), + (63, 64, 1), + (64, 64, 1), + (65, 64, 2), + (128, 64, 2), + (129, 64, 3), + (130, 64, 3), + (200, 64, 4), + (1, 1, 1), + (10, 1, 10), + (10, 8, 2), + (16, 8, 2), + (17, 8, 3), + ], +) +def test_infer_group_count_table(n, local, expect): + meta = _saxpy_meta() + meta["local_size_x"] = local + args = [n, 1.0, [0.0] * n, [0.0] * n] + assert infer_group_count_x(meta, args) == expect + + +def test_infer_group_count_from_n(): + meta = _saxpy_meta() + meta["local_size_x"] = 64 + args = [130, 1.0, [0.0] * 130, [0.0] * 130] + assert infer_group_count_x(meta, args) == 3 + + +def test_infer_group_count_exact_multiple(): + meta = _saxpy_meta() + meta["local_size_x"] = 64 + args = [128, 1.0, [0.0] * 128, [0.0] * 128] + assert infer_group_count_x(meta, args) == 2 + + +def test_infer_group_count_from_list_len_without_n(): + def k(x: list[float], y: list[float]) -> None: + pass + + meta = build_gpu_kernel_meta(k).to_dict() + meta["local_size_x"] = 64 + assert infer_group_count_x(meta, [[0.0] * 100, [0.0] * 50]) == 2 + + +def test_infer_group_count_empty_defaults_one(): + def k(a: float) -> None: + pass + + meta = build_gpu_kernel_meta(k).to_dict() + assert infer_group_count_x(meta, [1.0]) == 1 + + +def test_infer_group_count_n_zero_defaults_one(): + meta = _saxpy_meta() + assert infer_group_count_x(meta, [0, 1.0, [], []]) == 1 + + +def test_infer_group_count_n_negative_defaults_one(): + meta = _saxpy_meta() + assert infer_group_count_x(meta, [-5, 1.0, [], []]) == 1 + + +def test_infer_group_count_bad_local_size_falls_back_64(): + meta = _saxpy_meta() + meta["local_size_x"] = 0 + # n=130 / 64 => 3 + assert infer_group_count_x(meta, [130, 1.0, [0.0] * 130, [0.0] * 130]) == 3 + + +def test_infer_group_count_missing_local_size_defaults_64(): + meta = _saxpy_meta() + meta.pop("local_size_x", None) + assert infer_group_count_x(meta, [64, 1.0, [0.0] * 64, [0.0] * 64]) == 1 + + +def test_infer_group_count_ignores_non_n_int_named_differently(): + def k(count: int, x: list[float]) -> None: + pass + + meta = build_gpu_kernel_meta(k).to_dict() + meta["local_size_x"] = 64 + # No param named `n` -> use longest list (100) => ceil(100/64)=2 + # even though count is 1 + assert infer_group_count_x(meta, [1, [0.0] * 100]) == 2 + + +def test_infer_group_count_no_params_returns_one(): + assert infer_group_count_x({"local_size_x": 64}, []) == 1 + + +def test_infer_group_count_skips_non_dict_params(): + meta = { + "local_size_x": 64, + "params": ["bad", {"name": "n", "kind": "int"}], + } + # Index of `n` is 1 — args must be long enough for that slot. + assert infer_group_count_x(meta, [None, 130]) == 3 diff --git a/tests/unit/test_gpu_pack.py b/tests/unit/test_gpu_pack.py index 810272e..58ea917 100644 --- a/tests/unit/test_gpu_pack.py +++ b/tests/unit/test_gpu_pack.py @@ -12,14 +12,14 @@ import pytest from cthreads import gpu -from cthreads.gpu.errors import ( +from cthreads.gpu.frontend.errors import ( GpuInvalidArgument, GpuUseAfterDestroy, ) def _ext_gpu(): - return gpu._gpu + return gpu._ext_gpu_api._gpu def _require_gpu_testing(): @@ -48,7 +48,7 @@ def test_public_gpu_has_no_pack_roundtrip_exports(): def test_testing_submodule_absent_when_not_built(monkeypatch): - monkeypatch.setattr(gpu, "_gpu", None) + monkeypatch.setattr(gpu._ext_gpu_api, "_gpu", None) assert _ext_gpu() is None diff --git a/tests/unit/test_gpu_pipeline.py b/tests/unit/test_gpu_pipeline.py new file mode 100644 index 0000000..df02cd5 --- /dev/null +++ b/tests/unit/test_gpu_pipeline.py @@ -0,0 +1,460 @@ +""" +Full GPU pipeline tests: translate -> SPIR-V -> register -> launch via gpu(). + +Modular pieces are covered elsewhere; this file stresses end-to-end behavior +and edge cases across the public API. +""" + +from __future__ import annotations + +import math + +import pytest + +from helpers_gpu import prepare_module + +from cthreads.frontend.Registry import REGISTRY +from cthreads.gpu import GlobalIdx, Gpu, ThreadIdx, available, gpu, shutdown +from cthreads.gpu.compiler.translation.translate import translate_function_for_gpu + + +@pytest.fixture(autouse=True) +def _reset(): + prepare_mod = prepare_module() + REGISTRY.clear() + prepare_mod._gpu_prepared = False + yield + REGISTRY.clear() + try: + shutdown() + except Exception: + pass + prepare_mod._gpu_prepared = False + + +def test_pipeline_translate_source_is_shaderc_ready(): + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + r = translate_function_for_gpu(saxpy, compile_spirv=True) + assert r.spirv is not None + assert "layout(set = 0, binding = 0" in r.source + assert "void main()" in r.source + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +@pytest.mark.parametrize( + "n,a", + [ + (1, 1.0), + (2, -0.5), + (4, 2.0), + (8, 0.25), + (16, -2.0), + (31, 1.5), + (32, 0.0), + (33, 7.0), + (63, 0.5), + (64, -1.0), + (65, 3.25), + (96, 1.0), + (127, -3.0), + (128, 0.0), + (129, 2.5), + (200, math.pi), + (256, math.e), + (300, 1.125), + ], +) +def test_pipeline_saxpy_sizes(n, a): + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + x = [float(i + 1) for i in range(n)] + y = [float(i) for i in range(n)] + expect = [a * xi + yi for xi, yi in zip(x, y)] + gpu(saxpy, n, a, x, y).join() + assert y == pytest.approx(expect) + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +@pytest.mark.parametrize("n", [1, 2, 3, 7, 15, 16, 31, 32, 63, 64, 65, 100, 128, 200]) +def test_pipeline_int_list(n): + @Gpu + def add(n: int, x: list[int], y: list[int]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = x[i] + y[i] + + x = list(range(n)) + y = [10] * n + expect = [x[i] + 10 for i in range(n)] + gpu(add, n, x, y).join() + assert y == expect + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +@pytest.mark.parametrize("n", [1, 4, 17, 64, 100]) +def test_pipeline_bool_list(n): + @Gpu + def invert(n: int, flags: list[bool], ys: list[int]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + if flags[i]: + ys[i] = 1 + else: + ys[i] = 0 + + flags = [(i % 2) == 0 for i in range(n)] + ys = [9] * n + gpu(invert, n, flags, ys).join() + assert ys == [1 if f else 0 for f in flags] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_bool_scalar(): + @Gpu + def set_if(n: int, flag: bool, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + if flag: + y[i] = 1.0 + else: + y[i] = 0.0 + + y = [9.0, 9.0, 9.0] + gpu(set_if, 3, True, y).join() + assert y == [1.0, 1.0, 1.0] + gpu(set_if, 3, False, y).join() + assert y == [0.0, 0.0, 0.0] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_early_return_leaves_tail_untouched(): + @Gpu + def partial(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = 5.0 + + y = [0.0] * 10 + gpu(partial, 3, y).join() + assert y[:3] == [5.0, 5.0, 5.0] + assert y[3:] == [0.0] * 7 + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_two_kernels_same_process(): + @Gpu + def fill(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = 1.0 + + @Gpu + def double(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = y[i] * 2.0 + + y = [0.0, 0.0, 0.0, 0.0] + gpu(fill, 4, y).join() + gpu(double, 4, y).join() + assert y == [2.0, 2.0, 2.0, 2.0] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_inplace_identity_lists(): + @Gpu + def copy_x_to_y(n: int, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = x[i] + + x = [1.5, 2.5, 3.5] + y = [0.0, 0.0, 0.0] + x_id = id(x) + y_id = id(y) + gpu(copy_x_to_y, 3, x, y).join() + assert id(x) == x_id and id(y) == y_id + assert y == x + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_n_less_than_list_len_only_updates_prefix(): + @Gpu + def mark(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = 1.0 + + y = [0.0] * 8 + gpu(mark, 2, y).join() + assert y == [1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_three_lists(): + @Gpu + def madd(n: int, a: list[float], b: list[float], c: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + c[i] = a[i] * b[i] + c[i] + + n = 16 + a = [float(i) for i in range(n)] + b = [2.0] * n + c = [1.0] * n + expect = [a[i] * 2.0 + 1.0 for i in range(n)] + gpu(madd, n, a, b, c).join() + assert c == pytest.approx(expect) + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_local_temps_and_augassign(): + @Gpu + def scale_add(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + t: float = a * x[i] + t += y[i] + y[i] = t + + x = [1.0, 2.0, 3.0] + y = [10.0, 20.0, 30.0] + gpu(scale_add, 3, 2.0, x, y).join() + assert y == pytest.approx([12.0, 24.0, 36.0]) + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_for_range_reduction_into_slot(): + @Gpu + def sum_prefix(n: int, x: list[int], ys: list[int]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + s: int = 0 + for j in range(i + 1): + s += x[j] + ys[i] = s + + x = [1, 2, 3, 4] + ys = [0, 0, 0, 0] + gpu(sum_prefix, 4, x, ys).join() + assert ys == [1, 3, 6, 10] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_if_else_branch(): + @Gpu + def abs_copy(n: int, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + if x[i] < 0.0: + y[i] = 0.0 - x[i] + else: + y[i] = x[i] + + x = [-2.0, 0.0, 3.5] + y = [0.0, 0.0, 0.0] + gpu(abs_copy, 3, x, y).join() + assert y == pytest.approx([2.0, 0.0, 3.5]) + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_threadidx_local_only_first_lane_writes(): + """ThreadIdx.x == 0 within each workgroup; GlobalIdx still unique.""" + + @Gpu + def mark_lane0(n: int, y: list[int]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + lid: int = ThreadIdx.x + if lid == 0: + y[i] = 1 + else: + y[i] = 0 + + n = 130 + y = [9] * n + gpu(mark_lane0, n, y).join() + # Workgroup size 64: indices 0, 64, 128 are lane 0 in each group. + for i in range(n): + expect = 1 if (i % 64) == 0 else 0 + assert y[i] == expect, f"i={i}" + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_empty_n_no_write(): + @Gpu + def mark(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = 1.0 + + y = [0.0, 0.0, 0.0] + gpu(mark, 0, y).join() + assert y == [0.0, 0.0, 0.0] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_scalar_only_kernel_runs(): + @Gpu + def noop(n: int) -> None: + i: int = GlobalIdx.x + if i >= n: + return + + gpu(noop, 8).join() + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_lists_only_no_scalars(): + @Gpu + def copy_lists(x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= 4: + return + y[i] = x[i] + + x = [1.0, 2.0, 3.0, 4.0] + y = [0.0, 0.0, 0.0, 0.0] + gpu(copy_lists, x, y).join() + assert y == x + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_lists_with_dummy_n(): + """Workaround path: always pass `n` so binding 0 exists.""" + + @Gpu + def copy_n(n: int, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = x[i] + + x = [1.0, 2.0, 3.0, 4.0] + y = [0.0, 0.0, 0.0, 0.0] + gpu(copy_n, 4, x, y).join() + assert y == x + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +@pytest.mark.parametrize( + "ops", + [ + ("add", lambda a, b: a + b), + ("sub", lambda a, b: a - b), + ("mul", lambda a, b: a * b), + ], +) +def test_pipeline_binop_variants(ops): + name, py_op = ops + + if name == "add": + + @Gpu + def kern(n: int, x: list[float], y: list[float], z: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + z[i] = x[i] + y[i] + + elif name == "sub": + + @Gpu + def kern(n: int, x: list[float], y: list[float], z: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + z[i] = x[i] - y[i] + + else: + + @Gpu + def kern(n: int, x: list[float], y: list[float], z: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + z[i] = x[i] * y[i] + + x = [1.0, 2.0, 3.0, 4.0] + y = [4.0, 3.0, 2.0, 1.0] + z = [0.0] * 4 + gpu(kern, 4, x, y, z).join() + assert z == pytest.approx([py_op(a, b) for a, b in zip(x, y)]) + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_compare_chain_guard(): + @Gpu + def clamp_mark(n: int, lo: int, hi: int, y: list[int]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + if lo <= i < hi: + y[i] = 1 + else: + y[i] = 0 + + y = [9] * 10 + gpu(clamp_mark, 10, 3, 7, y).join() + assert y == [0, 0, 0, 1, 1, 1, 1, 0, 0, 0] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_many_launches_same_kernel(): + @Gpu + def add_k(n: int, k: float, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = y[i] + k + + y = [0.0] * 8 + for _ in range(10): + gpu(add_k, 8, 0.5, y).join() + assert y == pytest.approx([5.0] * 8) + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_pipeline_int_and_float_mixed_scalars(): + @Gpu + def mix(n: int, a: float, b: float, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + # i is int; promote via add with float literal (no float() call plugin yet) + t: float = a * b + t += 0.0 + 0.0 # keep float + if i == 0: + y[i] = t + elif i == 1: + y[i] = t + 1.0 + else: + y[i] = t + 2.0 + + y = [0.0, 0.0, 0.0] + gpu(mix, 3, 1.5, 2.0, y).join() + assert y == pytest.approx([3.0, 4.0, 5.0]) diff --git a/tests/unit/test_gpu_prepare.py b/tests/unit/test_gpu_prepare.py new file mode 100644 index 0000000..67fe1bf --- /dev/null +++ b/tests/unit/test_gpu_prepare.py @@ -0,0 +1,270 @@ +"""Unit / live tests for prepare() and gpu().""" + +from __future__ import annotations + +import pytest + +from helpers_gpu import prepare_module + +from cthreads.frontend.Registry import REGISTRY +from cthreads.gpu import GlobalIdx, Gpu, available, gpu, prepare, shutdown +from cthreads.gpu.frontend.errors import GPUNotAvailable +from cthreads.gpu.runtime import GpuJob + + +@pytest.fixture(autouse=True) +def _reset_gpu_prepare_state(): + prepare_mod = prepare_module() + REGISTRY.clear() + prepare_mod._gpu_prepared = False + yield + REGISTRY.clear() + # Isolation: shutdown releases ShaderCache; lib clears `_gpu_prepared`. + try: + shutdown() + except Exception: + pass + prepare_mod._gpu_prepared = False + + +def test_prepare_function_and_runtime_module_coexist(): + """Public API is the function; internals live on `cthreads.gpu.runtime`.""" + import cthreads.gpu.runtime as runtime_mod + from cthreads.gpu import prepare as prepare_fn + + assert callable(prepare_fn) + assert prepare_fn is runtime_mod.prepare + assert hasattr(runtime_mod, "_gpu_prepared") + assert prepare_module() is runtime_mod + + +def test_gpu_rejects_non_gpu_fn(): + def plain(n: int) -> None: + pass + + with pytest.raises(TypeError, match="@Gpu"): + gpu(plain, 1) + + +def test_gpu_rejects_non_callable(): + with pytest.raises(TypeError, match="callable"): + gpu(None) # type: ignore[arg-type] + + +def test_gpu_rejects_kwargs(): + if not available(): + pytest.skip("GPU not available") + + @Gpu + def k(n: int) -> None: + i: int = GlobalIdx.x + if i >= n: + return + + with pytest.raises(TypeError, match="keyword"): + gpu(k, 1, force=False, extra=1) + + +def test_gpu_rejects_arity_mismatch(): + if not available(): + pytest.skip("GPU not available") + + @Gpu + def k(n: int, a: float) -> None: + pass + + with pytest.raises(TypeError, match="expected 2"): + gpu(k, 1) + + +def test_prepare_raises_when_nothing_registered(): + if not available(): + pytest.skip("GPU not available") + with pytest.raises(RuntimeError, match="Nothing registered"): + prepare() + + +def test_prepare_unavailable(monkeypatch): + prepare_mod = prepare_module() + monkeypatch.setattr(prepare_mod._ext_gpu_api, "available", lambda: False) + with pytest.raises(GPUNotAvailable): + prepare() + + +def test_gpu_unavailable(monkeypatch): + prepare_mod = prepare_module() + monkeypatch.setattr(prepare_mod._ext_gpu_api, "available", lambda: False) + + def fake_gpu_fn(n: int) -> None: + pass + + fake_gpu_fn.__gpu__ = True # type: ignore[attr-defined] + with pytest.raises(GPUNotAvailable): + gpu(fake_gpu_fn, 1) + + +def test_gpu_job_result_is_none(): + job = GpuJob( + type( + "R", + (), + { + "start": lambda self: None, + "done": lambda self: True, + "join": lambda self: None, + }, + )() + ) + assert job.result() is None + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_live_prepare_and_gpu_saxpy(): + prepare_mod = prepare_module() + + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + x = [1.0, 2.0, 3.0, 4.0] + y = [10.0, 20.0, 30.0, 40.0] + a = 2.0 + expect = [a * xi + yi for xi, yi in zip(x, y)] + + info = prepare() + assert "rewritten" in info + assert prepare_mod._gpu_prepared is True + + job = gpu(saxpy, len(x), a, x, y) + assert isinstance(job, GpuJob) + job.join() + assert y == expect + assert job.done() + assert job.result() is None + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_live_gpu_auto_prepare(): + @Gpu + def fill(n: int, ys: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + ys[i] = 1.0 + + ys = [0.0, 0.0, 0.0, 0.0] + prepare_module()._gpu_prepared = False + gpu(fill, 4, ys).join() + assert ys == [1.0, 1.0, 1.0, 1.0] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_lib_reserved_param_name_out(): + """Live path: reserved `out` fails at GLSL compile with a clear hint.""" + + @Gpu + def fill(n: int, out: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + out[i] = 1.0 + + ys = [0.0, 0.0] + with pytest.raises(RuntimeError, match="reserved words|GLSL compile failed"): + gpu(fill, 2, ys).join() + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_live_gpu_second_launch_same_kernel(): + @Gpu + def add1(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = y[i] + 1.0 + + y = [0.0, 0.0, 0.0] + gpu(add1, 3, y).join() + gpu(add1, 3, y).join() + assert y == [2.0, 2.0, 2.0] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_live_larger_than_one_workgroup(): + @Gpu + def scale(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + + n = 200 + x = [float(i) for i in range(n)] + y = [0.0] * n + gpu(scale, n, 3.0, x, y).join() + assert y == [3.0 * float(i) for i in range(n)] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_live_force_reprepare(): + @Gpu + def k(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = 7.0 + + y = [0.0, 0.0] + gpu(k, 2, y).join() + y2 = [0.0, 0.0] + gpu(k, 2, y2, force=True).join() + assert y2 == [7.0, 7.0] + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_compile_sets_prepared_flag(): + from cthreads.gpu import compile + + prepare_mod = prepare_module() + + @Gpu + def k(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = 1.0 + + prepare_mod._gpu_prepared = False + info = compile() + assert prepare_mod._gpu_prepared is True + assert "rewritten" in info + + +@pytest.mark.skipif(not available(), reason="Vulkan GPU not available") +def test_shutdown_then_gpu_reregisters(): + """ + shutdown() releases ShaderCache; next gpu() must rewalk registry and work. + """ + prepare_mod = prepare_module() + + @Gpu + def fill(n: int, y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = 3.0 + + y = [0.0, 0.0] + gpu(fill, 2, y).join() + assert y == [3.0, 3.0] + + shutdown() + assert prepare_mod._gpu_prepared is False + + y2 = [0.0, 0.0] + gpu(fill, 2, y2).join() + assert y2 == [3.0, 3.0] + assert prepare_mod._gpu_prepared is True diff --git a/tests/unit/test_gpu_reserved_names.py b/tests/unit/test_gpu_reserved_names.py new file mode 100644 index 0000000..0a0f2b0 --- /dev/null +++ b/tests/unit/test_gpu_reserved_names.py @@ -0,0 +1,73 @@ +"""Unit tests for GLSL reserved-identifier / keyword collisions.""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu.compiler.translation.spirv import compile_glsl_to_spirv +from cthreads.gpu.compiler.translation.translate import translate_function_for_gpu + + +def _compiler_available() -> bool: + try: + compile_glsl_to_spirv( + "#version 450\nlayout(local_size_x = 1) in;\nvoid main() {}\n" + ) + return True + except RuntimeError: + return False + + +@pytest.mark.skipif(not _compiler_available(), reason="no GLSL compiler") +@pytest.mark.parametrize( + "param", + [ + "out", + "uniform", + "buffer", + "flat", + "smooth", + "shared", + ], +) +def test_reserved_list_param_names_fail_with_hint(param, tmp_module): + """No keyword denylist: glslang fails; we wrap with an identifier hint.""" + mod = tmp_module( + f""" + from cthreads.gpu import GlobalIdx + + def k(n: int, {param}: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + {param}[i] = 1.0 + """, + name=f"reserved_{param}", + ) + with pytest.raises(RuntimeError, match="reserved words|GLSL compile failed"): + translate_function_for_gpu(mod.k, compile_spirv=True) + + +@pytest.mark.skipif(not _compiler_available(), reason="no GLSL compiler") +def test_safe_param_names_compile(tmp_module): + mod = tmp_module( + """ + from cthreads.gpu import GlobalIdx + + def k(n: int, ys: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + ys[i] = 1.0 + """, + name="safe_param_names", + ) + r = translate_function_for_gpu(mod.k, compile_spirv=True) + assert r.spirv is not None + assert "} ys;" in r.source + + +@pytest.mark.skipif(not _compiler_available(), reason="no GLSL compiler") +def test_compile_error_includes_hint(): + with pytest.raises(RuntimeError, match="Hint:.*reserved"): + compile_glsl_to_spirv("#version 450\nvoid main() { not_a_type x; }\n") diff --git a/tests/unit/test_gpu_shader.py b/tests/unit/test_gpu_shader.py index bb69cc2..29d0bac 100644 --- a/tests/unit/test_gpu_shader.py +++ b/tests/unit/test_gpu_shader.py @@ -1,9 +1,8 @@ """ -GPU shader / launch smoke (test-only _ext.gpu.testing). +GPU shader / launch tests. -Covers create_entry, ShaderCache, update_descriptors, and launch_gpu_kernel -via smoke_launch_saxpy (join writeback). Live checks skip when CTHREADS_GPU -is off or no Vulkan device is available. +Substrate smokes stay on `_ext.gpu.testing`. Launch + join use the product +path: `_ext.gpu.launch_gpu_kernel` / `GpuJob` (via `_ext_gpu_api`). """ from __future__ import annotations @@ -11,11 +10,13 @@ import pytest from cthreads import gpu -from cthreads.gpu.errors import GpuInvalidArgument +from cthreads.gpu import _ext_gpu_api +from cthreads.gpu.frontend.errors import GpuInvalidArgument +from cthreads.gpu.gpu_kernel_meta import build_gpu_kernel_meta def _ext_gpu(): - return gpu._gpu + return gpu._ext_gpu_api._gpu def _require_gpu_testing(): @@ -37,9 +38,18 @@ def test_public_gpu_has_no_shader_smoke_exports(): assert not hasattr(gpu, "smoke_create_entry") assert not hasattr(gpu, "smoke_update_descriptors") assert not hasattr(gpu, "smoke_launch_saxpy") + assert not hasattr(gpu, "register_smoke_saxpy") assert not hasattr(gpu, "testing") +def test_ext_gpu_exports_launch_api(): + ext = _ext_gpu() + if ext is None: + pytest.skip("cthreads built without CTHREADS_GPU (_ext.gpu missing)") + assert hasattr(ext, "launch_gpu_kernel") + assert hasattr(ext, "GpuJob") + + def test_live_smoke_create_entry(): testing = _require_gpu_testing() try: @@ -64,13 +74,42 @@ def test_live_smoke_cache_register_and_get(): gpu.shutdown() -def test_live_smoke_launch_saxpy(): +def test_live_launch_saxpy_product_path(): + """ + Register smoke SPIR-V (test-only), then launch/join via product bindings. + """ testing = _require_gpu_testing() - if not hasattr(testing, "smoke_launch_saxpy"): - pytest.skip("rebuild with latest gpu testing (smoke_launch_saxpy)") + if not hasattr(testing, "register_smoke_saxpy"): + pytest.skip("rebuild with register_smoke_saxpy") + if not hasattr(_ext_gpu(), "launch_gpu_kernel"): + pytest.skip("rebuild with product launch_gpu_kernel") + + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + try: - testing.smoke_launch_saxpy() + symbol = testing.register_smoke_saxpy() + meta = build_gpu_kernel_meta(saxpy, symbol=symbol).to_dict() + meta["group_count_x"] = 1 + meta["group_count_y"] = 1 + meta["group_count_z"] = 1 + + n = 4 + a = 2.0 + x = [1.0, 2.0, 3.0, 4.0] + y = [10.0, 20.0, 30.0, 40.0] + expect = [a * xi + yi for xi, yi in zip(x, y)] + + job = _ext_gpu_api.launch_gpu_kernel(meta, [n, a, x, y]) + job.join() + assert y == expect + assert job.done() finally: + if hasattr(testing, "clear_shader_cache"): + try: + testing.clear_shader_cache() + except Exception: + pass gpu.shutdown() diff --git a/tests/unit/test_gpu_signature.py b/tests/unit/test_gpu_signature.py new file mode 100644 index 0000000..2df9323 --- /dev/null +++ b/tests/unit/test_gpu_signature.py @@ -0,0 +1,214 @@ +"""Unit tests for GpuSignature preamble emission.""" + +from __future__ import annotations + +import pytest + +from cthreads.compiler.translation.Source import Source +from cthreads.gpu.compiler.translation.Signature import GpuSignature +from cthreads.gpu.compiler.translation.context import GpuTranslationContext +from cthreads.gpu.gpu_kernel_meta import build_gpu_kernel_meta +from cthreads.types import PyBool, PyFloat, PyInt, PyList + + +def _sig(fn, local_size_x: int = 64): + ctx = GpuTranslationContext(fn=fn, local_size_x=local_size_x) + result = GpuSignature.translate(Source.parse_function(fn), ctx) + return ctx, result + + +def test_saxpy_preamble_matches_meta_and_smoke_shape(): + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + pass + + meta = build_gpu_kernel_meta(saxpy) + ctx, sig = _sig(saxpy) + assert sig.binding_count == meta.binding_count == 3 + assert sig.scalar_bytes == meta.scalar_bytes == 8 + assert sig.scalar_fields == [("n", "int"), ("a", "float")] + assert sig.list_fields == [(1, "x", "float"), (2, "y", "float")] + assert "layout(local_size_x = 64)" in sig.preamble + assert "binding = 0" in sig.preamble + assert "int n;" in sig.preamble and "float a;" in sig.preamble + assert "binding = 1" in sig.preamble and "float data[];" in sig.preamble + assert "} x;" in sig.preamble and "} y;" in sig.preamble + assert "n" in ctx.scalar_params and "a" in ctx.scalar_params + assert "x" in ctx.list_params and "y" in ctx.list_params + assert isinstance(ctx.symbols["n"], PyInt) + assert isinstance(ctx.symbols["x"], PyList) + + +def test_lists_only_no_binding_zero_scalars(): + def k(x: list[int], y: list[int]) -> None: + pass + + _, sig = _sig(k) + assert sig.binding_count == 2 + assert sig.scalar_bytes == 0 + assert sig.scalar_fields == [] + assert "binding = 0" in sig.preamble and "binding = 1" in sig.preamble + assert "buffer Scalars" not in sig.preamble + assert sig.list_fields == [(0, "x", "int"), (1, "y", "int")] + + +def test_scalars_only_binding_zero(): + def k(a: int, b: float, c: bool) -> None: + pass + + _, sig = _sig(k) + assert sig.binding_count == 1 + assert sig.scalar_bytes == 12 + assert "bool c;" in sig.preamble + assert "data[]" not in sig.preamble + + +@pytest.mark.parametrize("local", [1, 8, 32, 64, 128, 256]) +def test_custom_local_size(local): + def k(n: int) -> None: + pass + + _, sig = _sig(k, local_size_x=local) + assert f"layout(local_size_x = {local})" in sig.preamble + assert sig.local_size_x == local + + +def test_param_order_preserved_in_scalar_block(): + def k(b: float, a: int) -> None: + pass + + _, sig = _sig(k) + assert sig.scalar_fields == [("b", "float"), ("a", "int")] + pos_b = sig.preamble.index("float b;") + pos_a = sig.preamble.index("int a;") + assert pos_b < pos_a + + +def test_list_binding_order_follows_params(): + def k(n: int, z: list[float], a: list[int], b: list[bool]) -> None: + pass + + _, sig = _sig(k) + assert sig.list_fields == [ + (1, "z", "float"), + (2, "a", "int"), + (3, "b", "bool"), + ] + assert sig.binding_count == 4 + + +def test_missing_annotation_raises(): + def k(n: int, a) -> None: # type: ignore[no-untyped-def] + pass + + with pytest.raises(TypeError, match="type annotation"): + _sig(k) + + +def test_vararg_rejected(): + def k(*args: int) -> None: + pass + + with pytest.raises(TypeError, match=r"\*args"): + _sig(k) + + +def test_kwargs_rejected(): + def k(**kwargs: int) -> None: + pass + + with pytest.raises(TypeError, match=r"\*\*kwargs"): + _sig(k) + + +def test_kwonly_rejected(): + def k(*, n: int) -> None: + pass + + with pytest.raises(TypeError, match="kw-only"): + _sig(k) + + +def test_non_none_return_rejected(): + def k(n: int) -> int: + return n + + with pytest.raises(TypeError, match="return must be None"): + _sig(k) + + +def test_none_return_ok(): + def k(n: int) -> None: + pass + + _, sig = _sig(k) + assert sig.func_name == "k" + + +def test_unsupported_dict_param(): + def k(d: dict[str, int]) -> None: + pass + + with pytest.raises(TypeError): + _sig(k) + + +def test_unsupported_str_param(): + def k(s: str) -> None: + pass + + with pytest.raises(TypeError): + _sig(k) + + +def test_unsupported_nested_list(): + def k(x: list[list[int]]) -> None: + pass + + with pytest.raises(TypeError): + _sig(k) + + +def test_list_block_name_capitalized(): + def k(values: list[float]) -> None: + pass + + _, sig = _sig(k) + assert "buffer Values" in sig.preamble + assert "} values;" in sig.preamble + + +def test_single_list_binding_zero(): + def k(xs: list[int]) -> None: + pass + + _, sig = _sig(k) + assert sig.binding_count == 1 + assert "binding = 0" in sig.preamble + assert "buffer Scalars" not in sig.preamble + + +def test_bool_list_elem(): + def k(flags: list[bool]) -> None: + pass + + _, sig = _sig(k) + assert "bool data[];" in sig.preamble + + +def test_ctx_symbols_typed(): + def k(n: int, flag: bool, a: float, xs: list[int]) -> None: + pass + + ctx, _ = _sig(k) + assert isinstance(ctx.symbols["n"], PyInt) + assert isinstance(ctx.symbols["flag"], PyBool) + assert isinstance(ctx.symbols["a"], PyFloat) + assert isinstance(ctx.symbols["xs"], PyList) + + +def test_std430_mentioned_on_buffers(): + def k(n: int, x: list[float]) -> None: + pass + + _, sig = _sig(k) + assert sig.preamble.count("std430") >= 2 diff --git a/tests/unit/test_gpu_spirv.py b/tests/unit/test_gpu_spirv.py new file mode 100644 index 0000000..f8f965a --- /dev/null +++ b/tests/unit/test_gpu_spirv.py @@ -0,0 +1,117 @@ +"""Unit tests for GLSL -> SPIR-V (native glslang or glslc fallback).""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu import GlobalIdx +from cthreads.gpu.compiler.translation.spirv import compile_glsl_to_spirv +from cthreads.gpu.compiler.translation.translate import translate_function_for_gpu + +_MIN_COMP = """#version 450 +layout(local_size_x = 64) in; +layout(set = 0, binding = 0, std430) buffer Scalars { int n; } scalars; +void main() { + int i = int(gl_GlobalInvocationID.x); + if (i >= scalars.n) { return; } +} +""" + + +def _compiler_available() -> bool: + try: + compile_glsl_to_spirv(_MIN_COMP) + return True + except RuntimeError: + return False + + +pytestmark = pytest.mark.skipif( + not _compiler_available(), + reason="no native compile_glsl / glslc available", +) + + +def test_compile_min_comp_magic_and_alignment(): + data = compile_glsl_to_spirv(_MIN_COMP) + assert len(data) >= 20 + assert len(data) % 4 == 0 + assert data[:4] == b"\x03\x02\x23\x07" + + +def test_compile_empty_raises(): + with pytest.raises(Exception): + compile_glsl_to_spirv("") + + +def test_compile_bad_glsl_raises(): + with pytest.raises(Exception): + compile_glsl_to_spirv("#version 450\nvoid main() { not_a_type x; }\n") + + +def test_compile_rejects_uint_to_int_without_cast(): + bad = """#version 450 +layout(local_size_x = 1) in; +void main() { int i = gl_GlobalInvocationID.x; } +""" + with pytest.raises(Exception): + compile_glsl_to_spirv(bad) + + +def test_compile_with_cast_ok(): + ok = """#version 450 +layout(local_size_x = 1) in; +void main() { int i = int(gl_GlobalInvocationID.x); } +""" + data = compile_glsl_to_spirv(ok) + assert data[:4] == b"\x03\x02\x23\x07" + + +@pytest.mark.parametrize("local", [1, 8, 64, 256]) +def test_compile_local_sizes(local): + src = f"""#version 450 +layout(local_size_x = {local}) in; +void main() {{}} +""" + data = compile_glsl_to_spirv(src) + assert len(data) % 4 == 0 + + +def test_translate_compile_spirv_flag(): + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + r0 = translate_function_for_gpu(saxpy, compile_spirv=False) + assert r0.spirv is None + r1 = translate_function_for_gpu(saxpy, compile_spirv=True) + assert r1.spirv is not None + assert r1.spirv[:4] == b"\x03\x02\x23\x07" + assert "#version 450" in r1.source + assert "void main()" in r1.source + + +def test_translate_compile_many_kernels(): + def add(n: int, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = x[i] + y[i] + + def fill(n: int, y: list[int]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = i + + for fn in (add, fill): + r = translate_function_for_gpu(fn, compile_spirv=True) + assert r.spirv is not None + assert r.spirv[:4] == b"\x03\x02\x23\x07" + + +def test_compile_missing_version_raises(): + with pytest.raises(Exception): + compile_glsl_to_spirv("void main() {}\n") diff --git a/tests/unit/test_gpu_syntax.py b/tests/unit/test_gpu_syntax.py new file mode 100644 index 0000000..fb617ef --- /dev/null +++ b/tests/unit/test_gpu_syntax.py @@ -0,0 +1,701 @@ +"""Unit tests for GPU Syntax leaf translators and index AttrPlugin.""" + +from __future__ import annotations + +import ast + +import pytest + +from cthreads.compiler.translation.Source import Source +from cthreads.gpu import BlockDim, BlockIdx, GlobalIdx, GridDim, ThreadIdx +from cthreads.gpu.compiler.translation.Signature import GpuSignature +from cthreads.gpu.compiler.translation.context import GpuTranslationContext +from cthreads.gpu.compiler.translation.syntax.Syntax import GpuSyntax +from cthreads.types import PyFloat, PyInt, PyList + + +def _ctx_for(fn): + ctx = GpuTranslationContext(fn=fn) + GpuSignature.translate(Source.parse_function(fn), ctx) + return ctx + + +def _expr(src: str, ctx: GpuTranslationContext) -> str: + tree = ast.parse(src, mode="eval") + assert isinstance(tree, ast.Expression) + return GpuSyntax.expr(tree.body, ctx) + + +def _stmt(src: str, ctx: GpuTranslationContext) -> list[str]: + tree = ast.parse(src) + assert len(tree.body) == 1 + return GpuSyntax.stmt(tree.body[0], ctx) + + +def _inject_indexes(fn): + fn.__globals__.update( + { + "GlobalIdx": GlobalIdx, + "ThreadIdx": ThreadIdx, + "BlockIdx": BlockIdx, + "BlockDim": BlockDim, + "GridDim": GridDim, + } + ) + + +# --- Name ------------------------------------------------------------------- + + +def test_name_scalar_rewrites_to_scalars_block(): + def k(n: int, a: float, x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + assert _expr("n", ctx) == "scalars.n" + assert _expr("a", ctx) == "scalars.a" + + +def test_name_list_stays_bare(): + def k(x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + assert _expr("x", ctx) == "x" + + +def test_name_local_bare(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + ctx.symbols["i"] = PyInt() + assert _expr("i", ctx) == "i" + + +def test_name_unknown_raises(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="unknown name"): + _expr("missing", ctx) + + +def test_name_self_raises(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="self"): + _expr("self", ctx) + + +# --- Index ------------------------------------------------------------------ + + +def test_index_list_to_data(): + def k(x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + ctx.symbols["i"] = PyInt() + assert _expr("x[i]", ctx) == "(x.data[i])" + + +def test_index_nested_expr_index(): + def k(x: list[float], n: int) -> None: + pass + + ctx = _ctx_for(k) + assert _expr("x[n]", ctx) == "(x.data[scalars.n])" + + +def test_index_rejects_slice(): + def k(x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="slice"): + _expr("x[1:2]", ctx) + + +def test_index_rejects_non_list(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="list parameters"): + _expr("n[0]", ctx) + + +def test_index_literal(): + def k(x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + assert _expr("x[0]", ctx) == "(x.data[0])" + + +# --- Op --------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "src, needle", + [ + ("a + b", "+"), + ("a - b", "-"), + ("a * b", "*"), + ("a / b", "/"), + ("a // b", "/"), + ("a % b", "%"), + ("a << b", "<<"), + ("a >> b", ">>"), + ("a | b", "|"), + ("a ^ b", "^"), + ("a & b", "&"), + ], +) +def test_binop_table(src, needle): + def k(a: int, b: int) -> None: + pass + + ctx = _ctx_for(k) + out = _expr(src, ctx) + assert needle in out + assert "scalars.a" in out and "scalars.b" in out + + +def test_binop_mul_add(): + def k(a: float, x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + ctx.symbols["i"] = PyInt() + out = _expr("a * x[i] + x[i]", ctx) + assert "scalars.a" in out + assert "x.data[i]" in out + assert "*" in out and "+" in out + + +@pytest.mark.parametrize( + "src, op", + [ + ("a == b", "=="), + ("a != b", "!="), + ("a < b", "<"), + ("a <= b", "<="), + ("a > b", ">"), + ("a >= b", ">="), + ], +) +def test_compare_ops(src, op): + def k(a: int, b: int) -> None: + pass + + ctx = _ctx_for(k) + assert op in _expr(src, ctx) + + +def test_compare_and_bool(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + ctx.symbols["i"] = PyInt() + out = _expr("i >= n and True", ctx) + assert ">=" in out and "&&" in out + assert "scalars.n" in out + + +def test_bool_or(): + def k(a: bool, b: bool) -> None: + pass + + ctx = _ctx_for(k) + assert "||" in _expr("a or b", ctx) + + +def test_unary_not_neg(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + assert "!" in _expr("not n", ctx) + assert "-" in _expr("-n", ctx) + assert "+" in _expr("+n", ctx) + assert "~" in _expr("~n", ctx) + + +def test_pow_rejected(): + def k(a: float) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match=r"\*\*|pow"): + _expr("a ** 2", ctx) + + +def test_chained_compare(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + ctx.symbols["i"] = PyInt() + out = _expr("0 <= i < n", ctx) + assert "&&" in out + + +def test_unsupported_call_raises(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="unsupported call"): + _expr("abs(n)", ctx) + + +def test_float_call_rejected(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="unsupported call"): + _expr("float(n)", ctx) + + +# --- Assign ----------------------------------------------------------------- + + +def test_assign_list_element(): + def k(a: float, x: list[float], y: list[float]) -> None: + pass + + ctx = _ctx_for(k) + ctx.symbols["i"] = PyInt() + lines = _stmt("y[i] = a * x[i] + y[i]", ctx) + assert len(lines) == 1 + assert "y.data[i]" in lines[0] + assert "scalars.a" in lines[0] + + +def test_ann_assign_local(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt("i: int = n", ctx) + assert lines == [" int i = scalars.n;"] + assert isinstance(ctx.symbols["i"], PyInt) + assert "i" not in ctx.scalar_params + + +def test_ann_assign_from_global_idx(): + def k(n: int) -> None: + pass + + _inject_indexes(k) + ctx = _ctx_for(k) + lines = _stmt("i: int = GlobalIdx.x", ctx) + assert "int i =" in lines[0] + assert "gl_GlobalInvocationID.x" in lines[0] + + +def test_ann_assign_uninitialized(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + # Python allows `i: int` without value via AnnAssign value=None in AST — + # constructing via exec-style source needs a value; use explicit None-ish + # by building AST. + tree = ast.parse("i: int") + assert isinstance(tree.body[0], ast.AnnAssign) + lines = GpuSyntax.stmt(tree.body[0], ctx) + assert lines == [" int i;"] + + +def test_ann_assign_redeclare_raises(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + _stmt("i: int = 0", ctx) + with pytest.raises(TypeError, match="redeclaration"): + _stmt("i: int = 1", ctx) + + +def test_assign_unknown_name_raises(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="unknown name"): + _stmt("i = n", ctx) + + +def test_assign_bare_list_rejected(): + def k(x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="list parameter"): + _stmt("x = x", ctx) + + +def test_assign_multi_target_rejected(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + _stmt("i: int = 0", ctx) + _stmt("j: int = 0", ctx) + with pytest.raises(TypeError, match="single-target"): + _stmt("i = j = n", ctx) + + +def test_aug_assign(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + _stmt("i: int = 0", ctx) + lines = _stmt("i += n", ctx) + assert lines == [" i += scalars.n;"] + + +@pytest.mark.parametrize( + "src, op", + [ + ("i += n", "+="), + ("i -= n", "-="), + ("i *= n", "*="), + ("i %= n", "%="), + ("i &= n", "&="), + ("i |= n", "|="), + ("i ^= n", "^="), + ("i <<= n", "<<="), + ("i >>= n", ">>="), + ], +) +def test_aug_assign_ops(src, op): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + _stmt("i: int = 0", ctx) + assert op in _stmt(src, ctx)[0] + + +def test_aug_assign_pow_rejected(): + def k(a: float) -> None: + pass + + ctx = _ctx_for(k) + _stmt("t: float = a", ctx) + with pytest.raises(TypeError, match=r"\*\*|pow"): + _stmt("t **= 2", ctx) + + +def test_aug_assign_list_elem(): + def k(y: list[float]) -> None: + pass + + ctx = _ctx_for(k) + ctx.symbols["i"] = PyInt() + lines = _stmt("y[i] += 1.0", ctx) + assert "y.data[i]" in lines[0] + assert "+=" in lines[0] + + +def test_ann_assign_local_list_rejected(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="local list"): + _stmt("xs: list[float] = []", ctx) + + +def test_slice_assign_rejected(): + def k(x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="slice"): + _stmt("x[1:2] = x[0:1]", ctx) + + +# --- Flow ------------------------------------------------------------------- + + +def test_if_return(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + ctx.symbols["i"] = PyInt() + lines = _stmt("if i >= n:\n return", ctx) + assert lines[0].startswith(" if ") + assert any("return;" in L for L in lines) + + +def test_return_always_bare(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + assert _stmt("return", ctx) == [" return;"] + assert _stmt("return 1", ctx) == [" return;"] + + +def test_pass_break_continue(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + assert _stmt("pass", ctx) == [] + assert _stmt("break", ctx) == [" break;"] + assert _stmt("continue", ctx) == [" continue;"] + + +@pytest.mark.parametrize( + "src, needles", + [ + ("for i in range(n):\n pass", ["for (int i = 0;", "scalars.n", "i += 1"]), + ( + "for i in range(1, n):\n pass", + ["for (int i = 1;", "scalars.n", "i += 1"], + ), + ( + "for i in range(0, n, 2):\n pass", + ["for (int i = 0;", "scalars.n", "i += 2"], + ), + ], +) +def test_for_range_forms(src, needles): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt(src, ctx) + joined = "\n".join(lines) + for needle in needles: + assert needle in joined + assert "i" not in ctx.symbols + + +def test_for_range(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt("for i in range(n):\n pass", ctx) + assert "for (int i = 0;" in lines[0] + assert "scalars.n" in lines[0] + assert "i" not in ctx.symbols + + +def test_for_list_rejected(): + def k(x: list[float]) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="range"): + _stmt("for v in x:\n pass", ctx) + + +def test_for_else_rejected(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="for/else"): + _stmt("for i in range(n):\n pass\nelse:\n pass", ctx) + + +def test_while_else_rejected(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="while/else"): + _stmt("while n > 0:\n break\nelse:\n pass", ctx) + + +def test_for_rebind_rejected(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + _stmt("i: int = 0", ctx) + with pytest.raises(TypeError, match="rebinds"): + _stmt("for i in range(n):\n pass", ctx) + + +def test_for_range_zero_args_rejected(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="range"): + _stmt("for i in range():\n pass", ctx) + + +def test_for_range_four_args_rejected(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="range"): + _stmt("for i in range(0, n, 1, 2):\n pass", ctx) + + +def test_while(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt("while n > 0:\n break", ctx) + assert lines[0].startswith(" while ") + + +def test_if_else(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt("if n > 0:\n pass\nelse:\n return", ctx) + assert any("else" in L for L in lines) + + +def test_if_elif_lowers_as_nested_else(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt( + "if n > 0:\n return\nelif n < 0:\n return\nelse:\n pass", + ctx, + ) + joined = "\n".join(lines) + assert "if (" in joined + assert "else" in joined + + +def test_docstring_expr_ignored(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + assert _stmt('"doc"', ctx) == [] + + +def test_unsupported_expr_stmt_comment(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt("n", ctx) + assert lines[0].startswith(" // unsupported") + + +def test_unsupported_stmt_comment(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + tree = ast.parse("raise RuntimeError()") + lines = GpuSyntax.stmt(tree.body[0], ctx) + assert "unsupported statement" in lines[0] + + +# --- Index builtins / AttrPlugin -------------------------------------------- + + +@pytest.mark.parametrize( + "expr, glsl", + [ + ("GlobalIdx.x", "int(gl_GlobalInvocationID.x)"), + ("GlobalIdx.y", "int(gl_GlobalInvocationID.y)"), + ("GlobalIdx.z", "int(gl_GlobalInvocationID.z)"), + ("ThreadIdx.x", "int(gl_LocalInvocationID.x)"), + ("ThreadIdx.y", "int(gl_LocalInvocationID.y)"), + ("ThreadIdx.z", "int(gl_LocalInvocationID.z)"), + ("BlockIdx.x", "int(gl_WorkGroupID.x)"), + ("BlockIdx.y", "int(gl_WorkGroupID.y)"), + ("BlockIdx.z", "int(gl_WorkGroupID.z)"), + ("BlockDim.x", "int(gl_WorkGroupSize.x)"), + ("BlockDim.y", "int(gl_WorkGroupSize.y)"), + ("BlockDim.z", "int(gl_WorkGroupSize.z)"), + ("GridDim.x", "int(gl_NumWorkGroups.x)"), + ("GridDim.y", "int(gl_NumWorkGroups.y)"), + ("GridDim.z", "int(gl_NumWorkGroups.z)"), + ], +) +def test_index_builtins(expr, glsl): + def k(n: int) -> None: + pass + + _inject_indexes(k) + ctx = _ctx_for(k) + assert _expr(expr, ctx) == glsl + + +def test_module_qualified_global_idx(): + import cthreads.gpu as gpu_mod + + def k(n: int) -> None: + pass + + k.__globals__["gpu"] = gpu_mod + ctx = _ctx_for(k) + assert _expr("gpu.GlobalIdx.x", ctx) == "int(gl_GlobalInvocationID.x)" + + +def test_unknown_attr_raises(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="unsupported attribute"): + _expr("n.imag", ctx) + + +def test_index_builtin_bad_axis_raises(): + def k(n: int) -> None: + pass + + _inject_indexes(k) + ctx = _ctx_for(k) + with pytest.raises(TypeError, match="unsupported attribute"): + _expr("GlobalIdx.w", ctx) + + +def test_literal_constants(): + def k(n: int) -> None: + pass + + ctx = _ctx_for(k) + assert _expr("2", ctx) == "2" + assert _expr("True", ctx) == "true" + assert _expr("False", ctx) == "false" + assert _expr("1.5", ctx) == "1.5" + assert _expr("0.0", ctx) == "0.0" + + +def test_nested_if_in_for(): + def k(n: int, y: list[int]) -> None: + pass + + ctx = _ctx_for(k) + lines = _stmt( + "for i in range(n):\n" + " if i > 0:\n" + " y[i] = i\n", + ctx, + ) + joined = "\n".join(lines) + assert "for (int i = 0;" in joined + assert "if (" in joined + assert "y.data[i]" in joined diff --git a/tests/unit/test_gpu_translate.py b/tests/unit/test_gpu_translate.py new file mode 100644 index 0000000..99a735f --- /dev/null +++ b/tests/unit/test_gpu_translate.py @@ -0,0 +1,135 @@ +"""Unit tests for translate_function_for_gpu (no device required).""" + +from __future__ import annotations + +import pytest + +from cthreads.gpu import BlockIdx, GlobalIdx, ThreadIdx +from cthreads.gpu.compiler.translation.translate import translate_function_for_gpu +from cthreads.gpu.gpu_kernel_meta import build_gpu_kernel_meta + + +def test_translate_saxpy_source_shape(): + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + r = translate_function_for_gpu(saxpy) + meta = build_gpu_kernel_meta(saxpy) + assert r.func_name == "saxpy" + assert r.binding_count == meta.binding_count + assert r.scalar_bytes == meta.scalar_bytes + assert r.source.startswith("#version 450") + assert "gl_GlobalInvocationID.x" in r.source + assert "y.data[" in r.source + assert r.spirv is None + + +@pytest.mark.parametrize("local", [1, 32, 64, 128]) +def test_translate_custom_local_size(local): + def k(n: int) -> None: + i: int = GlobalIdx.x + if i >= n: + return + + r = translate_function_for_gpu(k, local_size_x=local) + assert f"layout(local_size_x = {local})" in r.source + assert r.local_size_x == local + + +def test_translate_rejects_bad_signature(): + def bad(x: dict[str, int]) -> None: + pass + + with pytest.raises(TypeError): + translate_function_for_gpu(bad) + + +def test_translate_pass_only_body(): + def k(n: int) -> None: + pass + + r = translate_function_for_gpu(k) + assert "void main()" in r.source + assert "binding = 0" in r.source + + +def test_translate_lists_only(): + def k(x: list[int], y: list[int]) -> None: + i: int = GlobalIdx.x + y[i] = x[i] + + r = translate_function_for_gpu(k) + assert "buffer Scalars" not in r.source + assert "binding = 0" in r.source + assert "binding = 1" in r.source + assert r.scalar_bytes == 0 + assert r.binding_count == 2 + + +def test_translate_for_range_and_while(): + def k(n: int, y: list[int]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + s: int = 0 + for j in range(n): + s += j + while s > 0: + s -= 1 + break + y[i] = s + + r = translate_function_for_gpu(k) + assert "for (int j =" in r.source + assert "while (" in r.source + + +def test_translate_thread_and_block_idx(): + def k(n: int, y: list[int]) -> None: + i: int = GlobalIdx.x + t: int = ThreadIdx.x + b: int = BlockIdx.x + if i >= n: + return + y[i] = t + b + + r = translate_function_for_gpu(k) + assert "gl_LocalInvocationID.x" in r.source + assert "gl_WorkGroupID.x" in r.source + + +def test_translate_bool_and_or(): + def k(n: int, flag: bool, y: list[int]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + if flag and i > 0 or i < 0: + y[i] = 1 + else: + y[i] = 0 + + r = translate_function_for_gpu(k) + assert "&&" in r.source and "||" in r.source + + +def test_translate_pow_rejected(): + def k(a: float, y: list[float]) -> None: + i: int = GlobalIdx.x + y[i] = a ** 2.0 + + with pytest.raises(TypeError, match=r"\*\*|pow"): + translate_function_for_gpu(k) + + +def test_translate_result_fields(): + def k(n: int) -> None: + pass + + r = translate_function_for_gpu(k) + assert r.func_name == "k" + assert isinstance(r.source, str) + assert r.local_size_x == 64 + assert r.spirv is None From 0ec8c6086b2db17257bf64a36df8f7f16434d833 Mon Sep 17 00:00:00 2001 From: T-Karu-smaecs Date: Mon, 14 Sep 2026 17:12:19 +0200 Subject: [PATCH 3/4] skip gpu tests on gh actions --- tests/helpers_gpu.py | 26 ++++++++++++++++++++++++++ tests/unit/test_gpu_pipeline.py | 6 +++++- tests/unit/test_gpu_reserved_names.py | 18 ++++++------------ tests/unit/test_gpu_spirv.py | 14 ++++---------- 4 files changed, 41 insertions(+), 23 deletions(-) diff --git a/tests/helpers_gpu.py b/tests/helpers_gpu.py index ee5d5ca..6078775 100644 --- a/tests/helpers_gpu.py +++ b/tests/helpers_gpu.py @@ -4,6 +4,7 @@ from __future__ import annotations +import os from types import ModuleType import cthreads.gpu.runtime as runtime_mod @@ -14,3 +15,28 @@ def prepare_module() -> ModuleType: Return `cthreads.gpu.runtime` (holds `_gpu_prepared` / prepare / gpu). """ return runtime_mod + + +def on_github_actions() -> bool: + """True when running under GitHub Actions CI.""" + return os.environ.get("GITHUB_ACTIONS", "").lower() == "true" + + +def glsl_compiler_available() -> bool: + """ + True when in-process compile_glsl or glslc can compile a tiny compute shader. + + Always False on GitHub Actions for now (CI builds without CTHREADS_GPU / + glslang). Broad except: probe must never fail collection. + """ + if on_github_actions(): + return False + try: + from cthreads.gpu.compiler.translation.spirv import compile_glsl_to_spirv + + compile_glsl_to_spirv( + "#version 450\nlayout(local_size_x = 1) in;\nvoid main() {}\n" + ) + return True + except Exception: + return False diff --git a/tests/unit/test_gpu_pipeline.py b/tests/unit/test_gpu_pipeline.py index df02cd5..8a74c72 100644 --- a/tests/unit/test_gpu_pipeline.py +++ b/tests/unit/test_gpu_pipeline.py @@ -11,7 +11,7 @@ import pytest -from helpers_gpu import prepare_module +from helpers_gpu import glsl_compiler_available, prepare_module from cthreads.frontend.Registry import REGISTRY from cthreads.gpu import GlobalIdx, Gpu, ThreadIdx, available, gpu, shutdown @@ -32,6 +32,10 @@ def _reset(): prepare_mod._gpu_prepared = False +@pytest.mark.skipif( + not glsl_compiler_available(), + reason="no GLSL compiler (skipped on GitHub Actions / CPU-only builds)", +) def test_pipeline_translate_source_is_shaderc_ready(): def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: i: int = GlobalIdx.x diff --git a/tests/unit/test_gpu_reserved_names.py b/tests/unit/test_gpu_reserved_names.py index 0a0f2b0..485ac46 100644 --- a/tests/unit/test_gpu_reserved_names.py +++ b/tests/unit/test_gpu_reserved_names.py @@ -4,21 +4,17 @@ import pytest +from helpers_gpu import glsl_compiler_available + from cthreads.gpu.compiler.translation.spirv import compile_glsl_to_spirv from cthreads.gpu.compiler.translation.translate import translate_function_for_gpu - -def _compiler_available() -> bool: - try: - compile_glsl_to_spirv( - "#version 450\nlayout(local_size_x = 1) in;\nvoid main() {}\n" - ) - return True - except RuntimeError: - return False +pytestmark = pytest.mark.skipif( + not glsl_compiler_available(), + reason="no GLSL compiler (skipped on GitHub Actions / CPU-only builds)", +) -@pytest.mark.skipif(not _compiler_available(), reason="no GLSL compiler") @pytest.mark.parametrize( "param", [ @@ -48,7 +44,6 @@ def k(n: int, {param}: list[float]) -> None: translate_function_for_gpu(mod.k, compile_spirv=True) -@pytest.mark.skipif(not _compiler_available(), reason="no GLSL compiler") def test_safe_param_names_compile(tmp_module): mod = tmp_module( """ @@ -67,7 +62,6 @@ def k(n: int, ys: list[float]) -> None: assert "} ys;" in r.source -@pytest.mark.skipif(not _compiler_available(), reason="no GLSL compiler") def test_compile_error_includes_hint(): with pytest.raises(RuntimeError, match="Hint:.*reserved"): compile_glsl_to_spirv("#version 450\nvoid main() { not_a_type x; }\n") diff --git a/tests/unit/test_gpu_spirv.py b/tests/unit/test_gpu_spirv.py index f8f965a..ccad9a1 100644 --- a/tests/unit/test_gpu_spirv.py +++ b/tests/unit/test_gpu_spirv.py @@ -4,6 +4,8 @@ import pytest +from helpers_gpu import glsl_compiler_available + from cthreads.gpu import GlobalIdx from cthreads.gpu.compiler.translation.spirv import compile_glsl_to_spirv from cthreads.gpu.compiler.translation.translate import translate_function_for_gpu @@ -18,17 +20,9 @@ """ -def _compiler_available() -> bool: - try: - compile_glsl_to_spirv(_MIN_COMP) - return True - except RuntimeError: - return False - - pytestmark = pytest.mark.skipif( - not _compiler_available(), - reason="no native compile_glsl / glslc available", + not glsl_compiler_available(), + reason="no GLSL compiler (skipped on GitHub Actions / CPU-only builds)", ) From 76293327e50da2bf5dafdbc37e1c544aa8b01d36 Mon Sep 17 00:00:00 2001 From: T-Karu-smaecs Date: Tue, 15 Sep 2026 21:31:55 +0200 Subject: [PATCH 4/4] Add GpuArena residency and sort-grid GPU SPH. Introduce process-wide GpuState and Python GpuArena so bound lists stay on device across launches. Support join(download=False), borrow resident buffers in launch, and skip GLSL compile tests on GitHub Actions. --- src/cthreads/cpp/CMakeLists.txt | 1 + src/cthreads/cpp/bindings/gpu_module.cpp | 132 +++++++- src/cthreads/cpp/bindings/gpu_module.hpp | 2 +- src/cthreads/cpp/gpu/headers/module.hpp | 14 +- src/cthreads/cpp/gpu/headers/state.hpp | 245 ++++++++++++++ src/cthreads/cpp/gpu/impl/context.cpp | 5 +- src/cthreads/cpp/gpu/impl/module.cpp | 105 +++++- src/cthreads/cpp/gpu/impl/state.cpp | 200 ++++++++++++ src/cthreads/python/cthreads/gpu/__init__.py | 2 + .../python/cthreads/gpu/_ext_gpu_api.py | 16 + src/cthreads/python/cthreads/gpu/arena.py | 309 ++++++++++++++++++ .../translation/plugins/math_calls.py | 34 +- src/cthreads/python/cthreads/gpu/runtime.py | 71 +++- tests/unit/test_gpu_arena.py | 100 ++++++ 14 files changed, 1202 insertions(+), 34 deletions(-) create mode 100644 src/cthreads/cpp/gpu/headers/state.hpp create mode 100644 src/cthreads/cpp/gpu/impl/state.cpp create mode 100644 src/cthreads/python/cthreads/gpu/arena.py create mode 100644 tests/unit/test_gpu_arena.py diff --git a/src/cthreads/cpp/CMakeLists.txt b/src/cthreads/cpp/CMakeLists.txt index 3b06479..35e540a 100644 --- a/src/cthreads/cpp/CMakeLists.txt +++ b/src/cthreads/cpp/CMakeLists.txt @@ -172,6 +172,7 @@ if(CTHREADS_GPU) "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/memory.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/pack.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/shader_cache.cpp" + "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/state.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/shader.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/descriptors.cpp" "${CMAKE_CURRENT_SOURCE_DIR}/gpu/impl/module.cpp" diff --git a/src/cthreads/cpp/bindings/gpu_module.cpp b/src/cthreads/cpp/bindings/gpu_module.cpp index bc9e577..fe28d21 100644 --- a/src/cthreads/cpp/bindings/gpu_module.cpp +++ b/src/cthreads/cpp/bindings/gpu_module.cpp @@ -7,16 +7,20 @@ #include "../gpu/headers/context.hpp" #include "../gpu/headers/compile_glsl.hpp" +#include "../gpu/headers/memory.hpp" #include "../gpu/headers/module.hpp" #include "../gpu/headers/shader_cache.hpp" +#include "../gpu/headers/state.hpp" #include +#include #include #include #include #include #include +#include #include namespace py = pybind11; @@ -63,10 +67,12 @@ void bind_gpu(py::module_& parent) { ) .def( "join", - [](cthreads::gpu::SpawnedGpuKernel& self) { - self.join(cthreads::gpu::context()); + [](cthreads::gpu::SpawnedGpuKernel& self, bool download) { + self.join(cthreads::gpu::context(), download); }, - "Wait for the GPU fence, download ref lists, release inflight state." + py::arg("download") = true, + "Wait for the GPU fence; if download=True, write ref lists back, then " + "release inflight state." ) .def( "done", @@ -141,6 +147,126 @@ void bind_gpu(py::module_& parent) { "Requires the kernel symbol to already be in ShaderCache." ); + // Named device-local buffer registry (singleton). Not a public DeviceBuffer: + // Python only sees names + sizes; VkBuffer stays inside the registry. + py::class_< + cthreads::gpu::memory::GpuState, + std::unique_ptr< + cthreads::gpu::memory::GpuState, + py::nodelete>>(g, "GpuState") + .def_static( + "instance", + []() -> cthreads::gpu::memory::GpuState& { + return cthreads::gpu::memory::GpuState::getInstance(); + }, + py::return_value_policy::reference, + "Process-wide GpuState singleton." + ) + .def( + "contains", + &cthreads::gpu::memory::GpuState::contains, + py::arg("name"), + "True if name is already registered." + ) + .def( + "size", + &cthreads::gpu::memory::GpuState::size, + "Number of registered buffers." + ) + .def( + "names", + &cthreads::gpu::memory::GpuState::names, + "Snapshot of registered names (order is not meaningful)." + ) + .def( + "add", + [](cthreads::gpu::memory::GpuState& self, + const std::string& name, + std::uint64_t nbytes) { + cthreads::gpu::init(); + auto buffer = cthreads::gpu::memory::create_buffer( + cthreads::gpu::context(), + static_cast(nbytes), + cthreads::gpu::memory::BufferKind::DeviceLocal + ); + self.add(name, std::move(buffer)); + }, + py::arg("name"), + py::arg("nbytes"), + "Allocate a device-local buffer of nbytes and register it under name. " + "Duplicate or empty names raise." + ) + .def( + "remove", + [](cthreads::gpu::memory::GpuState& self, const std::string& name) { + self.remove(cthreads::gpu::context(), name); + }, + py::arg("name"), + "Destroy and unregister name. Raises if unknown or in_use." + ) + .def( + "buffer_size", + [](cthreads::gpu::memory::GpuState& self, const std::string& name) { + return static_cast(self.get(name).size); + }, + py::arg("name"), + "Byte size of the buffer registered under name." + ) + .def( + "is_in_use", + &cthreads::gpu::memory::GpuState::is_in_use, + py::arg("name"), + "True if name is checked out for a launch." + ) + .def( + "mark_in_use", + &cthreads::gpu::memory::GpuState::mark_in_use, + py::arg("name"), + "Check out name for a launch. Raises if unknown or already in_use." + ) + .def( + "release_in_use", + &cthreads::gpu::memory::GpuState::release_in_use, + py::arg("name"), + "Clear in_use after a launch finishes. Raises if unknown or not in_use." + ) + .def( + "upload", + [](cthreads::gpu::memory::GpuState& self, + const std::string& name, + const py::bytes& data) { + cthreads::gpu::init(); + const std::string raw = data; + self.upload( + cthreads::gpu::context(), + name, + raw.empty() ? nullptr : raw.data(), + static_cast(raw.size()) + ); + }, + py::arg("name"), + py::arg("data"), + "Upload host bytes into a registered device-local buffer (H2D)." + ) + .def( + "download", + [](cthreads::gpu::memory::GpuState& self, const std::string& name) { + cthreads::gpu::init(); + const std::uint64_t nbytes = static_cast( + self.get(name).size); + std::string raw(static_cast(nbytes), '\0'); + self.download( + cthreads::gpu::context(), + name, + raw.empty() ? nullptr : raw.data(), + static_cast(nbytes) + ); + return py::bytes(raw); + }, + py::arg("name"), + "Download a registered device-local buffer into bytes (D2H)." + ); + // Test-only pack round-trips live in a separate submodule / translation unit // so product bindings stay small. Not re-exported by cthreads.gpu. bind_gpu_testing(g); diff --git a/src/cthreads/cpp/bindings/gpu_module.hpp b/src/cthreads/cpp/bindings/gpu_module.hpp index d66e9d5..a2490c3 100644 --- a/src/cthreads/cpp/bindings/gpu_module.hpp +++ b/src/cthreads/cpp/bindings/gpu_module.hpp @@ -6,6 +6,6 @@ namespace py = pybind11; /** * Register ``cthreads._ext.gpu`` (probe API, launch_gpu_kernel / GpuJob, - * and test-only ``testing`` submodule). + * GpuState singleton, and test-only ``testing`` submodule). */ void bind_gpu(py::module_& parent); diff --git a/src/cthreads/cpp/gpu/headers/module.hpp b/src/cthreads/cpp/gpu/headers/module.hpp index e5dfa20..dba01bd 100644 --- a/src/cthreads/cpp/gpu/headers/module.hpp +++ b/src/cthreads/cpp/gpu/headers/module.hpp @@ -82,6 +82,12 @@ struct SpawnedGpuKernel { // Ref list slots only; value scalars are not written back. std::vector writeback_lists; + // Parallel to pack.container_slots: true => destroy_buffer on release. + // False => borrowed from GpuState; handles cleared without destroy. + std::vector container_owned; + // GpuState names marked in_use for this launch; released in release_inflight. + std::vector resident_names; + bool finished = false; std::mutex done_mu; std::condition_variable done_cv; @@ -102,13 +108,15 @@ struct SpawnedGpuKernel { void start(); /** - * Wait until the GPU fence signals, then download/writeback and release - * inflight GPU objects. Rethrows eptr if set. Idempotent after finished. + * Wait until the GPU fence signals, optionally download/writeback, then + * release inflight GPU objects. Rethrows eptr if set. Idempotent after finished. * * #### Parameters: * - context: Context& = same device that created pack / submitted work. + * - download: bool = if true (default), download ref lists into values_keep. + * If false, skip writeback (resident buffers stay device-authoritative). */ - void join(Context& context); + void join(Context& context, bool download = true); /** * Block until done_flag is set (join or failure path). Does not download. diff --git a/src/cthreads/cpp/gpu/headers/state.hpp b/src/cthreads/cpp/gpu/headers/state.hpp new file mode 100644 index 0000000..b35135f --- /dev/null +++ b/src/cthreads/cpp/gpu/headers/state.hpp @@ -0,0 +1,245 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include "memory.hpp" + +namespace cthreads::gpu { +struct Context; +} + +/** + * Process-wide registry of named device-local GpuBuffers. + * + * Owns GPU memory outside of a single launch so kernels can reuse the same + * buffers across dispatches. Names are unique: adding a duplicate name throws. + * Intended for a later Python binding of this singleton so host code can + * allocate, upload, launch, download, and free without leaking or double-owning + * Vulkan handles. + * + * This is not a public DeviceBuffer type. Python should see a controlled state + * / arena API that calls into this registry; VkBuffer stays inside _ext. + * + * #### Technical terms: + * - GpuBuffer: one contiguous device (or staging) byte region (see memory.hpp). + * - Context: process-wide Vulkan connection (device, queue, loaded entry points). + * - in_use: true while a launch has checked out this name; blocks remove and a + * second mark_in_use until release_in_use. + * - singleton: one process-wide instance via getInstance(); not copyable. + */ +namespace cthreads::gpu::memory { + +/** + * One named row in GpuState: owned buffer plus launch checkout flag. + * + * #### Fields: + * - buffer: GpuBuffer = owned Vulkan allocation (moved in on add). + * - in_use: bool = true while a kernel job holds this name for dispatch. + */ +struct GpuStateEntry { + GpuBuffer buffer{}; + bool in_use = false; +}; + +/** + * Singleton map of unique string names to owned GpuBuffers. + * + * Thread-safe: every public method locks an internal mutex. Call clear(context) + * from Context shutdown before destroying the logical device so VkBuffer / + * VkDeviceMemory handles are freed while the device is still alive. + * + * Writers / readers: any native or pybind caller. Duplicate names are rejected + * on add. Removing or clearing while in_use is rejected on remove; clear on + * shutdown still destroys (process teardown). + */ +class GpuState { +private: + std::unordered_map _entries; + mutable std::mutex _mutex; + + GpuState() = default; + ~GpuState(); + +public: + static GpuState& getInstance(); + + GpuState(const GpuState&) = delete; + GpuState& operator=(const GpuState&) = delete; + GpuState(GpuState&&) = delete; + GpuState& operator=(GpuState&&) = delete; + + /** + * Destroy every registered buffer and empty the map. + * + * Call from Context shutdown before destroying the logical device. + * + * #### Parameters: + * - context: Context& = device used to destroy buffers + */ + void clear(cthreads::gpu::Context& context); + + /** + * True if `name` is already registered. + * + * #### Parameters: + * - name: const string& = registry key + * + * #### Returns: + * - bool = true when an entry exists for name + */ + bool contains(const std::string& name) const; + + /** + * Number of registered buffers. + * + * #### Returns: + * - size_t = entry count + */ + size_t size() const; + + /** + * Snapshot of all registered names (order is not meaningful). + * + * #### Returns: + * - vector = copy of keys for introspection / pybind + */ + std::vector names() const; + + /** + * Take ownership of a buffer under a unique name. + * + * Moves `buffer` into the registry. The caller must not use or destroy the + * moved-from GpuBuffer afterward (handles are null after a successful add). + * + * #### Parameters: + * - name: const string& = unique key (must be non-empty and not already used) + * - buffer: GpuBuffer&& = owned allocation to store (typically DeviceLocal) + * + * #### Throws: + * - runtime_error = empty name, duplicate name, or empty buffer handles + */ + void add(const std::string& name, GpuBuffer&& buffer); + + /** + * Destroy and unregister one buffer by name. + * + * #### Parameters: + * - context: Context& = same device that created the buffer + * - name: const string& = registry key + * + * #### Throws: + * - runtime_error = unknown name, or entry is in_use (release first) + */ + void remove(cthreads::gpu::Context& context, const std::string& name); + + /** + * Mutable reference to the buffer stored under `name`. + * + * #### Parameters: + * - name: const string& = registry key + * + * #### Returns: + * - GpuBuffer& = owned buffer (valid until remove/clear) + * + * #### Throws: + * - runtime_error = unknown name + */ + GpuBuffer& get(const std::string& name); + + /** + * Const reference to the buffer stored under `name`. + * + * #### Parameters: + * - name: const string& = registry key + * + * #### Returns: + * - const GpuBuffer& = owned buffer (valid until remove/clear) + * + * #### Throws: + * - runtime_error = unknown name + */ + const GpuBuffer& get(const std::string& name) const; + + /** + * True if the named entry is checked out for a launch. + * + * #### Parameters: + * - name: const string& = registry key + * + * #### Returns: + * - bool = entry.in_use + * + * #### Throws: + * - runtime_error = unknown name + */ + bool is_in_use(const std::string& name) const; + + /** + * Mark a buffer as checked out so remove and a second mark fail. + * + * Call from the launch path before recording work that binds this buffer. + * Pair with release_in_use after the fence wait (or on launch failure). + * + * #### Parameters: + * - name: const string& = registry key + * + * #### Throws: + * - runtime_error = unknown name, or already in_use + */ + void mark_in_use(const std::string& name); + + /** + * Clear the in_use flag after a launch finishes (or aborts). + * + * #### Parameters: + * - name: const string& = registry key + * + * #### Throws: + * - runtime_error = unknown name, or entry was not in_use + */ + void release_in_use(const std::string& name); + + /** + * Upload host bytes into a registered device-local buffer (H2D). + * + * #### Parameters: + * - context: Context& = same device that created the buffer + * - name: const string& = registry key + * - data: const void* = host source bytes + * - size: VkDeviceSize = bytes to copy; must be > 0 and <= buffer.size + * + * #### Throws: + * - runtime_error = unknown name, in_use, bad size, or transfer failure + */ + void upload( + cthreads::gpu::Context& context, + const std::string& name, + const void* data, + VkDeviceSize size + ); + + /** + * Download a registered device-local buffer into host bytes (D2H). + * + * #### Parameters: + * - context: Context& = same device that created the buffer + * - name: const string& = registry key + * - data: void* = host destination bytes + * - size: VkDeviceSize = bytes to copy; must be > 0 and <= buffer.size + * + * #### Throws: + * - runtime_error = unknown name, in_use, bad size, or transfer failure + */ + void download( + cthreads::gpu::Context& context, + const std::string& name, + void* data, + VkDeviceSize size + ); +}; + +} // namespace cthreads::gpu::memory diff --git a/src/cthreads/cpp/gpu/impl/context.cpp b/src/cthreads/cpp/gpu/impl/context.cpp index 944765d..70d4b57 100644 --- a/src/cthreads/cpp/gpu/impl/context.cpp +++ b/src/cthreads/cpp/gpu/impl/context.cpp @@ -7,6 +7,7 @@ #include "../headers/memory.hpp" #include "../headers/shader_cache.hpp" +#include "../headers/state.hpp" #if defined(_WIN32) // Windows (32-bit or 64-bit) @@ -428,9 +429,11 @@ namespace { } void shutdown_unlocked(Context& c) { - // Children before parents: launch engine, transfer engine, shader cache, then device. + // Children before parents: launch engine, transfer engine, named + // GpuState buffers, shader cache, then device. shutdown_launch_engine(c); shutdown_transfer_engine(c); + memory::GpuState::getInstance().clear(c); shader::ShaderCache::getInstance().clear(c); // 1) release logical device diff --git a/src/cthreads/cpp/gpu/impl/module.cpp b/src/cthreads/cpp/gpu/impl/module.cpp index 0927d93..d1f7102 100644 --- a/src/cthreads/cpp/gpu/impl/module.cpp +++ b/src/cthreads/cpp/gpu/impl/module.cpp @@ -4,6 +4,8 @@ #include "../headers/shader_cache.hpp" #include "../headers/pack.hpp" #include "../headers/descriptors.hpp" +#include "../headers/state.hpp" +#include "../headers/memory.hpp" #include #include @@ -56,9 +58,27 @@ void release_inflight(Context& context, SpawnedGpuKernel& job) { pack::free_set(context, job.descriptor_pool, job.descriptor_set); } pack::destroy_pool(context, job.descriptor_pool); - // destroy the pack + + // Borrowed GpuState buffers: drop handles without destroy_buffer. + for (size_t i = 0; i < job.pack.container_slots.size(); ++i) { + const bool owned = + (i < job.container_owned.size()) ? (job.container_owned[i] != 0) : true; + if (!owned) { + job.pack.container_slots[i].buffer = memory::GpuBuffer{}; + } + } + for (const std::string& name : job.resident_names) { + try { + memory::GpuState::getInstance().release_in_use(name); + } catch (...) { + // Best-effort on teardown / double-release paths. + } + } + job.resident_names.clear(); + job.container_owned.clear(); + pack::destroy_gpu_pack(context, job.pack); - job.symbol.clear(); // clear the symbol + job.symbol.clear(); job.writeback_lists.clear(); job.values_keep.reset(); } @@ -293,7 +313,7 @@ bool SpawnedGpuKernel::done() { return done_flag; } -void SpawnedGpuKernel::join(Context& context) { +void SpawnedGpuKernel::join(Context& context, bool download) { if (finished) { if (eptr) { std::rethrow_exception(eptr); @@ -320,7 +340,8 @@ void SpawnedGpuKernel::join(Context& context) { } // Permanent list writeback path (Threadable/schema marshal is later). - if (!writeback_lists.empty()) { + // download=false: fence only; resident GpuState buffers stay authoritative. + if (download && !writeback_lists.empty()) { compute_to_transfer_barrier(context); writeback_ref_lists(context, *this); } @@ -512,6 +533,21 @@ std::shared_ptr launch_gpu_kernel( } } + // Optional residency: value_index -> GpuState name (Python GpuArena). + // Those list SSBOs are borrowed; skip create/upload for them. + std::unordered_map resident_by_value_index; + if (meta.contains("resident") && !meta["resident"].is_none()) { + py::dict res = meta["resident"].cast(); + for (auto item : res) { + const size_t value_index = + py::reinterpret_borrow(item.first).cast(); + const std::string state_name = + py::reinterpret_borrow(item.second) + .cast(); + resident_by_value_index.emplace(value_index, state_name); + } + } + // Job owns GPU objects from here on so failures can release_inflight. auto job = std::make_shared(); job->symbol = symbol; @@ -541,12 +577,53 @@ std::shared_ptr launch_gpu_kernel( } try { - // init the gpu pack (device-local scalar blob + one buffer per list) - job->pack = pack::create_gpu_pack( - context, - scalar_bytes, - container_specs - ); + // Build pack: owned scalars + per-list either create or borrow from GpuState. + job->container_owned.assign(container_plans.size(), 1); + if (scalar_bytes > 0) { + job->pack.scalar_buffer = memory::create_buffer( + context, + static_cast(scalar_bytes), + memory::BufferKind::DeviceLocal + ); + } + job->pack.container_slots.resize(container_plans.size()); + memory::GpuState& state = memory::GpuState::getInstance(); + + for (size_t c = 0; c < container_plans.size(); ++c) { + const ContainerSlotPlan& plan = container_plans[c]; + job->pack.container_slots[c].spec = + pack::ContainerSpec{plan.elem_bytes, plan.numel}; + if (plan.numel == 0) { + continue; + } + + auto res_it = resident_by_value_index.find(plan.value_index); + if (res_it != resident_by_value_index.end()) { + const std::string& state_name = res_it->second; + // Checkout before reading handles so remove cannot race. + state.mark_in_use(state_name); + job->resident_names.push_back(state_name); + memory::GpuBuffer& registered = state.get(state_name); + const VkDeviceSize need = + static_cast(plan.elem_bytes * plan.numel); + if (registered.size < need) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: resident buffer '" + + state_name + "' is too small for list arg"); + } + // Borrow handles; GpuState remains the owner. + job->pack.container_slots[c].buffer = registered; + job->container_owned[c] = 0; + continue; + } + + job->pack.container_slots[c].buffer = memory::create_buffer( + context, + static_cast(plan.elem_bytes * plan.numel), + memory::BufferKind::DeviceLocal + ); + job->container_owned[c] = 1; + } // Pack Python scalars into a host byte blob, then upload through staging. if (scalar_bytes > 0) { @@ -563,13 +640,13 @@ std::shared_ptr launch_gpu_kernel( context, job->pack, scalar_host.data(), scalar_bytes); } - // Upload each list container from ordered_values (pack slot order). + // Upload each non-resident list container from ordered_values. for (size_t c = 0; c < container_plans.size(); ++c) { const ContainerSlotPlan& plan = container_plans[c]; - if (plan.numel == 0) { - continue; // empty slot: no VkBuffer; update_descriptors still rejects empty for now + if (plan.numel == 0 || job->container_owned[c] == 0) { + continue; // empty or resident (already on device) } - py::list list_val = ordered_values[plan.value_index].cast(); // get the py side list that was passed in the kernel call + py::list list_val = ordered_values[plan.value_index].cast(); if (plan.elem_kind == "float") { std::vector host(plan.numel); for (size_t j = 0; j < plan.numel; ++j) { diff --git a/src/cthreads/cpp/gpu/impl/state.cpp b/src/cthreads/cpp/gpu/impl/state.cpp new file mode 100644 index 0000000..749dd8c --- /dev/null +++ b/src/cthreads/cpp/gpu/impl/state.cpp @@ -0,0 +1,200 @@ +#include "../headers/state.hpp" +#include "../headers/context.hpp" + +#include +#include + +namespace cthreads::gpu::memory { + +GpuState& GpuState::getInstance() { + static GpuState instance; + return instance; +} + +GpuState::~GpuState() { + // Static teardown order vs Context is undefined. Shutdown must clear first + // so Vulkan handles are already destroyed; only drop the map here. + std::lock_guard lock(_mutex); + _entries.clear(); +} + +bool GpuState::contains(const std::string& name) const { + std::lock_guard lock(_mutex); + return _entries.find(name) != _entries.end(); +} + +size_t GpuState::size() const { + std::lock_guard lock(_mutex); + return _entries.size(); +} + +std::vector GpuState::names() const { + std::lock_guard lock(_mutex); + std::vector out; + out.reserve(_entries.size()); + for (const auto& pair : _entries) { + out.push_back(pair.first); + } + return out; +} + +void GpuState::add(const std::string& name, GpuBuffer&& buffer) { + if (name.empty()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState.add requires a non-empty " + "name"); + } + if (buffer.buffer == VK_NULL_HANDLE || buffer.memory == VK_NULL_HANDLE || + buffer.size == 0) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState.add requires a non-empty " + "GpuBuffer"); + } + + std::lock_guard lock(_mutex); + if (_entries.find(name) != _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState name already registered: " + + name); + } + + GpuStateEntry entry{}; + // GpuBuffer has no custom move that clears handles, so copy then null the + // caller only after the map insert succeeds (otherwise we would leak). + entry.buffer = buffer; + entry.in_use = false; + _entries.emplace(name, std::move(entry)); + buffer = GpuBuffer{}; +} + +void GpuState::remove(Context& context, const std::string& name) { + std::lock_guard lock(_mutex); + auto it = _entries.find(name); + if (it == _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState unknown name: " + name); + } + if (it->second.in_use) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState cannot remove in-use " + "buffer: " + + name); + } + destroy_buffer(context, it->second.buffer); + _entries.erase(it); +} + +GpuBuffer& GpuState::get(const std::string& name) { + std::lock_guard lock(_mutex); + auto it = _entries.find(name); + if (it == _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState unknown name: " + name); + } + return it->second.buffer; +} + +const GpuBuffer& GpuState::get(const std::string& name) const { + std::lock_guard lock(_mutex); + auto it = _entries.find(name); + if (it == _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState unknown name: " + name); + } + return it->second.buffer; +} + +bool GpuState::is_in_use(const std::string& name) const { + std::lock_guard lock(_mutex); + auto it = _entries.find(name); + if (it == _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState unknown name: " + name); + } + return it->second.in_use; +} + +void GpuState::mark_in_use(const std::string& name) { + std::lock_guard lock(_mutex); + auto it = _entries.find(name); + if (it == _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState unknown name: " + name); + } + if (it->second.in_use) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState buffer already in use: " + + name); + } + it->second.in_use = true; +} + +void GpuState::release_in_use(const std::string& name) { + std::lock_guard lock(_mutex); + auto it = _entries.find(name); + if (it == _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState unknown name: " + name); + } + if (!it->second.in_use) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState buffer is not in use: " + + name); + } + it->second.in_use = false; +} + +void GpuState::upload( + Context& context, + const std::string& name, + const void* data, + VkDeviceSize size +) { + std::lock_guard lock(_mutex); + auto it = _entries.find(name); + if (it == _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState unknown name: " + name); + } + if (it->second.in_use) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState cannot upload in-use " + "buffer: " + + name); + } + // upload_buffer validates size / kind; keep the map lock so remove cannot race. + upload_buffer(context, it->second.buffer, data, size); +} + +void GpuState::download( + Context& context, + const std::string& name, + void* data, + VkDeviceSize size +) { + std::lock_guard lock(_mutex); + auto it = _entries.find(name); + if (it == _entries.end()) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState unknown name: " + name); + } + if (it->second.in_use) { + throw std::runtime_error( + "cthreads.gpu.GpuInvalidArgument: GpuState cannot download in-use " + "buffer: " + + name); + } + download_buffer(context, it->second.buffer, data, size); +} + +void GpuState::clear(Context& context) { + std::lock_guard lock(_mutex); + for (auto& pair : _entries) { + // Process teardown: destroy even if a launch left in_use set. + destroy_buffer(context, pair.second.buffer); + pair.second.in_use = false; + } + _entries.clear(); +} + +} // namespace cthreads::gpu::memory diff --git a/src/cthreads/python/cthreads/gpu/__init__.py b/src/cthreads/python/cthreads/gpu/__init__.py index 2455106..2d508f0 100644 --- a/src/cthreads/python/cthreads/gpu/__init__.py +++ b/src/cthreads/python/cthreads/gpu/__init__.py @@ -31,6 +31,7 @@ init, shutdown, ) +from .arena import GpuArena from .runtime import GpuJob, compile, gpu, prepare @@ -42,6 +43,7 @@ def __getattr__(name: str): __all__ = [ "Gpu", + "GpuArena", "GpuJob", "BlockDim", "BlockIdx", diff --git a/src/cthreads/python/cthreads/gpu/_ext_gpu_api.py b/src/cthreads/python/cthreads/gpu/_ext_gpu_api.py index eebc9f0..befb337 100644 --- a/src/cthreads/python/cthreads/gpu/_ext_gpu_api.py +++ b/src/cthreads/python/cthreads/gpu/_ext_gpu_api.py @@ -171,3 +171,19 @@ def register_shader(symbol: str, spirv: bytes, binding_count: int) -> None: - Exception = native create/insert failures (unmapped) """ _require_ext_gpu().register_shader(symbol, spirv, binding_count) + + +def gpu_state() -> Any: + """ + Return the process-wide native GpuState singleton. + + Named device-local buffers live here outside of a single launch. Does not + expose Vulkan handles; use `add(name, nbytes)` / `remove(name)` / etc. + + #### Returns + - Any = native `_ext.gpu.GpuState` + + #### Raises + - RuntimeError = `_ext.gpu` is not built + """ + return _require_ext_gpu().GpuState.instance() diff --git a/src/cthreads/python/cthreads/gpu/arena.py b/src/cthreads/python/cthreads/gpu/arena.py new file mode 100644 index 0000000..f5ec6b4 --- /dev/null +++ b/src/cthreads/python/cthreads/gpu/arena.py @@ -0,0 +1,309 @@ +""" +GpuArena: bind Python lists into process GpuState for launch reuse. + +Option B launch style: pass the same list objects to `gpu()` after bind. +Launch checks id + length; resident buffers skip alloc/upload. Host edits +between syncs are the caller's responsibility (no proxy yet). +""" + +from __future__ import annotations + +import struct +import threading +import uuid +from dataclasses import dataclass +from typing import Any, Iterator + +from . import _ext_gpu_api +from .frontend.errors import GPUNotAvailable, GpuInvalidArgument + +_ELEM_BYTES: dict[str, int] = { + "bool": 4, + "int": 4, + "float": 4, + "double": 8, +} + +# id(list) -> BoundSlot for every live arena bind (process-wide Option B lookup). +_ID_TO_SLOT: dict[int, "BoundSlot"] = {} +_ID_LOCK = threading.Lock() + + +@dataclass +class BoundSlot: + """One arena-bound Python list and its GpuState name.""" + + arena_id: str + name: str + state_name: str + host: list[Any] + numel: int + elem_kind: str + elem_bytes: int + + +def lookup_resident(value: Any) -> BoundSlot | None: + """ + Return the BoundSlot for a Python list if it is currently arena-bound. + + #### Args: + - value: Any = launch argument (typically a list) + + #### Returns + - BoundSlot | None = slot when id(value) is registered + """ + if not isinstance(value, list): + return None + with _ID_LOCK: + return _ID_TO_SLOT.get(id(value)) + + +def infer_elem_kind(values: list[Any]) -> str: + """ + Infer GPU list elem_kind from the first element (bool before int). + + #### Args: + - values: list[Any] = non-empty host list + + #### Returns + - str = "bool" | "int" | "float" | "double" + + #### Raises + - GpuInvalidArgument = empty list or unsupported element type + """ + if not values: + raise GpuInvalidArgument( + "GpuArena.bind: cannot infer elem type from an empty list " + "(pass a non-empty list)" + ) + sample: Any = values[0] + if isinstance(sample, bool): + return "bool" + if isinstance(sample, int): + return "int" + if isinstance(sample, float): + return "float" + raise GpuInvalidArgument( + f"GpuArena.bind: unsupported list element type {type(sample)!r}" + ) + + +def list_to_bytes(values: list[Any], elem_kind: str) -> bytes: + """Pack a Python list into std430-friendly host bytes.""" + if elem_kind == "float": + return struct.pack(f"{len(values)}f", *[float(v) for v in values]) + if elem_kind == "double": + return struct.pack(f"{len(values)}d", *[float(v) for v in values]) + if elem_kind == "int": + return struct.pack(f"{len(values)}i", *[int(v) for v in values]) + if elem_kind == "bool": + return struct.pack( + f"{len(values)}i", *[1 if bool(v) else 0 for v in values] + ) + raise GpuInvalidArgument(f"unsupported elem_kind: {elem_kind!r}") + + +def bytes_into_list(data: bytes, values: list[Any], elem_kind: str) -> None: + """Write downloaded bytes back into the same Python list object.""" + n: int = len(values) + if elem_kind == "float": + unpacked = struct.unpack(f"{n}f", data) + for i, v in enumerate(unpacked): + values[i] = float(v) + return + if elem_kind == "double": + unpacked = struct.unpack(f"{n}d", data) + for i, v in enumerate(unpacked): + values[i] = float(v) + return + if elem_kind == "int": + unpacked = struct.unpack(f"{n}i", data) + for i, v in enumerate(unpacked): + values[i] = int(v) + return + if elem_kind == "bool": + unpacked = struct.unpack(f"{n}i", data) + for i, v in enumerate(unpacked): + values[i] = bool(v) + return + raise GpuInvalidArgument(f"unsupported elem_kind: {elem_kind!r}") + + +class GpuArena: + """ + Session that keeps named list buffers resident in GpuState. + + #### Example: + ``py + with GpuArena() as arena: + arena.bind(x=x, y=y) + for _ in range(100): + gpu(saxpy, n, 2.0, x, y).join(download=False) + arena.sync() + `` + """ + + def __init__(self) -> None: + self._id: str = uuid.uuid4().hex + self._slots: dict[str, BoundSlot] = {} + self._released: bool = False + + def __enter__(self) -> "GpuArena": + return self + + def __exit__(self, *args: Any) -> None: + self.release() + + def bind(self, **named_lists: list[Any]) -> "GpuArena": + """ + Allocate/upload device buffers for the given lists (by kwarg name). + + Reuses GpuState entries when the same slot name is rebound with the + same length and elem kind; otherwise removes and recreates. + + #### Args: + - **named_lists: list[Any] = keyword slot name -> Python list object + + #### Returns + - GpuArena = this arena (for chaining) + + #### Raises + - GPUNotAvailable = GPU extension missing + - GpuInvalidArgument = bad args, empty lists, duplicate list ids + """ + if self._released: + raise GpuInvalidArgument("GpuArena.bind: arena already released") + if not named_lists: + raise GpuInvalidArgument("GpuArena.bind: expected at least one list") + if not _ext_gpu_api.available(): + raise GPUNotAvailable( + "GPU is not available (build with CTHREADS_GPU=ON and a Vulkan device)" + ) + + state = _ext_gpu_api.gpu_state() + _ext_gpu_api.init() + + for slot_name, host in named_lists.items(): + if not isinstance(host, list): + raise GpuInvalidArgument( + f"GpuArena.bind: {slot_name!r} must be a list, got {type(host)!r}" + ) + elem_kind: str = infer_elem_kind(host) + elem_bytes: int = _ELEM_BYTES[elem_kind] + numel: int = len(host) + nbytes: int = numel * elem_bytes + state_name: str = f"{self._id}/{slot_name}" + host_id: int = id(host) + + with _ID_LOCK: + existing = _ID_TO_SLOT.get(host_id) + if existing is not None and ( + existing.arena_id != self._id or existing.name != slot_name + ): + raise GpuInvalidArgument( + f"GpuArena.bind: list for {slot_name!r} is already bound " + f"as {existing.arena_id}/{existing.name}" + ) + + # Replace prior slot with the same name if shape changed. + prior = self._slots.get(slot_name) + if prior is not None: + if ( + prior.host is host + and prior.numel == numel + and prior.elem_kind == elem_kind + ): + state.upload(state_name, list_to_bytes(host, elem_kind)) + continue + self._unbind_slot(prior, state) + + if state.contains(state_name): + state.remove(state_name) + state.add(state_name, nbytes) + state.upload(state_name, list_to_bytes(host, elem_kind)) + + slot = BoundSlot( + arena_id=self._id, + name=slot_name, + state_name=state_name, + host=host, + numel=numel, + elem_kind=elem_kind, + elem_bytes=elem_bytes, + ) + self._slots[slot_name] = slot + with _ID_LOCK: + _ID_TO_SLOT[host_id] = slot + return self + + def sync(self, *names: str) -> None: + """ + Download resident buffers into the bound Python lists. + + #### Args: + - *names: str = optional slot names; default = all bound slots + + #### Raises + - GpuInvalidArgument = unknown name or arena released + """ + if self._released: + raise GpuInvalidArgument("GpuArena.sync: arena already released") + state = _ext_gpu_api.gpu_state() + targets: list[BoundSlot] + if names: + targets = [] + for name in names: + slot = self._slots.get(name) + if slot is None: + raise GpuInvalidArgument( + f"GpuArena.sync: unknown slot {name!r}" + ) + targets.append(slot) + else: + targets = list(self._slots.values()) + + for slot in targets: + if len(slot.host) != slot.numel: + raise GpuInvalidArgument( + f"GpuArena.sync: list for {slot.name!r} changed length " + f"(was {slot.numel}, now {len(slot.host)}); rebind" + ) + data: bytes = bytes(state.download(slot.state_name)) + bytes_into_list(data, slot.host, slot.elem_kind) + + def release(self) -> None: + """Destroy all GpuState buffers owned by this arena and drop id map entries.""" + if self._released: + return + state = None + try: + if _ext_gpu_api.available(): + state = _ext_gpu_api.gpu_state() + except Exception: + state = None + for slot in list(self._slots.values()): + self._unbind_slot(slot, state) + self._slots.clear() + self._released = True + + def _unbind_slot(self, slot: BoundSlot, state: Any) -> None: + with _ID_LOCK: + cur = _ID_TO_SLOT.get(id(slot.host)) + if cur is slot: + del _ID_TO_SLOT[id(slot.host)] + if state is not None and state.contains(slot.state_name): + try: + state.remove(slot.state_name) + except Exception: + pass + self._slots.pop(slot.name, None) + + def __contains__(self, name: str) -> bool: + return name in self._slots + + def names(self) -> list[str]: + """Registered slot names in this arena.""" + return list(self._slots.keys()) + + def __iter__(self) -> Iterator[str]: + return iter(self._slots) diff --git a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/math_calls.py b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/math_calls.py index 979cdd8..3362059 100644 --- a/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/math_calls.py +++ b/src/cthreads/python/cthreads/gpu/compiler/translation/plugins/math_calls.py @@ -1,5 +1,5 @@ """ -Minimal GLSL math CallPlugins for @Gpu (sqrt first — needed for SPH forces). +Minimal GLSL math CallPlugins for @Gpu (sqrt / floor / int — SPH grid + forces). """ from __future__ import annotations @@ -12,7 +12,7 @@ class MathCallPlugin(CallPlugin): """ - Lower `sqrt(x)` / `math.sqrt(x)` to GLSL `sqrt(...)`. + Lower a small set of math / cast calls to GLSL. """ def try_lower( @@ -26,17 +26,31 @@ def try_lower( if len(node.args) != 1: return None fn = node.func - is_sqrt = False + arg = translate_expr(node.args[0], ctx) + if isinstance(fn, ast.Name) and fn.id == "sqrt": - is_sqrt = True - elif ( + return f"sqrt({arg})" + if ( isinstance(fn, ast.Attribute) and fn.attr == "sqrt" and isinstance(fn.value, ast.Name) and fn.value.id == "math" ): - is_sqrt = True - if not is_sqrt: - return None - arg = translate_expr(node.args[0], ctx) - return f"sqrt({arg})" + return f"sqrt({arg})" + + # Truncate toward -inf (GLSL floor); used for cell indices. + if isinstance(fn, ast.Name) and fn.id == "floor": + return f"floor({arg})" + if ( + isinstance(fn, ast.Attribute) + and fn.attr == "floor" + and isinstance(fn.value, ast.Name) + and fn.value.id == "math" + ): + return f"floor({arg})" + + # Python int(x) on floats -> GLSL int(x) (trunc toward zero). + if isinstance(fn, ast.Name) and fn.id == "int": + return f"int({arg})" + + return None diff --git a/src/cthreads/python/cthreads/gpu/runtime.py b/src/cthreads/python/cthreads/gpu/runtime.py index 0c98c92..931b076 100644 --- a/src/cthreads/python/cthreads/gpu/runtime.py +++ b/src/cthreads/python/cthreads/gpu/runtime.py @@ -9,8 +9,9 @@ from ..job import Job from . import _ext_gpu_api +from .arena import lookup_resident from .compiler.orchestrator import GpuCompileSession -from .frontend.errors import GPUNotAvailable, _map_error +from .frontend.errors import GPUNotAvailable, GpuInvalidArgument, _map_error from .gpu_kernel_meta import build_gpu_kernel_meta from .gpu_marshal import infer_group_count_x, ordered_values_for_meta @@ -24,6 +25,31 @@ class GpuJob(Job): Job wrapper for native GpuJob handles (void kernels; `result()` is None). """ + def join(self, download: bool = True) -> None: + """ + Wait for the GPU fence; optionally write ref lists back into Python. + + #### Args: + - download: bool = if True (default), download ref lists on join. + If False, skip writeback (use GpuArena.sync for resident lists). + + #### Returns + - None + """ + if not self._started: + self.start() + raw_join = self._raw.join + try: + raw_join(download) + except TypeError: + # Older native builds without the download argument. + if download is False: + raise GpuInvalidArgument( + "GpuJob.join(download=False) requires a rebuild with " + "CTHREADS_GPU residency support" + ) from None + raw_join() + def result(self) -> None: """ GPU kernels are writeback-only; there is no scalar return value. @@ -76,6 +102,42 @@ def prepare(force: bool = False) -> dict[str, Any]: return compile(force=force) +def _resident_meta_for_args( + meta: dict[str, Any], ordered: list[Any] +) -> dict[int, str]: + """ + Map value_index -> GpuState name for arena-bound list args. + + Raises if a bound list length no longer matches the registered numel. + """ + params = meta.get("params") + if not isinstance(params, list): + return {} + resident: dict[int, str] = {} + for i, param in enumerate(params): + if not isinstance(param, dict) or param.get("kind") != "list": + continue + slot = lookup_resident(ordered[i]) + if slot is None: + continue + host = ordered[i] + if not isinstance(host, list): + continue + if len(host) != slot.numel: + raise GpuInvalidArgument( + f"gpu(): bound list {slot.name!r} length changed " + f"(was {slot.numel}, now {len(host)}); call arena.bind again" + ) + meta_kind = param.get("elem_kind") + if meta_kind is not None and meta_kind != slot.elem_kind: + raise GpuInvalidArgument( + f"gpu(): bound list {slot.name!r} elem_kind {slot.elem_kind!r} " + f"does not match kernel {meta_kind!r}" + ) + resident[i] = slot.state_name + return resident + + def gpu( fn: Callable[..., Any], *args: Any, @@ -86,7 +148,8 @@ def gpu( Launch a `@Gpu` kernel and return a joinable job handle. Ensures GPU compile/emit has run, then submits via `launch_gpu_kernel`. - List arguments are written back in place on `join()`. + List arguments are written back in place on `join()` unless + `join(download=False)` is used with GpuArena-resident lists. #### Args: - fn: Callable = `@Gpu`-decorated kernel @@ -150,6 +213,10 @@ def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: if meta.get("group_count_z") is None: meta["group_count_z"] = 1 + resident = _resident_meta_for_args(meta, ordered) + if resident: + meta["resident"] = resident + try: raw = _ext_gpu_api.launch_gpu_kernel(meta, ordered) except Exception as exc: diff --git a/tests/unit/test_gpu_arena.py b/tests/unit/test_gpu_arena.py new file mode 100644 index 0000000..ab0011c --- /dev/null +++ b/tests/unit/test_gpu_arena.py @@ -0,0 +1,100 @@ +"""GpuArena residency: bind once, relaunch without re-upload, sync download.""" + +from __future__ import annotations + +import pytest + +from helpers_gpu import prepare_module + +from cthreads.frontend.Registry import REGISTRY +from cthreads.gpu import ( + GlobalIdx, + Gpu, + GpuArena, + GpuInvalidArgument, + available, + gpu, + shutdown, +) + + +pytestmark = pytest.mark.skipif( + not available(), + reason="GPU / Vulkan not available", +) + + +@pytest.fixture(autouse=True) +def _reset(): + prepare_mod = prepare_module() + REGISTRY.clear() + prepare_mod._gpu_prepared = False + yield + REGISTRY.clear() + try: + shutdown() + except Exception: + pass + prepare_mod._gpu_prepared = False + + +def test_arena_bind_reuse_and_sync(): + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + x = [1.0, 2.0, 3.0, 4.0] + y = [10.0, 20.0, 30.0, 40.0] + with GpuArena() as arena: + arena.bind(x=x, y=y) + gpu(saxpy, len(x), 2.0, x, y).join(download=False) + gpu(saxpy, len(x), 2.0, x, y).join(download=False) + # Host lists unchanged until sync (device authoritative). + assert y == [10.0, 20.0, 30.0, 40.0] + arena.sync() + assert y == pytest.approx([14.0, 28.0, 42.0, 56.0]) + + +def test_arena_join_download_true_still_writeback(): + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + x = [1.0, 2.0] + y = [0.0, 0.0] + with GpuArena() as arena: + arena.bind(x=x, y=y) + gpu(saxpy, len(x), 3.0, x, y).join(download=True) + assert y == pytest.approx([3.0, 6.0]) + + +def test_arena_length_mismatch_raises(): + @Gpu + def saxpy(n: int, a: float, x: list[float], y: list[float]) -> None: + i: int = GlobalIdx.x + if i >= n: + return + y[i] = a * x[i] + y[i] + + x = [1.0, 2.0, 3.0] + y = [0.0, 0.0, 0.0] + with GpuArena() as arena: + arena.bind(x=x, y=y) + x.append(4.0) + with pytest.raises(GpuInvalidArgument, match="length changed"): + gpu(saxpy, len(x), 1.0, x, y) + + +def test_arena_duplicate_list_bind_raises(): + x = [1.0, 2.0] + with GpuArena() as a1: + a1.bind(x=x) + with GpuArena() as a2: + with pytest.raises(GpuInvalidArgument, match="already bound"): + a2.bind(x=x)