Fix GL Push Constant Layout

naga/src/back/glsl/mod.rs

@@ -309,6 +309,8 @@ pub struct ReflectionInfo {
  pub uniforms: crate::FastHashMap<Handle<crate::GlobalVariable>, String>,
  /// Mapping between names and attribute locations.
  pub varying: crate::FastHashMap<String, VaryingLocation>,
+ /// List of push constant items in the shader.
+ pub push_constant_items: Vec<PushConstantItem>,
 }
 
 /// Mapping between a texture and its sampler, if it exists.
@@ -328,6 +330,50 @@ pub struct TextureMapping {
  pub sampler: Option<Handle<crate::GlobalVariable>>,
 }
 
+/// All information to bind a single uniform value to the shader.
+///
+/// Push constants are emulated using traditional uniforms in OpenGL.
+///
+/// These are composed of a set of primatives (scalar, vector, matrix) that
+/// are given names. Because they are not backed by the concept of a buffer,
+/// we must do the work of calculating the offset of each primative in the
+/// push constant block.
+#[derive(Debug, Clone)]
+pub struct PushConstantItem {
+ /// GL uniform name for the item. This name is the same as if you were
+ /// to access it directly from a GLSL shader.
+ ///
+ /// The with the following example, the following names will be generated,
+ /// one name per GLSL uniform.
+ ///
+ /// ```glsl
+ /// struct InnerStruct {
+ /// value: f32,
+ /// }
+ ///
+ /// struct PushConstant {
+ /// InnerStruct inner;
+ /// vec4 array[2];
+ /// }
+ ///
+ /// uniform PushConstants _push_constant_binding_cs;
+ /// ```
+ ///
+ /// ```text
+ /// - _push_constant_binding_cs.inner.value
+ /// - _push_constant_binding_cs.array[0]
+ /// - _push_constant_binding_cs.array[1]
+ /// ```
+ ///
+ pub access_path: String,
+ /// Type of the uniform. This will only ever be a scalar, vector, or matrix.
+ pub ty: Handle<crate::Type>,
+ /// The offset in the push constant memory block this uniform maps to.
+ ///
+ /// The size of the uniform can be derived from the type.
+ pub offset: u32,
+}
+
 /// Helper structure that generates a number
 #[derive(Default)]
 struct IdGenerator(u32);
