Support push constants output in naga_oil composition strategy

Swoorup · Aug 10, 2024 · cf33819 · cf33819
1 parent 2334515
commit cf33819
Show file tree

Hide file tree

Showing 13 changed files with 1,627 additions and 1,410 deletions.
diff --git a/example/build.rs b/example/build.rs
@@ -1,8 +1,8 @@
 use miette::{IntoDiagnostic, Result};
 use wgsl_bindgen::qs::quote;
 use wgsl_bindgen::{
-  GlamWgslTypeMap, Regex, WgslBindgenOptionBuilder, WgslShaderSourceType,
-  WgslTypeSerializeStrategy,
+  GlamWgslTypeMap, Regex, WgslBindgenOptionBuilder, WgslShaderIrCapabilities,
+  WgslShaderSourceType, WgslTypeSerializeStrategy,
 };
 
 fn main() -> Result<()> {
@@ -13,6 +13,7 @@ fn main() -> Result<()> {
     .skip_hash_check(true)
     .serialization_strategy(WgslTypeSerializeStrategy::Bytemuck)
     .type_map(GlamWgslTypeMap)
+    .ir_capabilities(WgslShaderIrCapabilities::PUSH_CONSTANT)
     .override_struct_field_type(
       [("utils::types::VectorsU32", "a", quote!(crate::MyTwoU32))].map(Into::into),
     )

diff --git a/example/shaders/triangle.wgsl b/example/shaders/triangle.wgsl
@@ -31,6 +31,12 @@ fn vs_main(in: VertexInput) -> VertexOutput {
   return out;
 }
 
+struct PushConstants {
+    color_matrix: mat4x4<f32>
+}
+
+var<push_constant> constants: PushConstants;
+
 // wgsl outputs with pipeline constants are not supported.
 // https://github.com/gfx-rs/wgpu/blob/abba12ae4e5488b08d9e189fc37dab5e1755b443/naga/src/back/wgsl/writer.rs#L108-L113
 // override force_black: bool;
@@ -44,5 +50,5 @@ fn fs_main(in: VertexOutput) -> @location(0) vec4<f32> {
   // } else {
   //   return vec4(color * uniforms.color_rgb.rgb * scale, 1.0);
   // }
-  return vec4(color * uniforms.color_rgb.rgb, 1.0);
+  return constants.color_matrix * vec4(color * uniforms.color_rgb.rgb, 1.0);
 }
diff --git a/example/src/main.rs b/example/src/main.rs
@@ -57,8 +57,13 @@ impl State {
       .request_device(
         &wgpu::DeviceDescriptor {
           label: None,
-          required_features: wgpu::Features::TEXTURE_COMPRESSION_BC,
-          required_limits: wgpu::Limits::default(),
+          required_features: wgpu::Features::TEXTURE_COMPRESSION_BC
+            | wgpu::Features::PUSH_CONSTANTS,
+          required_limits: wgpu::Limits {
+            max_push_constant_size: 128,
+            ..Default::default()
+          },
+          memory_hints: Default::default(),
         },
         None,
       )
@@ -99,6 +104,7 @@ impl State {
       depth_stencil: None,
       multisample: wgpu::MultisampleState::default(),
       multiview: None,
+      cache: Default::default(),
     });
 
     // Create a gradient texture.
@@ -242,6 +248,17 @@ impl State {
 
     render_pass.set_pipeline(&self.pipeline);
 
+    // Push constant data also needs to follow alignment rules.
+    let push_constant = shader_bindings::triangle::PushConstants {
+      color_matrix: glam::Mat4::IDENTITY,
+    };
+
+    render_pass.set_push_constants(
+      wgpu::ShaderStages::VERTEX_FRAGMENT,
+      0,
+      &bytemuck::cast_slice(&[push_constant]),
+    );
+
     // Use this function to ensure all bind groups are set.
     self.bind_group0.set(&mut render_pass);
     self.bind_group1.set(&mut render_pass);