Allow obtaining custom implementation from wgpu api types (#7541)

Author: sagudev

examples/standalone/custom_backend/src/custom.rs

@@ -3,9 +3,9 @@ use std::pin::Pin;
 use std::sync::Arc;
 
 use wgpu::custom::{
-    AdapterInterface, DeviceInterface, DispatchAdapter, DispatchDevice, DispatchQueue,
-    DispatchShaderModule, DispatchSurface, InstanceInterface, QueueInterface, RequestAdapterFuture,
-    ShaderModuleInterface,
+    AdapterInterface, ComputePipelineInterface, DeviceInterface, DispatchAdapter, DispatchDevice,
+    DispatchQueue, DispatchShaderModule, DispatchSurface, InstanceInterface, QueueInterface,
+    RequestAdapterFuture, ShaderModuleInterface,
 };
 
 #[derive(Debug, Clone)]
@@ -163,9 +163,10 @@ impl DeviceInterface for CustomDevice {
 
     fn create_compute_pipeline(
         &self,
-        _desc: &wgpu::ComputePipelineDescriptor<'_>,
+        desc: &wgpu::ComputePipelineDescriptor<'_>,
     ) -> wgpu::custom::DispatchComputePipeline {
-        unimplemented!()
+        let module = desc.module.as_custom::<CustomShaderModule>().unwrap();
+        wgpu::custom::DispatchComputePipeline::custom(CustomComputePipeline(module.0.clone()))
     }
 
     unsafe fn create_pipeline_cache(
@@ -265,7 +266,7 @@ impl DeviceInterface for CustomDevice {
 }
 
 #[derive(Debug)]
-struct CustomShaderModule(Counter);
+pub struct CustomShaderModule(pub Counter);
 
 impl ShaderModuleInterface for CustomShaderModule {
     fn get_compilation_info(&self) -> Pin<Box<dyn wgpu::custom::ShaderCompilationInfoFuture>> {
@@ -346,3 +347,12 @@ impl QueueInterface for CustomQueue {
         unimplemented!()
     }
 }
+
+#[derive(Debug)]
+pub struct CustomComputePipeline(pub Counter);
+
+impl ComputePipelineInterface for CustomComputePipeline {
+    fn get_bind_group_layout(&self, _index: u32) -> wgpu::custom::DispatchBindGroupLayout {
+        unimplemented!()
+    }
+}

examples/standalone/custom_backend/src/main.rs

@@ -1,6 +1,6 @@
 use std::marker::PhantomData;
 
-use custom::Counter;
+use custom::{Counter, CustomShaderModule};
 use wgpu::{DeviceDescriptor, RequestAdapterOptions};
 
 mod custom;
@@ -31,12 +31,26 @@ async fn main() {
             .unwrap();
         assert_eq!(counter.count(), 5);
 
-        let _module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
+        let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
             label: Some("shader"),
             source: wgpu::ShaderSource::Dummy(PhantomData),
         });
 
+        let custom_module = module.as_custom::<CustomShaderModule>().unwrap();
+        assert_eq!(custom_module.0.count(), 6);
+        let _module_clone = module.clone();
         assert_eq!(counter.count(), 6);
+
+        let _pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
+            label: None,
+            layout: None,
+            module: &module,
+            entry_point: None,
+            compilation_options: Default::default(),
+            cache: None,
+        });
+
+        assert_eq!(counter.count(), 7);
     }
     assert_eq!(counter.count(), 1);
 }

wgpu/src/api/adapter.rs

@@ -133,6 +133,12 @@ impl Adapter {
         }
     }
 
