Add awaitable Python GPU operations and wheel publishing

.gitignore

@@ -19,6 +19,8 @@ third_party/emdawnwebgpu/*
 # cmake build directories
 out
 build
+dist/
+__pycache__/
 
 # clangd files
 .cache

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

CMakeLists.txt

@@ -81,16 +81,34 @@ 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
     GIT_TAG v3.0.4
     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
-                 $<TARGET_OBJECTS:gpu_cpp_web_bindings>)
-  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
+                   $<TARGET_OBJECTS:gpu_cpp_web_bindings>)
+    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=$<TARGET_FILE_DIR:gpu_cpp>")
+    ENVIRONMENT "PYTHONPATH=${CMAKE_CURRENT_BINARY_DIR}/python")
 endif()

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

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

bindings/python/gpu_cpp.cpp

@@ -5,6 +5,7 @@
 #include <pybind11/stl.h>
 
 #include <cstdint>
+#include <memory>
 #include <span>
 #include <string>
 #include <utility>
@@ -61,6 +62,12 @@ std::span<const T> input(const py::buffer_info &buffer) {
           static_cast<size_t>(buffer.size)};
 }
 
+struct PythonFuture {
+  std::shared_ptr<Context> 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 <typename T>
-void download(Context &context, const Tensor &tensor, const py::buffer_info &buffer) {
-  auto future = toCPU(
-      context, tensor,
-      std::span<T>(static_cast<T *>(buffer.ptr), static_cast<size_t>(buffer.size)));
-  py::gil_scoped_release release;
-  wait(context, future);
+PythonFuture download(std::shared_ptr<Context> context, const Tensor &tensor) {
+  auto values = std::make_shared<std::vector<T>>(size(tensor.shape));
+  auto owner = py::capsule(
+      new std::shared_ptr<std::vector<T>>(values), [](void *pointer) {
+        delete static_cast<std::shared_ptr<std::vector<T>> *>(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> context, const Tensor &tensor) {
   switch (tensor.type) {
-  case kf16: download<uint16_t>(context, tensor, buffer); break;
-  case kf32: download<float>(context, tensor, buffer); break;
-  case ki32: download<int32_t>(context, tensor, buffer); break;
+  case kf16: return download<uint16_t>(std::move(context), tensor);
+  case kf32: return download<float>(std::move(context), tensor);
+  case ki32: return download<int32_t>(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_<NumType>(module, "NumType")
@@ -130,13 +143,13 @@ PYBIND11_MODULE(gpu_cpp, module) {
       .value("i32", ki32)
       .export_values();
 
-  py::class_<Context>(module, "Context")
+  py::class_<Context, std::shared_ptr<Context>>(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<Context>(createContext(options));
+             return std::make_shared<Context>(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_<Binding>(module, "Binding");
   py::class_<Kernel>(module, "Kernel");
-  py::class_<PythonFuture>(module, "Future");
+  py::class_<PythonFuture>(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<py::ssize_t>{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> 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"));

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('<f', 3))
 
 dispatched = gpu.dispatch_kernel(context, kernel)
-gpu.wait(context, dispatched)
-np.testing.assert_array_equal(gpu.to_numpy(context, gpu_output), values * 3)
+gpu.wait(dispatched)
+np.testing.assert_array_equal(gpu.wait(gpu.to_numpy(context, gpu_output)), values * 3)
 assert gpu_output.shape == [12]
 assert gpu_output.dtype == gpu.f32
 
 # The tensor and compiled kernel remain reusable with new input data.
 updated = np.arange(12, dtype=np.float32)[::-1].copy()
 gpu.to_gpu(context, gpu_values, updated)
-dispatched = gpu.dispatch_kernel(context, kernel)
-gpu.wait(context, dispatched)
-np.testing.assert_array_equal(gpu.to_numpy(context, gpu_output), updated * 3)
 
-print('Python upload, reusable dispatch, and NumPy readback story passed')
+async def reuse_kernel():
+    await gpu.dispatch_kernel(context, kernel)
+    np.testing.assert_array_equal(await gpu.to_numpy(context, gpu_output), updated * 3)
+
+asyncio.run(reuse_kernel())
+print('Python blocking and async dispatch/readback story passed')