Custom Backend: No api to get custom buffer during create_bind_group.

[raphaelhetzel]

Hi, I'm trying to build a custom WGPU backend based on the example: https://github.com/gfx-rs/wgpu/blob/trunk/examples/standalone/custom_backend/src/custom.rs. The backend aims to support WASM outside a browser and is backed by a instance of wgpu running on the host.
What I'm essentially doing is mapping all things to objects held externally and just storing an unique id inside WASM, e.g., a buffer is represented as a u64 inside WASM, which maps to a core buffer outside of the VM.
That custom Implementation lives in my own crate (no fork of wgpu, just depending on wgpu).

However, there apears to be no way to get to my custom buffer (and therefore the id needed for serialization) during create_bind_group(&self, desc: &wgpu::BindGroupDescriptor<'_>) -> wgpu::custom::DispatchBindGroup in the DeviceInterface. The crate-internal implementations rely on https://github.com/gfx-rs/wgpu/blob/trunk/wgpu/src/api/buffer.rs#L178, followed by e.g. https://github.com/gfx-rs/wgpu/blob/trunk/wgpu/src/dispatch.rs#L581, but that is crate-only public. There appears to be https://github.com/gfx-rs/wgpu/blob/trunk/wgpu/src/dispatch.rs#L600, but I can't get to that as the inner buffer is crate-public.

Is there any way to do this using the existing public API or should the inner buffer be fully public to support this?

[sagudev]

First of all, thank you for trying custom backend. I mainly wrote for servo/vello, so we can use internal WebGPU objects of servo for vello, but it's great to see other uses.

That custom Implementation lives in my own crate (no fork of wgpu, just depending on wgpu).

That was exactly the point of supporting custom backend ❤.

I am not sure I understand your issue correctly, but you can implement your own custom buffer as:

#[derive(Debug)]
struct CustomBuffer(Counter);

impl wgpu::custom::BufferInterface for CustomBuffer {
    // ...
}

then to create it you must simply implement something like this on device:

fn create_buffer(&self, _desc: &wgpu::BufferDescriptor<'_>) -> wgpu::custom::DispatchBuffer {
        DispatchBuffer::custom(CustomBuffer(self.0.clone()))
    }

Then when your consumer will call create_buffer from your custom device it will already create custom buffer, that you should be able to use in BindGroupDescriptor.

I am not sure how exactly are you generating/storing ids, but instead of storing internal buffer (in hashmap of ids or whatever you have) you should be able to store just normal wgpu::Buffer.

[sagudev]

Or are you trying to access CustomBuffer that is inside wgpu::Buffer that is provided via wgpu::BindGroupDescriptor?

In that case I think we actually a have problem, and we would need something akin to as_hal methods to make it work. Let me try this.

[raphaelhetzel]

Thanks for the replies and the Custom API. I guess your second comment describes my problem pretty well.

The creation of a buffer is clear.
In my case this boils down to

#[derive(Debug)]
struct CustomBuffer {
    ident: u64,
}

where that id is then used for referencing the buffer living on the host:

extern "C" {
    fn webgpu_buffer_map_async(buffer_id: u64, mode: u64, start: u64, end: u64);
}

The problem arrives once I try to use the BindGroupDescriptor to create a bind group on the host (and thereby needing to pass a list of buffers to the host).

I guess code is easier here, this are the changes that I required to do what I want to do:
trunk...raphaelhetzel:wgpu:custom_fix

I am not sure if I am just using this wrong or whether a public API for this just does not exist yet (I guess your servo integration is able to just pass and maintain Rust structures, while I need to have u64 as resource identifiers (similar to what I assumed the wasm implementation does.))

[sagudev]

I guess your servo integration is able to just pass and maintain Rust structures

This integration is yet to be done. I only have smaller prototype of vello using custom backend, but that was written on different revision of custom backend.

I am not sure if I am just using this wrong or whether a public API for this just does not exist yet (I guess your servo integration is able to just pass and maintain Rust structures, while I need to have u64 as resource identifiers (similar to what I assumed the wasm implementation does.))

This is definitely missing API. I will fix this.

[raphaelhetzel]

Thanks, let me know if you need any help.
The workaround above was what was required to get my guest-code to compile (only compute, so no textures etc).
Haven't tested whether its actually working (still working on the host).

[sagudev]

Your workarounds were very educational, at least to know what is needed. I managed to get safer downcasting and hopefully nicer API. It still need some more polish, but I think it's generally ready: #7541

[raphaelhetzel]

Thanks for your changes, I now switched to your changes while debugging an early drop.

Your changes don't cover all situations yet. I guess they can be summarizes as "getting a dispatch entity instead of an api entity".
Here are the cases where i still need a workaround:

wgpu::custom::QueueInterface::write_buffer (buffer is a &wgpu::custom::DispatchBuffer)
wgpu::custom::QueueInterface::submit (command_buffers: &mut dyn Iterator<Item = wgpu::custom::DispatchCommandBuffer>
wgpu::custom::CommandEncoderInterface::wgpu::custom::CommandEncoderInterface  (similar to above, I just have dispatch buffers)
wgpu::custom::ComputePassInterface::set_pipeline (pipeline is a wgpu::custom::DispatchComputePipeline)
wgpu::custom::ComputePassInterface::set_bind_group (bind_group: Option<&wgpu::custom::DispatchBindGroup>)
wgpu::custom::ComputePassInterface::dispatch_workgroups_indirect (buffers again)

My code for reference (the edgeless_function_gpu crate/folder is the wasm-guest using the custom backend):
edgeless-project/nextless@6d74f9e

[sagudev]

Your changes don't cover all situations yet. I guess they can be summarizes as "getting a dispatch entity instead of an api entity".

Indeed, thanks for great description. I think this should do it: 3e758ee

[raphaelhetzel]

Yes that does indeed work, thanks.
Hope this can be merged soon.

One thing I forgot to mention / have currently just not used: The SubmissionIndex used in poll is currently not useable in a custom backend as the inner part is private. I am not sure if you want a custom submission index for that or whether that should just be made public for the sake of custom backends.

Other than that, I still need to properly disable all web-dependencies for wasm but not using the webgpu backend (currently just commented that out), but I guess that is only partially related to the custom backend and not really related to the api changes.

But with your changes and commenting out the web-dependencies (and some wasm specific methods), I can now use a implementation of the custom backend to run the compute hello world example inside wasmtime.

[sagudev]

One thing I forgot to mention / have currently just not used: The SubmissionIndex used in poll is currently not useable in a custom backend as the inner part is private. I am not sure if you want a custom submission index for that or whether that should just be made public for the sake of custom backends.

Indeed, I will open new issue about this.

Other than that, I still need to properly disable all web-dependencies for wasm but not using the webgpu backend (currently just commented that out), but I guess that is only partially related to the custom backend and not really related to the api changes.

But with your changes and commenting out the web-dependencies (and some wasm specific methods), I can now use a implementation of the custom backend to run the compute hello world example inside wasmtime.

This might be related to #7407, but it was ultimately scraped due to complexity (we want too keep same experience for existing wgpu users). Here WGPU custom backend users will continue to depend on core (at least on native). Do you have any commits I could look at to see what is needed?

[raphaelhetzel]

I think this is simpler than the native case. Here is my attempt to fix this: #7565