Reading vertex buffer, VERTEX_WRITABLE_STORAGE

[JamesSWilliams]

Hey,

I'm trying to implement a simple shader where I calculate a new color for a few points, later I want to read the buffer back.
Currently I'm failing at setting the buffer to read, write using "buffer/storage".

Binding("test", "buffer/storage", geometry.label, "VERTEX"),

Uncaught WGPU error (Unknown):
Entry { binding: 5, error: MissingFeatures(MissingFeatures(VERTEX_WRITABLE_STORAGE)) }
thread '<unnamed>' panicked at 'invalid bind group layout for bind group descriptor', src/device.rs:559:17
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
fatal runtime error: failed to initiate panic, error 5

Any help is appreciated!

[Korijn]

Can you share a minimal reproduction example?

[JamesSWilliams]

In this script I'm trying set visibility to read/write and try to write to it f_visibility, but I'm getting the above error:
I'm using a macbook pro.

import numpy as np
import pygfx as gfx
import wgpu
from pygfx import Material
from pygfx.objects import Points
from pygfx.renderers.wgpu import (
    Binding,
    RenderMask,
    WorldObjectShader,
    register_wgpu_render_function,
)
from pygfx.utils import Color, unpack_bitfield
from wgpu.gui.auto import WgpuCanvas, run


class TestPointMaterial(Material):
    uniform_type = dict(
        Material.uniform_type,
        mode="u4",
        size="f4",
        color="4xf4",
    )

    def __init__(
        self,
        size: float = 1,
        color=(1, 1, 1, 1),
        **kwargs
    ):
        super().__init__(**kwargs)

        self.size = size
        self.color = color

    def _wgpu_get_pick_info(self, pick_value):
        # This should match with the shader
        values = unpack_bitfield(pick_value, wobject_id=20, index=26)
        return {
            "vertex_index": values["index"],
        }

    @property
    def color(self):
        """The color of the points (if map is not set)."""
        return Color(self.uniform_buffer.data["color"])

    @color.setter
    def color(self, color):
        color = Color(color)
        self.uniform_buffer.data["color"] = color
        self.uniform_buffer.update_range(0, 1)
        self._store.color_is_transparent = color.a < 1

    @property
    def color_is_transparent(self):
        """Whether the color is (semi) transparent (i.e. not fully opaque)."""
        return self._store.color_is_transparent

    @property
    def size(self) -> float:
        """The size (diameter) of the points, in logical pixels."""
        return float(self.uniform_buffer.data["size"])

    @size.setter
    def size(self, size: float) -> None:
        self.uniform_buffer.data["size"] = size
        self.uniform_buffer.update_range(0, 1)


@register_wgpu_render_function(Points, TestPointMaterial)
class PointsShader(WorldObjectShader):

    type = "render"

    def get_bindings(self, wobject, shared):
        geometry = wobject.geometry
        material = wobject.material

        rbuffer = "buffer/read_only_storage"
        bindings = [
            Binding("u_stdinfo", "buffer/uniform", shared.uniform_buffer),
            Binding("u_wobject", "buffer/uniform", wobject.uniform_buffer),
            Binding("u_material", "buffer/uniform", material.uniform_buffer),
            Binding("s_positions", rbuffer, geometry.positions, "VERTEX"),
            Binding("f_visibility", "buffer/storage", geometry.visibility, "VERTEX"),
        ]

        bindings = dict(enumerate(bindings))
        self.define_bindings(0, bindings)

        return {
            0: bindings,
        }

    def get_pipeline_info(self, wobject, shared):
        return {
            "primitive_topology": wgpu.PrimitiveTopology.triangle_list,
            "cull_mode": wgpu.CullMode.none,
        }

    def get_render_info(self, wobject, shared):
        n = wobject.geometry.positions.nitems * 6

        render_mask = wobject.render_mask
        if not render_mask:
            render_mask = RenderMask.opaque

        return {
            "indices": (n, 1),
            "render_mask": render_mask,
        }

    def get_code(self) -> str:
        return (
            self.code_definitions()
            + self.code_common()
            + self.code_vertex()
            + self.code_fragment()
        )

    def code_vertex(self) -> str:
        return """

        struct VertexInput {
            @builtin(vertex_index) vertex_index : u32,
        };

        @vertex
        fn vs_main(in: VertexInput) -> Varyings {

            let index = i32(in.vertex_index);
            let i0 = index / 6;
            let sub_index = index % 6;

            let raw_pos = load_s_positions(i0);
            let world_pos = u_wobject.world_transform * vec4<f32>(raw_pos, 1.0);
            let ndc_pos = u_stdinfo.projection_transform * u_stdinfo.cam_transform * world_pos;

            var deltas = array<vec2<f32>, 6>(
                vec2<f32>(-1.0, -1.0),
                vec2<f32>(-1.0,  1.0),
                vec2<f32>( 1.0, -1.0),
                vec2<f32>(-1.0,  1.0),
                vec2<f32>( 1.0, -1.0),
                vec2<f32>( 1.0,  1.0),
            );

            // Need size here in vertex shader too
            let size = u_material.size;

            let aa_margin = 1.0;
            let delta_logical = deltas[sub_index] * (size + aa_margin);
            let delta_ndc = delta_logical * (1.0 / u_stdinfo.logical_size);

            var varyings: Varyings;
            varyings.position = vec4<f32>(ndc_pos.xy + delta_ndc * ndc_pos.w, ndc_pos.zw);
            varyings.world_pos = vec3<f32>(world_pos.xyz / world_pos.w);
            varyings.pointcoord = vec2<f32>(delta_logical);
            varyings.size = f32(size);

            // Picking
            varyings.pick_idx = u32(i0);

            // Do calculation and save to buffer
            f_visibility[i0] = 1.0;

            return varyings;
        }
        """

    def code_fragment(self) -> str:
        # Also see See https://github.com/vispy/vispy/blob/master/vispy/visuals/markers.py
        return """

        @fragment
        fn fs_main(varyings: Varyings) -> FragmentOutput {
            var final_color : vec4<f32>;

            let d = length(varyings.pointcoord);
            let aa_width = 1.0;

            let size = u_material.size;
            let color = u_material.color;


            if (d <= size - 0.5 * aa_width) {
                final_color = color;
            } else if (d <= size + 0.5 * aa_width) {
                let alpha1 = 0.5 + (size - d) / aa_width;
                let alpha2 = pow(alpha1, 2.0);  // this works better
                final_color = vec4<f32>(color.rgb, color.a * alpha2);
            } else {
                discard;
            }

            let physical_color = srgb2physical(final_color.rgb);
            let opacity = final_color.a * u_material.opacity;
            let out_color = vec4<f32>(physical_color, opacity);

            // Wrap up
            apply_clipping_planes(varyings.world_pos);
            var out = get_fragment_output(varyings.position.z, out_color);

            $$ if write_pick
            // The wobject-id must be 20 bits. In total it must not exceed 64 bits.
            out.pick = (
                pick_pack(u32(u_wobject.id), 20) +
                pick_pack(varyings.pick_idx, 26)
            );
            $$ endif

            return out;
        }
        """