+    #[cfg(custom)]
+    /// Returns custom implementation of adapter (if custom backend and is internally T)
+    pub fn as_custom<T: custom::AdapterInterface>(&self) -> Option<&T> {
+        self.inner.as_custom()
+    }
+
     #[cfg(custom)]
     /// Creates Adapter from custom implementation
     pub fn from_custom<T: custom::AdapterInterface>(adapter: T) -> Self {

wgpu/src/api/bind_group.rs

@@ -17,6 +17,14 @@ static_assertions::assert_impl_all!(BindGroup: Send, Sync);
 
 crate::cmp::impl_eq_ord_hash_proxy!(BindGroup => .inner);
 
+impl BindGroup {
+    #[cfg(custom)]
+    /// Returns custom implementation of BindGroup (if custom backend and is internally T)
+    pub fn as_custom<T: custom::BindGroupInterface>(&self) -> Option<&T> {
+        self.inner.as_custom()
+    }
+}
+
 /// Resource to be bound by a [`BindGroup`] for use with a pipeline.
 ///
 /// The pipeline’s [`BindGroupLayout`] must contain a matching [`BindingType`].

wgpu/src/api/bind_group_layout.rs

@@ -20,6 +20,14 @@ static_assertions::assert_impl_all!(BindGroupLayout: Send, Sync);
 
 crate::cmp::impl_eq_ord_hash_proxy!(BindGroupLayout => .inner);
 
+impl BindGroupLayout {
+    #[cfg(custom)]
+    /// Returns custom implementation of BindGroupLayout (if custom backend and is internally T)
+    pub fn as_custom<T: custom::BindGroupLayoutInterface>(&self) -> Option<&T> {
+        self.inner.as_custom()
+    }
+}
+
 /// Describes a [`BindGroupLayout`].
 ///
 /// For use with [`Device::create_bind_group_layout`].

wgpu/src/api/blas.rs

@@ -174,6 +174,12 @@ impl Blas {
             hal_blas_callback(None)
         }
     }
+
+    #[cfg(custom)]
+    /// Returns custom implementation of Blas (if custom backend and is internally T)
+    pub fn as_custom<T: crate::custom::BlasInterface>(&self) -> Option<&T> {
+        self.inner.as_custom()
+    }
 }
 
 /// Context version of [BlasTriangleGeometry].

wgpu/src/api/buffer.rs

@@ -386,6 +386,12 @@ impl Buffer {
     ) -> BufferViewMut<'_> {
         self.slice(bounds).get_mapped_range_mut()
     }
+
+    #[cfg(custom)]
+    /// Returns custom implementation of Buffer (if custom backend and is internally T)
+    pub fn as_custom<T: custom::BufferInterface>(&self) -> Option<&T> {
+        self.inner.as_custom()
+    }
 }
 
 /// A slice of a [`Buffer`], to be mapped, used for vertex or index data, or the like.

wgpu/src/api/command_buffer.rs

@@ -13,3 +13,11 @@ pub struct CommandBuffer {
 }
 #[cfg(send_sync)]
 static_assertions::assert_impl_all!(CommandBuffer: Send, Sync);
+
+impl CommandBuffer {
+    #[cfg(custom)]
+    /// Returns custom implementation of CommandBuffer (if custom backend and is internally T)
+    pub fn as_custom<T: custom::CommandBufferInterface>(&self) -> Option<&T> {
+        self.buffer.as_custom()
+    }
+}

wgpu/src/api/command_encoder.rs

@@ -262,6 +262,12 @@ impl CommandEncoder {
             hal_command_encoder_callback(None)
         }
     }
+
+    #[cfg(custom)]
+    /// Returns custom implementation of CommandEncoder (if custom backend and is internally T)
+    pub fn as_custom<T: custom::CommandEncoderInterface>(&self) -> Option<&T> {
+        self.inner.as_custom()
+    }
 }
 
 /// [`Features::TIMESTAMP_QUERY_INSIDE_ENCODERS`] must be enabled on the device in order to call these functions.

wgpu/src/api/compute_pass.rs

@@ -94,6 +94,12 @@ impl ComputePass<'_> {
         self.inner
             .dispatch_workgroups_indirect(&indirect_buffer.inner, indirect_offset);
     }
+
+    #[cfg(custom)]
+    /// Returns custom implementation of ComputePass (if custom backend and is internally T)
+    pub fn as_custom<T: custom::ComputePassInterface>(&self) -> Option<&T> {
+        self.inner.as_custom()
+    }
 }
 
 /// [`Features::PUSH_CONSTANTS`] must be enabled on the device in order to call these functions.