From 8cfdd8b3713ddd2e0b3d8442e95fff813de7eaea Mon Sep 17 00:00:00 2001 From: Jeremy Howard Date: Fri, 21 Aug 2026 13:34:57 +1000 Subject: [PATCH] Add awaitable Python GPU operations and wheel publishing --- .gitignore | 2 + CHANGELOG.md | 4 +- CMakeLists.txt | 48 +++++++++++++------ DEV.md | 14 +++++- Makefile | 16 ------- README.md | 22 ++++++--- bindings/python/gpu_cpp.cpp | 84 ++++++++++++++++++++------------- bindings/python/test_gpu_cpp.py | 16 ++++--- gpu.hpp | 74 ++++++++++++++++++++++++++--- pyproject.toml | 30 ++++++++++++ python/gpu_cpp/__init__.py | 1 + python/gpu_cpp/_async.py | 6 +++ tools/build_dawn.sh | 3 +- tools/build_python_wheel.sh | 13 +++++ 14 files changed, 248 insertions(+), 85 deletions(-) delete mode 100644 Makefile create mode 100644 pyproject.toml create mode 100644 python/gpu_cpp/__init__.py create mode 100644 python/gpu_cpp/_async.py create mode 100755 tools/build_python_wheel.sh diff --git a/.gitignore b/.gitignore index a6f10b7..37af17f 100644 --- a/.gitignore +++ b/.gitignore @@ -19,6 +19,8 @@ third_party/emdawnwebgpu/* # cmake build directories out build +dist/ +__pycache__/ # clangd files .cache diff --git a/CHANGELOG.md b/CHANGELOG.md index 2740924..afd1f69 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,13 +16,15 @@ browser target. Chrome story covering browser errors, context reuse, dispatch, and readback. - Reproducible scripts pinning Dawn, Emdawn, emsdk, SPIRV-Tools, and SPIRV-Headers to exact revisions. +- A self-contained macOS ARM64 Python wheel and direct PyPI publish workflow. ### Changed - WebGPU objects now use Dawn's generated C++ RAII facade and normal C++ value semantics. - The Python binding now follows the C++ API and preserves NumPy shape and - `float16`, `float32`, and `int32` dtypes. + `float16`, `float32`, and `int32` dtypes. Futures retain their context and + support blocking waits or native `asyncio` dispatch and NumPy readback. - Host-side f16 uses native `_Float16` where supported; the portable core treats f16 as two-byte IEEE 754 storage. - Examples and the build are consolidated under the root CMake project and the diff --git a/CMakeLists.txt b/CMakeLists.txt index 89643b7..6363cc5 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -81,7 +81,7 @@ target_link_libraries(gpucpp INTERFACE Dawn::webgpu) endif() if(GPUCPP_BUILD_PYTHON) - find_package(Python 3.9 COMPONENTS Interpreter Development.Module REQUIRED) + find_package(Python 3.10 COMPONENTS Interpreter Development.Module REQUIRED) include(FetchContent) FetchContent_Declare(pybind11 GIT_REPOSITORY https://github.com/pybind/pybind11.git @@ -89,8 +89,26 @@ if(GPUCPP_BUILD_PYTHON) GIT_SHALLOW TRUE) FetchContent_MakeAvailable(pybind11) - pybind11_add_module(gpu_cpp bindings/python/gpu_cpp.cpp) - target_link_libraries(gpu_cpp PRIVATE gpucpp) + set(GPUCPP_PYTHON_PACKAGE_DIR + "${CMAKE_CURRENT_BINARY_DIR}/python/gpu_cpp") + file(MAKE_DIRECTORY "${GPUCPP_PYTHON_PACKAGE_DIR}") + configure_file(python/gpu_cpp/__init__.py + "${GPUCPP_PYTHON_PACKAGE_DIR}/__init__.py" COPYONLY) + configure_file(python/gpu_cpp/_async.py + "${GPUCPP_PYTHON_PACKAGE_DIR}/_async.py" COPYONLY) + + pybind11_add_module(_gpu_cpp bindings/python/gpu_cpp.cpp) + target_link_libraries(_gpu_cpp PRIVATE gpucpp) + set_target_properties(_gpu_cpp PROPERTIES + LIBRARY_OUTPUT_DIRECTORY "${GPUCPP_PYTHON_PACKAGE_DIR}") + if(APPLE) + set_target_properties(_gpu_cpp PROPERTIES INSTALL_RPATH "@loader_path") + elseif(UNIX) + set_target_properties(_gpu_cpp PROPERTIES INSTALL_RPATH "$ORIGIN") + endif() + + install(TARGETS _gpu_cpp LIBRARY DESTINATION gpu_cpp) + install(FILES "${DAWN_LIBRARY}" DESTINATION gpu_cpp) endif() if(GPUCPP_BUILD_EXAMPLES) @@ -114,7 +132,7 @@ if(GPUCPP_BUILD_EXAMPLES) endif() -enable_testing() +include(CTest) if(EMSCRIPTEN) add_library(gpu_cpp_web_bindings OBJECT bindings/web/gpu_cpp.cpp) target_link_libraries(gpu_cpp_web_bindings PRIVATE gpucpp) @@ -126,14 +144,16 @@ if(EMSCRIPTEN) "-sMODULARIZE=1" "-sEXPORT_ES6=1" "-sEXPORT_NAME=createGpuCpp" "-sENVIRONMENT=web" "-sALLOW_MEMORY_GROWTH=1") - add_executable(web_binding_test tests/web_binding_test.cpp - $) - set_target_properties(web_binding_test PROPERTIES SUFFIX ".html") - target_link_libraries(web_binding_test PRIVATE gpucpp) - target_link_options(web_binding_test PRIVATE "--bind" "--emrun" - "-sMODULARIZE=1" "-sEXPORT_ES6=1" - "-sALLOW_MEMORY_GROWTH=1") -else() + if(BUILD_TESTING) + add_executable(web_binding_test tests/web_binding_test.cpp + $) + set_target_properties(web_binding_test PROPERTIES SUFFIX ".html") + target_link_libraries(web_binding_test PRIVATE gpucpp) + target_link_options(web_binding_test PRIVATE "--bind" "--emrun" + "-sMODULARIZE=1" "-sEXPORT_ES6=1" + "-sALLOW_MEMORY_GROWTH=1") + endif() +elseif(BUILD_TESTING) set(TEST_SPIRV "${CMAKE_CURRENT_BINARY_DIR}/write42.spv") add_custom_command( OUTPUT "${TEST_SPIRV}" @@ -146,9 +166,9 @@ else() target_link_libraries(gpu_test PRIVATE gpucpp) add_test(NAME gpu COMMAND gpu_test "${TEST_SPIRV}") endif() -if(GPUCPP_BUILD_PYTHON) +if(BUILD_TESTING AND GPUCPP_BUILD_PYTHON) add_test(NAME python COMMAND "${Python_EXECUTABLE}" "${CMAKE_CURRENT_SOURCE_DIR}/bindings/python/test_gpu_cpp.py") set_tests_properties(python PROPERTIES - ENVIRONMENT "PYTHONPATH=$") + ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_BINARY_DIR}/python") endif() diff --git a/DEV.md b/DEV.md index f6b445c..d3d4d1e 100644 --- a/DEV.md +++ b/DEV.md @@ -57,7 +57,14 @@ reviewable. Avoid adding tests that merely restate Dawn validation; add a new story only when gpu.cpp itself owns meaningful behavior. The Python story covers NumPy upload/readback, explicit binding access, -uniform parameters, and reuse of a compiled kernel with updated tensor data. +uniform parameters, reuse of a compiled kernel, blocking waits, and native +`asyncio` dispatch and readback. + +`tools/build_python_wheel.sh` packages the staged macOS 12 Dawn runtime into a +self-contained ARM64 wheel and checks its metadata. Publish that fresh wheel +with `twine upload dist/*.whl`, using the developer's standard `~/.pypirc` +credentials. Publishing from other platforms belongs in CI rather than this +local release path. `./test --web` cross-compiles the same public core and Embind API against Emdawnwebgpu, opens the result with `emrun`, and checks rejected invalid WGSL, @@ -80,6 +87,11 @@ they write into `Context` error state, which foreground API calls surface. Native waits pump `ProcessEvents`; browser waits call `WaitAny`, allowing JSPI to suspend Wasm while JavaScript and WebGPU make progress. +Python futures retain shared ownership of their context. Their awaiter polls +`ProcessEvents` on the asyncio thread rather than moving Dawn work to an +executor thread. Async readback uses callback-owned storage, so abandoning a +future cannot leave Dawn writing into a freed NumPy buffer. + The browser Embind layer owns one persistent `Context`. Each `run()` creates short-lived tensors and a kernel from JavaScript typed arrays, rejects overlapping calls, and copies `readWrite` results back in place. Wasm diff --git a/Makefile b/Makefile deleted file mode 100644 index 5ee7d49..0000000 --- a/Makefile +++ /dev/null @@ -1,16 +0,0 @@ -.PHONY: all build dawn run test - -all: build - -build: - cmake -S . -B build -G Ninja -DCMAKE_BUILD_TYPE=Release - cmake --build build - -dawn: - tools/build_dawn.sh - -run: build - ./build/hello_gpu - -test: - ./test diff --git a/README.md b/README.md index 02c5e0d..d211ff1 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ driver; Mesa's Vulkan driver is sufficient for development and CI. ./test # configure, build, and run the native GPU stories ./test --rebuild-web # first browser setup ./test --web # build and run the browser story in Chrome -make run # run the hello-world example +./build/hello_gpu # run the hello-world example after building ``` `tools/build_dawn.sh` checks out exact revisions, builds a monolithic shared @@ -134,9 +134,10 @@ See [DEV.md](DEV.md) for the dependency layout and update process, and ## Python -The optional pybind11 module is built by default when gpu.cpp is the top-level -CMake project. It accepts C-contiguous NumPy arrays and preserves tensor shape -and `float16`, `float32`, or `int32` dtype information: +Install the self-contained macOS ARM64 wheel from PyPI with `pip install +gpu-cpp`. The optional pybind11 module is also built by default when gpu.cpp is +the top-level CMake project. It accepts C-contiguous NumPy arrays and preserves +tensor shape and `float16`, `float32`, or `int32` dtype information: ```python import numpy as np @@ -150,8 +151,17 @@ kernel = gpu.create_kernel( context, gpu.WGSL(source, workgroup_size=[4, 1, 1]), [gpu.read(gpu_values), gpu.read_write(gpu_output)]) dispatched = gpu.dispatch_kernel(context, kernel) -gpu.wait(context, dispatched) -result = gpu.to_numpy(context, gpu_output) +gpu.wait(dispatched) +result = gpu.wait(gpu.to_numpy(context, gpu_output)) +``` + +Futures own their context and support both blocking scripts and native +`asyncio`/notebook execution. Readback futures also own their destination array +and return it on completion: + +```python +await gpu.dispatch_kernel(context, kernel) +result = await gpu.to_numpy(context, gpu_output) ``` Set `GPUCPP_BUILD_PYTHON=OFF` when embedding gpu.cpp in a CMake project that diff --git a/bindings/python/gpu_cpp.cpp b/bindings/python/gpu_cpp.cpp index e292097..f07fac2 100644 --- a/bindings/python/gpu_cpp.cpp +++ b/bindings/python/gpu_cpp.cpp @@ -5,6 +5,7 @@ #include #include +#include #include #include #include @@ -61,6 +62,12 @@ std::span input(const py::buffer_info &buffer) { static_cast(buffer.size)}; } +struct PythonFuture { + std::shared_ptr context; + gpu::Future value; + py::object result; +}; + Tensor tensor(Context &context, const py::array &array) { requireContiguous(array); const auto type = numType(array); @@ -91,37 +98,43 @@ void upload(Context &context, const Tensor &tensor, const py::array &array) { } template -void download(Context &context, const Tensor &tensor, const py::buffer_info &buffer) { - auto future = toCPU( - context, tensor, - std::span(static_cast(buffer.ptr), static_cast(buffer.size))); - py::gil_scoped_release release; - wait(context, future); +PythonFuture download(std::shared_ptr context, const Tensor &tensor) { + auto values = std::make_shared>(size(tensor.shape)); + auto owner = py::capsule( + new std::shared_ptr>(values), [](void *pointer) { + delete static_cast> *>(pointer); + }); + py::array result(dtype(tensor.type), dimensions(tensor.shape), values->data(), + owner); + auto future = toCPU(*context, tensor, values); + return {std::move(context), std::move(future), std::move(result)}; } -py::array numpy(Context &context, const Tensor &tensor) { - py::array result(dtype(tensor.type), dimensions(tensor.shape)); - const auto buffer = result.request(); +PythonFuture numpy(std::shared_ptr context, const Tensor &tensor) { switch (tensor.type) { - case kf16: download(context, tensor, buffer); break; - case kf32: download(context, tensor, buffer); break; - case ki32: download(context, tensor, buffer); break; + case kf16: return download(std::move(context), tensor); + case kf32: return download(std::move(context), tensor); + case ki32: return download(std::move(context), tensor); } - return result; + throw std::invalid_argument("unknown numeric type"); } -struct PythonFuture { - explicit PythonFuture(gpu::Future value) : value(std::move(value)) {} - PythonFuture(const PythonFuture &) = delete; - PythonFuture &operator=(const PythonFuture &) = delete; - PythonFuture(PythonFuture &&) = default; - PythonFuture &operator=(PythonFuture &&) = default; - gpu::Future value; -}; +bool pollFuture(PythonFuture &future) { + py::gil_scoped_release release; + return poll(*future.context, future.value); +} + +py::object waitFuture(PythonFuture &future) { + { + py::gil_scoped_release release; + wait(*future.context, future.value); + } + return future.result; +} } // namespace -PYBIND11_MODULE(gpu_cpp, module) { +PYBIND11_MODULE(_gpu_cpp, module) { module.doc() = "Native WebGPU compute with gpu.cpp"; py::enum_(module, "NumType") @@ -130,13 +143,13 @@ PYBIND11_MODULE(gpu_cpp, module) { .value("i32", ki32) .export_values(); - py::class_(module, "Context") + py::class_>(module, "Context") .def(py::init([](bool shaderF16, bool spirv) { ContextOptions options{.enableSPIRV = spirv}; if (shaderF16) options.requiredFeatures = {wgpu::FeatureName::ShaderF16}; py::gil_scoped_release release; - return std::make_unique(createContext(options)); + return std::make_shared(createContext(options)); }), py::kw_only(), py::arg("shader_f16") = false, py::arg("spirv") = false); @@ -175,7 +188,14 @@ PYBIND11_MODULE(gpu_cpp, module) { .def_readonly("dtype", &Tensor::type); py::class_(module, "Binding"); py::class_(module, "Kernel"); - py::class_(module, "Future"); + py::class_(module, "Future") + .def("_poll", &pollFuture) + .def("_result", [](PythonFuture &future) { return future.result; }) + .def("__await__", [](py::object self) { + return py::module_::import("gpu_cpp._async") + .attr("wait")(std::move(self)) + .attr("__await__")(); + }); module.def("create_tensor", [](Context &context, @@ -202,13 +222,13 @@ PYBIND11_MODULE(gpu_cpp, module) { py::arg("workgroups") = std::vector{1, 1, 1}, py::arg("parameters") = py::bytes()); module.def("dispatch_kernel", - [](Context &context, const Kernel &kernel) { - return PythonFuture(dispatchKernel(context, kernel)); - }); - module.def("wait", [](Context &context, PythonFuture &future) { - py::gil_scoped_release release; - wait(context, future.value); - }); + [](std::shared_ptr context, const Kernel &kernel) { + auto future = dispatchKernel(*context, kernel); + return PythonFuture{std::move(context), std::move(future), + py::none()}; + }, + py::arg("context"), py::arg("kernel")); + module.def("wait", &waitFuture, py::arg("future")); module.def("to_gpu", &upload, py::arg("context"), py::arg("tensor"), py::arg("array")); module.def("to_numpy", &numpy, py::arg("context"), py::arg("tensor")); diff --git a/bindings/python/test_gpu_cpp.py b/bindings/python/test_gpu_cpp.py index 431c10c..4bcfd46 100644 --- a/bindings/python/test_gpu_cpp.py +++ b/bindings/python/test_gpu_cpp.py @@ -1,4 +1,4 @@ -import struct +import asyncio, struct import numpy as np @@ -30,16 +30,18 @@ kernel = gpu.create_kernel(context, shader, bindings, workgroups=[3, 1, 1], parameters=struct.pack(' #include #include +#include #include #include #include @@ -265,8 +266,8 @@ inline constexpr wgpu::CallbackMode persistentMode() { } template -T await(const wgpu::Instance &instance, std::future &result, - wgpu::Future event) { +void waitReady(const wgpu::Instance &instance, std::future &result, + wgpu::Future event) { #if defined(__EMSCRIPTEN__) if (instance.WaitAny(event, UINT64_MAX) != wgpu::WaitStatus::Success) throw std::runtime_error("could not wait for WebGPU operation"); @@ -278,6 +279,12 @@ T await(const wgpu::Instance &instance, std::future &result, std::this_thread::sleep_for(std::chrono::microseconds(100)); } #endif +} + +template +T await(const wgpu::Instance &instance, std::future &result, + wgpu::Future event) { + waitReady(instance, result, event); return result.get(); } @@ -309,6 +316,21 @@ struct Context { class Future { std::future result; wgpu::Future event; + std::exception_ptr failure; + bool completed = false; + + void finish(Context &context) { + if (!completed) { + try { + result.get(); + context.check(); + } catch (...) { + failure = std::current_exception(); + } + completed = true; + } + if (failure) std::rethrow_exception(failure); + } public: Future() = default; @@ -319,6 +341,7 @@ class Future { Future(Future &&) = default; Future &operator=(Future &&) = default; + friend bool poll(Context &, Future &); friend void wait(Context &, Future &); }; @@ -679,12 +702,36 @@ inline Future dispatchKernel(Context &context, const Kernel &kernel) { } inline void wait(Context &context, Future &future) { - detail::await(context.instance, future.result, future.event); - context.check(); + if (!future.completed) + detail::waitReady(context.instance, future.result, future.event); + future.finish(context); +} + +inline bool poll(Context &context, Future &future) { + if (future.completed) { + future.finish(context); + return true; + } +#if defined(__EMSCRIPTEN__) + const auto status = context.instance.WaitAny(future.event, 0); + if (status == wgpu::WaitStatus::TimedOut) return false; + if (status != wgpu::WaitStatus::Success) + throw std::runtime_error("could not poll WebGPU operation"); +#else + context.instance.ProcessEvents(); +#endif + if (future.result.wait_for(std::chrono::milliseconds(0)) != + std::future_status::ready) + return false; + future.finish(context); + return true; } +namespace detail { + template -Future toCPU(Context &context, const Tensor &tensor, std::span output) { +Future toCPU(Context &context, const Tensor &tensor, std::span output, + std::shared_ptr lifetime) { if (sizeof(T) != sizeBytes(tensor.type)) throw std::invalid_argument("host and tensor element sizes differ"); if (output.size_bytes() != tensor.bytes) @@ -706,8 +753,8 @@ Future toCPU(Context &context, const Tensor &tensor, std::span output) { auto event = readback.MapAsync( wgpu::MapMode::Read, 0, tensor.bytes, detail::completionMode(), - [promise, readback, output](wgpu::MapAsyncStatus status, - wgpu::StringView message) { + [promise, readback, output, lifetime = std::move(lifetime)]( + wgpu::MapAsyncStatus status, wgpu::StringView message) { if (status != wgpu::MapAsyncStatus::Success) { promise->set_exception(std::make_exception_ptr(std::runtime_error( "WebGPU readback failed: " + detail::string(message)))); @@ -722,6 +769,19 @@ Future toCPU(Context &context, const Tensor &tensor, std::span output) { return {std::move(future), event}; } +} // namespace detail + +template +Future toCPU(Context &context, const Tensor &tensor, std::span output) { + return detail::toCPU(context, tensor, output, {}); +} + +template +Future toCPU(Context &context, const Tensor &tensor, + const std::shared_ptr> &output) { + return detail::toCPU(context, tensor, std::span(*output), output); +} + template Future toCPU(Context &context, const Tensor &tensor, std::vector &output) { return toCPU(context, tensor, std::span(output)); diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..9061b3c --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["scikit-build-core>=1.0"] +build-backend = "scikit_build_core.build" + +[project] +name = "gpu-cpp" +version = "0.2.0" +description = "A small C++20 and Python interface for GPU compute through WebGPU" +readme = "README.md" +requires-python = ">=3.10" +license = "Apache-2.0" +license-files = ["LICENSE"] +authors = [{name = "Answer.AI"}] +dependencies = ["numpy"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Programming Language :: C++", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3 :: Only", + "Topic :: Scientific/Engineering", +] + +[project.urls] +Repository = "https://github.com/AnswerDotAI/gpu.cpp" +Issues = "https://github.com/AnswerDotAI/gpu.cpp/issues" + +[tool.scikit-build] +cmake.build-type = "Release" +cmake.args = ["-DGPUCPP_BUILD_EXAMPLES=OFF", "-DBUILD_TESTING=OFF"] +wheel.packages = ["python/gpu_cpp"] diff --git a/python/gpu_cpp/__init__.py b/python/gpu_cpp/__init__.py new file mode 100644 index 0000000..3bf0665 --- /dev/null +++ b/python/gpu_cpp/__init__.py @@ -0,0 +1 @@ +from ._gpu_cpp import * diff --git a/python/gpu_cpp/_async.py b/python/gpu_cpp/_async.py new file mode 100644 index 0000000..367b469 --- /dev/null +++ b/python/gpu_cpp/_async.py @@ -0,0 +1,6 @@ +import asyncio + + +async def wait(future): + while not future._poll(): await asyncio.sleep(0.001) + return future._result() diff --git a/tools/build_dawn.sh b/tools/build_dawn.sh index 4e9a9ae..be89383 100755 --- a/tools/build_dawn.sh +++ b/tools/build_dawn.sh @@ -13,7 +13,8 @@ readonly STAGE="$ROOT/third_party/dawn" backend_flags=(-DDAWN_ENABLE_NULL=OFF) case "$(uname -s)" in - Darwin) backend_flags+=(-DDAWN_ENABLE_METAL=ON -DDAWN_ENABLE_VULKAN=OFF) ;; + Darwin) backend_flags+=(-DDAWN_ENABLE_METAL=ON -DDAWN_ENABLE_VULKAN=OFF + -DCMAKE_OSX_DEPLOYMENT_TARGET="${MACOSX_DEPLOYMENT_TARGET:-12.0}") ;; Linux) backend_flags+=(-DDAWN_ENABLE_METAL=OFF -DDAWN_ENABLE_VULKAN=ON) ;; *) echo "gpu.cpp supports Dawn on macOS and Linux" >&2; exit 1 ;; esac diff --git a/tools/build_python_wheel.sh b/tools/build_python_wheel.sh new file mode 100755 index 0000000..ef45956 --- /dev/null +++ b/tools/build_python_wheel.sh @@ -0,0 +1,13 @@ +#!/usr/bin/env bash +set -euo pipefail + +readonly ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) + +if [[ $(uname -s) != Darwin || $(uname -m) != arm64 ]]; then + echo "Local Python publishing currently supports macOS ARM64" >&2 + exit 1 +fi + +export MACOSX_DEPLOYMENT_TARGET=${MACOSX_DEPLOYMENT_TARGET:-12.0} +uv build --wheel --clear --out-dir "$ROOT/dist" "$ROOT" +twine check "$ROOT"/dist/*.whl