canvas = WgpuCanvas()
renderer = gfx.WgpuRenderer(canvas)

scene = gfx.Scene()

positions = np.random.normal(0, 0.5, (100, 3)).astype(np.float32)
visibility = np.random.rand(100).astype(np.float32) 
geometry = gfx.Geometry(positions=positions, visibility=visibility)

material = TestPointMaterial(size=10)
points = gfx.Points(geometry, material)
scene.add(points)

scene.add(
    gfx.Background(None, gfx.BackgroundMaterial((0.2, 0.0, 0, 1), (0, 0.0, 0.2, 1))),
)

camera = gfx.NDCCamera()


if __name__ == "__main__":
    canvas.request_draw(lambda: renderer.render(scene, camera))
    run()

[JamesSWilliams]

Also I'm new to webgpu, so not sure if this even makes sense what I'm doing.
With the amount of points is pretty slow to move them around. Could the buffer be reused by a compute shader?
The goal is to only retrive the point data from the GPU on save.

[almarklein]

Thanks for the example, that helps a lot narrowing this down! 🙏

The error indicates that using a buffer with "storage" binding type requires the VERTEX_WRITABLE_STORAGE feature. I tried adding that feature in the adapter.request_device() call:

pygfx/pygfx/renderers/wgpu/_shared.py

Lines 54 to 56 in 5d46ef0

	self._device = self.adapter.request_device(
	required_features=[], required_limits={}
	)

And then your example works!

So a quick fix for you could be to change your local pygfx by adding "vertex_writable_storage" to the required_features.

For the long run we should decide how we want to enable functionality like this.

[almarklein]

I created #569 to track the more general issue of enabling wgpu extensions.

[almarklein]

With the amount of points is pretty slow to move them around. Could the buffer be reused by a compute shader?
The goal is to only retrive the point data from the GPU on save.

If you download the buffer to CPU memory on each draw, that can indeed be slow. So if you only need it on a safe, I'd do the downloading only when saving.

It should indeed be possible to use the buffer in a compute shader.

[mrkbac]

I've created a small hack to download vertex buffers form the GPU here: https://github.com/mrkbac/pygfx/tree/buffer-download.
I'm not sure how to signal that the download is completed, currently i have a small callback, but this looks rather hacky.

geom.positions._download_done_cb = lambda: print("DOWNLOAD DONE")
geom.positions.download_range()

[almarklein]

Note that the lazy upload is there to 1) make the Buffer agnostic about wgpu, 2) more efficient uploading because all changes since the last draw can be uploaded in one go.

This is probably not needed for downloading, so you may be able to get away with placing this code somewhere in your application code:

wgpu_buffer = buffer._wgpu_object
assert wgpu_buffer is not None

device = get_shared().device
bytes_per_item = buffer.nbytes // buffer.nitems
download_data = device.queue.read_buffer(wgpu_buffer)

[mrkbac]

Thanks for the tips:

I've got it running without my fork, just have to patch the required_features and ensure_wgpu_object_patched

shard = gfx.renderers.wgpu.get_shared()
shard._device = shard.adapter.request_device(
    required_features=["VERTEX_WRITABLE_STORAGE"], required_limits={},
)

def ensure_wgpu_object_patched(resource):
    if resource._wgpu_object is not None:
        return resource._wgpu_object

    if isinstance(resource, gfx.Buffer):
        resource._wgpu_usage |= wgpu.BufferUsage.COPY_SRC
    return gfx.renderers.wgpu._update.ensure_wgpu_object(resource)

gfx.renderers.wgpu._pipeline.ensure_wgpu_object = ensure_wgpu_object_patched

[almarklein]

We don't have an API to specify the buffer usage yet, in part because we aim to keep the main API free of wgpu. But you can just do your_buffer._wgpu_usage |= wgpu.BufferUsage.COPY_SRC ... It's a bit of a hack but it could work for now. I created #574 to look into a proper way to allow doing this.

[almarklein]

We've narrowed this down to two issues, one of which is fixed (being able to set wgpu features), and the other (setting buffer usage) can be resolved by setting a private attribute, and we have an issue to make this less hacky on the long term.