Implement GPUComputePassEncoder functions

Implement the `dispatch`, `endPass`, `setBindGroup`, `setPipeline` functions of `GPUComputePassEncoder`.

8 files changed, 101 insertions, 23 deletions

diff --git a/components/script/dom/bindings/trace.rs b/components/script/dom/bindings/trace.rs
index c481cf8b2bb..3fb93b98dfa 100644
--- a/components/script/dom/bindings/trace.rs
+++ b/components/script/dom/bindings/trace.rs
@@ -152,8 +152,8 @@ use tendril::{StrTendril, TendrilSink};
 use time::{Duration, Timespec, Tm};
 use uuid::Uuid;
 use webgpu::{
-    WebGPU, WebGPUAdapter, WebGPUBindGroup, WebGPUBindGroupLayout, WebGPUBuffer,
-    WebGPUCommandBuffer, WebGPUCommandEncoder, WebGPUComputePipeline, WebGPUDevice,
+    wgpu::command::RawPass, WebGPU, WebGPUAdapter, WebGPUBindGroup, WebGPUBindGroupLayout,
+    WebGPUBuffer, WebGPUCommandBuffer, WebGPUCommandEncoder, WebGPUComputePipeline, WebGPUDevice,
     WebGPUPipelineLayout, WebGPUQueue, WebGPUShaderModule,
 };
 use webrender_api::{DocumentId, ImageKey};
@@ -542,7 +542,7 @@ unsafe_no_jsmanaged_fields!(WebGPUShaderModule);
 unsafe_no_jsmanaged_fields!(WebGPUCommandBuffer);
 unsafe_no_jsmanaged_fields!(WebGPUCommandEncoder);
 unsafe_no_jsmanaged_fields!(WebGPUDevice);
-unsafe_no_jsmanaged_fields!(webgpu::wgpu::command::RawPass);
+unsafe_no_jsmanaged_fields!(RefCell<Option<RawPass>>);
 unsafe_no_jsmanaged_fields!(GPUBufferState);
 unsafe_no_jsmanaged_fields!(WebXRSwapChainId);
 unsafe_no_jsmanaged_fields!(MediaList);
diff --git a/components/script/dom/gpubindgroup.rs b/components/script/dom/gpubindgroup.rs
index c14ca7a23b5..96f2d8db59b 100644
--- a/components/script/dom/gpubindgroup.rs
+++ b/components/script/dom/gpubindgroup.rs
@@ -45,6 +45,12 @@ impl GPUBindGroup {
     }
 }
 
