From 265d990cbe75c9b35cace68da67f0468170d05a6 Mon Sep 17 00:00:00 2001 From: Raph Levien Date: Fri, 10 Apr 2020 16:27:21 -0700 Subject: [PATCH] [derive] Add layout Add a layout pass to the struct derive logic. --- piet-gpu-derive/src/glsl.rs | 196 ++++++++++++++++++++++++++++++++++ piet-gpu-derive/src/layout.rs | 177 ++++++++++++++++++++++++++++++ piet-gpu-derive/src/lib.rs | 14 ++- piet-gpu-derive/src/parse.rs | 32 ++++-- piet-gpu-hal/src/lib.rs | 1 - piet-gpu-hal/src/vulkan.rs | 51 ++++----- 6 files changed, 429 insertions(+), 42 deletions(-) create mode 100644 piet-gpu-derive/src/glsl.rs create mode 100644 piet-gpu-derive/src/layout.rs diff --git a/piet-gpu-derive/src/glsl.rs b/piet-gpu-derive/src/glsl.rs new file mode 100644 index 0000000..0fc3a19 --- /dev/null +++ b/piet-gpu-derive/src/glsl.rs @@ -0,0 +1,196 @@ +//! Generation of GLSL struct definitions and accessor functions. + +use std::fmt::Write; +use std::ops::Deref; + +use crate::layout::{LayoutModule, LayoutType, LayoutTypeDef}; +use crate::parse::{GpuScalar, GpuType}; + +pub fn gen_glsl(module: &LayoutModule) -> String { + let mut r = String::new(); + writeln!(&mut r, "// Code auto-generated by piet-gpu-derive\n").unwrap(); + // Note: GLSL needs definitions before uses. We could do a topological sort here, + // but easiest for now to just require that in spec. + for name in &module.def_names { + 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); + } + } + 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::Enum(_)) => gen_enum_read(&mut r, &module.name, &name), + } + } + r +} + +fn gen_refdef(r: &mut String, name: &str) { + writeln!(r, "struct {}Ref {{", name).unwrap(); + writeln!(r, " uint offset;").unwrap(); + writeln!(r, "}};\n").unwrap(); +} + +fn gen_struct_def(r: &mut String, name: &str, size: usize, 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(); + + writeln!(r, "#define {}_size {}\n", name, size).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(); + } + writeln!(r, " {} s;", name).unwrap(); + for (name, offset, ty) in fields { + writeln!(r, " s.{} = {};", name, gen_extract(*offset, &ty.ty)).unwrap(); + } + writeln!(r, " return s;").unwrap(); + writeln!(r, "}}\n").unwrap(); +} + +fn gen_enum_read(r: &mut String, bufname: &str, name: &str) { + writeln!(r, "uint {}_tag({}Ref ref) {{", name, name).unwrap(); + writeln!(r, " return {}[ref.offset >> 2];", bufname).unwrap(); + writeln!(r, "}}\n").unwrap(); +} + +fn gen_extract(offset: usize, ty: &GpuType) -> String { + match ty { + GpuType::Scalar(scalar) => gen_extract_scalar(offset, scalar), + GpuType::Vector(scalar, size) => { + let mut r = glsl_type(ty); + r.push_str("("); + for i in 0..*size { + if i != 0 { + r.push_str(", "); + } + let el_offset = offset + i * scalar.size(); + r.push_str(&gen_extract_scalar(el_offset, scalar)); + } + r.push_str(")"); + r + } + GpuType::InlineStruct(name) => format!( + "{}_read({}Ref({}))", + name, + name, + simplified_add("ref.offset", offset) + ), + GpuType::Ref(inner) => { + if let GpuType::InlineStruct(name) = inner.deref() { + format!( + "{}Ref({})", + name, + gen_extract_scalar(offset, &GpuScalar::U32) + ) + } else { + panic!("only know how to deal with Ref of struct") + } + } + } +} + +fn gen_extract_scalar(offset: usize, ty: &GpuScalar) -> String { + match ty { + GpuScalar::F32 => format!("uintBitsToFloat(raw{})", offset / 4), + GpuScalar::U8 | GpuScalar::U16 | GpuScalar::U32 => extract_ubits(offset, ty.size()), + GpuScalar::I8 | GpuScalar::I16 | GpuScalar::I32 => extract_ibits(offset, ty.size()), + } +} + +fn extract_ubits(offset: usize, nbytes: usize) -> String { + if nbytes == 4 { + return format!("raw{}", offset / 4); + } + let mask = (1 << (nbytes * 8)) - 1; + if offset % 4 == 0 { + format!("raw{} & 0x{:x}", offset / 4, mask) + } else if offset % 4 + nbytes == 4 { + format!("raw{} >> {}", offset / 4, (offset % 4) * 8) + } else { + format!("(raw{} >> {}) & 0x{:x}", offset / 4, (offset % 4) * 8, mask) + } +} + +fn extract_ibits(offset: usize, nbytes: usize) -> String { + if nbytes == 4 { + return format!("int(raw{})", offset / 4); + } + if offset % 4 + nbytes == 4 { + format!("int(raw{}) >> {}", offset / 4, (offset % 4) * 8) + } else { + format!( + "(int(raw{}) << {}) >> {}", + offset / 4, + (3 - offset % 4) * 8, + (4 - nbytes) * 8 + ) + } +} + +fn glsl_type(ty: &GpuType) -> String { + match ty { + GpuType::Scalar(scalar) => glsl_scalar(scalar).into(), + GpuType::Vector(scalar, size) => { + if *size == 1 { + glsl_scalar(scalar).into() + } else { + format!("{}{}", glsl_vecname(scalar), size) + } + } + GpuType::InlineStruct(name) => name.clone(), + GpuType::Ref(inner) => { + if let GpuType::InlineStruct(name) = inner.deref() { + format!("{}Ref", name) + } else { + panic!("only know how to deal with Ref of struct") + } + } + } +} + +// GLSL type that can contain the scalar value. +fn glsl_scalar(s: &GpuScalar) -> &'static str { + match s { + GpuScalar::F32 => "float", + GpuScalar::I8 | GpuScalar::I16 | GpuScalar::I32 => "int", + GpuScalar::U8 | GpuScalar::U16 | GpuScalar::U32 => "uint", + } +} + +fn glsl_vecname(s: &GpuScalar) -> &'static str { + match s { + GpuScalar::F32 => "vec", + GpuScalar::I8 | GpuScalar::I16 | GpuScalar::I32 => "ivec", + 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 { + String::from(var_name) + } else { + format!("{} + {}", var_name, c) + } +} diff --git a/piet-gpu-derive/src/layout.rs b/piet-gpu-derive/src/layout.rs new file mode 100644 index 0000000..25cf38c --- /dev/null +++ b/piet-gpu-derive/src/layout.rs @@ -0,0 +1,177 @@ +//! Logic for layout of structures in memory. + +use std::collections::{HashMap, HashSet}; + +use crate::parse::{GpuModule, GpuType, GpuTypeDef}; + +#[derive(Clone)] +pub struct LayoutType { + size: Size, + pub ty: GpuType, +} + +#[derive(Clone)] +pub enum LayoutTypeDef { + /// Name, offset, field type. Make a separate struct? + Struct(Vec<(String, usize, LayoutType)>), + Enum(Vec<(String, Vec<(usize, LayoutType)>)>), +} + +pub struct LayoutModule { + pub name: String, + pub def_names: Vec, + pub defs: HashMap, +} + +struct LayoutSession<'a> { + enum_variants: HashSet, + orig_defs: HashMap, + defs: HashMap, +} + +#[derive(Clone, Copy)] +pub struct Size { + pub size: usize, + alignment: usize, +} + +impl LayoutType { + fn from_gpu(ty: &GpuType, session: &mut LayoutSession) -> LayoutType { + let size = session.get_size(ty); + LayoutType { + size, + ty: ty.clone(), + } + } +} + +impl LayoutTypeDef { + // Maybe have a type representing the tuple? + fn from_gpu(def: &GpuTypeDef, session: &mut LayoutSession) -> (Size, LayoutTypeDef) { + match def { + GpuTypeDef::Struct(_name, fields) => { + let mut offset = 0; + let mut result = Vec::new(); + for field in fields { + let layout_ty = LayoutType::from_gpu(&field.1, session); + offset += align_padding(offset, layout_ty.size.alignment); + let size = layout_ty.size.size; + result.push((field.0.clone(), offset, layout_ty)); + offset += size; + } + offset += align_padding(offset, 4); + let size = Size::new_struct(offset); + (size, LayoutTypeDef::Struct(result)) + } + GpuTypeDef::Enum(en) => { + let mut result = Vec::new(); + 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 + }; + for field in &variant.1 { + let layout_ty = LayoutType::from_gpu(field, session); + offset += align_padding(offset, layout_ty.size.alignment); + let size = layout_ty.size.size; + r2.push((offset, layout_ty)); + offset += size; + } + max_offset = max_offset.max(offset); + result.push((variant.0.clone(), r2)); + } + max_offset += align_padding(max_offset, 4); + let size = Size::new_struct(max_offset); + (size, LayoutTypeDef::Enum(result)) + } + } + } +} + +impl LayoutModule { + pub fn from_gpu(module: &GpuModule) -> LayoutModule { + let def_names = module + .defs + .iter() + .map(|def| def.name().to_owned()) + .collect::>(); + let mut session = LayoutSession::new(module); + for def in &module.defs { + let _ = session.layout_def(def.name()); + } + LayoutModule { + name: module.name.clone(), + def_names, + defs: session.defs, + } + } +} + +impl<'a> LayoutSession<'a> { + fn new(module: &GpuModule) -> LayoutSession { + let mut orig_defs = HashMap::new(); + let mut enum_variants = HashSet::new(); + for def in &module.defs { + orig_defs.insert(def.name().to_owned(), def.clone()); + if let GpuTypeDef::Enum(en) = def { + for variant in &en.variants { + if let Some(GpuType::InlineStruct(name)) = variant.1.first() { + enum_variants.insert(name.clone()); + } + } + } + } + LayoutSession { + enum_variants, + orig_defs, + defs: HashMap::new(), + } + } + + /// Do layout of one def. + /// + /// This might be called recursively. + /// Note: expect stack overflow for circular dependencies. + fn layout_def(&mut self, name: &str) -> Size { + if let Some(def) = self.defs.get(name) { + return def.0; + } + let def = self.orig_defs.get(name).unwrap(); + let layout = LayoutTypeDef::from_gpu(def, self); + let size = layout.0; + self.defs.insert(name.to_owned(), layout); + size + } + + fn get_size(&mut self, ty: &GpuType) -> Size { + match ty { + GpuType::Scalar(scalar) => Size::new(scalar.size()), + GpuType::Vector(scalar, len) => Size::new(scalar.size() * len), + GpuType::Ref(_) => Size::new(4), + GpuType::InlineStruct(name) => self.layout_def(name), + } + } + + fn is_enum_variant(&self, name: &str) -> bool { + self.enum_variants.contains(name) + } +} + +impl Size { + fn new(size: usize) -> Size { + let alignment = if size < 4 { 1 } else { 4 }; + Size { size, alignment } + } + + fn new_struct(size: usize) -> Size { + let alignment = 4; + Size { size, alignment } + } +} + +fn align_padding(offset: usize, alignment: usize) -> usize { + offset.wrapping_neg() & (alignment - 1) +} diff --git a/piet-gpu-derive/src/lib.rs b/piet-gpu-derive/src/lib.rs index 31568f3..fd086d5 100644 --- a/piet-gpu-derive/src/lib.rs +++ b/piet-gpu-derive/src/lib.rs @@ -1,9 +1,12 @@ +mod glsl; +mod layout; mod parse; use proc_macro::TokenStream; -use quote::quote; +use quote::{format_ident, quote}; use syn::parse_macro_input; +use layout::LayoutModule; use parse::GpuModule; #[proc_macro] @@ -11,6 +14,13 @@ pub fn piet_gpu(input: TokenStream) -> TokenStream { let input = parse_macro_input!(input as syn::ItemMod); //println!("input: {:#?}", input); let module = GpuModule::from_syn(&input).unwrap(); - let expanded = quote! {}; + let layout = LayoutModule::from_gpu(&module); + let glsl = glsl::gen_glsl(&layout); + let gen_gpu_fn = format_ident!("gen_gpu_{}", layout.name); + let expanded = quote! { + fn #gen_gpu_fn() -> String { + #glsl.into() + } + }; expanded.into() } diff --git a/piet-gpu-derive/src/parse.rs b/piet-gpu-derive/src/parse.rs index c335006..ee3ddbd 100644 --- a/piet-gpu-derive/src/parse.rs +++ b/piet-gpu-derive/src/parse.rs @@ -2,11 +2,9 @@ extern crate proc_macro; -use std::collections::HashSet; - use syn::{ - Expr, ExprLit, Fields, FieldsNamed, FieldsUnnamed, GenericArgument, ItemEnum, ItemStruct, - Lit, PathArguments, TypeArray, TypePath, + Expr, ExprLit, Fields, FieldsNamed, FieldsUnnamed, GenericArgument, ItemEnum, ItemStruct, Lit, + PathArguments, TypeArray, TypePath, }; /// A scalar that can be represented in a packed data structure. @@ -60,6 +58,18 @@ impl GpuScalar { _ => None, }) } + + /// Size of scalar type. + /// + /// This is arguably a concern at the layout level, not syntax, but it's here because + /// it's not likely to be variable, so reduces the total number of types. + pub fn size(self) -> usize { + match self { + GpuScalar::F32 | GpuScalar::I32 | GpuScalar::U32 => 4, + GpuScalar::I8 | GpuScalar::U8 => 1, + GpuScalar::I16 | GpuScalar::U16 => 2, + } + } } fn ty_as_single_ident(ty: &syn::Type) -> Option { @@ -166,7 +176,14 @@ impl GpuTypeDef { Err("unknown item".into()) } } - } + } + + pub fn name(&self) -> &str { + match self { + GpuTypeDef::Struct(name, _) => name, + GpuTypeDef::Enum(en) => &en.name, + } + } } impl GpuModule { @@ -179,10 +196,7 @@ impl GpuModule { defs.push(def); } } - Ok(GpuModule { - name, - defs, - }) + Ok(GpuModule { name, defs }) } } diff --git a/piet-gpu-hal/src/lib.rs b/piet-gpu-hal/src/lib.rs index 07ba686..e4eb3b5 100644 --- a/piet-gpu-hal/src/lib.rs +++ b/piet-gpu-hal/src/lib.rs @@ -2,7 +2,6 @@ /// /// This abstraction is inspired by gfx-hal, but is specialized to the needs of piet-gpu. /// In time, it may go away and be replaced by either gfx-hal or wgpu. - pub mod vulkan; /// This isn't great but is expedient. diff --git a/piet-gpu-hal/src/vulkan.rs b/piet-gpu-hal/src/vulkan.rs index 6d6f4d0..192ac29 100644 --- a/piet-gpu-hal/src/vulkan.rs +++ b/piet-gpu-hal/src/vulkan.rs @@ -74,17 +74,14 @@ impl VkInstance { None, )?; - Ok(VkInstance { - entry, - instance, - }) + Ok(VkInstance { entry, instance }) } } /// Create a device from the instance, suitable for compute. - /// + /// /// # Safety - /// + /// /// The caller is responsible for making sure that the instance outlives the device. /// We could enforce that, for example having an `Arc` of the raw instance, but for /// now keep things simple. @@ -127,11 +124,7 @@ impl crate::Device for VkDevice { type Pipeline = Pipeline; type MemFlags = MemFlags; - fn create_buffer( - &self, - size: u64, - mem_flags: MemFlags, - ) -> Result { + fn create_buffer(&self, size: u64, mem_flags: MemFlags) -> Result { unsafe { let device = &self.device.device; let buffer = device.create_buffer( @@ -197,20 +190,22 @@ impl crate::Device for VkDevice { None, )?; - let pipeline = device.create_compute_pipelines( - vk::PipelineCache::null(), - &[vk::ComputePipelineCreateInfo::builder() - .stage( - vk::PipelineShaderStageCreateInfo::builder() - .stage(vk::ShaderStageFlags::COMPUTE) - .module(compute_shader_module) - .name(&entry_name) - .build(), - ) - .layout(pipeline_layout) - .build()], - None, - ).map_err(|(_pipeline, err)| err)?[0]; + let pipeline = device + .create_compute_pipelines( + vk::PipelineCache::null(), + &[vk::ComputePipelineCreateInfo::builder() + .stage( + vk::PipelineShaderStageCreateInfo::builder() + .stage(vk::ShaderStageFlags::COMPUTE) + .module(compute_shader_module) + .name(&entry_name) + .build(), + ) + .layout(pipeline_layout) + .build()], + None, + ) + .map_err(|(_pipeline, err)| err)?[0]; Ok(Pipeline { pipeline, pipeline_layout, @@ -333,11 +328,7 @@ impl crate::Device for VkDevice { Ok(()) } - unsafe fn write_buffer( - &self, - buffer: &Buffer, - contents: &[T], - ) -> Result<(), Error> { + unsafe fn write_buffer(&self, buffer: &Buffer, contents: &[T]) -> Result<(), Error> { let device = &self.device.device; let buf = device.map_memory( buffer.buffer_memory,