webgpu 0.0.1__py3-none-any.whl

webgpu/main.py

@@ -0,0 +1,177 @@
+"""Main file for the webgpu example, creates a small 2d mesh and renders it using WebGPU"""
+
+import urllib.parse
+
+import js
+import ngsolve as ngs
+from netgen.occ import unit_cube
+from netgen.geom2d import unit_square
+from pyodide.ffi import create_proxy
+
+from .gpu import WebGPU, init_webgpu
+from .mesh import *
+from .lic import LineIntegralConvolutionRenderObject
+
+from .webgpu_api import BackendType, RenderPassEncoder, ColorTargetState
+
+# s = ColorTargetState()
+# print("BackendType", BackendType)
+# print("ColorTargetState", s)
+
+
+def f(encoder: RenderPassEncoder):
+    return
+
+
+gpu: WebGPU = None
+mesh_object: RenderObject = None
+elements_object = None
+point_number_object = None
+
+cf = None
+render_function = None
+
+
+async def main():
+    global gpu, mesh_object, cf, render_function
+
+    gpu = await init_webgpu(js.document.getElementById("canvas"))
+    # print("DEVICE", dir(gpu.native_device))
+
+    point_number_object = None
+
+    if 1:
+        from ngsolve.meshes import MakeStructured3DMesh
+
+        # create new ngsolve mesh and evaluate arbitrary function on it
+        # mesh = ngs.Mesh(unit_cube.GenerateMesh(maxh=0.2))
+        # mesh = MakeStructured3DMesh(True, 10, prism=True)
+
+        import netgen.occ as occ
+        from netgen.meshing import IdentificationType
+
+        idtype = IdentificationType.CLOSESURFACES
+        inner = occ.Box((0, 0, 0), (1, 1, 1))
+        trafo = occ.gp_Trsf().Scale(inner.center, 1.1)
+        outer = trafo(inner)
+
+        # inner.Identify(outer, "", idtype, trafo)
+        shape = occ.Glue([outer - inner, inner])
+
+        geo = occ.OCCGeometry(shape)
+        mesh = geo.GenerateMesh(maxh=0.3)
+
+        # mesh = unit_square.GenerateMesh(maxh=0.3)
+        mesh = ngs.Mesh(mesh)
+
+        order = 3
+        cf = cf or ngs.sin(10 * ngs.x) * ngs.sin(10 * ngs.y)
+        # cf = ngs.x
+        data = MeshData(mesh, cf, order)
+        gpu.u_function.min = -1
+        gpu.u_function.max = 1
+    else:
+        # use compute shader to create a unit_square mesh
+        # but has always P1 and 'x' hard-coded as function
+        query = urllib.parse.parse_qs(js.location.search[1:])
+        N = 10
+        N = int(query.get("n", [N])[0])
+        data = create_testing_square_mesh(gpu, N)
+        gpu.u_function.min = 0
+        gpu.u_function.max = 1
+
+    # lic = LineIntegralConvolutionRenderObject(gpu, 1000, 800)
+    # print("LIC", lic)
+
+    # mesh_object = MeshRenderObject(gpu, data)
+    # mesh_object = MeshRenderObjectIndexed(gpu, data) # function values are wrong, due to ngsolve vertex numbering order
+    # mesh_object = MeshRenderObjectDeferred(
+    #     gpu, data
+    # )  # function values are wrong, due to ngsolve vertex numbering order
+    point_number_object = PointNumbersRenderObject(gpu, data, font_size=16)
+    elements_object = Mesh3dElementsRenderObject(gpu, data)
+
+    t_last = 0
+    fps = 0
+    frame_counter = 0
+    params = pyodide.ffi.to_js({"shrink": 0.3})
+
+    def render(time):
+        # this is the render function, it's called for every frame
+        if not isinstance(time, float):
+            time = 0
+
+        nonlocal t_last, fps, frame_counter
+        print("params", params.shrink)
+        dt = time - t_last
+        t_last = time
+        frame_counter += 1
+        print(f"frame time {dt:.2f} ms")
+
+        gpu.u_mesh.shrink = params.shrink
+
+        # copy camera position etc. to GPU
+        gpu.update_uniforms()
+
+        command_encoder = gpu.device.createCommandEncoder()
+
+        if mesh_object is not None:
+            mesh_object.render(command_encoder)
+
+        if elements_object is not None:
+            elements_object.render(command_encoder)
+
+        if point_number_object is not None:
+            point_number_object.render(command_encoder)
+
+        gpu.device.queue.submit([command_encoder.finish()])
+        if frame_counter < 20:
+            js.requestAnimationFrame(render_function)
+
+    render_function = create_proxy(render)
+    gpu.input_handler._update_uniforms()
+    gpu.input_handler.render_function = render_function
+
+    render_function.request_id = js.requestAnimationFrame(render_function)
+
+    try:
+        gui = js.gui
+
+        gui.reset(recursive=True)
+
+        folder = js.window.folder
+        folder.reset()
+        folder.add(params, "shrink", 0.1, 1.0)
+        gui.onChange(render_function)
+    except Exception as e:
+        print(e)
+
+
+def cleanup():
+    print("cleanup")
+    global gpu, mesh_object
+    if "gpu" in globals():
+        del gpu
+    if "mesh_object" in globals():
+        del mesh_object
+
+
+async def user_function(data):
+    code, expr = data
+    import base64
+    import marshal
+    import types
+
+    code = base64.b64decode(code.encode("utf-8"))
+    code = marshal.loads(code)
+    func = types.FunctionType(code, globals(), "user_function")
+    func(expr)
+
+
+async def reload(*args, **kwargs):
+    print("reload")
+    cleanup()
+    reload_package("webgpu")
+    from webgpu.main import main
+
+    await main()