+impl GPUBindGroup {
+    pub fn id(&self) -> &WebGPUBindGroup {
+        &self.bind_group
+    }
+}
+
 impl GPUBindGroupMethods for GPUBindGroup {
     /// https://gpuweb.github.io/gpuweb/#dom-gpuobjectbase-label
     fn GetLabel(&self) -> Option<DOMString> {
diff --git a/components/script/dom/gpucommandencoder.rs b/components/script/dom/gpucommandencoder.rs
index b0086ec0704..03b9373f3ab 100644
--- a/components/script/dom/gpucommandencoder.rs
+++ b/components/script/dom/gpucommandencoder.rs
@@ -17,7 +17,7 @@ use crate::dom::gpucomputepassencoder::GPUComputePassEncoder;
 use dom_struct::dom_struct;
 use ipc_channel::ipc;
 use std::collections::HashSet;
-use webgpu::{wgpu::command::RawPass, WebGPU, WebGPUCommandEncoder, WebGPURequest};
+use webgpu::{WebGPU, WebGPUCommandEncoder, WebGPURequest};
 
 #[dom_struct]
 pub struct GPUCommandEncoder {
@@ -69,11 +69,7 @@ impl GPUCommandEncoderMethods for GPUCommandEncoder {
         &self,
         _descriptor: &GPUComputePassDescriptor,
     ) -> DomRoot<GPUComputePassEncoder> {
-        GPUComputePassEncoder::new(
-            &self.global(),
-            self.channel.clone(),
-            RawPass::new_compute(self.encoder.0),
-        )
+        GPUComputePassEncoder::new(&self.global(), self.channel.clone(), self.encoder)
     }
 
     /// https://gpuweb.github.io/gpuweb/#dom-gpucommandencoder-copybuffertobuffer
diff --git a/components/script/dom/gpucomputepassencoder.rs b/components/script/dom/gpucomputepassencoder.rs
index ecdd4407a8d..f6116c88f9d 100644
--- a/components/script/dom/gpucomputepassencoder.rs
+++ b/components/script/dom/gpucomputepassencoder.rs
@@ -10,8 +10,20 @@ use crate::dom::bindings::reflector::{reflect_dom_object, Reflector};
 use crate::dom::bindings::root::DomRoot;
 use crate::dom::bindings::str::DOMString;
 use crate::dom::globalscope::GlobalScope;
+use crate::dom::gpubindgroup::GPUBindGroup;
+use crate::dom::gpucomputepipeline::GPUComputePipeline;
 use dom_struct::dom_struct;
-use webgpu::{wgpu::command::RawPass, WebGPU};
+use std::cell::RefCell;
+use webgpu::{
+    wgpu::command::{
+        compute_ffi::{
+            wgpu_compute_pass_dispatch, wgpu_compute_pass_set_bind_group,
+            wgpu_compute_pass_set_pipeline,
+        },
+        RawPass,
+    },
+    WebGPU, WebGPUCommandEncoder, WebGPURequest,
+};
 
 #[dom_struct]
 pub struct GPUComputePassEncoder {
@@ -20,26 +32,26 @@ pub struct GPUComputePassEncoder {
     channel: WebGPU,
     label: DomRefCell<Option<DOMString>>,
     #[ignore_malloc_size_of = "defined in wgpu-core"]
-    pass: RawPass,
+    raw_pass: RefCell<Option<RawPass>>,
 }
 
 impl GPUComputePassEncoder {
-    pub fn new_inherited(channel: WebGPU, pass: RawPass) -> GPUComputePassEncoder {
+    fn new_inherited(channel: WebGPU, parent: WebGPUCommandEncoder) -> GPUComputePassEncoder {
         GPUComputePassEncoder {
             channel,
             reflector_: Reflector::new(),
             label: DomRefCell::new(None),
-            pass,
+            raw_pass: RefCell::new(Some(RawPass::new_compute(parent.0))),
         }
     }
 
     pub fn new(
         global: &GlobalScope,
         channel: WebGPU,
-        pass: RawPass,
+        parent: WebGPUCommandEncoder,
     ) -> DomRoot<GPUComputePassEncoder> {
         reflect_dom_object(
-            Box::new(GPUComputePassEncoder::new_inherited(channel, pass)),
+            Box::new(GPUComputePassEncoder::new_inherited(channel, parent)),
             global,
             GPUComputePassEncoderBinding::Wrap,
         )
@@ -56,4 +68,49 @@ impl GPUComputePassEncoderMethods for GPUComputePassEncoder {
     fn SetLabel(&self, value: Option<DOMString>) {
         *self.label.borrow_mut() = value;
     }
+
+    #[allow(unsafe_code)]
+    /// https://gpuweb.github.io/gpuweb/#dom-gpucomputepassencoder-dispatch
+    fn Dispatch(&self, x: u32, y: u32, z: u32) {
+        if let Some(raw_pass) = self.raw_pass.borrow_mut().as_mut() {
+            unsafe { wgpu_compute_pass_dispatch(raw_pass, x, y, z) };
+        }
+    }
+
+    #[allow(unsafe_code)]
+    /// https://gpuweb.github.io/gpuweb/#dom-gpurenderpassencoder-endpass
+    fn EndPass(&self) {
+        if let Some(raw_pass) = self.raw_pass.borrow_mut().take() {
+            let (pass_data, id) = unsafe { raw_pass.finish_compute() };
+
+            self.channel
+                .0
+                .send(WebGPURequest::RunComputePass(id, pass_data))
+                .unwrap();
+        }
+    }
+
+    #[allow(unsafe_code)]
+    /// https://gpuweb.github.io/gpuweb/#dom-gpuprogrammablepassencoder-setbindgroup
+    fn SetBindGroup(&self, index: u32, bind_group: &GPUBindGroup, dynamic_offsets: Vec<u32>) {
+        if let Some(raw_pass) = self.raw_pass.borrow_mut().as_mut() {
+            unsafe {
+                wgpu_compute_pass_set_bind_group(
+                    raw_pass,
+                    index,
+                    bind_group.id().0,
+                    dynamic_offsets.as_ptr(),
+                    dynamic_offsets.len(),
+                )
+            };
+        }
+    }
+
+    #[allow(unsafe_code)]
+    /// https://gpuweb.github.io/gpuweb/#dom-gpucomputepassencoder-setpipeline
+    fn SetPipeline(&self, pipeline: &GPUComputePipeline) {
+        if let Some(raw_pass) = self.raw_pass.borrow_mut().as_mut() {
+            unsafe { wgpu_compute_pass_set_pipeline(raw_pass, pipeline.id().0) };
+        }
+    }
 }
diff --git a/components/script/dom/gpucomputepipeline.rs b/components/script/dom/gpucomputepipeline.rs
index 32c0338a7e9..5d7620b4556 100644
--- a/components/script/dom/gpucomputepipeline.rs
+++ b/components/script/dom/gpucomputepipeline.rs
@@ -42,6 +42,12 @@ impl GPUComputePipeline {
     }
 }
 
+impl GPUComputePipeline {
+    pub fn id(&self) -> &WebGPUComputePipeline {
+        &self.compute_pipeline
+    }
+}
+
 impl GPUComputePipelineMethods for GPUComputePipeline {
     /// https://gpuweb.github.io/gpuweb/#dom-gpuobjectbase-label
     fn GetLabel(&self) -> Option<DOMString> {
diff --git a/components/script/dom/webidls/GPUComputePassEncoder.webidl b/components/script/dom/webidls/GPUComputePassEncoder.webidl
index 0c03436a7e4..6d0423a3fec 100644
--- a/components/script/dom/webidls/GPUComputePassEncoder.webidl
+++ b/components/script/dom/webidls/GPUComputePassEncoder.webidl
@@ -5,11 +5,13 @@
 // https://gpuweb.github.io/gpuweb/#gpucomputepassencoder
 [Exposed=(Window, DedicatedWorker), Serializable, Pref="dom.webgpu.enabled"]
 interface GPUComputePassEncoder {
-    // void setPipeline(GPUComputePipeline pipeline);
-    // void dispatch(unsigned long x, optional unsigned long y = 1, optional unsigned long z = 1);
+    void setPipeline(GPUComputePipeline pipeline);
+    void dispatch(GPUSize32 x, optional GPUSize32 y = 1, optional GPUSize32 z = 1);
     // void dispatchIndirect(GPUBuffer indirectBuffer, GPUBufferSize indirectOffset);
 
-    // void endPass();
+    void endPass();
 };
 GPUComputePassEncoder includes GPUObjectBase;
 GPUComputePassEncoder includes GPUProgrammablePassEncoder;
+
+typedef [EnforceRange] unsigned long GPUSize32;
diff --git a/components/script/dom/webidls/GPUProgrammablePassEncoder.webidl b/components/script/dom/webidls/GPUProgrammablePassEncoder.webidl
index 2a44e806065..71bfdcb163e 100644
--- a/components/script/dom/webidls/GPUProgrammablePassEncoder.webidl
+++ b/components/script/dom/webidls/GPUProgrammablePassEncoder.webidl
@@ -5,15 +5,18 @@
 // https://gpuweb.github.io/gpuweb/#gpuprogrammablepassencoder
 [Exposed=(Window, DedicatedWorker)]
 interface mixin GPUProgrammablePassEncoder {
-    // void setBindGroup(unsigned long index, GPUBindGroup bindGroup,
-    //                   optional sequence<unsigned long> dynamicOffsets = []);
+    void setBindGroup(GPUIndex32 index, GPUBindGroup bindGroup,
+                      optional sequence<GPUBufferDynamicOffset> dynamicOffsets = []);
 
-    // void setBindGroup(unsigned long index, GPUBindGroup bindGroup,
+    // void setBindGroup(GPUIndex32 index, GPUBindGroup bindGroup,
     //                   Uint32Array dynamicOffsetsData,
-    //                   unsigned long long dynamicOffsetsDataStart,
-    //                   unsigned long long dynamicOffsetsDataLength);
+    //                   GPUSize64 dynamicOffsetsDataStart,
+    //                   GPUSize64 dynamicOffsetsDataLength);
 
     // void pushDebugGroup(DOMString groupLabel);
     // void popDebugGroup();
     // void insertDebugMarker(DOMString markerLabel);
 };
+
+typedef [EnforceRange] unsigned long GPUBufferDynamicOffset;
+typedef [EnforceRange] unsigned long GPUIndex32;
diff --git a/components/webgpu/lib.rs b/components/webgpu/lib.rs
index d7259dba795..ae26972e073 100644
--- a/components/webgpu/lib.rs
+++ b/components/webgpu/lib.rs
@@ -103,6 +103,7 @@ pub enum WebGPURequest {
         // wgpu::command::CommandBufferDescriptor,
     ),
     Submit(wgpu::id::QueueId, Vec<wgpu::id::CommandBufferId>),
+    RunComputePass(wgpu::id::CommandEncoderId, Vec<u8>),
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize)]
@@ -418,6 +419,13 @@ impl WGPU {
                         &command_buffer_ids
                     ));
                 },
+                WebGPURequest::RunComputePass(command_encoder_id, raw_data) => {
+                    let global = &self.global;
+                    gfx_select!(command_encoder_id => global.command_encoder_run_compute_pass(
+                        command_encoder_id,
+                        &raw_data
+                    ));
+                },
                 WebGPURequest::Exit(sender) => {
                     self.deinit();
                     if let Err(e) = sender.send(()) {