diff --git a/piet-gpu-derive/src/glsl.rs b/piet-gpu-derive/src/glsl.rs index 0fc3a19..5164179 100644 --- a/piet-gpu-derive/src/glsl.rs +++ b/piet-gpu-derive/src/glsl.rs @@ -15,18 +15,32 @@ pub fn gen_glsl(module: &LayoutModule) -> String { gen_refdef(&mut r, &name); } for name in &module.def_names { - let def = module.defs.get(name).unwrap(); - if let (size, LayoutTypeDef::Struct(fields)) = def { - gen_struct_def(&mut r, &name, size.size, fields); + match module.defs.get(name).unwrap() { + (size, LayoutTypeDef::Struct(fields)) => { + gen_struct_def(&mut r, name, fields); + gen_item_def(&mut r, name, size.size); + } + (size, LayoutTypeDef::Enum(en)) => { + gen_enum_def(&mut r, name, en); + gen_item_def(&mut r, name, size.size); + } } } for name in &module.def_names { let def = module.defs.get(name).unwrap(); match def { - (size, LayoutTypeDef::Struct(fields)) => { - gen_struct_read(&mut r, &module.name, &name, size.size, fields) + (_size, LayoutTypeDef::Struct(fields)) => { + gen_struct_read(&mut r, &module.name, &name, fields); + if module.gpu_write { + gen_struct_write(&mut r, &module.name, &name, fields); + } + } + (_size, LayoutTypeDef::Enum(en)) => { + gen_enum_read(&mut r, &module.name, &name, en); + if module.gpu_write { + gen_enum_write(&mut r, &module.name, &name, en); + } } - (_size, LayoutTypeDef::Enum(_)) => gen_enum_read(&mut r, &module.name, &name), } } r @@ -38,28 +52,50 @@ fn gen_refdef(r: &mut String, name: &str) { writeln!(r, "}};\n").unwrap(); } -fn gen_struct_def(r: &mut String, name: &str, size: usize, fields: &[(String, usize, LayoutType)]) { +fn gen_struct_def(r: &mut String, name: &str, fields: &[(String, usize, LayoutType)]) { writeln!(r, "struct {} {{", name).unwrap(); for (name, _offset, ty) in fields { writeln!(r, " {} {};", glsl_type(&ty.ty), name).unwrap(); } writeln!(r, "}};\n").unwrap(); +} +fn gen_enum_def(r: &mut String, name: &str, variants: &[(String, Vec<(usize, LayoutType)>)]) { + for (i, (var_name, _payload)) in variants.iter().enumerate() { + writeln!(r, "#define {}_{} {}", name, var_name, i).unwrap(); + } +} + +fn gen_item_def(r: &mut String, name: &str, size: usize) { writeln!(r, "#define {}_size {}\n", name, size).unwrap(); + writeln!( + r, + "{}Ref {}_index({}Ref ref, uint index) {{", + name, name, name + ) + .unwrap(); + writeln!( + r, + " return {}Ref(ref.offset + index * {}_size);", + name, name + ) + .unwrap(); + writeln!(r, "}}\n").unwrap(); } fn gen_struct_read( r: &mut String, bufname: &str, name: &str, - size: usize, fields: &[(String, usize, LayoutType)], ) { writeln!(r, "{} {}_read({}Ref ref) {{", name, name, name).unwrap(); writeln!(r, " uint ix = ref.offset >> 2;").unwrap(); - for i in 0..(size / 4) { - // TODO: don't generate raw reads for inline structs - writeln!(r, " uint raw{} = {}[ix + {}];", i, bufname, i).unwrap(); + let coverage = crate::layout::struct_coverage(fields, false); + for (i, fields) in coverage.iter().enumerate() { + if !fields.is_empty() { + writeln!(r, " uint raw{} = {}[ix + {}];", i, bufname, i).unwrap(); + } } writeln!(r, " {} s;", name).unwrap(); for (name, offset, ty) in fields { @@ -69,10 +105,35 @@ fn gen_struct_read( writeln!(r, "}}\n").unwrap(); } -fn gen_enum_read(r: &mut String, bufname: &str, name: &str) { +fn gen_enum_read( + r: &mut String, + bufname: &str, + name: &str, + variants: &[(String, Vec<(usize, LayoutType)>)], +) { writeln!(r, "uint {}_tag({}Ref ref) {{", name, name).unwrap(); writeln!(r, " return {}[ref.offset >> 2];", bufname).unwrap(); writeln!(r, "}}\n").unwrap(); + for (var_name, payload) in variants { + if payload.len() == 1 { + if let GpuType::InlineStruct(structname) = &payload[0].1.ty { + writeln!( + r, + "{} {}_{}_read({}Ref ref) {{", + structname, name, var_name, name + ) + .unwrap(); + writeln!( + r, + " return {}_read({}Ref(ref.offset + {}));", + structname, structname, payload[0].0 + ) + .unwrap(); + writeln!(r, "}}\n").unwrap(); + } + } + // TODO: support for variants that aren't one struct. + } } fn gen_extract(offset: usize, ty: &GpuType) -> String { @@ -141,14 +202,146 @@ fn extract_ibits(offset: usize, nbytes: usize) -> String { format!("int(raw{}) >> {}", offset / 4, (offset % 4) * 8) } else { format!( - "(int(raw{}) << {}) >> {}", + "int(raw{} << {}) >> {}", offset / 4, - (3 - offset % 4) * 8, + ((4 - nbytes) - offset % 4) * 8, (4 - nbytes) * 8 ) } } +// Writing + +fn gen_struct_write( + r: &mut String, + bufname: &str, + name: &str, + fields: &[(String, usize, LayoutType)], +) { + writeln!(r, "void {}_write({}Ref ref, {} s) {{", name, name, name).unwrap(); + let coverage = crate::layout::struct_coverage(fields, true); + for (i, field_ixs) in coverage.iter().enumerate() { + let mut pieces = Vec::new(); + for field_ix in field_ixs { + let (name, offset, ty) = &fields[*field_ix]; + match &ty.ty { + GpuType::Scalar(scalar) => { + let inner = format!("s.{}", name); + pieces.push(gen_pack_bits_scalar(scalar, *offset, &inner)); + } + GpuType::Vector(scalar, len) => { + let size = scalar.size(); + let ix_lo = (i * 4 - offset) / size; + let ix_hi = ((4 + i * 4 - offset) / size).min(*len); + for ix in ix_lo..ix_hi { + let scalar_offset = offset + ix * size; + let inner = format!("s.{}.{}", name, &"xyzw"[ix..ix + 1]); + pieces.push(gen_pack_bits_scalar(scalar, scalar_offset, &inner)); + } + } + GpuType::InlineStruct(structname) => { + writeln!( + r, + " {}_write({}Ref({}), s.{});", + structname, + structname, + simplified_add("ref.offset", *offset), + name + ) + .unwrap(); + } + GpuType::Ref(_) => pieces.push(format!("s.{}.offset", name)), + } + } + if !pieces.is_empty() { + write!(r, " {}[{}] = ", bufname, i).unwrap(); + for (j, piece) in pieces.iter().enumerate() { + if j != 0 { + write!(r, " | ").unwrap(); + } + write!(r, "{}", piece).unwrap(); + } + writeln!(r, ";").unwrap(); + } + } + writeln!(r, "}}\n").unwrap(); +} + +fn gen_pack_bits_scalar(ty: &GpuScalar, offset: usize, inner: &str) -> String { + let shift = (offset % 4) * 8; + let bits = match ty { + GpuScalar::F32 => format!("floatBitsToUint({})", inner), + // Note: this doesn't mask small unsigned int types; the caller is + // responsible for making sure they don't overflow. + GpuScalar::U8 | GpuScalar::U16 | GpuScalar::U32 => inner.into(), + GpuScalar::I8 => { + if shift == 24 { + format!("uint({})", inner) + } else { + format!("(uint({}) & 0xff)", inner) + } + } + GpuScalar::I16 => { + if shift == 16 { + format!("uint({})", inner) + } else { + format!("(uint({}) & 0xffff)", inner) + } + } + GpuScalar::I32 => format!("uint({})", inner), + }; + if shift == 0 { + bits + } else { + format!("({} << {})", bits, shift) + } +} + +fn gen_enum_write( + r: &mut String, + bufname: &str, + name: &str, + variants: &[(String, Vec<(usize, LayoutType)>)], +) { + for (var_name, payload) in variants { + if payload.is_empty() { + writeln!(r, "void {}_{}_write({}Ref ref) {{", name, var_name, name).unwrap(); + writeln!( + r, + " {}[ref.offset >> 2] = {}_{};", + bufname, name, var_name + ) + .unwrap(); + writeln!(r, "}}\n").unwrap(); + } else if payload.len() == 1 { + if let GpuType::InlineStruct(structname) = &payload[0].1.ty { + writeln!( + r, + "void {}_{}_write({}Ref ref, {} s) {{", + name, var_name, name, structname + ) + .unwrap(); + writeln!( + r, + " {}[ref.offset >> 2] = {}_{};", + bufname, name, var_name + ) + .unwrap(); + writeln!( + r, + " {}_write({}Ref(ref.offset + {}), s);", + structname, structname, payload[0].0 + ) + .unwrap(); + writeln!(r, "}}\n").unwrap(); + } + } + // TODO: support for variants that aren't one struct. + } +} + +// Utility functions + fn glsl_type(ty: &GpuType) -> String { match ty { GpuType::Scalar(scalar) => glsl_scalar(scalar).into(), @@ -186,6 +379,7 @@ fn glsl_vecname(s: &GpuScalar) -> &'static str { GpuScalar::U8 | GpuScalar::U16 | GpuScalar::U32 => "uvec", } } + /// If `c = 0`, return `"var_name"`, else `"var_name + c"` fn simplified_add(var_name: &str, c: usize) -> String { if c == 0 { diff --git a/piet-gpu-derive/src/layout.rs b/piet-gpu-derive/src/layout.rs index 25cf38c..aef0048 100644 --- a/piet-gpu-derive/src/layout.rs +++ b/piet-gpu-derive/src/layout.rs @@ -1,5 +1,11 @@ //! Logic for layout of structures in memory. +// This is fairly simple now, but there are some extensions that are likely: +// * Addition of f16 types +// + These will probably have 2-byte alignments to support `packHalf2x16` +// * 1 byte tag values (so small struct fields can be packed along with tag) +// * (Possibly) reordering for better packing + use std::collections::{HashMap, HashSet}; use crate::parse::{GpuModule, GpuType, GpuTypeDef}; @@ -21,6 +27,16 @@ pub struct LayoutModule { pub name: String, pub def_names: Vec, pub defs: HashMap, + enum_variants: HashSet, + + /// Generate shader code to write the module. + /// + /// This is derived from the presence of the `gpu_write` attribute in the source module. + pub gpu_write: bool, + /// Generate Rust code to encode the module. + /// + /// This is derived from the presence of the `rust_encode` attribute in the source module. + pub rust_encode: bool, } struct LayoutSession<'a> { @@ -50,6 +66,13 @@ impl LayoutTypeDef { fn from_gpu(def: &GpuTypeDef, session: &mut LayoutSession) -> (Size, LayoutTypeDef) { match def { GpuTypeDef::Struct(_name, fields) => { + // TODO: We want to be able to pack enums more tightly, in particular + // other struct fields along with the enum tag. Structs in that category + // (first field has an alignment < 4, serve as enum variant) will have a + // different layout. This is why we're tracking `is_enum_variant`. + // + // But it's a bit of YAGNI for now; we're currently reserving 4 bytes for + // the tag, so structure layout doesn't care. let mut offset = 0; let mut result = Vec::new(); for field in fields { @@ -68,11 +91,7 @@ impl LayoutTypeDef { let mut max_offset = 0; for variant in &en.variants { let mut r2 = Vec::new(); - let mut offset = if session.is_enum_variant(&en.name) { - 4 - } else { - 0 - }; + let mut offset = 4; for field in &variant.1 { let layout_ty = LayoutType::from_gpu(field, session); offset += align_padding(offset, layout_ty.size.alignment); @@ -102,12 +121,22 @@ impl LayoutModule { for def in &module.defs { let _ = session.layout_def(def.name()); } + let gpu_write = module.attrs.contains("gpu_write"); + let rust_encode = module.attrs.contains("rust_encode"); LayoutModule { name: module.name.clone(), + gpu_write, + rust_encode, def_names, + enum_variants: session.enum_variants, defs: session.defs, } } + + #[allow(unused)] + pub fn is_enum_variant(&self, name: &str) -> bool { + self.enum_variants.contains(name) + } } impl<'a> LayoutSession<'a> { @@ -155,14 +184,52 @@ impl<'a> LayoutSession<'a> { } } + #[allow(unused)] fn is_enum_variant(&self, name: &str) -> bool { self.enum_variants.contains(name) } } +/// Compute coverage of fields. +/// +/// Each element of the result represents a list of fields for one 4-byte chunk of +/// the struct layout. Inline structs are only included if requested. +pub fn struct_coverage( + fields: &[(String, usize, LayoutType)], + include_inline: bool, +) -> Vec> { + let mut result: Vec> = Vec::new(); + for (i, (_name, offset, ty)) in fields.iter().enumerate() { + let size = match ty.ty { + GpuType::Scalar(scalar) => scalar.size(), + GpuType::Vector(scalar, len) => scalar.size() * len, + GpuType::Ref(_) => 4, + GpuType::InlineStruct(_) => { + if include_inline { + 4 + } else { + 0 + } + } + }; + if size > 0 { + for ix in (offset / 4)..(offset + size + 3) / 4 { + if ix >= result.len() { + result.resize_with(ix + 1, Default::default); + } + result[ix].push(i); + } + } + } + result +} + impl Size { fn new(size: usize) -> Size { - let alignment = if size < 4 { 1 } else { 4 }; + // Note: there is special case we could do better: + // `(u8, u16, u8)`, where the alignment could be 1. However, + // this case can also be solved by reordering. + let alignment = size.min(4); Size { size, alignment } } diff --git a/piet-gpu-derive/src/parse.rs b/piet-gpu-derive/src/parse.rs index ee3ddbd..8e51bab 100644 --- a/piet-gpu-derive/src/parse.rs +++ b/piet-gpu-derive/src/parse.rs @@ -2,6 +2,8 @@ extern crate proc_macro; +use std::collections::HashSet; + use syn::{ Expr, ExprLit, Fields, FieldsNamed, FieldsUnnamed, GenericArgument, ItemEnum, ItemStruct, Lit, PathArguments, TypeArray, TypePath, @@ -42,6 +44,7 @@ pub enum GpuTypeDef { pub struct GpuModule { pub name: String, + pub attrs: HashSet, pub defs: Vec, } @@ -72,22 +75,6 @@ impl GpuScalar { } } -fn ty_as_single_ident(ty: &syn::Type) -> Option { - if let syn::Type::Path(TypePath { - path: syn::Path { segments, .. }, - .. - }) = ty - { - if segments.len() == 1 { - let seg = &segments[0]; - if seg.arguments == PathArguments::None { - return Some(seg.ident.to_string()); - } - } - } - None -} - impl GpuType { fn from_syn(ty: &syn::Type) -> Result { //println!("gputype {:#?}", ty); @@ -189,6 +176,12 @@ impl GpuTypeDef { impl GpuModule { pub fn from_syn(module: &syn::ItemMod) -> Result { let name = module.ident.to_string(); + let mut attrs = HashSet::new(); + for attr in &module.attrs { + if let Some(id) = path_as_single_ident(&attr.path) { + attrs.insert(id.to_owned()); + } + } let mut defs = Vec::new(); if let Some((_brace, items)) = &module.content { for item in items { @@ -196,7 +189,25 @@ impl GpuModule { defs.push(def); } } - Ok(GpuModule { name, defs }) + Ok(GpuModule { name, attrs, defs }) + } +} + +fn path_as_single_ident(path: &syn::Path) -> Option { + if path.segments.len() == 1 { + let seg = &path.segments[0]; + if seg.arguments == PathArguments::None { + return Some(seg.ident.to_string()); + } + } + None +} + +fn ty_as_single_ident(ty: &syn::Type) -> Option { + if let syn::Type::Path(TypePath { path, .. }) = ty { + path_as_single_ident(path) + } else { + None } }