webgpu/platform.py

@@ -0,0 +1,129 @@
+"""Platform specific code, currenty there are two possibilities:
+
+1. Running in a browser with Pyodide
+   Webgpu calls are done directly in the browser using the Pyodide interface provided by the "js" module.
+
+2. Running in a Python environment with a websocket connection to a browser
+   Webgpu calls are transferred via a websocket connection to the browser environment (using the webgpu.link.websocket module)
+"""
+
+from collections.abc import Mapping
+
+is_pyodide = False
+create_proxy = None
+destroy_proxy = None
+js = None
+websocket_server = None
+link = None
+
+try:
+    import js as pyodide_js
+    import pyodide.ffi
+    from pyodide.ffi import JsPromise, JsProxy, create_proxy
+
+    def destroy_proxy(proxy):
+        proxy.destroy()
+
+    is_pyodide = True
+
+    def _default_converter(value, a, b):
+        from .webgpu_api import BaseWebGPUHandle, BaseWebGPUObject
+
+        if isinstance(value, BaseWebGPUHandle):
+            return pyodide.ffi.to_js(value.handle)
+        if isinstance(value, BaseWebGPUObject):
+            return value.__dict__
+
+    def _convert(d):
+        from .webgpu_api import BaseWebGPUHandle, BaseWebGPUObject
+
+        if d is None:
+            return None
+        if isinstance(d, BaseWebGPUHandle):
+            return d.handle
+        if isinstance(d, BaseWebGPUObject):
+            return _convert(d.__dict__) if d.__dict__ else None
+        if isinstance(d, Mapping):
+            if not d:
+                return None
+            ret = {}
+            for key in d:
+                value = _convert(d[key])
+                if value is not None:
+                    ret[key] = value
+            return ret
+
+        if isinstance(d, list):
+            return [_convert(value) for value in d]
+
+        return d
+
+    def toJS(value):
+        value = _convert(value)
+        ret = pyodide.ffi.to_js(
+            value,
+            dict_converter=pyodide_js.Object.fromEntries,
+            default_converter=_default_converter,
+            create_pyproxies=False,
+        )
+        return ret
+
+except ImportError:
+    pass
+
+if not is_pyodide:
+    from .link.proxy import Proxy as JsProxy
+    from .link.websocket import WebsocketLinkServer
+
+    toJS = lambda x: x
+
+    class JsPromise:
+        pass
+
+
+if is_pyodide:
+    import json
+
+    from .link.base import LinkBase
+
+    LinkBase.register_serializer(JsProxy, lambda v: json.loads(pyodide_js.JSON.stringify(v)))
+
+
+def init(before_wait_for_connection=None):
+    global js, create_proxy, destroy_proxy, websocket_server, link
+    if is_pyodide or js is not None:
+        return
+
+    websocket_server = WebsocketLinkServer()
+    create_proxy = websocket_server.create_proxy
+    destroy_proxy = websocket_server.destroy_proxy
+    link = websocket_server
+
+    if before_wait_for_connection:
+        before_wait_for_connection(websocket_server)
+
+    websocket_server.wait_for_connection()
+    js = websocket_server.get(None, None)
+
+    from .link.base import LinkBase
+    from .webgpu_api import BaseWebGPUHandle, BaseWebGPUObject
+
+    LinkBase.register_serializer(BaseWebGPUHandle, lambda v: v.handle)
+    LinkBase.register_serializer(BaseWebGPUObject, lambda v: v.__dict__ or None)
+
+    websocket_server._start_handling_messages.set()
+
+
+def init_pyodide(link_):
+    global link
+    link = link_
+    print("init pyodide with link")
+    global js
+    js = link.get(None, None)
+    print("JS:", js)
+
+    from .link.base import LinkBase
+    from .webgpu_api import BaseWebGPUHandle, BaseWebGPUObject
+
+    LinkBase.register_serializer(BaseWebGPUHandle, lambda v: v.handle)
+    LinkBase.register_serializer(BaseWebGPUObject, lambda v: v.__dict__ or None)

