Promise / Future API with .then() (#739)

* WIP Promise / Future API with .then()

* Fix leftover dict default

* update all methods

* The _sync classes are implemented in the base classes

* Can remove some proxy methods in _api.py now

* Adjust a test in codegen

* Add docs

* wip

* wip

* cleaner

* tweaks

* remove accidentally added

* Move implementation of promise to separate module

* codegen

* Change chaining

* Also keep working async without a loop

* tweaks

* more docs and tweaks

* add catch method

* add to classes, so we can override behaviour in backends

* Improve how GPUPromise is documented

* Tweak signature

* Add tests

* fix names

* fix more names

* Fix tests

* cleanup after self-review

* restore cube example

* tweak loop in classes.py

* add back incremental sleeping

* codegen

* fix typo

* tweak example

* fix typo

Author: almarklein

codegen/apipatcher.py

@@ -38,11 +38,13 @@ def patch_base_api(code):
     idl = get_idl_parser()
 
     # Write __all__
+    extra_public_classes = ["GPUPromise"]
+    all_public_classes = [*idl.classes.keys(), *extra_public_classes]
     part1, found_all, part2 = code.partition("\n__all__ =")
     if found_all:
         part2 = part2.split("]", 1)[-1]
         line = "\n__all__ = ["
-        line += ", ".join(f'"{name}"' for name in sorted(idl.classes.keys()))
+        line += ", ".join(f'"{name}"' for name in sorted(all_public_classes))
         line += "]"
         code = part1 + line + part2
 
@@ -158,14 +160,22 @@ def patch_classes(self):
         for classname, i1, i2 in self.iter_classes():
             seen_classes.add(classname)
             self._apidiffs = set()
+            pre_lines = "\n".join(self.lines[i1 - 3 : i1])
+            self._apidiffs_from_lines(pre_lines, classname)
             if self.class_is_known(classname):
+                if "@apidiff.add" in pre_lines:
+                    print(f"ERROR: apidiff.add for known {classname}")
+                elif "@apidiff.hide" in pre_lines:
+                    pass  # continue as normal
                 old_line = self.lines[i1]
                 new_line = self.get_class_def(classname)
                 if old_line != new_line:
                     fixme_line = "# FIXME: was " + old_line.split("class ", 1)[-1]
                     self.replace_line(i1, f"{fixme_line}\n{new_line}")
                 self.patch_properties(classname, i1 + 1, i2)
                 self.patch_methods(classname, i1 + 1, i2)
+            elif "@apidiff.add" in pre_lines:
+                pass
             else:
                 msg = f"unknown api: class {classname}"
                 self.insert_line(i1, "# FIXME: " + msg)
@@ -422,14 +432,15 @@ def get_property_def(self, classname, propname) -> str:
             print(f"Error resolving type for {classname}.{propname}: {err}")
             prop_type = None
 
+        if prop_type and propname.endswith("_async"):
+            prop_type = f"GPUPromise[{prop_type}]"
+
         line = (
             "def "
             + to_snake_case(propname)
             + "(self)"
             + f"{f' -> {prop_type}' if prop_type else ''}:"
         )
-        if propname.endswith("_async"):
-            line = "async " + line
         return "    " + line
 
     def get_method_def(self, classname, methodname) -> str:
@@ -439,8 +450,6 @@ def get_method_def(self, classname, methodname) -> str:
 
         # Construct preamble
         preamble = "def " + to_snake_case(methodname) + "("
-        if methodname.endswith("_async"):
-            preamble = "async " + preamble
 
         # Get arg names and types
         idl_line = functions[name_idl]
@@ -455,6 +464,8 @@ def get_method_def(self, classname, methodname) -> str:
             return_type = None
         if return_type:
             return_type = self.idl.resolve_type(return_type)
+            if methodname.endswith("_async"):
+                return_type = f"GPUPromise[{return_type}]"
 
         # If one arg that is a dict, flatten dict to kwargs
         if len(args) == 1 and args[0].typename.endswith(
@@ -601,6 +612,8 @@ def __init__(self, base_api_code):
                 pre_lines = "\n".join(p1.lines[j1 - 3 : j1])
                 if "@apidiff.hide" in pre_lines:
                     continue  # method (currently) not part of our API
+                if methodname.endswith("_sync"):
+                    continue  # the base class implements _sync versions (using promise.sync_wait())
                 body = "\n".join(p1.lines[j1 + 1 : j2 + 1])
                 must_overload = "raise NotImplementedError()" in body
                 methods[methodname] = p1.lines[j1], must_overload
@@ -690,7 +703,7 @@ def apply(self, code):
                 defer_func_name = "_" + methodname
                 defer_line_starts = (
                     f"return self.{defer_func_name[:-7]}",
-                    f"awaitable = self.{defer_func_name[:-7]}",
+                    f"promise = self.{defer_func_name[:-7]}",
                 )
                 this_method_defers = any(
                     line.strip().startswith(defer_line_starts)

codegen/tests/test_codegen_result.py

@@ -1,17 +1,19 @@
 """Test some aspects of the generated code."""
 
 from codegen.files import read_file
+from codegen.utils import format_code
 
 
 def test_async_methods_and_props():
-    # Test that only and all async methods are suffixed with '_async'
+    # Test that async methods return a promise
 
     for fname in ["_classes.py", "backends/wgpu_native/_api.py"]:
-        code = read_file(fname)
+        code = format_code(read_file(fname), singleline=True)
         for line in code.splitlines():
             line = line.strip()
             if line.startswith("def "):
-                assert not line.endswith("_async"), line
-            elif line.startswith("async def "):
-                name = line.split("def", 1)[1].split("(")[0].strip()
-                assert name.endswith("_async"), line
+                res_type = line.split("->")[-1].strip()
+                if "_async(" in line:
+                    assert res_type.startswith("GPUPromise")
+                else:
+                    assert "GPUPromise" not in line

docs/_templates/wgpu_class_layout.rst

@@ -2,6 +2,17 @@
 
 .. currentmodule:: {{ module }}
 
+{% if objname == "GPUPromise" %}
+
 .. autoclass:: {{ objname }}
     :members:
+    :inherited-members:
     :show-inheritance:
+
+{% else %}
+
+.. autoclass:: {{ objname }}
+    :members:
+    :show-inheritance:
+
+{% endif %}

docs/wgpu.rst

@@ -65,6 +65,12 @@ come in two flafours:
   not part of the WebGPU spec, and as a consequence, code that uses this method
   is less portable (to e.g. pyodide/pyscript).
 
+The async methods return a :class:`GPUPromise`, which resolves to the actual result. You can wait for it to resolve in three ways:
+
+* In async code, use ``await promise``.
+* In sync code, use ``promise.then(callback)`` to register a callback that is executed when the promise resolves.
+* In sync code, you can use ``promise.sync_wait()``. This is similar to the ``_sync()`` flavour mentioned above (it makes your code less portable).
+
 
 Canvas API
 ----------
@@ -241,6 +247,7 @@ List of GPU classes
     ~GPUPipelineBase
     ~GPUPipelineError
     ~GPUPipelineLayout
+    ~GPUPromise
     ~GPUQuerySet
     ~GPUQueue
     ~GPURenderBundle

examples/cube.py

@@ -34,7 +34,8 @@ def setup_drawing_sync(
 
     adapter = wgpu.gpu.request_adapter_sync(power_preference=power_preference)
     device = adapter.request_device_sync(
-        required_limits=limits, label="Cube Example device"
+        label="Cube Example device",
+        required_limits=limits,
     )
 
     pipeline_layout, uniform_buffer, bind_group = create_pipeline_layout(device)
@@ -43,7 +44,7 @@ def setup_drawing_sync(
     render_pipeline = device.create_render_pipeline(**pipeline_kwargs)
 
     return get_draw_function(
-        canvas, device, render_pipeline, uniform_buffer, bind_group, asynchronous=False
+        canvas, device, render_pipeline, uniform_buffer, bind_group
     )
 
 
@@ -53,10 +54,10 @@ async def setup_drawing_async(canvas, limits=None):
     The given canvas must implement WgpuCanvasInterface, but nothing more.
     Returns the draw function.
     """
-
     adapter = await wgpu.gpu.request_adapter_async(power_preference="high-performance")
+
     device = await adapter.request_device_async(
-        required_limits=limits, label="Cube Example async device"
+        label="Cube Example async device", required_limits=limits
     )
 
     pipeline_layout, uniform_buffer, bind_group = create_pipeline_layout(device)
@@ -65,7 +66,19 @@ async def setup_drawing_async(canvas, limits=None):
     render_pipeline = await device.create_render_pipeline_async(**pipeline_kwargs)
 
     return get_draw_function(
-        canvas, device, render_pipeline, uniform_buffer, bind_group, asynchronous=True
+        canvas, device, render_pipeline, uniform_buffer, bind_group
+    )
+
+
+def get_drawing_func(canvas, device):
+    pipeline_layout, uniform_buffer, bind_group = create_pipeline_layout(device)
+    pipeline_kwargs = get_render_pipeline_kwargs(canvas, device, pipeline_layout)
+
+    render_pipeline = device.create_render_pipeline(**pipeline_kwargs)
+    # render_pipeline = device.create_render_pipeline(**pipeline_kwargs)
+
+    return get_draw_function(
+        canvas, device, render_pipeline, uniform_buffer, bind_group
     )
 
 
@@ -242,8 +255,6 @@ def get_draw_function(
     render_pipeline: wgpu.GPURenderPipeline,
     uniform_buffer: wgpu.GPUBuffer,
     bind_group: wgpu.GPUBindGroup,
-    *,
-    asynchronous: bool,
 ):
     # Create vertex buffer, and upload data
     vertex_buffer = device.create_buffer_with_data(
@@ -288,48 +299,8 @@ def update_transform():
         )
         uniform_data["transform"] = rot2 @ rot1 @ ortho
 
-    def upload_uniform_buffer_sync():
-        if True:
-            tmp_buffer = uniform_buffer.copy_buffer
-            tmp_buffer.map_sync(wgpu.MapMode.WRITE)
-            tmp_buffer.write_mapped(uniform_data)
-            tmp_buffer.unmap()
-        else:
-            tmp_buffer = device.create_buffer_with_data(
-                data=uniform_data, usage=wgpu.BufferUsage.COPY_SRC
-            )
-        command_encoder = device.create_command_encoder(
-            label="Cube Example uniform buffer upload command encoder"
-        )
-        command_encoder.copy_buffer_to_buffer(
-            tmp_buffer, 0, uniform_buffer, 0, uniform_data.nbytes
-        )
-        device.queue.submit(
-            [
-                command_encoder.finish(
-                    label="Cube Example uniform buffer upload command buffer"
-                )
-            ]
-        )
-
-    async def upload_uniform_buffer_async():
-        tmp_buffer = uniform_buffer.copy_buffer
-        await tmp_buffer.map_async(wgpu.MapMode.WRITE)
-        tmp_buffer.write_mapped(uniform_data)
-        tmp_buffer.unmap()
-        command_encoder = device.create_command_encoder(
-            label="Cube Example uniform buffer upload async command encoder"
-        )
-        command_encoder.copy_buffer_to_buffer(
-            tmp_buffer, 0, uniform_buffer, 0, uniform_data.nbytes
-        )
-        device.queue.submit(
-            [
-                command_encoder.finish(
-                    label="Cube Example uniform buffer upload async command buffer"
-                )
-            ]
-        )
+    def upload_uniform_buffer():
+        device.queue.write_buffer(uniform_buffer, 0, uniform_data)
 
     def draw_frame():
         current_texture_view: wgpu.GPUTextureView = (
@@ -367,20 +338,12 @@ def draw_frame():
             [command_encoder.finish(label="Cube Example render pass command buffer")]
         )
 
-    def draw_frame_sync():
+    def draw_func():
         update_transform()
-        upload_uniform_buffer_sync()
+        upload_uniform_buffer()
         draw_frame()
 
-    async def draw_frame_async():
-        update_transform()
-        await upload_uniform_buffer_async()
-        draw_frame()
-
-    if asynchronous:
-        return draw_frame_async
-    else:
-        return draw_frame_sync
+    return draw_func
 
 
 # %% WGSL
@@ -509,6 +472,7 @@ async def draw_frame_async():
 for a in wgpu.gpu.enumerate_adapters_sync():
     print(a.summary)
 
+
 if __name__ == "__main__":
     canvas = RenderCanvas(
         size=(640, 480),
@@ -517,6 +481,19 @@ async def draw_frame_async():
         max_fps=60,
         vsync=True,
     )
-    draw_frame = setup_drawing_sync(canvas)
-    canvas.request_draw(draw_frame)
+
+    # Pick one
+
+    if True:
+        # Async
+        @loop.add_task
+        async def init():
+            draw_frame = await setup_drawing_async(canvas)
+            canvas.request_draw(draw_frame)
+    else:
+        # Sync
+        draw_frame = setup_drawing_sync(canvas)
+        canvas.request_draw(draw_frame)
+
+    # loop.add_task(poller)
     loop.run()

tests/test_api.py

@@ -22,8 +22,10 @@ def test_basic_api():
 
     code1 = wgpu.GPU.request_adapter_sync.__code__
     code2 = wgpu.GPU.request_adapter_async.__code__
+    varnames1 = set(code1.co_varnames) - {"gpu", "promise", "loop"}
+    varnames2 = set(code2.co_varnames) - {"gpu", "promise", "loop"}
     # nargs1 = code1.co_argcount + code1.co_kwonlyargcount
-    assert code1.co_varnames == code2.co_varnames
+    assert varnames1 == varnames2
 
     assert repr(wgpu.classes.GPU()).startswith(
         "<wgpu.GPU "
@@ -98,7 +100,7 @@ def test_enums_and_flags_and_structs():
 
 def test_base_wgpu_api():
     # Fake a device and an adapter
-    adapter = wgpu.GPUAdapter(None, set(), {}, wgpu.GPUAdapterInfo({}))
+    adapter = wgpu.GPUAdapter(None, set(), {}, wgpu.GPUAdapterInfo({}), None)
     queue = wgpu.GPUQueue("", None, None)
     device = wgpu.GPUDevice("device08", -1, adapter, {42, 43}, {}, queue)