Retrieve Handle<ShaderStorageBuffer> from storage buffer added to compute worker

[BrianWiz]

I would like to add a storage to a compute worker and then retrieve the handle to it (Handle<ShaderStorageBuffer>), so that I can use it as a uniform on a material.

I've been able to accomplish below by manually setting up the pipelines outside of the plugin. For context I am creating a terrain (heightmap) editor.

So, I have a compute shader that writes heightmap heights to a storage buffer when a user does something (such as sculpting up a mountain).

I have an extended standard material that has the handle of that same storage buffer as one of its uniforms, and it reads from the storage buffer in the vertex shader to set the mesh's vertex positions. Resulting in a mesh that visually reacts to changes made on the heightmap.

This gives me really fast real-time editing of terrain with noise and other computations on the GPU. I've already tried writing to the storage on the CPU side and it's too slow, so simplifying it to not use compute shaders unfortunately is a non-option for me.

Anyway, in bevy_app_compute I see I can do worker.get_buffer("some_buffer") which gives me a Buffer, but I don't know if I can even get the handle from that. If I can it's not obvious.

Tl;dr I need to retrieve the handle to a storage buffer before spawning the terrain mesh with a newly created material.

Happy to implement this myself, but if you have thoughts on it I would love to hear it.

Edit:

Research notes and stream of thought (expect possible misleading info, lots of code and concepts I'm not familiar with)

- update 1: in my implementation I created the buffer on the main app by creating a ShaderStorageBuffer which is prepared as a RenderAsset. This then gets added to Assets<ShaderStorageBuffer>, this gives me the handle. After quickly scanning thru the bevy_app_compute codebase, it appears they're being created in a different way that by-passes this step, thus there is no asset handle involved.

- update 2: Possibly the solution is harder than expected. So I think the "Bevy way" works something like this:

• create your ShaderStorageBuffer, add to assets, retrieve handle and apply to new material asset, spawn mesh with material assigned
• Bevy then extracts the ShaderStorageBuffer into the RenderApp, I think this involves copying data across CPU threads once, honestly I'm a bit iffy on the details, but there's a step it takes from the main app to the render app.
• Bevy will eventually call prepare_asset() and create a GpuShaderStorageBuffer which contains the Buffer, and that's when they call render_device.create_buffer_with_data(), then wgpu takes it from there, eventually uploading the data to the GPU.
• I don't yet understand how the #[storage] attribute on a material takes the asset handle and sets up the right relationships needed for wgpu. I think this is probably where the answer will be, maybe we don't need to have an asset handle.

- update 3: I wonder if it would make more sense to have the plugin create the buffer asset and store the asset handle on the compute worker. In which case, I should look into how the plugin actually instructs wgpu/bevy to create and run the compute pass. I would imagine that's where you'd retrieve the GpuShaderStorageBuffer from the handle. Alternatively, I should also look into the material side of things, maybe with a bit more manual setup of the material it can point to the correct data, removing the need to modify the plugin. Because the more I look at the plugin the more I realize it may be a huge refactor to use the storage assets, because I believe ComputeWorker::build() is called on the RenderApp (edit: I think it actually is ran on the Main app so that is good news).

- update 4: I've been adding this feature to the plugin and I came across a bit of a code smell with how you define which vars when calling builder.add_pass(). When dispatch is eventually called it looks up the vars and grabs the buffer from that. The issue now is that i'm adding a feature to store/access storage buffers using asset handles. And so in dispatch it needs to look up from both the buffershashmap and the newly added storage_buffer_assets hashmap. It's just a bit weird. I'll think about and propose maybe a cleaner version, but probably should be done after any PRs I make for this feature.

[BrianWiz]

Draft PR: #23