webgpu/render_object.py

@@ -0,0 +1,155 @@
+import uuid
+from typing import Callable
+
+from .camera import Camera
+from .canvas import Canvas
+from .light import Light
+from .utils import BaseBinding, create_bind_group, get_device, is_pyodide
+from .webgpu_api import (
+    CommandEncoder,
+    CompareFunction,
+    DepthStencilState,
+    Device,
+    FragmentState,
+    PrimitiveState,
+    PrimitiveTopology,
+    VertexState,
+)
+
+
+class RenderOptions:
+    viewport: tuple[int, int, int, int, float, float]
+    canvas: Canvas
+
+    def __init__(self, canvas):
+        self.canvas = canvas
+        self.light = Light(self.device)
+        self.camera = Camera(canvas)
+
+    @property
+    def device(self) -> Device:
+        return self.canvas.device
+
+    def update_buffers(self):
+        self.camera._update_uniforms()
+
+    def get_bindings(self):
+        return [
+            *self.light.get_bindings(),
+            *self.camera.get_bindings(),
+        ]
+
+    def begin_render_pass(self, command_encoder: CommandEncoder, **kwargs):
+        load_op = command_encoder.getLoadOp()
+
+        render_pass_encoder = command_encoder.beginRenderPass(
+            self.canvas.color_attachments(load_op),
+            self.canvas.depth_stencil_attachment(load_op),
+            **kwargs,
+        )
+
+        render_pass_encoder.setViewport(0, 0, self.canvas.width, self.canvas.height, 0.0, 1.0)
+
+        return render_pass_encoder
+
+
+class BaseRenderObject:
+    options: RenderOptions
+    label: str = ""
+    _timestamp: float = -1
+    active: bool = True
+
+    def __init__(self, label=None):
+        if label is None:
+            self.label = self.__class__.__name__
+        else:
+            self.label = label
+
+    def get_bounding_box(self) -> tuple[list[float], list[float]] | None:
+        return ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0])
+
+    def update(self, timestamp):
+        if timestamp == self._timestamp:
+            return
+        self._timestamp = timestamp
+        self.create_render_pipeline()
+
+    @property
+    def device(self) -> Device:
+        return get_device()
+
+    @property
+    def canvas(self) -> Canvas:
+        return self.options.canvas
+
+    def create_render_pipeline(self) -> None:
+        raise NotImplementedError
+
+    def render(self, encoder: CommandEncoder):
+        raise NotImplementedError
+
+    def get_bindings(self) -> list[BaseBinding]:
+        raise NotImplementedError
+
+    def get_shader_code(self) -> str:
+        raise NotImplementedError
+
+    def add_options_to_gui(self, gui):
+        pass
+
+
+class MultipleRenderObject(BaseRenderObject):
+    def __init__(self, render_objects):
+        self.render_objects = render_objects
+
+    def update(self, timestamp):
+        for r in self.render_objects:
+            r.options = self.options
+            r.update(timestamp=timestamp)
+
+    def redraw(self, timestamp=None):
+        for r in self.render_objects:
+            r.redraw(timestamp=timestamp)
+
+    def render(self, encoder):
+        for r in self.render_objects:
+            r.render(encoder)
+
+
+class RenderObject(BaseRenderObject):
+    """Base class for render objects"""
+
+    n_vertices: int = 0
+    n_instances: int = 1
+    topology: PrimitiveTopology = PrimitiveTopology.triangle_list
+    depthBias: int = 0
+    vertex_entry_point: str = "vertex_main"
+    fragment_entry_point: str = "fragment_main"
+
+    def create_render_pipeline(self) -> None:
+        shader_module = self.device.createShaderModule(self.get_shader_code())
+        layout, self.group = create_bind_group(self.device, self.get_bindings())
+        self.pipeline = self.device.createRenderPipeline(
+            self.device.createPipelineLayout([layout]),
+            vertex=VertexState(module=shader_module, entryPoint=self.vertex_entry_point),
+            fragment=FragmentState(
+                module=shader_module,
+                entryPoint=self.fragment_entry_point,
+                targets=[self.options.canvas.color_target],
+            ),
+            primitive=PrimitiveState(topology=self.topology),
+            depthStencil=DepthStencilState(
+                format=self.options.canvas.depth_format,
+                depthWriteEnabled=True,
+                depthCompare=CompareFunction.less,
+                depthBias=self.depthBias,
+            ),
+            multisample=self.options.canvas.multisample,
+        )
+
+    def render(self, encoder: CommandEncoder) -> None:
+        render_pass = self.options.begin_render_pass(encoder)
+        render_pass.setPipeline(self.pipeline)
+        render_pass.setBindGroup(0, self.group)
+        render_pass.draw(self.n_vertices, self.n_instances)
+        render_pass.end()

