From 7e912b71e3f30bc74285b3954d2dfded2c969334 Mon Sep 17 00:00:00 2001 From: Kane Rogers Date: Sat, 23 Sep 2023 04:02:17 +1000 Subject: [PATCH 1/2] Move to shader defined argument buffers --- Cargo.lock | 1 - Cargo.toml | 22 +- wgpu-hal/src/metal/command.rs | 373 ++++++++++++++++++++-------------- wgpu-hal/src/metal/device.rs | 145 ++++++++++--- wgpu-hal/src/metal/mod.rs | 37 +++- 5 files changed, 386 insertions(+), 192 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index ba1a403628..e057b04a6a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1603,7 +1603,6 @@ dependencies = [ [[package]] name = "naga" version = "0.13.0" -source = "git+https://github.com/gfx-rs/naga?rev=cc87b8f9eb30bb55d0735b89d3df3e099e1a6e7c#cc87b8f9eb30bb55d0735b89d3df3e099e1a6e7c" dependencies = [ "bit-set", "bitflags 2.4.0", diff --git a/Cargo.toml b/Cargo.toml index 55c6048b86..5f7d7924db 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,7 +21,7 @@ default-members = [ "wgpu-hal", "wgpu-info", "wgpu-types", - "tests" + "tests", ] [workspace.package] @@ -81,7 +81,11 @@ noise = "0.8" obj = "0.10" # parking_lot 0.12 switches from `winapi` to `windows`; permit either parking_lot = ">=0.11,<0.13" -pico-args = { version = "0.5.0", features = ["eq-separator", "short-space-opt", "combined-flags"] } +pico-args = { version = "0.5.0", features = [ + "eq-separator", + "short-space-opt", + "combined-flags", +] } png = "0.17.10" pollster = "0.3" profiling = { version = "1", default-features = false } @@ -96,9 +100,9 @@ thiserror = "1" wgpu = { version = "0.17.0", path = "./wgpu" } wgpu-core = { version = "0.17.0", path = "./wgpu-core" } wgpu-example = { version = "0.17.0", path = "./examples/common" } -wgpu-test = { version = "0.17", path = "./tests"} +wgpu-test = { version = "0.17", path = "./tests" } wgpu-types = { version = "0.17.0", path = "./wgpu-types" } -winit = { version = "0.28.6", features = [ "android-native-activity" ] } +winit = { version = "0.28.6", features = ["android-native-activity"] } # Metal dependencies block = "0.1" @@ -114,7 +118,11 @@ android_system_properties = "0.1.1" # DX dependencies bit-set = "0.5" -gpu-allocator = { version = "0.21", default_features = false, features = ["d3d12", "windows", "public-winapi"] } +gpu-allocator = { version = "0.21", default_features = false, features = [ + "d3d12", + "windows", + "public-winapi", +] } d3d12 = "0.7.0" range-alloc = "0.1" winapi = "0.3" @@ -145,7 +153,7 @@ tokio = "1.32.0" termcolor = "1.2.0" [patch."https://github.com/gfx-rs/naga"] -#naga = { path = "../naga" } +naga = { path = "../naga" } [patch."https://github.com/zakarumych/gpu-descriptor"] #gpu-descriptor = { path = "../gpu-descriptor/gpu-descriptor" } @@ -154,7 +162,7 @@ termcolor = "1.2.0" #gpu-alloc = { path = "../gpu-alloc/gpu-alloc" } [patch.crates-io] -#naga = { path = "../naga" } +naga = { path = "../naga" } #glow = { path = "../glow" } #d3d12 = { path = "../d3d12-rs" } #metal = { path = "../metal-rs" } diff --git a/wgpu-hal/src/metal/command.rs b/wgpu-hal/src/metal/command.rs index cc737fd228..9553fb8bb6 100644 --- a/wgpu-hal/src/metal/command.rs +++ b/wgpu-hal/src/metal/command.rs @@ -1,3 +1,5 @@ +use metal::MTLResourceUsage; + use super::{conv, AsNative}; use crate::CommandEncoder as _; use std::{borrow::Cow, mem, ops::Range}; @@ -499,161 +501,221 @@ impl crate::CommandEncoder for super::CommandEncoder { ) { let bg_info = &layout.bind_group_infos[group_index as usize]; + // TODO iterate through shader stages if let Some(ref encoder) = self.state.render { - let mut changes_sizes_buffer = false; - for index in 0..group.counters.vs.buffers { - let buf = &group.buffers[index as usize]; - let mut offset = buf.offset; - if let Some(dyn_index) = buf.dynamic_index { - offset += dynamic_offsets[dyn_index as usize] as wgt::BufferAddress; - } - encoder.set_vertex_buffer( - (bg_info.base_resource_indices.vs.buffers + index) as u64, - Some(buf.ptr.as_native()), - offset, - ); - if let Some(size) = buf.binding_size { - let br = naga::ResourceBinding { - group: group_index, - binding: buf.binding_location, - }; - self.state.storage_buffer_length_map.insert(br, size); - changes_sizes_buffer = true; - } - } - if changes_sizes_buffer { - if let Some((index, sizes)) = self.state.make_sizes_buffer_update( - naga::ShaderStage::Vertex, - &mut self.temp.binding_sizes, - ) { - encoder.set_vertex_bytes( - index as _, - (sizes.len() * WORD_SIZE) as u64, - sizes.as_ptr() as _, - ); - } - } - - changes_sizes_buffer = false; - for index in 0..group.counters.fs.buffers { - let buf = &group.buffers[(group.counters.vs.buffers + index) as usize]; - let mut offset = buf.offset; - if let Some(dyn_index) = buf.dynamic_index { - offset += dynamic_offsets[dyn_index as usize] as wgt::BufferAddress; - } - encoder.set_fragment_buffer( - (bg_info.base_resource_indices.fs.buffers + index) as u64, - Some(buf.ptr.as_native()), - offset, - ); - if let Some(size) = buf.binding_size { - let br = naga::ResourceBinding { - group: group_index, - binding: buf.binding_location, - }; - self.state.storage_buffer_length_map.insert(br, size); - changes_sizes_buffer = true; - } - } - if changes_sizes_buffer { - if let Some((index, sizes)) = self.state.make_sizes_buffer_update( - naga::ShaderStage::Fragment, - &mut self.temp.binding_sizes, - ) { - encoder.set_fragment_bytes( - index as _, - (sizes.len() * WORD_SIZE) as u64, - sizes.as_ptr() as _, - ); - } - } - - for index in 0..group.counters.vs.samplers { - let res = group.samplers[index as usize]; - encoder.set_vertex_sampler_state( - (bg_info.base_resource_indices.vs.samplers + index) as u64, - Some(res.as_native()), - ); - } - for index in 0..group.counters.fs.samplers { - let res = group.samplers[(group.counters.vs.samplers + index) as usize]; - encoder.set_fragment_sampler_state( - (bg_info.base_resource_indices.fs.samplers + index) as u64, - Some(res.as_native()), - ); - } - - for index in 0..group.counters.vs.textures { - let res = group.textures[index as usize]; - encoder.set_vertex_texture( - (bg_info.base_resource_indices.vs.textures + index) as u64, - Some(res.as_native()), - ); - } - for index in 0..group.counters.fs.textures { - let res = group.textures[(group.counters.vs.textures + index) as usize]; - encoder.set_fragment_texture( - (bg_info.base_resource_indices.fs.textures + index) as u64, - Some(res.as_native()), - ); - } - } - - if let Some(ref encoder) = self.state.compute { - let index_base = super::ResourceData { - buffers: group.counters.vs.buffers + group.counters.fs.buffers, - samplers: group.counters.vs.samplers + group.counters.fs.samplers, - textures: group.counters.vs.textures + group.counters.fs.textures, - }; - - let mut changes_sizes_buffer = false; - for index in 0..group.counters.cs.buffers { - let buf = &group.buffers[(index_base.buffers + index) as usize]; - let mut offset = buf.offset; - if let Some(dyn_index) = buf.dynamic_index { - offset += dynamic_offsets[dyn_index as usize] as wgt::BufferAddress; - } - encoder.set_buffer( - (bg_info.base_resource_indices.cs.buffers + index) as u64, - Some(buf.ptr.as_native()), - offset, - ); - if let Some(size) = buf.binding_size { - let br = naga::ResourceBinding { - group: group_index, - binding: buf.binding_location, - }; - self.state.storage_buffer_length_map.insert(br, size); - changes_sizes_buffer = true; - } - } - if changes_sizes_buffer { - if let Some((index, sizes)) = self.state.make_sizes_buffer_update( - naga::ShaderStage::Compute, - &mut self.temp.binding_sizes, - ) { - encoder.set_bytes( - index as _, - (sizes.len() * WORD_SIZE) as u64, - sizes.as_ptr() as _, - ); + for entry in &group.bindings { + for (stage, stage_info) in [ + (wgt::ShaderStages::VERTEX, &self.state.stage_infos.vs), + (wgt::ShaderStages::FRAGMENT, &self.state.stage_infos.fs), + ] { + if entry.visibility.contains(stage) { + // TODO: It *should* be impossible to have bindings but no argument buffer. + let Some(ref argument_buffer) = + stage_info.argument_buffer.as_ref() else { + log::warn!("Attempted to set bind group but no argument buffer found for vertex stage!"); + continue; + }; + + let index = argument_buffer.entries[&naga::ResourceBinding { + group: group_index, + binding: entry.binding, + }]; + + match &entry.resource { + super::MetalBindGroupResource::Buffer(buffers) => { + for (buffer_index, buffer) in buffers.iter().enumerate() { + let mut offset = buffer.offset; + if let Some(dyn_index) = buffer.dynamic_index { + offset += dynamic_offsets[dyn_index as usize] as wgt::BufferAddress; + } + let native = buffer.ptr.as_native(); + encoder.use_resource(native, MTLResourceUsage::Read); + argument_buffer.encoder.set_buffer( + (index as usize + buffer_index) as _, + native, + offset, + ) + } + } + super::MetalBindGroupResource::Texture(textures) => { + // PERF: potentially needless vec allocation? + for (texture_index, texture) in textures.iter().enumerate() { + encoder + .use_resource(texture.as_native(), MTLResourceUsage::Read); + argument_buffer.encoder.set_texture( + (index as usize + texture_index) as _, + texture.as_native(), + ) + } + } + super::MetalBindGroupResource::Sampler(samplers) => { + for (sampler_index, sampler) in samplers.iter().enumerate() { + // encoder + // .use_resource(sampler.as_native(), MTLResourceUsage::Read); + argument_buffer.encoder.set_sampler_state( + (index as usize + sampler_index) as _, + sampler.as_native(), + ) + } + } + } + } } } - for index in 0..group.counters.cs.samplers { - let res = group.samplers[(index_base.samplers + index) as usize]; - encoder.set_sampler_state( - (bg_info.base_resource_indices.cs.samplers + index) as u64, - Some(res.as_native()), - ); - } - for index in 0..group.counters.cs.textures { - let res = group.textures[(index_base.textures + index) as usize]; - encoder.set_texture( - (bg_info.base_resource_indices.cs.textures + index) as u64, - Some(res.as_native()), - ); - } + // TODO: Handle dynamic offsets + // let mut changes_sizes_buffer = false; + // for index in 0..group.counters.vs.buffers { + // let buf = &group.buffers[index as usize]; + // let mut offset = buf.offset; + // if let Some(dyn_index) = buf.dynamic_index { + // offset += dynamic_offsets[dyn_index as usize] as wgt::BufferAddress; + // } + // encoder.set_vertex_buffer( + // (bg_info.base_resource_indices.vs.buffers + index) as u64, + // Some(buf.ptr.as_native()), + // offset, + // ); + // if let Some(size) = buf.binding_size { + // let br = naga::ResourceBinding { + // group: group_index, + // binding: buf.binding_location, + // }; + // self.state.storage_buffer_length_map.insert(br, size); + // changes_sizes_buffer = true; + // } + // } + // if changes_sizes_buffer { + // if let Some((index, sizes)) = self.state.make_sizes_buffer_update( + // naga::ShaderStage::Vertex, + // &mut self.temp.binding_sizes, + // ) { + // encoder.set_vertex_bytes( + // index as _, + // (sizes.len() * WORD_SIZE) as u64, + // sizes.as_ptr() as _, + // ); + // } + // } + + // changes_sizes_buffer = false; + // for index in 0..group.counters.fs.buffers { + // let buf = &group.buffers[(group.counters.vs.buffers + index) as usize]; + // let mut offset = buf.offset; + // if let Some(dyn_index) = buf.dynamic_index { + // offset += dynamic_offsets[dyn_index as usize] as wgt::BufferAddress; + // } + // encoder.set_fragment_buffer( + // (bg_info.base_resource_indices.fs.buffers + index) as u64, + // Some(buf.ptr.as_native()), + // offset, + // ); + // if let Some(size) = buf.binding_size { + // let br = naga::ResourceBinding { + // group: group_index, + // binding: buf.binding_location, + // }; + // self.state.storage_buffer_length_map.insert(br, size); + // changes_sizes_buffer = true; + // } + // } + // if changes_sizes_buffer { + // if let Some((index, sizes)) = self.state.make_sizes_buffer_update( + // naga::ShaderStage::Fragment, + // &mut self.temp.binding_sizes, + // ) { + // encoder.set_fragment_bytes( + // index as _, + // (sizes.len() * WORD_SIZE) as u64, + // sizes.as_ptr() as _, + // ); + // } + // } + + // for index in 0..group.counters.vs.samplers { + // let res = group.samplers[index as usize]; + // encoder.set_vertex_sampler_state( + // (bg_info.base_resource_indices.vs.samplers + index) as u64, + // Some(res.as_native()), + // ); + // } + // for index in 0..group.counters.fs.samplers { + // let res = group.samplers[(group.counters.vs.samplers + index) as usize]; + // encoder.set_fragment_sampler_state( + // (bg_info.base_resource_indices.fs.samplers + index) as u64, + // Some(res.as_native()), + // ); + // } + + // for index in 0..group.counters.vs.textures { + // let res = group.textures[index as usize]; + // encoder.set_vertex_texture( + // (bg_info.base_resource_indices.vs.textures + index) as u64, + // Some(res.as_native()), + // ); + // } + // for index in 0..group.counters.fs.textures { + // let res = group.textures[(group.counters.vs.textures + index) as usize]; + // encoder.set_fragment_texture( + // (bg_info.base_resource_indices.fs.textures + index) as u64, + // Some(res.as_native()), + // ); + // } + + // let index_base = super::ResourceData { + // buffers: group.counters.vs.buffers + group.counters.fs.buffers, + // samplers: group.counters.vs.samplers + group.counters.fs.samplers, + // textures: group.counters.vs.textures + group.counters.fs.textures, + // }; + + // let mut changes_sizes_buffer = false; + // for index in 0..group.counters.cs.buffers { + // let buf = &group.buffers[(index_base.buffers + index) as usize]; + // let mut offset = buf.offset; + // if let Some(dyn_index) = buf.dynamic_index { + // offset += dynamic_offsets[dyn_index as usize] as wgt::BufferAddress; + // } + // encoder.set_buffer( + // (bg_info.base_resource_indices.cs.buffers + index) as u64, + // Some(buf.ptr.as_native()), + // offset, + // ); + // if let Some(size) = buf.binding_size { + // let br = naga::ResourceBinding { + // group: group_index, + // binding: buf.binding_location, + // }; + // self.state.storage_buffer_length_map.insert(br, size); + // changes_sizes_buffer = true; + // } + // } + // if changes_sizes_buffer { + // if let Some((index, sizes)) = self.state.make_sizes_buffer_update( + // naga::ShaderStage::Compute, + // &mut self.temp.binding_sizes, + // ) { + // encoder.set_bytes( + // index as _, + // (sizes.len() * WORD_SIZE) as u64, + // sizes.as_ptr() as _, + // ); + // } + // } + + // for index in 0..group.counters.cs.samplers { + // let res = group.samplers[(index_base.samplers + index) as usize]; + // encoder.set_sampler_state( + // (bg_info.base_resource_indices.cs.samplers + index) as u64, + // Some(res.as_native()), + // ); + // } + // for index in 0..group.counters.cs.textures { + // let res = group.textures[(index_base.textures + index) as usize]; + // encoder.set_texture( + // (bg_info.base_resource_indices.cs.textures + index) as u64, + // Some(res.as_native()), + // ); + // } } } @@ -761,6 +823,15 @@ impl crate::CommandEncoder for super::CommandEncoder { ); } } + + if let Some(ref vs_argument_buffer) = self.state.stage_infos.vs.argument_buffer { + // TODO this is probably overriding the buffer with, you know. The vertices. + encoder.set_vertex_buffer(0, Some(&vs_argument_buffer.buffer), 0); + } + + if let Some(ref fs_argument_buffer) = self.state.stage_infos.fs.argument_buffer { + encoder.set_fragment_buffer(0, Some(&fs_argument_buffer.buffer), 0); + } } unsafe fn set_index_buffer<'a>( diff --git a/wgpu-hal/src/metal/device.rs b/wgpu-hal/src/metal/device.rs index 475332b76d..2d2c40ad67 100644 --- a/wgpu-hal/src/metal/device.rs +++ b/wgpu-hal/src/metal/device.rs @@ -1,3 +1,4 @@ +use naga::{FastHashMap, ShaderStage}; use parking_lot::Mutex; use std::{ num::NonZeroU32, @@ -6,8 +7,8 @@ use std::{ thread, time, }; -use super::conv; -use crate::auxil::map_naga_stage; +use super::{conv, ArgumentBuffer}; +use crate::{auxil::map_naga_stage, RenderPipelineDescriptor}; type DeviceResult = Result; @@ -28,6 +29,7 @@ struct CompiledShader { sized_bindings: Vec, immutable_buffer_mask: usize, + argument_buffer_entries: FastHashMap, } fn create_stencil_desc( @@ -114,7 +116,7 @@ impl super::Device { }, }; - let (source, info) = naga::back::msl::write_string( + let (source, mut translation_info) = naga::back::msl::write_string( module, &stage.module.naga.info, &options, @@ -152,10 +154,13 @@ impl super::Device { .position(|ep| ep.stage == naga_stage && ep.name == stage.entry_point) .ok_or(crate::PipelineError::EntryPoint(naga_stage))?; let ep = &module.entry_points[ep_index]; - let ep_name = info.entry_point_names[ep_index] - .as_ref() + let ep_info = translation_info + .entry_point_info + .swap_remove(ep_index) .map_err(|e| crate::PipelineError::Linkage(stage_bit, format!("{}", e)))?; + let ep_name = &ep_info.name; + let wg_size = metal::MTLSize { width: ep.workgroup_size[0] as _, height: ep.workgroup_size[1] as _, @@ -168,14 +173,15 @@ impl super::Device { })?; // collect sizes indices, immutable buffers, and work group memory sizes - let ep_info = &stage.module.naga.info.get_entry_point(ep_index); + let function_info = &stage.module.naga.info.get_entry_point(ep_index); let mut wg_memory_sizes = Vec::new(); let mut sized_bindings = Vec::new(); let mut immutable_buffer_mask = 0; + for (var_handle, var) in module.global_variables.iter() { match var.space { naga::AddressSpace::WorkGroup => { - if !ep_info[var_handle].is_empty() { + if !function_info[var_handle].is_empty() { let size = module.types[var.ty].inner.size(module.to_ctx()); wg_memory_sizes.push(size); } @@ -193,7 +199,7 @@ impl super::Device { }; // check for an immutable buffer - if !ep_info[var_handle].is_empty() && !storage_access_store { + if !function_info[var_handle].is_empty() && !storage_access_store { let slot = ep_resources.resources[&br].buffer.unwrap(); immutable_buffer_mask |= 1 << slot; } @@ -222,6 +228,7 @@ impl super::Device { wg_memory_sizes, sized_bindings, immutable_buffer_mask, + argument_buffer_entries: ep_info.argument_buffer_entries, }) } @@ -271,6 +278,50 @@ impl super::Device { pub fn raw_device(&self) -> &Mutex { &self.shared.device } + + fn create_argument_buffer( + &self, + argument_buffer_entries: FastHashMap, + function: &metal::Function, + ) -> Option> { + // Check to see if this shader needs an argument buffer. + if argument_buffer_entries.is_empty() { + return None; + } + // How many bind groups does this entry point use? + // if shader.argument_buffers_used.is_empty() { + // return Vec::new(); + // } + + // let mut buffers = Vec::new(); + // for id in shader.argument_buffers_used.iter().copied() { + let encoder = function.new_argument_encoder(0); + + let size = encoder.encoded_length(); + if size == 0 { + return None; + } + + let id = 0; + + log::debug!("Creating argument buffer at index {id} of size: {size}"); + let buffer = self + .shared + .device + .lock() + .new_buffer(size, metal::MTLResourceOptions::empty()); + // TODO (KR): more meaningful label + buffer.set_label(&format!("ArgumentBufferGroup{id}")); + log::info!("Created argument buffer: {:?}", buffer); + encoder.set_argument_buffer(&buffer, 0); + Some(Arc::new(ArgumentBuffer { + id, + encoder, + buffer, + entries: argument_buffer_entries, + })) + // } + } } impl crate::Device for super::Device { @@ -475,6 +526,8 @@ impl crate::Device for super::Device { descriptor.set_lod_min_clamp(desc.lod_clamp.start); descriptor.set_lod_max_clamp(desc.lod_clamp.end); + descriptor.set_support_argument_buffers(true); + if let Some(fun) = desc.compare { descriptor.set_compare_function(conv::map_compare_function(fun)); } @@ -705,30 +758,35 @@ impl crate::Device for super::Device { for (&stage, counter) in super::NAGA_STAGES.iter().zip(bg.counters.iter_mut()) { let stage_bit = map_naga_stage(stage); let mut dynamic_offsets_count = 0u32; - for (entry, layout) in desc.entries.iter().zip(desc.layout.entries.iter()) { - let size = layout.count.map_or(1, |c| c.get()); + for (bind_group_entry, layout_entry) in + desc.entries.iter().zip(desc.layout.entries.iter()) + { + let size = layout_entry.count.map_or(1, |c| c.get()); + let visibility = layout_entry.visibility; + let binding = bind_group_entry.binding; if let wgt::BindingType::Buffer { has_dynamic_offset: true, .. - } = layout.ty + } = layout_entry.ty { dynamic_offsets_count += size; } - if !layout.visibility.contains(stage_bit) { + if !visibility.contains(stage_bit) { continue; } - match layout.ty { + match layout_entry.ty { wgt::BindingType::Buffer { ty, has_dynamic_offset, .. } => { - let start = entry.resource_index as usize; + let start = bind_group_entry.resource_index as usize; let end = start + size as usize; - bg.buffers - .extend(desc.buffers[start..end].iter().map(|source| { - // Given the restrictions on `BufferBinding::offset`, - // this should never be `None`. + // Given the restrictions on `BufferBinding::offset`, + // this should never be `None`. + let buffers = desc.buffers[start..end] + .iter() + .map(|source| { let remaining_size = wgt::BufferSize::new(source.buffer.size - source.offset); let binding_size = match ty { @@ -746,26 +804,44 @@ impl crate::Device for super::Device { None }, binding_size, - binding_location: layout.binding, + binding_location: layout_entry.binding, } - })); + }) + .collect(); + + bg.bindings.push(super::MetalBinding { + binding, + visibility, + resource: super::MetalBindGroupResource::Buffer(buffers), + }); counter.buffers += 1; } wgt::BindingType::Sampler { .. } => { - let start = entry.resource_index as usize; + let start = bind_group_entry.resource_index as usize; let end = start + size as usize; - bg.samplers - .extend(desc.samplers[start..end].iter().map(|samp| samp.as_raw())); + let samplers = desc.samplers[start..end] + .iter() + .map(|samp| samp.as_raw()) + .collect(); + bg.bindings.push(super::MetalBinding { + binding, + visibility, + resource: super::MetalBindGroupResource::Sampler(samplers), + }); counter.samplers += size; } wgt::BindingType::Texture { .. } | wgt::BindingType::StorageTexture { .. } => { - let start = entry.resource_index as usize; + let start = bind_group_entry.resource_index as usize; let end = start + size as usize; - bg.textures.extend( - desc.textures[start..end] - .iter() - .map(|tex| tex.view.as_raw()), - ); + let textures = desc.textures[start..end] + .iter() + .map(|tex| tex.view.as_raw()) + .collect(); + bg.bindings.push(super::MetalBinding { + binding, + visibility, + resource: super::MetalBindGroupResource::Texture(textures), + }); counter.textures += size; } } @@ -796,7 +872,7 @@ impl crate::Device for super::Device { unsafe fn create_render_pipeline( &self, - desc: &crate::RenderPipelineDescriptor, + desc: &RenderPipelineDescriptor, ) -> Result { objc::rc::autoreleasepool(|| { let descriptor = metal::RenderPipelineDescriptor::new(); @@ -822,6 +898,9 @@ impl crate::Device for super::Device { naga::ShaderStage::Vertex, )?; + let argument_buffer = + self.create_argument_buffer(vs.argument_buffer_entries, &vs.function); + descriptor.set_vertex_function(Some(&vs.function)); if self.shared.private_caps.supports_mutability { Self::set_buffers_mutability( @@ -834,6 +913,7 @@ impl crate::Device for super::Device { push_constants: desc.layout.push_constants_infos.vs, sizes_slot: desc.layout.per_stage_map.vs.sizes_buffer, sized_bindings: vs.sized_bindings, + argument_buffer, }; (vs.library, info) @@ -849,6 +929,9 @@ impl crate::Device for super::Device { naga::ShaderStage::Fragment, )?; + let argument_buffer = + self.create_argument_buffer(fs.argument_buffer_entries, &fs.function); + descriptor.set_fragment_function(Some(&fs.function)); if self.shared.private_caps.supports_mutability { Self::set_buffers_mutability( @@ -861,6 +944,7 @@ impl crate::Device for super::Device { push_constants: desc.layout.push_constants_infos.fs, sizes_slot: desc.layout.per_stage_map.fs.sizes_buffer, sized_bindings: fs.sized_bindings, + argument_buffer, }; (Some(fs.library), Some(info)) @@ -1053,6 +1137,7 @@ impl crate::Device for super::Device { push_constants: desc.layout.push_constants_infos.cs, sizes_slot: desc.layout.per_stage_map.cs.sizes_buffer, sized_bindings: cs.sized_bindings, + argument_buffer: None, }; if let Some(name) = desc.label { diff --git a/wgpu-hal/src/metal/mod.rs b/wgpu-hal/src/metal/mod.rs index 76f57002ff..5bd7d352c2 100644 --- a/wgpu-hal/src/metal/mod.rs +++ b/wgpu-hal/src/metal/mod.rs @@ -34,8 +34,11 @@ use std::{ use arrayvec::ArrayVec; use metal::foreign_types::ForeignTypeRef as _; +use naga::{FastHashMap, ResourceBinding}; use parking_lot::Mutex; +use crate::BindGroupEntry; + #[derive(Clone)] pub struct Api; @@ -629,9 +632,24 @@ struct BufferResource { #[derive(Debug, Default)] pub struct BindGroup { counters: MultiStageResourceCounters, - buffers: Vec, - samplers: Vec, - textures: Vec, + // buffers: Vec, + // samplers: Vec, + // textures: Vec, + bindings: Vec, +} + +#[derive(Debug)] +pub struct MetalBinding { + binding: u32, + visibility: wgt::ShaderStages, + resource: MetalBindGroupResource, +} + +#[derive(Debug)] +enum MetalBindGroupResource { + Buffer(Vec), + Texture(Vec), + Sampler(Vec), } unsafe impl Send for BindGroup {} @@ -656,6 +674,9 @@ struct PipelineStageInfo { /// /// See [`device::CompiledShader::sized_bindings`] for more details. sized_bindings: Vec, + + /// An argument buffer, if necessary + argument_buffer: Option>, } impl PipelineStageInfo { @@ -663,6 +684,7 @@ impl PipelineStageInfo { self.push_constants = None; self.sizes_slot = None; self.sized_bindings.clear(); + self.argument_buffer = None; } fn assign_from(&mut self, other: &Self) { @@ -670,6 +692,7 @@ impl PipelineStageInfo { self.sizes_slot = other.sizes_slot; self.sized_bindings.clear(); self.sized_bindings.extend_from_slice(&other.sized_bindings); + self.argument_buffer = other.argument_buffer.clone(); } } @@ -816,3 +839,11 @@ pub struct CommandBuffer { unsafe impl Send for CommandBuffer {} unsafe impl Sync for CommandBuffer {} + +#[derive(Debug, Clone)] +struct ArgumentBuffer { + id: u32, + encoder: metal::ArgumentEncoder, + buffer: metal::Buffer, + entries: FastHashMap, +} From f7b06fe44bc0cd1ef5656cee73fcc2e2d9816421 Mon Sep 17 00:00:00 2001 From: Kane Rogers Date: Mon, 9 Oct 2023 17:24:04 +1100 Subject: [PATCH 2/2] Further work on argument buffers --- wgpu-hal/src/metal/command.rs | 69 ++++++++++++++++++----------------- wgpu-hal/src/metal/device.rs | 49 ++++++++++++++----------- 2 files changed, 63 insertions(+), 55 deletions(-) diff --git a/wgpu-hal/src/metal/command.rs b/wgpu-hal/src/metal/command.rs index 9553fb8bb6..ab6b23a6be 100644 --- a/wgpu-hal/src/metal/command.rs +++ b/wgpu-hal/src/metal/command.rs @@ -494,28 +494,39 @@ impl crate::CommandEncoder for super::CommandEncoder { unsafe fn set_bind_group( &mut self, - layout: &super::PipelineLayout, + _layout: &super::PipelineLayout, group_index: u32, group: &super::BindGroup, dynamic_offsets: &[wgt::DynamicOffset], ) { - let bg_info = &layout.bind_group_infos[group_index as usize]; - // TODO iterate through shader stages if let Some(ref encoder) = self.state.render { - for entry in &group.bindings { - for (stage, stage_info) in [ - (wgt::ShaderStages::VERTEX, &self.state.stage_infos.vs), - (wgt::ShaderStages::FRAGMENT, &self.state.stage_infos.fs), - ] { - if entry.visibility.contains(stage) { - // TODO: It *should* be impossible to have bindings but no argument buffer. - let Some(ref argument_buffer) = - stage_info.argument_buffer.as_ref() else { - log::warn!("Attempted to set bind group but no argument buffer found for vertex stage!"); - continue; - }; + for (stage, stage_info) in [ + (wgt::ShaderStages::VERTEX, &self.state.stage_infos.vs), + (wgt::ShaderStages::FRAGMENT, &self.state.stage_infos.fs), + ] { + // TODO: It *should* be impossible to have bindings but no argument buffer. + let Some(ref argument_buffer) = + stage_info.argument_buffer.as_ref() else { + log::warn!("Attempted to set bind group but no argument buffer found for {stage:?}!"); + continue; + }; + match stage { + wgt::ShaderStages::VERTEX => { + encoder.set_vertex_buffer(0, Some(&argument_buffer.buffer), 0); + } + wgt::ShaderStages::FRAGMENT => { + encoder.set_fragment_buffer(0, Some(&argument_buffer.buffer), 0); + } + wgt::ShaderStages::COMPUTE => { + todo!("Support argument buffers on compute"); + } + _ => unreachable!(), + } + + for entry in &group.bindings { + if entry.visibility.contains(stage) { let index = argument_buffer.entries[&naga::ResourceBinding { group: group_index, binding: entry.binding, @@ -525,20 +536,23 @@ impl crate::CommandEncoder for super::CommandEncoder { super::MetalBindGroupResource::Buffer(buffers) => { for (buffer_index, buffer) in buffers.iter().enumerate() { let mut offset = buffer.offset; + let index = index as usize + buffer_index; if let Some(dyn_index) = buffer.dynamic_index { - offset += dynamic_offsets[dyn_index as usize] as wgt::BufferAddress; + offset += dynamic_offsets[dyn_index as usize] + as wgt::BufferAddress; + log::trace!( + "index {index}, dyn_index {dyn_index}, offset {offset}" + ); } let native = buffer.ptr.as_native(); + // TODO: use resource_at encoder.use_resource(native, MTLResourceUsage::Read); - argument_buffer.encoder.set_buffer( - (index as usize + buffer_index) as _, - native, - offset, - ) + argument_buffer + .encoder + .set_buffer(index as _, native, offset) } } super::MetalBindGroupResource::Texture(textures) => { - // PERF: potentially needless vec allocation? for (texture_index, texture) in textures.iter().enumerate() { encoder .use_resource(texture.as_native(), MTLResourceUsage::Read); @@ -550,8 +564,6 @@ impl crate::CommandEncoder for super::CommandEncoder { } super::MetalBindGroupResource::Sampler(samplers) => { for (sampler_index, sampler) in samplers.iter().enumerate() { - // encoder - // .use_resource(sampler.as_native(), MTLResourceUsage::Read); argument_buffer.encoder.set_sampler_state( (index as usize + sampler_index) as _, sampler.as_native(), @@ -823,15 +835,6 @@ impl crate::CommandEncoder for super::CommandEncoder { ); } } - - if let Some(ref vs_argument_buffer) = self.state.stage_infos.vs.argument_buffer { - // TODO this is probably overriding the buffer with, you know. The vertices. - encoder.set_vertex_buffer(0, Some(&vs_argument_buffer.buffer), 0); - } - - if let Some(ref fs_argument_buffer) = self.state.stage_infos.fs.argument_buffer { - encoder.set_fragment_buffer(0, Some(&fs_argument_buffer.buffer), 0); - } } unsafe fn set_index_buffer<'a>( diff --git a/wgpu-hal/src/metal/device.rs b/wgpu-hal/src/metal/device.rs index 2d2c40ad67..8bd7a69b45 100644 --- a/wgpu-hal/src/metal/device.rs +++ b/wgpu-hal/src/metal/device.rs @@ -288,13 +288,7 @@ impl super::Device { if argument_buffer_entries.is_empty() { return None; } - // How many bind groups does this entry point use? - // if shader.argument_buffers_used.is_empty() { - // return Vec::new(); - // } - // let mut buffers = Vec::new(); - // for id in shader.argument_buffers_used.iter().copied() { let encoder = function.new_argument_encoder(0); let size = encoder.encoded_length(); @@ -310,8 +304,7 @@ impl super::Device { .device .lock() .new_buffer(size, metal::MTLResourceOptions::empty()); - // TODO (KR): more meaningful label - buffer.set_label(&format!("ArgumentBufferGroup{id}")); + buffer.set_label(&format!("{} Argument Buffer", function.name())); log::info!("Created argument buffer: {:?}", buffer); encoder.set_argument_buffer(&buffer, 0); Some(Arc::new(ArgumentBuffer { @@ -320,7 +313,6 @@ impl super::Device { buffer, entries: argument_buffer_entries, })) - // } } } @@ -898,16 +890,23 @@ impl crate::Device for super::Device { naga::ShaderStage::Vertex, )?; + descriptor.set_vertex_function(Some(&vs.function)); let argument_buffer = self.create_argument_buffer(vs.argument_buffer_entries, &vs.function); - descriptor.set_vertex_function(Some(&vs.function)); - if self.shared.private_caps.supports_mutability { - Self::set_buffers_mutability( - descriptor.vertex_buffers().unwrap(), - vs.immutable_buffer_mask, - ); - } + descriptor + .vertex_buffers() + .unwrap() + .object_at(0) + .unwrap() + .set_mutability(metal::MTLMutability::Mutable); + + // if self.shared.private_caps.supports_mutability { + // Self::set_buffers_mutability( + // descriptor.vertex_buffers().unwrap(), + // vs.immutable_buffer_mask, + // ); + // } let info = super::PipelineStageInfo { push_constants: desc.layout.push_constants_infos.vs, @@ -931,14 +930,20 @@ impl crate::Device for super::Device { let argument_buffer = self.create_argument_buffer(fs.argument_buffer_entries, &fs.function); + descriptor + .vertex_buffers() + .unwrap() + .object_at(0) + .unwrap() + .set_mutability(metal::MTLMutability::Mutable); descriptor.set_fragment_function(Some(&fs.function)); - if self.shared.private_caps.supports_mutability { - Self::set_buffers_mutability( - descriptor.fragment_buffers().unwrap(), - fs.immutable_buffer_mask, - ); - } + // if self.shared.private_caps.supports_mutability { + // Self::set_buffers_mutability( + // descriptor.fragment_buffers().unwrap(), + // fs.immutable_buffer_mask, + // ); + // } let info = super::PipelineStageInfo { push_constants: desc.layout.push_constants_infos.fs,