hal/gl: Allow push constants trough emulation

Cargo.toml

@@ -18,6 +18,7 @@ default-members = ["wgpu", "wgpu-hal", "wgpu-info"]
 
 [patch."https://github.com/gfx-rs/naga"]
 #naga = { path = "../naga" }
+naga = { git = "https://github.com/JCapucho/naga", branch = "glsl-out-push-constants-v2" }
 
 [patch."https://github.com/zakarumych/gpu-descriptor"]
 #gpu-descriptor = { path = "../gpu-descriptor/gpu-descriptor" }

wgpu-hal/src/gles/adapter.rs

@@ -286,7 +286,8 @@ impl super::Adapter {
 
         let mut features = wgt::Features::empty()
             | wgt::Features::TEXTURE_ADAPTER_SPECIFIC_FORMAT_FEATURES
-            | wgt::Features::CLEAR_TEXTURE;
+            | wgt::Features::CLEAR_TEXTURE
+            | wgt::Features::PUSH_CONSTANTS;
         features.set(
             wgt::Features::ADDRESS_MODE_CLAMP_TO_BORDER | wgt::Features::ADDRESS_MODE_CLAMP_TO_ZERO,
             extensions.contains("GL_EXT_texture_border_clamp"),
@@ -399,7 +400,7 @@ impl super::Adapter {
             } else {
                 !0
             },
-            max_push_constant_size: 0,
+            max_push_constant_size: super::MAX_PUSH_CONSTANTS as u32 * 4,
             min_uniform_buffer_offset_alignment,
             min_storage_buffer_offset_alignment,
             max_inter_stage_shader_components: gl.get_parameter_i32(glow::MAX_VARYING_COMPONENTS)

wgpu-hal/src/gles/command.rs

@@ -28,6 +28,7 @@ pub(super) struct State {
     has_pass_label: bool,
     instance_vbuf_mask: usize,
     dirty_vbuf_mask: usize,
+    push_offset_to_uniform: ArrayVec<super::UniformDesc, { super::MAX_PUSH_CONSTANTS }>,
 }
 
 impl super::CommandBuffer {
@@ -43,6 +44,21 @@ impl super::CommandBuffer {
         self.data_bytes.extend(marker.as_bytes());
         start..self.data_bytes.len() as u32
     }
+
+    fn add_push_constant_data(&mut self, data: &[u32]) -> Range<u32> {
+        let data_raw = unsafe {
+            std::slice::from_raw_parts(
+                data.as_ptr() as *const _,
+                data.len() * mem::size_of::<u32>(),
+            )
+        };
+        let start = self.data_bytes.len();
+        assert!(start < u32::MAX as usize);
+        self.data_bytes.extend_from_slice(data_raw);
+        let end = self.data_bytes.len();
+        assert!(end < u32::MAX as usize);
+        (start as u32)..(end as u32)
+    }
 }
 
 impl super::CommandEncoder {
@@ -148,8 +164,10 @@ impl super::CommandEncoder {
     fn set_pipeline_inner(&mut self, inner: &super::PipelineInner) {
         self.cmd_buffer.commands.push(C::SetProgram(inner.program));
 
-        //TODO: push constants
-        let _ = &inner.uniforms;
+        self.state.push_offset_to_uniform.clear();
+        self.state
+            .push_offset_to_uniform
+            .extend(inner.uniforms.iter().cloned());
 
         // rebind textures, if needed
         let mut dirty_textures = 0u32;
@@ -603,10 +621,25 @@ impl crate::CommandEncoder<super::Api> for super::CommandEncoder {
         &mut self,
         _layout: &super::PipelineLayout,
         _stages: wgt::ShaderStages,
-        _offset: u32,
-        _data: &[u32],
+        start_offset: u32,
+        data: &[u32],
     ) {
-        unimplemented!()
+        let range = self.cmd_buffer.add_push_constant_data(data);
+
+        let end = start_offset + data.len() as u32 * 4;
+        let mut offset = start_offset;
+        while offset < end {
+            let uniform = self.state.push_offset_to_uniform[offset as usize / 4].clone();
+            let size = uniform.size;
+            if uniform.location.is_none() {
+                panic!("No uniform for push constant");
+            }
+            self.cmd_buffer.commands.push(C::SetPushConstants {
+                uniform,
+                offset: range.start + offset,
+            });
+            offset += size;
+        }
     }
 
     unsafe fn insert_debug_marker(&mut self, label: &str) {

wgpu-hal/src/gles/conv.rs

@@ -375,3 +375,59 @@ pub(super) fn map_storage_access(access: wgt::StorageTextureAccess) -> u32 {
         wgt::StorageTextureAccess::ReadWrite => glow::READ_WRITE,
     }
 }
+
+pub(super) fn is_sampler(glsl_uniform_type: u32) -> bool {
+    match glsl_uniform_type {
+        glow::INT_SAMPLER_1D
+        | glow::INT_SAMPLER_1D_ARRAY
+        | glow::INT_SAMPLER_2D
+        | glow::INT_SAMPLER_2D_ARRAY
+        | glow::INT_SAMPLER_2D_MULTISAMPLE
+        | glow::INT_SAMPLER_2D_MULTISAMPLE_ARRAY
+        | glow::INT_SAMPLER_2D_RECT
+        | glow::INT_SAMPLER_3D
+        | glow::INT_SAMPLER_CUBE
+        | glow::INT_SAMPLER_CUBE_MAP_ARRAY
+        | glow::UNSIGNED_INT_SAMPLER_1D
+        | glow::UNSIGNED_INT_SAMPLER_1D_ARRAY
+        | glow::UNSIGNED_INT_SAMPLER_2D
+        | glow::UNSIGNED_INT_SAMPLER_2D_ARRAY
+        | glow::UNSIGNED_INT_SAMPLER_2D_MULTISAMPLE
+        | glow::UNSIGNED_INT_SAMPLER_2D_MULTISAMPLE_ARRAY
+        | glow::UNSIGNED_INT_SAMPLER_2D_RECT
+        | glow::UNSIGNED_INT_SAMPLER_3D
+        | glow::UNSIGNED_INT_SAMPLER_CUBE
+        | glow::UNSIGNED_INT_SAMPLER_CUBE_MAP_ARRAY
+        | glow::SAMPLER_1D
+        | glow::SAMPLER_1D_SHADOW
+        | glow::SAMPLER_1D_ARRAY
+        | glow::SAMPLER_1D_ARRAY_SHADOW
+        | glow::SAMPLER_2D
+        | glow::SAMPLER_2D_SHADOW
+        | glow::SAMPLER_2D_ARRAY
+        | glow::SAMPLER_2D_ARRAY_SHADOW
+        | glow::SAMPLER_2D_MULTISAMPLE
+        | glow::SAMPLER_2D_MULTISAMPLE_ARRAY
+        | glow::SAMPLER_2D_RECT
+        | glow::SAMPLER_2D_RECT_SHADOW
+        | glow::SAMPLER_3D
+        | glow::SAMPLER_CUBE
+        | glow::SAMPLER_CUBE_MAP_ARRAY
+        | glow::SAMPLER_CUBE_MAP_ARRAY_SHADOW
+        | glow::SAMPLER_CUBE_SHADOW => true,
+        _ => false,
+    }
+}
+
+pub(super) fn uniform_byte_size(glsl_uniform_type: u32) -> u32 {
+    match glsl_uniform_type {
+        glow::FLOAT | glow::INT => 4,
+        glow::FLOAT_VEC2 | glow::INT_VEC2 => 8,
+        glow::FLOAT_VEC3 | glow::INT_VEC3 => 12,
+        glow::FLOAT_VEC4 | glow::INT_VEC4 => 16,
+        glow::FLOAT_MAT2 => 16,
+        glow::FLOAT_MAT3 => 36,
+        glow::FLOAT_MAT4 => 64,
+        _ => panic!("Unsupported uniform datatype!"),
+    }
+}

wgpu-hal/src/gles/device.rs

@@ -272,30 +272,35 @@ impl super::Device {
             }
         }
 
-        let uniforms = {
-            let count = gl.get_active_uniforms(program);
-            let mut offset = 0;
-            let mut uniforms = Vec::new();
-
-            for uniform in 0..count {
-                let glow::ActiveUniform { size, utype, name } =
-                    gl.get_active_uniform(program, uniform).unwrap();
-
-                if let Some(location) = gl.get_uniform_location(program, &name) {
-                    // Sampler2D won't show up in UniformLocation and the only other uniforms
-                    // should be push constants
-                    uniforms.push(super::UniformDesc {
-                        location,
-                        offset,
-                        utype,
-                    });
+        let mut uniforms: [super::UniformDesc; super::MAX_PUSH_CONSTANTS] = Default::default();
+        let count = gl.get_active_uniforms(program);
+        let mut offset = 0;
 
-                    offset += size as u32;
-                }
+        for uniform in 0..count {
+            let glow::ActiveUniform { utype, name, .. } =
+                gl.get_active_uniform(program, uniform).unwrap();
+
+            if conv::is_sampler(utype) {
+                continue;
             }
 
-            uniforms.into_boxed_slice()
-        };
+            if let Some(location) = gl.get_uniform_location(program, &name) {
+                if uniforms[offset / 4].location.is_some() {
+                    panic!("Offset already occupied")
+                }
+
+                // `size` will always be 1 so we need to guess the real size from the type
+                let uniform_size = conv::uniform_byte_size(utype);
+
+                uniforms[offset / 4] = super::UniformDesc {
+                    location: Some(location),
+                    size: uniform_size,
+                    utype,
+                };
+
+                offset += uniform_size as usize;
+            }
+        }
 
         Ok(super::PipelineInner {
             program,
@@ -858,7 +863,6 @@ impl crate::Device<super::Api> for super::Device {
                 version: self.shared.shading_language_version,
                 writer_flags,
                 binding_map,
-                push_constant_binding: 0, //TODO?
             },
         })
     }

wgpu-hal/src/gles/mod.rs

@@ -88,6 +88,7 @@ const MAX_TEXTURE_SLOTS: usize = 16;
 const MAX_SAMPLERS: usize = 16;
 const MAX_VERTEX_ATTRIBUTES: usize = 16;
 const ZERO_BUFFER_SIZE: usize = 256 << 10;
+const MAX_PUSH_CONSTANTS: usize = 16;
 
 impl crate::Api for Api {
     type Instance = Instance;
@@ -392,22 +393,27 @@ struct VertexBufferDesc {
     stride: u32,
 }
 
-#[allow(unused)]
-#[derive(Clone)]
+#[derive(Clone, Debug, Default)]
 struct UniformDesc {
-    location: glow::UniformLocation,
-    offset: u32,
+    location: Option<glow::UniformLocation>,
+    size: u32,
     utype: u32,
 }
 
+// Safe: WASM doesn't have threads
+#[cfg(target_arch = "wasm32")]
+unsafe impl Sync for UniformDesc {}
+#[cfg(target_arch = "wasm32")]
+unsafe impl Send for UniformDesc {}
+
 /// For each texture in the pipeline layout, store the index of the only
 /// sampler (in this layout) that the texture is used with.
 type SamplerBindMap = [Option<u8>; MAX_TEXTURE_SLOTS];
 
 struct PipelineInner {
     program: glow::Program,
     sampler_map: SamplerBindMap,
-    uniforms: Box<[UniformDesc]>,
+    uniforms: [UniformDesc; MAX_PUSH_CONSTANTS],
 }
 
 #[derive(Clone, Debug)]
@@ -716,6 +722,11 @@ enum Command {
     InsertDebugMarker(Range<u32>),
     PushDebugGroup(Range<u32>),
     PopDebugGroup,
+    SetPushConstants {
+        uniform: UniformDesc,
+        /// Offset from the start of the `data_bytes`
+        offset: u32,
+    },
 }
 
 #[derive(Default)]

wgpu-hal/src/gles/queue.rs

@@ -1067,6 +1067,70 @@ impl super::Queue {
                 #[cfg(not(target_arch = "wasm32"))]
                 gl.pop_debug_group();
             }
+            C::SetPushConstants {
+                ref uniform,
+                offset,
+            } => {
+                fn get_data<T>(data: &[u8], offset: u32) -> &[T] {
+                    let raw = &data[(offset as usize)..];
+                    unsafe {
+                        slice::from_raw_parts(
+                            raw.as_ptr() as *const _,
+                            raw.len() / mem::size_of::<T>(),
+                        )
+                    }
+                }
+
+                let location = uniform.location.as_ref();
+
+                match uniform.utype {
+                    glow::FLOAT => {
+                        let data = get_data::<f32>(data_bytes, offset)[0];
+                        gl.uniform_1_f32(location, data);
+                    }
+                    glow::FLOAT_VEC2 => {
+                        let data = get_data::<[f32; 2]>(data_bytes, offset)[0];
+                        gl.uniform_2_f32_slice(location, &data);
+                    }
+                    glow::FLOAT_VEC3 => {
+                        let data = get_data::<[f32; 3]>(data_bytes, offset)[0];
+                        gl.uniform_3_f32_slice(location, &data);
+                    }
+                    glow::FLOAT_VEC4 => {
+                        let data = get_data::<[f32; 4]>(data_bytes, offset)[0];
+                        gl.uniform_4_f32_slice(location, &data);
+                    }
+                    glow::INT => {
+                        let data = get_data::<i32>(data_bytes, offset)[0];
+                        gl.uniform_1_i32(location, data);
+                    }
+                    glow::INT_VEC2 => {
+                        let data = get_data::<[i32; 2]>(data_bytes, offset)[0];
+                        gl.uniform_2_i32_slice(location, &data);
+                    }
+                    glow::INT_VEC3 => {
+                        let data = get_data::<[i32; 3]>(data_bytes, offset)[0];
+                        gl.uniform_3_i32_slice(location, &data);
+                    }
+                    glow::INT_VEC4 => {
+                        let data = get_data::<[i32; 4]>(data_bytes, offset)[0];
+                        gl.uniform_4_i32_slice(location, &data);
+                    }
+                    glow::FLOAT_MAT2 => {
+                        let data = get_data::<[f32; 4]>(data_bytes, offset)[0];
+                        gl.uniform_matrix_2_f32_slice(location, false, &data);
+                    }
+                    glow::FLOAT_MAT3 => {
+                        let data = get_data::<[f32; 9]>(data_bytes, offset)[0];
+                        gl.uniform_matrix_3_f32_slice(location, false, &data);
+                    }
+                    glow::FLOAT_MAT4 => {
+                        let data = get_data::<[f32; 16]>(data_bytes, offset)[0];
+                        gl.uniform_matrix_4_f32_slice(location, false, &data);
+                    }
+                    _ => panic!("Unsupported uniform datatype!"),
+                }
+            }
         }
     }
 }