webgpu/scene.py

@@ -0,0 +1,201 @@
+import math
+import time
+from threading import Timer
+
+from . import platform
+from .canvas import Canvas
+from .render_object import BaseRenderObject, RenderOptions
+from .utils import is_pyodide, max_bounding_box
+from .webgpu_api import *
+
+_TARGET_FPS = 60
+
+
+def debounce(render_function):
+    # Render only once every 1/_TARGET_FPS seconds
+    def debounced(*args, **kwargs):
+        if debounced.timer is not None:
+            # we already have a render scheduled, so do nothing
+            return
+
+        def f():
+            # clear the timer, so we can schedule a new one with the next function call
+            t = time.time()
+            render_function(*args, **kwargs)
+            debounced.timer = None
+            debounced.t_last = t
+
+        t_wait = max(1 / _TARGET_FPS - (time.time() - debounced.t_last), 0)
+        debounced.timer = Timer(t_wait, f)
+        debounced.timer.start()
+
+    debounced.timer = None
+    debounced.t_last = time.time()
+    return debounced
+
+
+class Scene:
+    canvas: Canvas = None
+    render_objects: list[BaseRenderObject]
+    options: RenderOptions
+    gui: object = None
+
+    def __init__(
+        self,
+        render_objects: list[BaseRenderObject],
+        id: str | None = None,
+        canvas: Canvas | None = None,
+    ):
+        if id is None:
+            import uuid
+
+            id = str(uuid.uuid4())
+
+        self._id = id
+        self.render_objects = render_objects
+
+        if is_pyodide:
+            _scenes_by_id[id] = self
+            if canvas is not None:
+                self.init(canvas)
+
+        self.t_last = 0
+
+    def __repr__(self):
+        return ""
+
+    @property
+    def id(self) -> str:
+        return self._id
+
+    @property
+    def device(self) -> Device:
+        return self.canvas.device
+
+    def init(self, canvas):
+        self.canvas = canvas
+        self.options = RenderOptions(self.canvas)
+
+        timestamp = time.time()
+        for obj in self.render_objects:
+            obj.options = self.options
+            obj.update(timestamp=timestamp)
+
+        pmin, pmax = max_bounding_box([o.get_bounding_box() for o in self.render_objects])
+        camera = self.options.camera
+        camera.transform._center = 0.5 * (pmin + pmax)
+
+        def norm(v):
+            return math.sqrt(v[0] ** 2 + v[1] ** 2 + v[2] ** 2)
+
+        camera.transform._scale = 2 / norm(pmax - pmin)
+        if not (pmin[2] == 0 and pmax[2] == 0):
+            camera.transform.rotate(270, 0)
+            camera.transform.rotate(0, -20)
+            camera.transform.rotate(20, 0)
+
+        self._js_render = platform.create_proxy(self._render_direct)
+        camera.register_callbacks(canvas.input_handler, self.render)
+        self.options.update_buffers()
+        if is_pyodide:
+            _scenes_by_id[self.id] = self
+
+        canvas.on_resize(self.render)
+
+    def redraw(self):
+        import time
+
+        ts = time.time()
+        for obj in self.render_objects:
+            obj.redraw(timestamp=ts)
+
+        self.render()
+        return
+
+        if is_pyodide:
+            import js
+
+            js.requestAnimationFrame(self._js_render)
+        else:
+            # TODO: check if we are in a jupyter kernel
+            from .jupyter import run_code_in_pyodide
+
+            ts = time.time()
+            for obj in self.render_objects:
+                obj.redraw(timestamp=ts)
+
+            run_code_in_pyodide(f"import webgpu.scene; webgpu.scene.redraw_scene('{self.id}')")
+
+    def _render(self):
+        platform.js.requestAnimationFrame(self._js_render)
+
+    def _render_direct(self, t=0):
+        encoder = self.device.createCommandEncoder()
+
+        for obj in self.render_objects:
+            if obj.active:
+                obj.render(encoder)
+
+        encoder.copyTextureToTexture(
+            TexelCopyTextureInfo(self.canvas.target_texture),
+            TexelCopyTextureInfo(self.canvas.context.getCurrentTexture()),
+            [self.canvas.width, self.canvas.height, 1],
+        )
+        self.device.queue.submit([encoder.finish()])
+
+    @debounce
+    def render(self, t=0):
+        # self.canvas.resize()
+
+        if is_pyodide:
+            self._render()
+            return
+        # print("render")
+        # print("canvas", self.canvas.canvas)
+        # from . import proxy
+        # proxy.js.console.log("canvas", self.canvas.canvas)
+        # print("canvas size ", self.canvas.canvas.width, self.canvas.canvas.height)
+        # print(
+        #     "texture size",
+        #     self.canvas.target_texture.width,
+        #     self.canvas.target_texture.height,
+        # )
+        encoder = self.device.createCommandEncoder()
+
+        for obj in self.render_objects:
+            if obj.active:
+                obj.render(encoder)
+
+        self.device.queue.submit([encoder.finish()])
+
+        if not is_pyodide:
+            platform.js.patchedRequestAnimationFrame(
+                self.canvas.device.handle,
+                self.canvas.context,
+                self.canvas.target_texture,
+            )
+
+    def cleanup(self):
+        for obj in self.render_objects:
+            obj.options = None
+
+        self.options.camera.unregister_callbacks(self.canvas.input_handler)
+        self.options.camera._render_function = None
+        self.canvas.input_handler.unregister_callbacks()
+        platform.destroy_proxy(self._js_render)
+        del self._js_render
+        self.canvas._on_resize_callbacks.remove(self.render)
+        self.canvas = None
+
+        if is_pyodide:
+            del _scenes_by_id[self.id]
+
+
+if is_pyodide:
+    _scenes_by_id: dict[str, Scene] = {}
+
+    def get_scene(id: str) -> Scene:
+        return _scenes_by_id[id]
+
+    def redraw_scene(id: str):
+        get_scene(id).redraw()

webgpu/shaders/camera.wgsl

@@ -0,0 +1,21 @@
+struct CameraUniforms {
+  model_view: mat4x4<f32>,
+  model_view_projection: mat4x4<f32>,
+  normal_mat: mat4x4<f32>,
+  aspect: f32,
+
+  padding0: u32,
+  padding1: u32,
+  padding2: u32,
+};
+
+@group(0) @binding(0) var<uniform> u_camera : CameraUniforms;
+
+fn cameraMapPoint(p: vec3f) -> vec4f {
+    return u_camera.model_view_projection * vec4<f32>(p, 1.0);
+}
+
+fn cameraMapNormal(n: vec3f) -> vec4f {
+    return u_camera.normal_mat * vec4(n, 1.0);
+}
+