From 0db71153adbef63235eb36e3e91f5265f3255c30 Mon Sep 17 00:00:00 2001 From: Chad Brokaw Date: Wed, 25 Jan 2023 14:27:14 -0500 Subject: [PATCH 01/12] Playing with shader permutations and AOT compilation --- Cargo.toml | 1 + shader/permutations | 3 + vello_shaders/Cargo.toml | 17 +++ vello_shaders/build.rs | 99 ++++++++++++++ vello_shaders/src/compile/mod.rs | 159 ++++++++++++++++++++++ vello_shaders/src/compile/msl.rs | 49 +++++++ vello_shaders/src/compile/permutations.rs | 41 ++++++ vello_shaders/src/compile/preprocess.rs | 159 ++++++++++++++++++++++ vello_shaders/src/lib.rs | 30 ++++ vello_shaders/src/types.rs | 30 ++++ 10 files changed, 588 insertions(+) create mode 100644 shader/permutations create mode 100644 vello_shaders/Cargo.toml create mode 100644 vello_shaders/build.rs create mode 100644 vello_shaders/src/compile/mod.rs create mode 100644 vello_shaders/src/compile/msl.rs create mode 100644 vello_shaders/src/compile/permutations.rs create mode 100644 vello_shaders/src/compile/preprocess.rs create mode 100644 vello_shaders/src/lib.rs create mode 100644 vello_shaders/src/types.rs diff --git a/Cargo.toml b/Cargo.toml index 397ece2..e682fe3 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -9,6 +9,7 @@ members = [ "examples/with_bevy", "examples/run_wasm", "examples/scenes", + "vello_shaders", ] [workspace.package] diff --git a/shader/permutations b/shader/permutations new file mode 100644 index 0000000..beb1e48 --- /dev/null +++ b/shader/permutations @@ -0,0 +1,3 @@ +pathtag_scan ++ pathtag_scan_large ++ pathtag_scan_small: small diff --git a/vello_shaders/Cargo.toml b/vello_shaders/Cargo.toml new file mode 100644 index 0000000..abd7dc0 --- /dev/null +++ b/vello_shaders/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "vello_shaders" +version = "0.1.0" +edition = "2021" + +[features] +default = ["compile", "wgsl", "msl"] +compile = ["naga"] +wgsl = [] +msl = [] + +[dependencies] +naga = { git = "https://github.com/gfx-rs/naga", features = ["wgsl-in", "msl-out", "validate"], optional = true } + +[build-dependencies] +naga = { git = "https://github.com/gfx-rs/naga", features = ["wgsl-in", "msl-out", "validate"] } + diff --git a/vello_shaders/build.rs b/vello_shaders/build.rs new file mode 100644 index 0000000..e447292 --- /dev/null +++ b/vello_shaders/build.rs @@ -0,0 +1,99 @@ +#[path = "src/compile/mod.rs"] +mod compile; +#[path = "src/types.rs"] +mod types; + +use std::collections::HashMap; +use std::env; +use std::fmt::Write; +use std::path::Path; + +use compile::ShaderInfo; + +fn main() { + let out_dir = env::var_os("OUT_DIR").unwrap(); + let dest_path = Path::new(&out_dir).join("shaders.rs"); + let shaders = compile::ShaderInfo::from_dir("../shader"); + let mut buf = String::default(); + write_types(&mut buf, &shaders).unwrap(); + if cfg!(feature = "wgsl") { + write_shaders(&mut buf, "wgsl", &shaders, |info| { + info.source.as_bytes().to_owned() + }) + .unwrap(); + } + if cfg!(feature = "msl") { + write_shaders(&mut buf, "msl", &shaders, |info| { + compile::msl::translate(info).unwrap().as_bytes().to_owned() + }) + .unwrap(); + } + std::fs::write(&dest_path, &buf).unwrap(); + println!("cargo:rerun-if-changed=../shaders"); +} + +fn write_types( + buf: &mut String, + shaders: &HashMap, +) -> Result<(), std::fmt::Error> { + writeln!(buf, "pub struct Shaders<'a> {{")?; + for (name, _) in shaders { + writeln!(buf, " pub {name}: ComputeShader<'a>,")?; + } + writeln!(buf, "}}")?; + writeln!(buf, "pub struct Pipelines {{")?; + for (name, _) in shaders { + writeln!(buf, " pub {name}: T,")?; + } + writeln!(buf, "}}")?; + writeln!(buf, "impl Pipelines {{")?; + writeln!(buf, " pub fn from_shaders>(shaders: &Shaders, device: &H::Device, host: &mut H) -> Result {{")?; + writeln!(buf, " Ok(Self {{")?; + for (name, _) in shaders { + writeln!( + buf, + " {name}: host.new_compute_pipeline(device, &shaders.{name})?," + )?; + } + writeln!(buf, " }})")?; + writeln!(buf, " }}")?; + writeln!(buf, "}}")?; + Ok(()) +} + +fn write_shaders( + buf: &mut String, + mod_name: &str, + shaders: &HashMap, + translate: impl Fn(&ShaderInfo) -> Vec, +) -> Result<(), std::fmt::Error> { + writeln!(buf, "pub mod {mod_name} {{")?; + writeln!(buf, " use super::*;")?; + writeln!(buf, " use BindType::*;")?; + writeln!(buf, " pub const SHADERS: Shaders<'static> = Shaders {{")?; + for (name, info) in shaders { + let bind_tys = info + .bindings + .iter() + .map(|binding| binding.ty) + .collect::>(); + let source = translate(info); + writeln!(buf, " {name}: ComputeShader {{")?; + writeln!(buf, " name: Cow::Borrowed({:?}),", name)?; + writeln!( + buf, + " code: Cow::Borrowed(&{:?}),", + source.as_slice() + )?; + writeln!( + buf, + " workgroup_size: {:?},", + info.workgroup_size + )?; + writeln!(buf, " bindings: Cow::Borrowed(&{:?}),", bind_tys)?; + writeln!(buf, " }},")?; + } + writeln!(buf, " }};")?; + writeln!(buf, "}}")?; + Ok(()) +} diff --git a/vello_shaders/src/compile/mod.rs b/vello_shaders/src/compile/mod.rs new file mode 100644 index 0000000..c995b0d --- /dev/null +++ b/vello_shaders/src/compile/mod.rs @@ -0,0 +1,159 @@ +use naga::{ + front::wgsl, + valid::{Capabilities, ModuleInfo, ValidationError, ValidationFlags}, + AddressSpace, ImageClass, Module, StorageAccess, WithSpan, +}; + +use std::{ + collections::{HashMap, HashSet}, + path::Path, +}; + +pub mod permutations; +pub mod preprocess; + +pub mod msl; + +use crate::types::{BindType, BindingInfo}; + +#[derive(Debug)] +pub enum Error { + Parse(wgsl::ParseError), + Validate(WithSpan), + EntryPointNotFound, +} + +impl From for Error { + fn from(e: wgsl::ParseError) -> Self { + Self::Parse(e) + } +} + +impl From> for Error { + fn from(e: WithSpan) -> Self { + Self::Validate(e) + } +} + +#[derive(Debug)] +pub struct ShaderInfo { + pub source: String, + pub module: Module, + pub module_info: ModuleInfo, + pub workgroup_size: [u32; 3], + pub bindings: Vec, +} + +impl ShaderInfo { + pub fn new(source: String, entry_point: &str) -> Result { + let module = wgsl::parse_str(&source)?; + let module_info = naga::valid::Validator::new( + ValidationFlags::all() & !ValidationFlags::CONTROL_FLOW_UNIFORMITY, + Capabilities::all(), + ) + .validate(&module)?; + let (entry_index, entry) = module + .entry_points + .iter() + .enumerate() + .find(|(_, entry)| entry.name.as_str() == entry_point) + .ok_or(Error::EntryPointNotFound)?; + let mut bindings = vec![]; + let entry_info = module_info.get_entry_point(entry_index); + for (var_handle, var) in module.global_variables.iter() { + if entry_info[var_handle].is_empty() { + continue; + } + let Some(binding) = &var.binding else { + continue; + }; + let mut resource = BindingInfo { + name: var.name.clone(), + location: (binding.group, binding.binding), + ty: BindType::Buffer, + }; + let binding_ty = match module.types[var.ty].inner { + naga::TypeInner::BindingArray { base, .. } => &module.types[base].inner, + ref ty => ty, + }; + if let naga::TypeInner::Image { class, .. } = &binding_ty { + resource.ty = BindType::ImageRead; + if let ImageClass::Storage { access, .. } = class { + if access.contains(StorageAccess::STORE) { + resource.ty = BindType::Image; + } + } + } else { + resource.ty = BindType::BufReadOnly; + match var.space { + AddressSpace::Storage { access } => { + if access.contains(StorageAccess::STORE) { + resource.ty = BindType::Buffer; + } + } + AddressSpace::Uniform => { + resource.ty = BindType::Uniform; + } + _ => {} + } + } + bindings.push(resource); + } + bindings.sort_by_key(|res| res.location); + let workgroup_size = entry.workgroup_size; + Ok(ShaderInfo { + source, + module, + module_info, + workgroup_size, + bindings, + }) + } + + pub fn from_dir(shader_dir: impl AsRef) -> HashMap { + use std::fs; + let shader_dir = shader_dir.as_ref(); + let permutation_map = if let Ok(permutations_source) = + std::fs::read_to_string(shader_dir.join("permutations")) + { + permutations::parse(&permutations_source) + } else { + Default::default() + }; + println!("{:?}", permutation_map); + let imports = preprocess::get_imports(shader_dir); + let mut info = HashMap::default(); + let mut defines = HashSet::default(); + defines.insert("full".to_string()); + for entry in shader_dir + .read_dir() + .expect("Can read shader import directory") + { + let entry = entry.expect("Can continue reading shader import directory"); + if entry.file_type().unwrap().is_file() { + let file_name = entry.file_name(); + if let Some(name) = file_name.to_str() { + let suffix = ".wgsl"; + if let Some(shader_name) = name.strip_suffix(suffix) { + let contents = fs::read_to_string(shader_dir.join(&file_name)) + .expect("Could read shader {shader_name} contents"); + if let Some(permutations) = permutation_map.get(shader_name) { + for permutation in permutations { + let mut defines = defines.clone(); + defines.extend(permutation.defines.iter().cloned()); + let source = preprocess::preprocess(&contents, &defines, &imports); + let shader_info = Self::new(source.clone(), "main").unwrap(); + info.insert(permutation.name.clone(), shader_info); + } + } else { + let source = preprocess::preprocess(&contents, &defines, &imports); + let shader_info = Self::new(source.clone(), "main").unwrap(); + info.insert(shader_name.to_string(), shader_info); + } + } + } + } + } + info + } +} diff --git a/vello_shaders/src/compile/msl.rs b/vello_shaders/src/compile/msl.rs new file mode 100644 index 0000000..04d432d --- /dev/null +++ b/vello_shaders/src/compile/msl.rs @@ -0,0 +1,49 @@ +use naga::back::msl; + +use super::{BindType, ShaderInfo}; + +pub fn translate(shader: &ShaderInfo) -> Result { + let mut map = msl::PerStageMap::default(); + let mut buffer_index = 0u8; + let mut image_index = 0u8; + let mut binding_map = msl::BindingMap::default(); + for resource in &shader.bindings { + let binding = naga::ResourceBinding { + group: resource.location.0, + binding: resource.location.1, + }; + let mut target = msl::BindTarget::default(); + match resource.ty { + BindType::Buffer | BindType::BufReadOnly | BindType::Uniform => { + target.buffer = Some(buffer_index); + buffer_index += 1; + } + BindType::Image | BindType::ImageRead => { + target.texture = Some(image_index); + image_index += 1; + } + } + target.mutable = resource.ty.is_mutable(); + binding_map.insert(binding, target); + } + map.cs = msl::PerStageResources { + resources: binding_map, + push_constant_buffer: None, + sizes_buffer: Some(30), + }; + let options = msl::Options { + lang_version: (2, 0), + per_stage_map: map, + inline_samplers: vec![], + spirv_cross_compatibility: false, + fake_missing_bindings: false, + bounds_check_policies: naga::proc::BoundsCheckPolicies::default(), + }; + let (source, _) = msl::write_string( + &shader.module, + &shader.module_info, + &options, + &msl::PipelineOptions::default(), + )?; + Ok(source) +} diff --git a/vello_shaders/src/compile/permutations.rs b/vello_shaders/src/compile/permutations.rs new file mode 100644 index 0000000..5dbb138 --- /dev/null +++ b/vello_shaders/src/compile/permutations.rs @@ -0,0 +1,41 @@ +use std::collections::HashMap; + +#[derive(Debug)] +pub struct Permutation { + /// The new name for the permutation + pub name: String, + /// Set of defines to apply for the permutation + pub defines: Vec, +} + +pub fn parse(source: &str) -> HashMap> { + let mut map: HashMap> = Default::default(); + let mut current_source: Option = None; + for line in source.lines() { + let line = line.trim(); + if line.is_empty() || line.starts_with('#') { + continue; + } + if let Some(line) = line.strip_prefix('+') { + if let Some(source) = ¤t_source { + let mut parts = line.split(':').map(|s| s.trim()); + let Some(name) = parts.next() else { + continue; + }; + let mut defines = vec![]; + if let Some(define_list) = parts.next() { + defines.extend(define_list.split(' ').map(|s| s.trim().to_string())); + } + map.entry(source.to_string()) + .or_default() + .push(Permutation { + name: name.to_string(), + defines, + }); + } + } else { + current_source = Some(line.to_string()); + } + } + map +} diff --git a/vello_shaders/src/compile/preprocess.rs b/vello_shaders/src/compile/preprocess.rs new file mode 100644 index 0000000..95b8a4e --- /dev/null +++ b/vello_shaders/src/compile/preprocess.rs @@ -0,0 +1,159 @@ +use std::{ + collections::{HashMap, HashSet}, + fs, + path::Path, + vec, +}; + +pub fn get_imports(shader_dir: &Path) -> HashMap { + let mut imports = HashMap::new(); + let imports_dir = shader_dir.join("shared"); + for entry in imports_dir + .read_dir() + .expect("Can read shader import directory") + { + let entry = entry.expect("Can continue reading shader import directory"); + if entry.file_type().unwrap().is_file() { + let file_name = entry.file_name(); + if let Some(name) = file_name.to_str() { + let suffix = ".wgsl"; + if let Some(import_name) = name.strip_suffix(suffix) { + let contents = fs::read_to_string(imports_dir.join(&file_name)) + .expect("Could read shader {import_name} contents"); + imports.insert(import_name.to_owned(), contents); + } + } + } + } + imports +} + +pub struct StackItem { + active: bool, + else_passed: bool, +} + +pub fn preprocess( + input: &str, + defines: &HashSet, + imports: &HashMap, +) -> String { + let mut output = String::with_capacity(input.len()); + let mut stack = vec![]; + 'all_lines: for (line_number, mut line) in input.lines().enumerate() { + loop { + if line.is_empty() { + break; + } + let hash_index = line.find('#'); + let comment_index = line.find("//"); + let hash_index = match (hash_index, comment_index) { + (Some(hash_index), None) => hash_index, + (Some(hash_index), Some(comment_index)) if hash_index < comment_index => hash_index, + // Add this line to the output - all directives are commented out or there are no directives + _ => break, + }; + let directive_start = &line[hash_index + '#'.len_utf8()..]; + let directive_len = directive_start + // The first character which can't be part of the directive name marks the end of the directive + // In practise this should always be whitespace, but in theory a 'unit' directive + // could be added + .find(|c: char| !c.is_alphanumeric()) + .unwrap_or(directive_start.len()); + let directive = &directive_start[..directive_len]; + let directive_is_at_start = line.trim_start().starts_with('#'); + + match directive { + if_item @ ("ifdef" | "ifndef" | "else" | "endif") if !directive_is_at_start => { + eprintln!("#{if_item} directives must be the first non_whitespace items on their line, ignoring (line {line_number})"); + break; + } + def_test @ ("ifdef" | "ifndef") => { + let def = directive_start[directive_len..].trim(); + let exists = defines.contains(def); + let mode = def_test == "ifdef"; + stack.push(StackItem { + active: mode == exists, + else_passed: false, + }); + // Don't add this line to the output; instead process the next line + continue 'all_lines; + } + "else" => { + let item = stack.last_mut(); + if let Some(item) = item { + if item.else_passed { + eprintln!("Second else for same ifdef/ifndef (line {line_number}); ignoring second else") + } else { + item.else_passed = true; + item.active = !item.active; + } + } + let remainder = directive_start[directive_len..].trim(); + if !remainder.is_empty() { + eprintln!("#else directives don't take an argument. `{remainder}` will not be in output (line {line_number})"); + } + // Don't add this line to the output; it should be empty (see warning above) + continue 'all_lines; + } + "endif" => { + if stack.pop().is_none() { + eprintln!("Mismatched endif (line {line_number})"); + } + let remainder = directive_start[directive_len..].trim(); + if !remainder.is_empty() { + eprintln!("#endif directives don't take an argument. `{remainder}` will not be in output (line {line_number})"); + } + // Don't add this line to the output; it should be empty (see warning above) + continue 'all_lines; + } + "import" => { + output.push_str(&line[..hash_index]); + let directive_end = &directive_start[directive_len..]; + let import_name_start = if let Some(import_name_start) = + directive_end.find(|c: char| !c.is_whitespace()) + { + import_name_start + } else { + eprintln!("#import needs a non_whitespace argument (line {line_number})"); + continue 'all_lines; + }; + let import_name_start = &directive_end[import_name_start..]; + let import_name_end_index = import_name_start + // The first character which can't be part of the import name marks the end of the import + .find(|c: char| !(c == '_' || c.is_alphanumeric())) + .unwrap_or(import_name_start.len()); + let import_name = &import_name_start[..import_name_end_index]; + line = &import_name_start[import_name_end_index..]; + let import = imports.get(import_name); + if let Some(import) = import { + // In theory, we can cache this until the top item of the stack changes + // However, in practise there will only ever be at most 2 stack items, so it's reasonable to just recompute it every time + if stack.iter().all(|item| item.active) { + output.push_str(&preprocess(import, defines, imports)); + } + } else { + eprintln!("Unknown import `{import_name}` (line {line_number})"); + } + continue; + } + val => { + eprintln!("Unknown preprocessor directive `{val}` (line {line_number})"); + } + } + } + if stack.iter().all(|item| item.active) { + // Naga does not yet recognize `const` but web does not allow global `let`. We + // use `let` in our canonical sources to satisfy wgsl-analyzer but replace with + // `const` when targeting web. + if line.starts_with("let ") { + output.push_str("const"); + output.push_str(&line[3..]); + } else { + output.push_str(line); + } + output.push('\n'); + } + } + output +} diff --git a/vello_shaders/src/lib.rs b/vello_shaders/src/lib.rs new file mode 100644 index 0000000..ad71439 --- /dev/null +++ b/vello_shaders/src/lib.rs @@ -0,0 +1,30 @@ +mod types; + +#[cfg(feature = "compile")] +pub mod compile; + +pub use types::{BindType, BindingInfo}; + +use std::borrow::Cow; + +#[derive(Clone, Debug)] +pub struct ComputeShader<'a> { + pub name: Cow<'a, str>, + pub code: Cow<'a, [u8]>, + pub workgroup_size: [u32; 3], + pub bindings: Cow<'a, [BindType]>, +} + +pub trait PipelineHost { + type Device; + type ComputePipeline; + type Error; + + fn new_compute_pipeline( + &mut self, + device: &Self::Device, + shader: &ComputeShader, + ) -> Result; +} + +include!(concat!(env!("OUT_DIR"), "/shaders.rs")); diff --git a/vello_shaders/src/types.rs b/vello_shaders/src/types.rs new file mode 100644 index 0000000..376ec01 --- /dev/null +++ b/vello_shaders/src/types.rs @@ -0,0 +1,30 @@ +//! Types that are shared between the main crate and build. + +/// The type of resource that will be bound to a slot in a shader. +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +pub enum BindType { + /// A storage buffer with read/write access. + Buffer, + /// A storage buffer with read only access. + BufReadOnly, + /// A small storage buffer to be used as uniforms. + Uniform, + /// A storage image. + Image, + /// A storage image with read only access. + ImageRead, + // TODO: Sampler, maybe others +} + +impl BindType { + pub fn is_mutable(self) -> bool { + matches!(self, Self::Buffer | Self::Image) + } +} + +#[derive(Clone, Debug)] +pub struct BindingInfo { + pub name: Option, + pub location: (u32, u32), + pub ty: BindType, +} From a5434569b6eebcd75b6e0254b5f0c37931678360 Mon Sep 17 00:00:00 2001 From: Chad Brokaw Date: Wed, 25 Jan 2023 15:36:19 -0500 Subject: [PATCH 02/12] fix build rerun path --- vello_shaders/build.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vello_shaders/build.rs b/vello_shaders/build.rs index e447292..39efbfa 100644 --- a/vello_shaders/build.rs +++ b/vello_shaders/build.rs @@ -29,7 +29,7 @@ fn main() { .unwrap(); } std::fs::write(&dest_path, &buf).unwrap(); - println!("cargo:rerun-if-changed=../shaders"); + println!("cargo:rerun-if-changed=../shader"); } fn write_types( From a532eacf7b30676b156c99ba43f9dbb5c194cd5f Mon Sep 17 00:00:00 2001 From: Chad Brokaw Date: Wed, 25 Jan 2023 16:02:57 -0500 Subject: [PATCH 03/12] Replace HashMap with sorted Vec so builds are deterministic --- vello_shaders/build.rs | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/vello_shaders/build.rs b/vello_shaders/build.rs index 39efbfa..4c0c20e 100644 --- a/vello_shaders/build.rs +++ b/vello_shaders/build.rs @@ -13,7 +13,10 @@ use compile::ShaderInfo; fn main() { let out_dir = env::var_os("OUT_DIR").unwrap(); let dest_path = Path::new(&out_dir).join("shaders.rs"); - let shaders = compile::ShaderInfo::from_dir("../shader"); + let mut shaders = compile::ShaderInfo::from_dir("../shader"); + // Drop the HashMap and sort by name so that we get deterministic order. + let mut shaders = shaders.drain().collect::>(); + shaders.sort_by(|x, y| x.0.cmp(&y.0)); let mut buf = String::default(); write_types(&mut buf, &shaders).unwrap(); if cfg!(feature = "wgsl") { @@ -34,7 +37,7 @@ fn main() { fn write_types( buf: &mut String, - shaders: &HashMap, + shaders: &[(String, ShaderInfo)], ) -> Result<(), std::fmt::Error> { writeln!(buf, "pub struct Shaders<'a> {{")?; for (name, _) in shaders { @@ -64,7 +67,7 @@ fn write_types( fn write_shaders( buf: &mut String, mod_name: &str, - shaders: &HashMap, + shaders: &[(String, ShaderInfo)], translate: impl Fn(&ShaderInfo) -> Vec, ) -> Result<(), std::fmt::Error> { writeln!(buf, "pub mod {mod_name} {{")?; From 87803cc8a6ebec596d7c00697aa5e649fb341b14 Mon Sep 17 00:00:00 2001 From: Chad Brokaw Date: Wed, 25 Jan 2023 16:05:01 -0500 Subject: [PATCH 04/12] fmt --- vello_shaders/build.rs | 6 +----- vello_shaders/src/compile/mod.rs | 4 ++-- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/vello_shaders/build.rs b/vello_shaders/build.rs index 4c0c20e..aa13e97 100644 --- a/vello_shaders/build.rs +++ b/vello_shaders/build.rs @@ -3,7 +3,6 @@ mod compile; #[path = "src/types.rs"] mod types; -use std::collections::HashMap; use std::env; use std::fmt::Write; use std::path::Path; @@ -35,10 +34,7 @@ fn main() { println!("cargo:rerun-if-changed=../shader"); } -fn write_types( - buf: &mut String, - shaders: &[(String, ShaderInfo)], -) -> Result<(), std::fmt::Error> { +fn write_types(buf: &mut String, shaders: &[(String, ShaderInfo)]) -> Result<(), std::fmt::Error> { writeln!(buf, "pub struct Shaders<'a> {{")?; for (name, _) in shaders { writeln!(buf, " pub {name}: ComputeShader<'a>,")?; diff --git a/vello_shaders/src/compile/mod.rs b/vello_shaders/src/compile/mod.rs index c995b0d..c1815a4 100644 --- a/vello_shaders/src/compile/mod.rs +++ b/vello_shaders/src/compile/mod.rs @@ -65,8 +65,8 @@ impl ShaderInfo { continue; } let Some(binding) = &var.binding else { - continue; - }; + continue; + }; let mut resource = BindingInfo { name: var.name.clone(), location: (binding.group, binding.binding), From 14ab8d90ae8d5226c1b9627b7a1d7e3206e17479 Mon Sep 17 00:00:00 2001 From: Arman Uguray Date: Wed, 29 Mar 2023 10:09:27 -0700 Subject: [PATCH 05/12] [vello_shaders] Use the per_entry_point_map option for MSL --- vello_shaders/src/compile/msl.rs | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/vello_shaders/src/compile/msl.rs b/vello_shaders/src/compile/msl.rs index 04d432d..7fa8740 100644 --- a/vello_shaders/src/compile/msl.rs +++ b/vello_shaders/src/compile/msl.rs @@ -3,7 +3,7 @@ use naga::back::msl; use super::{BindType, ShaderInfo}; pub fn translate(shader: &ShaderInfo) -> Result { - let mut map = msl::PerStageMap::default(); + let mut map = msl::EntryPointResourceMap::default(); let mut buffer_index = 0u8; let mut image_index = 0u8; let mut binding_map = msl::BindingMap::default(); @@ -26,18 +26,22 @@ pub fn translate(shader: &ShaderInfo) -> Result { target.mutable = resource.ty.is_mutable(); binding_map.insert(binding, target); } - map.cs = msl::PerStageResources { - resources: binding_map, - push_constant_buffer: None, - sizes_buffer: Some(30), - }; + map.insert( + "main".to_string(), + msl::EntryPointResources { + resources: binding_map, + push_constant_buffer: None, + sizes_buffer: Some(30), + }, + ); let options = msl::Options { lang_version: (2, 0), - per_stage_map: map, + per_entry_point_map: map, inline_samplers: vec![], spirv_cross_compatibility: false, fake_missing_bindings: false, bounds_check_policies: naga::proc::BoundsCheckPolicies::default(), + zero_initialize_workgroup_memory: false, }; let (source, _) = msl::write_string( &shader.module, From b52ef32c9030fd82f0dce09985d4412361b4ecd8 Mon Sep 17 00:00:00 2001 From: Arman Uguray Date: Wed, 29 Mar 2023 10:12:23 -0700 Subject: [PATCH 06/12] [vello_shaders] Use thiserror::Error for the library Error type --- vello_shaders/Cargo.toml | 4 ++- vello_shaders/src/compile/mod.rs | 43 ++++++++++++++------------------ 2 files changed, 22 insertions(+), 25 deletions(-) diff --git a/vello_shaders/Cargo.toml b/vello_shaders/Cargo.toml index abd7dc0..72217a8 100644 --- a/vello_shaders/Cargo.toml +++ b/vello_shaders/Cargo.toml @@ -5,13 +5,15 @@ edition = "2021" [features] default = ["compile", "wgsl", "msl"] -compile = ["naga"] +compile = ["naga", "thiserror"] wgsl = [] msl = [] [dependencies] naga = { git = "https://github.com/gfx-rs/naga", features = ["wgsl-in", "msl-out", "validate"], optional = true } +thiserror = { version = "1.0.40", optional = true } [build-dependencies] naga = { git = "https://github.com/gfx-rs/naga", features = ["wgsl-in", "msl-out", "validate"] } +thiserror = "1.0.40" diff --git a/vello_shaders/src/compile/mod.rs b/vello_shaders/src/compile/mod.rs index c1815a4..6d88667 100644 --- a/vello_shaders/src/compile/mod.rs +++ b/vello_shaders/src/compile/mod.rs @@ -1,12 +1,14 @@ -use naga::{ - front::wgsl, - valid::{Capabilities, ModuleInfo, ValidationError, ValidationFlags}, - AddressSpace, ImageClass, Module, StorageAccess, WithSpan, -}; - -use std::{ - collections::{HashMap, HashSet}, - path::Path, +use { + naga::{ + front::wgsl, + valid::{Capabilities, ModuleInfo, ValidationError, ValidationFlags}, + AddressSpace, ImageClass, Module, StorageAccess, WithSpan, + }, + std::{ + collections::{HashMap, HashSet}, + path::Path, + }, + thiserror::Error, }; pub mod permutations; @@ -16,25 +18,18 @@ pub mod msl; use crate::types::{BindType, BindingInfo}; -#[derive(Debug)] +#[derive(Error, Debug)] pub enum Error { - Parse(wgsl::ParseError), - Validate(WithSpan), + #[error("failed to parse shader: {0}")] + Parse(#[from] wgsl::ParseError), + + #[error("failed to validate shader: {0}")] + Validate(#[from] WithSpan), + + #[error("missing entry point function")] EntryPointNotFound, } -impl From for Error { - fn from(e: wgsl::ParseError) -> Self { - Self::Parse(e) - } -} - -impl From> for Error { - fn from(e: WithSpan) -> Self { - Self::Validate(e) - } -} - #[derive(Debug)] pub struct ShaderInfo { pub source: String, From 4f445c2e0aaeac571fcdc56f997b29231f3f0b4a Mon Sep 17 00:00:00 2001 From: Arman Uguray Date: Wed, 29 Mar 2023 10:27:16 -0700 Subject: [PATCH 07/12] [vello_shaders] Provide workgroup shared memory allocation sizes Naga traslates workgroup variable declarations to threadgroup address-space entry-point parameters when generating MSL. Metal API validation requires that the memory sizes for these parameters be set explicitly by calling setThreadgroupMemoryLength:index on the MTLComputeCommandEncoder. The crate now calculates the required memory size for global workgroup variables that are accessed by the entry point and provides them alongside the binding list. This is abstracted separately from the binding list. While the current usage that we're aware of is limited to Metal, this information is being provided as part of the generic ComputeShader type instead of a MSL-specific type, as the information itself is computed from the parsed WGSL IR and not specific to Metal. --- vello_shaders/build.rs | 6 ++++ vello_shaders/src/compile/mod.rs | 52 ++++++++++++++++++++++++++++---- vello_shaders/src/lib.rs | 3 +- vello_shaders/src/types.rs | 7 +++++ 4 files changed, 61 insertions(+), 7 deletions(-) diff --git a/vello_shaders/build.rs b/vello_shaders/build.rs index aa13e97..65ea159 100644 --- a/vello_shaders/build.rs +++ b/vello_shaders/build.rs @@ -76,6 +76,7 @@ fn write_shaders( .iter() .map(|binding| binding.ty) .collect::>(); + let wg_bufs = &info.workgroup_buffers; let source = translate(info); writeln!(buf, " {name}: ComputeShader {{")?; writeln!(buf, " name: Cow::Borrowed({:?}),", name)?; @@ -90,6 +91,11 @@ fn write_shaders( info.workgroup_size )?; writeln!(buf, " bindings: Cow::Borrowed(&{:?}),", bind_tys)?; + writeln!( + buf, + " workgroup_buffers: Cow::Borrowed(&{:?}),", + wg_bufs + )?; writeln!(buf, " }},")?; } writeln!(buf, " }};")?; diff --git a/vello_shaders/src/compile/mod.rs b/vello_shaders/src/compile/mod.rs index 6d88667..c218be6 100644 --- a/vello_shaders/src/compile/mod.rs +++ b/vello_shaders/src/compile/mod.rs @@ -2,7 +2,8 @@ use { naga::{ front::wgsl, valid::{Capabilities, ModuleInfo, ValidationError, ValidationFlags}, - AddressSpace, ImageClass, Module, StorageAccess, WithSpan, + AddressSpace, ArraySize, ConstantInner, ImageClass, Module, ScalarValue, StorageAccess, + WithSpan, }, std::{ collections::{HashMap, HashSet}, @@ -16,7 +17,7 @@ pub mod preprocess; pub mod msl; -use crate::types::{BindType, BindingInfo}; +use crate::types::{BindType, BindingInfo, WorkgroupBufferInfo}; #[derive(Error, Debug)] pub enum Error { @@ -37,6 +38,7 @@ pub struct ShaderInfo { pub module_info: ModuleInfo, pub workgroup_size: [u32; 3], pub bindings: Vec, + pub workgroup_buffers: Vec, } impl ShaderInfo { @@ -54,12 +56,53 @@ impl ShaderInfo { .find(|(_, entry)| entry.name.as_str() == entry_point) .ok_or(Error::EntryPointNotFound)?; let mut bindings = vec![]; + let mut workgroup_buffers = vec![]; + let mut wg_buffer_idx = 0; let entry_info = module_info.get_entry_point(entry_index); for (var_handle, var) in module.global_variables.iter() { if entry_info[var_handle].is_empty() { continue; } + let binding_ty = match module.types[var.ty].inner { + naga::TypeInner::BindingArray { base, .. } => &module.types[base].inner, + ref ty => ty, + }; let Some(binding) = &var.binding else { + if var.space == AddressSpace::WorkGroup { + let index = wg_buffer_idx; + wg_buffer_idx += 1; + let size_in_bytes = match binding_ty { + naga::TypeInner::Array { + size: ArraySize::Constant(const_handle), + stride, + .. + } => { + let size: u32 = match module.constants[*const_handle].inner { + ConstantInner::Scalar { value, width: _ } => match value { + ScalarValue::Uint(value) => value.try_into().unwrap(), + ScalarValue::Sint(value) => value.try_into().unwrap(), + _ => continue, + }, + ConstantInner::Composite { .. } => continue, + }; + size * stride + }, + naga::TypeInner::Struct { span, .. } => *span, + naga::TypeInner::Scalar { width, ..} => *width as u32, + naga::TypeInner::Vector { width, ..} => *width as u32, + naga::TypeInner::Matrix { width, ..} => *width as u32, + naga::TypeInner::Atomic { width, ..} => *width as u32, + _ => { + // Not a valid workgroup variable type. At least not one that is used + // in our shaders. + continue; + } + }; + workgroup_buffers.push(WorkgroupBufferInfo { + size_in_bytes, + index, + }); + } continue; }; let mut resource = BindingInfo { @@ -67,10 +110,6 @@ impl ShaderInfo { location: (binding.group, binding.binding), ty: BindType::Buffer, }; - let binding_ty = match module.types[var.ty].inner { - naga::TypeInner::BindingArray { base, .. } => &module.types[base].inner, - ref ty => ty, - }; if let naga::TypeInner::Image { class, .. } = &binding_ty { resource.ty = BindType::ImageRead; if let ImageClass::Storage { access, .. } = class { @@ -102,6 +141,7 @@ impl ShaderInfo { module_info, workgroup_size, bindings, + workgroup_buffers, }) } diff --git a/vello_shaders/src/lib.rs b/vello_shaders/src/lib.rs index ad71439..ded1984 100644 --- a/vello_shaders/src/lib.rs +++ b/vello_shaders/src/lib.rs @@ -3,7 +3,7 @@ mod types; #[cfg(feature = "compile")] pub mod compile; -pub use types::{BindType, BindingInfo}; +pub use types::{BindType, BindingInfo, WorkgroupBufferInfo}; use std::borrow::Cow; @@ -13,6 +13,7 @@ pub struct ComputeShader<'a> { pub code: Cow<'a, [u8]>, pub workgroup_size: [u32; 3], pub bindings: Cow<'a, [BindType]>, + pub workgroup_buffers: Cow<'a, [WorkgroupBufferInfo]>, } pub trait PipelineHost { diff --git a/vello_shaders/src/types.rs b/vello_shaders/src/types.rs index 376ec01..a9db3fe 100644 --- a/vello_shaders/src/types.rs +++ b/vello_shaders/src/types.rs @@ -28,3 +28,10 @@ pub struct BindingInfo { pub location: (u32, u32), pub ty: BindType, } + +#[derive(Clone, Debug)] +pub struct WorkgroupBufferInfo { + pub size_in_bytes: u32, + /// The order in which the workgroup variable is declared in the shader module. + pub index: u32, +} From eb1d4be36ad46c0a544f077175b95a391c48dbf8 Mon Sep 17 00:00:00 2001 From: Arman Uguray Date: Wed, 29 Mar 2023 11:20:15 -0700 Subject: [PATCH 08/12] [vello_shaders] Pin naga dependency --- vello_shaders/Cargo.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vello_shaders/Cargo.toml b/vello_shaders/Cargo.toml index 72217a8..d3969ef 100644 --- a/vello_shaders/Cargo.toml +++ b/vello_shaders/Cargo.toml @@ -10,10 +10,10 @@ wgsl = [] msl = [] [dependencies] -naga = { git = "https://github.com/gfx-rs/naga", features = ["wgsl-in", "msl-out", "validate"], optional = true } +naga = { git = "https://github.com/gfx-rs/naga", rev = "53d62b9", features = ["wgsl-in", "msl-out", "validate"], optional = true } thiserror = { version = "1.0.40", optional = true } [build-dependencies] -naga = { git = "https://github.com/gfx-rs/naga", features = ["wgsl-in", "msl-out", "validate"] } +naga = { git = "https://github.com/gfx-rs/naga", rev = "53d62b9", features = ["wgsl-in", "msl-out", "validate"] } thiserror = "1.0.40" From 7a99ae5a9eef2c3481580b3b3e0b900ed6b58a8c Mon Sep 17 00:00:00 2001 From: Arman Uguray Date: Wed, 29 Mar 2023 11:20:38 -0700 Subject: [PATCH 09/12] [vello_shaders] Add copyright headers --- vello_shaders/build.rs | 3 +++ vello_shaders/src/compile/mod.rs | 3 +++ vello_shaders/src/compile/msl.rs | 3 +++ vello_shaders/src/compile/permutations.rs | 3 +++ vello_shaders/src/compile/preprocess.rs | 3 +++ vello_shaders/src/lib.rs | 3 +++ vello_shaders/src/types.rs | 3 +++ 7 files changed, 21 insertions(+) diff --git a/vello_shaders/build.rs b/vello_shaders/build.rs index 65ea159..0856974 100644 --- a/vello_shaders/build.rs +++ b/vello_shaders/build.rs @@ -1,3 +1,6 @@ +// Copyright 2023 The Vello authors +// SPDX-License-Identifier: Apache-2.0 OR MIT + #[path = "src/compile/mod.rs"] mod compile; #[path = "src/types.rs"] diff --git a/vello_shaders/src/compile/mod.rs b/vello_shaders/src/compile/mod.rs index c218be6..74cafda 100644 --- a/vello_shaders/src/compile/mod.rs +++ b/vello_shaders/src/compile/mod.rs @@ -1,3 +1,6 @@ +// Copyright 2023 The Vello authors +// SPDX-License-Identifier: Apache-2.0 OR MIT + use { naga::{ front::wgsl, diff --git a/vello_shaders/src/compile/msl.rs b/vello_shaders/src/compile/msl.rs index 7fa8740..5f7b831 100644 --- a/vello_shaders/src/compile/msl.rs +++ b/vello_shaders/src/compile/msl.rs @@ -1,3 +1,6 @@ +// Copyright 2023 The Vello authors +// SPDX-License-Identifier: Apache-2.0 OR MIT + use naga::back::msl; use super::{BindType, ShaderInfo}; diff --git a/vello_shaders/src/compile/permutations.rs b/vello_shaders/src/compile/permutations.rs index 5dbb138..f85f667 100644 --- a/vello_shaders/src/compile/permutations.rs +++ b/vello_shaders/src/compile/permutations.rs @@ -1,3 +1,6 @@ +// Copyright 2023 The Vello authors +// SPDX-License-Identifier: Apache-2.0 OR MIT + use std::collections::HashMap; #[derive(Debug)] diff --git a/vello_shaders/src/compile/preprocess.rs b/vello_shaders/src/compile/preprocess.rs index 95b8a4e..917f83f 100644 --- a/vello_shaders/src/compile/preprocess.rs +++ b/vello_shaders/src/compile/preprocess.rs @@ -1,3 +1,6 @@ +// Copyright 2023 The Vello authors +// SPDX-License-Identifier: Apache-2.0 OR MIT + use std::{ collections::{HashMap, HashSet}, fs, diff --git a/vello_shaders/src/lib.rs b/vello_shaders/src/lib.rs index ded1984..66ac937 100644 --- a/vello_shaders/src/lib.rs +++ b/vello_shaders/src/lib.rs @@ -1,3 +1,6 @@ +// Copyright 2023 The Vello authors +// SPDX-License-Identifier: Apache-2.0 OR MIT + mod types; #[cfg(feature = "compile")] diff --git a/vello_shaders/src/types.rs b/vello_shaders/src/types.rs index a9db3fe..548f7e9 100644 --- a/vello_shaders/src/types.rs +++ b/vello_shaders/src/types.rs @@ -1,3 +1,6 @@ +// Copyright 2023 The Vello authors +// SPDX-License-Identifier: Apache-2.0 OR MIT + //! Types that are shared between the main crate and build. /// The type of resource that will be bound to a slot in a shader. From d2b41d628b9279018f0f844853631aeb4ce81b35 Mon Sep 17 00:00:00 2001 From: Arman Uguray Date: Wed, 29 Mar 2023 11:20:56 -0700 Subject: [PATCH 10/12] [vello_shaders] Add a README --- vello_shaders/README.md | 7 +++++++ 1 file changed, 7 insertions(+) create mode 100644 vello_shaders/README.md diff --git a/vello_shaders/README.md b/vello_shaders/README.md new file mode 100644 index 0000000..c0cf1af --- /dev/null +++ b/vello_shaders/README.md @@ -0,0 +1,7 @@ +The `vello_shaders` crate provides a utility library to integrate the Vello shader modules into any +renderer project. The create provides the necessary metadata to construct the individual compute +pipelines on any GPU API while leaving the responsibility of all API interactions (such as +resource management and command encoding) up to the client. + +The shaders can be pre-compiled to any target shading language at build time based on feature flags. +Currently only WGSL and Metal Shading Language are supported. From 64020a3f77a188f43fd98bbfcd20c036a84d39b0 Mon Sep 17 00:00:00 2001 From: Arman Uguray Date: Wed, 29 Mar 2023 11:23:06 -0700 Subject: [PATCH 11/12] [vello_shaders] Fix typo in README --- vello_shaders/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vello_shaders/README.md b/vello_shaders/README.md index c0cf1af..9bec768 100644 --- a/vello_shaders/README.md +++ b/vello_shaders/README.md @@ -1,5 +1,5 @@ The `vello_shaders` crate provides a utility library to integrate the Vello shader modules into any -renderer project. The create provides the necessary metadata to construct the individual compute +renderer project. The crate provides the necessary metadata to construct the individual compute pipelines on any GPU API while leaving the responsibility of all API interactions (such as resource management and command encoding) up to the client. From 0a6a6e2c21ce3b0cb7b133d26dffe843b8222d49 Mon Sep 17 00:00:00 2001 From: Arman Uguray Date: Wed, 29 Mar 2023 12:24:28 -0700 Subject: [PATCH 12/12] [vello_shaders] Move vello_shaders to crates/shaders --- Cargo.toml | 3 ++- {vello_shaders => crates/shaders}/Cargo.toml | 0 {vello_shaders => crates/shaders}/README.md | 0 {vello_shaders => crates/shaders}/build.rs | 2 +- {vello_shaders => crates/shaders}/src/compile/mod.rs | 0 {vello_shaders => crates/shaders}/src/compile/msl.rs | 0 {vello_shaders => crates/shaders}/src/compile/permutations.rs | 0 {vello_shaders => crates/shaders}/src/compile/preprocess.rs | 0 {vello_shaders => crates/shaders}/src/lib.rs | 0 {vello_shaders => crates/shaders}/src/types.rs | 0 10 files changed, 3 insertions(+), 2 deletions(-) rename {vello_shaders => crates/shaders}/Cargo.toml (100%) rename {vello_shaders => crates/shaders}/README.md (100%) rename {vello_shaders => crates/shaders}/build.rs (98%) rename {vello_shaders => crates/shaders}/src/compile/mod.rs (100%) rename {vello_shaders => crates/shaders}/src/compile/msl.rs (100%) rename {vello_shaders => crates/shaders}/src/compile/permutations.rs (100%) rename {vello_shaders => crates/shaders}/src/compile/preprocess.rs (100%) rename {vello_shaders => crates/shaders}/src/lib.rs (100%) rename {vello_shaders => crates/shaders}/src/types.rs (100%) diff --git a/Cargo.toml b/Cargo.toml index e682fe3..5873f00 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -2,6 +2,8 @@ resolver = "2" members = [ + "crates/shaders", + "integrations/vello_svg", "examples/headless", @@ -9,7 +11,6 @@ members = [ "examples/with_bevy", "examples/run_wasm", "examples/scenes", - "vello_shaders", ] [workspace.package] diff --git a/vello_shaders/Cargo.toml b/crates/shaders/Cargo.toml similarity index 100% rename from vello_shaders/Cargo.toml rename to crates/shaders/Cargo.toml diff --git a/vello_shaders/README.md b/crates/shaders/README.md similarity index 100% rename from vello_shaders/README.md rename to crates/shaders/README.md diff --git a/vello_shaders/build.rs b/crates/shaders/build.rs similarity index 98% rename from vello_shaders/build.rs rename to crates/shaders/build.rs index 0856974..7e5a952 100644 --- a/vello_shaders/build.rs +++ b/crates/shaders/build.rs @@ -15,7 +15,7 @@ use compile::ShaderInfo; fn main() { let out_dir = env::var_os("OUT_DIR").unwrap(); let dest_path = Path::new(&out_dir).join("shaders.rs"); - let mut shaders = compile::ShaderInfo::from_dir("../shader"); + let mut shaders = compile::ShaderInfo::from_dir("../../shader"); // Drop the HashMap and sort by name so that we get deterministic order. let mut shaders = shaders.drain().collect::>(); shaders.sort_by(|x, y| x.0.cmp(&y.0)); diff --git a/vello_shaders/src/compile/mod.rs b/crates/shaders/src/compile/mod.rs similarity index 100% rename from vello_shaders/src/compile/mod.rs rename to crates/shaders/src/compile/mod.rs diff --git a/vello_shaders/src/compile/msl.rs b/crates/shaders/src/compile/msl.rs similarity index 100% rename from vello_shaders/src/compile/msl.rs rename to crates/shaders/src/compile/msl.rs diff --git a/vello_shaders/src/compile/permutations.rs b/crates/shaders/src/compile/permutations.rs similarity index 100% rename from vello_shaders/src/compile/permutations.rs rename to crates/shaders/src/compile/permutations.rs diff --git a/vello_shaders/src/compile/preprocess.rs b/crates/shaders/src/compile/preprocess.rs similarity index 100% rename from vello_shaders/src/compile/preprocess.rs rename to crates/shaders/src/compile/preprocess.rs diff --git a/vello_shaders/src/lib.rs b/crates/shaders/src/lib.rs similarity index 100% rename from vello_shaders/src/lib.rs rename to crates/shaders/src/lib.rs diff --git a/vello_shaders/src/types.rs b/crates/shaders/src/types.rs similarity index 100% rename from vello_shaders/src/types.rs rename to crates/shaders/src/types.rs