webgpu: Support pipeline-overridable constants (#33291)

* Impl pipeline constants

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* More relaxed lifetimes

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* Update expectations

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

* Replace convert function with `From` implementation

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

---------

Signed-off-by: sagudev <16504129+sagudev@users.noreply.github.com>

components/script/dom/gpuconvert.rs

@@ -6,8 +6,10 @@ use std::borrow::Cow;
 use std::num::NonZeroU64;
 
 use webgpu::wgc::command as wgpu_com;
+use webgpu::wgc::pipeline::ProgrammableStageDescriptor;
 use webgpu::wgt::{self, AstcBlock, AstcChannel};
 
+use super::bindings::codegen::Bindings::WebGPUBinding::GPUProgrammableStage;
 use super::bindings::error::Error;
 use crate::dom::bindings::codegen::Bindings::WebGPUBinding::{
     GPUAddressMode, GPUBindGroupLayoutEntry, GPUBlendComponent, GPUBlendFactor, GPUBlendOperation,
@@ -473,7 +475,7 @@ pub fn convert_ic_texture(
     })
 }
 
-pub fn convert_label(parent: &GPUObjectDescriptorBase) -> Option<Cow<'static, str>> {
+pub fn convert_label<'a>(parent: &GPUObjectDescriptorBase) -> Option<Cow<'a, str>> {
     if parent.label.is_empty() {
         None
     } else {
@@ -579,3 +581,23 @@ pub fn convert_color(color: &GPUColor) -> Fallible<wgt::Color> {
         }),
     }
 }
+
+impl<'a> From<&GPUProgrammableStage> for ProgrammableStageDescriptor<'a> {
+    fn from(stage: &GPUProgrammableStage) -> Self {
+        Self {
+            module: stage.module.id().0,
+            entry_point: stage
+                .entryPoint
+                .as_ref()
+                .map(|ep| Cow::Owned(ep.to_string())),
+            constants: Cow::Owned(
+                stage
+                    .constants
+                    .as_ref()
+                    .map(|records| records.iter().map(|(k, v)| (k.0.clone(), **v)).collect())
+                    .unwrap_or_default(),
+            ),
+            zero_initialize_workgroup_memory: true,
+        }
+    }
+}

components/script/dom/gpudevice.rs

@@ -6,7 +6,6 @@
 
 use std::borrow::Cow;
 use std::cell::Cell;
-use std::collections::HashMap;
 use std::rc::Rc;
 
 use dom_struct::dom_struct;
@@ -246,12 +245,12 @@ impl GPUDevice {
         }
     }
 
-    fn parse_render_pipeline(
+    fn parse_render_pipeline<'a>(
         &self,
         descriptor: &GPURenderPipelineDescriptor,
     ) -> Fallible<(
         Option<(PipelineLayoutId, Vec<BindGroupLayoutId>)>,
-        RenderPipelineDescriptor<'static>,
+        RenderPipelineDescriptor<'a>,
     )> {
         let (layout, implicit_ids, _) = self.get_pipeline_layout_data(&descriptor.parent.layout);
 
@@ -260,17 +259,7 @@ impl GPUDevice {
             layout,
             cache: None,
             vertex: wgpu_pipe::VertexState {
-                stage: wgpu_pipe::ProgrammableStageDescriptor {
-                    module: descriptor.vertex.parent.module.id().0,
-                    entry_point: descriptor
-                        .vertex
-                        .parent
-                        .entryPoint
-                        .as_ref()
-                        .map(|ep| Cow::Owned(ep.to_string())),
-                    constants: Cow::Owned(HashMap::new()),
-                    zero_initialize_workgroup_memory: true,
-                },
+                stage: (&descriptor.vertex.parent).into(),
                 buffers: Cow::Owned(
                     descriptor
                         .vertex
@@ -302,16 +291,7 @@ impl GPUDevice {
                 .as_ref()
                 .map(|stage| -> Fallible<wgpu_pipe::FragmentState> {
                     Ok(wgpu_pipe::FragmentState {
-                        stage: wgpu_pipe::ProgrammableStageDescriptor {
-                            module: stage.parent.module.id().0,
-                            entry_point: stage
-                                .parent
-                                .entryPoint
-                                .as_ref()
-                                .map(|ep| Cow::Owned(ep.to_string())),
-                            constants: Cow::Owned(HashMap::new()),
-                            zero_initialize_workgroup_memory: true,
-                        },
+                        stage: (&stage.parent).into(),
                         targets: Cow::Owned(
                             stage
                                 .targets
@@ -670,16 +650,7 @@ impl GPUDeviceMethods for GPUDevice {
         let desc = wgpu_pipe::ComputePipelineDescriptor {
             label: convert_label(&descriptor.parent.parent),
             layout,
-            stage: wgpu_pipe::ProgrammableStageDescriptor {
-                module: descriptor.compute.module.id().0,
-                entry_point: descriptor
-                    .compute
-                    .entryPoint
-                    .as_ref()
-                    .map(|ep| Cow::Owned(ep.to_string())),
-                constants: Cow::Owned(HashMap::new()),
-                zero_initialize_workgroup_memory: true,
-            },
+            stage: (&descriptor.compute).into(),
             cache: None,
         };
 
@@ -721,16 +692,7 @@ impl GPUDeviceMethods for GPUDevice {
         let desc = wgpu_pipe::ComputePipelineDescriptor {
             label: convert_label(&descriptor.parent.parent),
             layout,
-            stage: wgpu_pipe::ProgrammableStageDescriptor {
-                module: descriptor.compute.module.id().0,
-                entry_point: descriptor
-                    .compute
-                    .entryPoint
-                    .as_ref()
-                    .map(|ep| Cow::Owned(ep.to_string())),
-                constants: Cow::Owned(HashMap::new()),
-                zero_initialize_workgroup_memory: true,
-            },
+            stage: (&descriptor.compute).into(),
             cache: None,
         };
 

components/script/dom/webidls/WebGPU.webidl

@@ -625,8 +625,11 @@ interface mixin GPUPipelineBase {
 dictionary GPUProgrammableStage {
     required GPUShaderModule module;
     USVString entryPoint;
+    record<USVString, GPUPipelineConstantValue> constants;
 };
 
+typedef double GPUPipelineConstantValue; // May represent WGSL's bool, f32, i32, u32, and f16 if enabled.
+
 [Exposed=(Window, DedicatedWorker), Serializable, Pref="dom.webgpu.enabled"]
 interface GPUComputePipeline {
 };