@@ -1264,16 +1310,19 @@ impl<'a, W: Write> Writer<'a, W> {
  handle: Handle<crate::GlobalVariable>,
  global: &crate::GlobalVariable,
  ) -> String {
- match global.binding {
- Some(ref br) => {
+ match (&global.binding, global.space) {
+ (&Some(ref br), _) => {
  format!(
  "_group_{}_binding_{}_{}",
  br.group,
  br.binding,
  self.entry_point.stage.to_str()
  )
  }
- None => self.names[&NameKey::GlobalVariable(handle)].clone(),
+ (&None, crate::AddressSpace::PushConstant) => {
+ format!("_push_constant_binding_{}", self.entry_point.stage.to_str())
+ }
+ (&None, _) => self.names[&NameKey::GlobalVariable(handle)].clone(),
  }
  }
 
@@ -1283,15 +1332,20 @@ impl<'a, W: Write> Writer<'a, W> {
  handle: Handle<crate::GlobalVariable>,
  global: &crate::GlobalVariable,
  ) -> BackendResult {
- match global.binding {
- Some(ref br) => write!(
+ match (&global.binding, global.space) {
+ (&Some(ref br), _) => write!(
  self.out,
  "_group_{}_binding_{}_{}",
  br.group,
  br.binding,
  self.entry_point.stage.to_str()
  )?,
- None => write!(
+ (&None, crate::AddressSpace::PushConstant) => write!(
+ self.out,
+ "_push_constant_binding_{}",
+ self.entry_point.stage.to_str()
+ )?,
+ (&None, _) => write!(
  self.out,
  "{}",
  &self.names[&NameKey::GlobalVariable(handle)]
@@ -4069,6 +4123,7 @@ impl<'a, W: Write> Writer<'a, W> {
  }
  }
 
+ let mut push_constant_info = None;
  for (handle, var) in self.module.global_variables.iter() {
  if info[handle].is_empty() {
  continue;
@@ -4093,17 +4148,105 @@ impl<'a, W: Write> Writer<'a, W> {
  let name = self.reflection_names_globals[&handle].clone();
  uniforms.insert(handle, name);
  }
+ crate::AddressSpace::PushConstant => {
+ let name = self.reflection_names_globals[&handle].clone();
+ push_constant_info = Some((name, var.ty));
+ }
  _ => (),
  },
  }
  }
 
+ let mut push_constant_segments = Vec::new();
+ let mut push_constant_items = vec![];
+
+ if let Some((name, ty)) = push_constant_info {
+ // We don't have a layouter available to us, so we need to create one.
+ //
+ // This is potentially a bit wasteful, but the set of types in the program
+ // shouldn't be too large.
+ let mut layouter = crate::proc::Layouter::default();
+ layouter.update(self.module.to_ctx()).unwrap();
+
+ // We start with the name of the binding itself.
+ push_constant_segments.push(name);
+
+ // We then recursively collect all the uniform fields of the push constant.
+ self.collect_push_constant_items(
+ ty,
+ &mut push_constant_segments,
+ &layouter,
+ &mut 0,
+ &mut push_constant_items,
+ );
+ }
+
  Ok(ReflectionInfo {
  texture_mapping,
  uniforms,
  varying: mem::take(&mut self.varying),
+ push_constant_items,
  })
  }
+
+ fn collect_push_constant_items(
+ &mut self,
+ ty: Handle<crate::Type>,
+ segments: &mut Vec<String>,
+ layouter: &crate::proc::Layouter,
+ offset: &mut u32,
+ items: &mut Vec<PushConstantItem>,
+ ) {
+ // At this point in the recursion, `segments` contains the path
+ // needed to access `ty` from the root.
+
+ let layout = &layouter[ty];
+ *offset = layout.alignment.round_up(*offset);
+ match self.module.types[ty].inner {
+ // All these types map directly to GL uniforms.
+ TypeInner::Scalar { .. } | TypeInner::Vector { .. } | TypeInner::Matrix { .. } => {
+ // Build the full name, by combining all current segments.
+ let name: String = segments.iter().map(String::as_str).collect();
+ items.push(PushConstantItem {
+ access_path: name,
+ offset: *offset,
+ ty,
+ });
+ *offset += layout.size;
+ }
+ // Arrays are recursed into.
+ TypeInner::Array { base, size, .. } => {
+ let crate::ArraySize::Constant(count) = size else {
+ unreachable!("Cannot have dynamic arrays in push constants");
+ };
+
+ for i in 0..count.get() {
+ // Add the array accessor and recurse.
+ segments.push(format!("[{}]", i));
+ self.collect_push_constant_items(base, segments, layouter, offset, items);
+ segments.pop();
+ }
+
+ // Ensure the stride is kept by rounding up to the alignment.
+ *offset = layout.alignment.round_up(*offset)
+ }
+ TypeInner::Struct { ref members, .. } => {
+ for (index, member) in members.iter().enumerate() {
+ // Add struct accessor and recurse.
+ segments.push(format!(
+ ".{}",
+ self.names[&NameKey::StructMember(ty, index as u32)]
+ ));
+ self.collect_push_constant_items(member.ty, segments, layouter, offset, items);
+ segments.pop();
+ }
+
+ // Ensure ending padding is kept by rounding up to the alignment.
+ *offset = layout.alignment.round_up(*offset)
+ }
+ _ => unreachable!(),
+ }
+ }
 }
 
 /// Structure returned by [`glsl_scalar`]

naga/tests/out/glsl/push-constants.main.Fragment.glsl

@@ -9,14 +9,14 @@ struct PushConstants {
 struct FragmentIn {
  vec4 color;
 };
-uniform PushConstants pc;
+uniform PushConstants _push_constant_binding_fs;
 
 layout(location = 0) smooth in vec4 _vs2fs_location0;
 layout(location = 0) out vec4 _fs2p_location0;
 
 void main() {
  FragmentIn in_ = FragmentIn(_vs2fs_location0);
- float _e4 = pc.multiplier;
+ float _e4 = _push_constant_binding_fs.multiplier;
  _fs2p_location0 = (in_.color * _e4);
  return;
 }

naga/tests/out/glsl/push-constants.vert_main.Vertex.glsl

@@ -9,14 +9,14 @@ struct PushConstants {
 struct FragmentIn {
  vec4 color;
 };
-uniform PushConstants pc;
+uniform PushConstants _push_constant_binding_vs;
 
 layout(location = 0) in vec2 _p2vs_location0;
 
 void main() {
  vec2 pos = _p2vs_location0;
  uint vi = uint(gl_VertexID);
- float _e5 = pc.multiplier;
+ float _e5 = _push_constant_binding_vs.multiplier;
  gl_Position = vec4(((float(vi) * _e5) * pos), 0.0, 1.0);
  return;
 }

tests/src/image.rs

@@ -625,12 +625,16 @@ impl ReadbackBuffers {
  buffer_zero && stencil_buffer_zero
  }
 
- pub fn check_buffer_contents(&self, device: &Device, expected_data: &[u8]) -> bool {
- let result = self
- .retrieve_buffer(device, &self.buffer, self.buffer_aspect())
- .iter()
- .eq(expected_data.iter());
+ pub fn assert_buffer_contents(&self, device: &Device, expected_data: &[u8]) {
+ let result_buffer = self.retrieve_buffer(device, &self.buffer, self.buffer_aspect());
+ assert!(
+ result_buffer.len() >= expected_data.len(),
+ "Result buffer ({}) smaller than expected buffer ({})",
+ result_buffer.len(),
+ expected_data.len()
+ );
+ let result_buffer = &result_buffer[..expected_data.len()];
+ assert_eq!(result_buffer, expected_data);
  self.buffer.unmap();
- result
  }
 }

tests/tests/gpu.rs

@@ -1,4 +1,5 @@
 mod regression {
+ mod issue_3349;
  mod issue_3457;
  mod issue_4024;
  mod issue_4122;
@@ -19,6 +20,7 @@ mod occlusion_query;
 mod partially_bounded_arrays;
 mod pipeline;
 mod poll;
+mod push_constants;
 mod query_set;
 mod queue_transfer;
 mod resource_descriptor_accessor;

tests/tests/partially_bounded_arrays/mod.rs

@@ -97,9 +97,6 @@ static PARTIALLY_BOUNDED_ARRAY: GpuTestConfiguration = GpuTestConfiguration::new
 
  ctx.queue.submit(Some(encoder.finish()));
 
- assert!(
- readback_buffers
- .check_buffer_contents(device, bytemuck::bytes_of(&[4.0f32, 3.0, 2.0, 1.0])),
- "texture storage values are incorrect!"
- );
+ readback_buffers
+ .assert_buffer_contents(device, bytemuck::bytes_of(&[4.0f32, 3.0, 2.0, 1.0]));
  });