shame

shame copied to clipboard

TypeLayout Rework

This PR reworks TypeLayout for two purposes:

• Centralize all layout calculations into one place - the type_layout module.
• Prepare for adding #20 support .
• Prepare for #21.

New traits and types

Here is a walkthrough of all new traits and the most important types

/// Represents all types that can be layed out in memory. The center of all layout calculations.
enum LayoutableType { ... }
/// Implemented for all non-bool shame gpu types and automatically derived by `#[derive(GpuLayout)]`.
trait Layoutable { fn layoutable_type() -> LayoutableType }
/// Implemented for all sized non-bool shame gpu types and automatically derived
/// by `#[derive(GpuLayout)]` when possible. `SizedType` is a variant of `LayoutableType`.
trait LayoutableSized { fn layoutable_type_sized() -> SizedType }

/// TypeLayout is what it used to be, however instead of mainly being constructed from
/// `StoreType`s, it's now mainly constructed from `LayoutableType`s.
struct TypeLayout { ... }
/// Represents type layout rules. Either packed or wgsl storage/uniform layout rules.
enum Repr { Storage, Uniform, Packed }
/// When creating a `TypeLayout` from a `LayoutableType` a `Repr` is required, so the
/// `TypeLayout` knows how to lay out the `LayoutableType` in memory.
let layout = TypeLayout::new_layout_for(layoutable_type, Repr::Storage);

/// However, there is a problem with
let layout = TypeLayout::new_layout_for(layoutable_type, Repr::Uniform);
/// because while the produced `TypeLayout` represents a layout that follows
/// all uniform layout rules, it may not be representable as a wgsl type.
/// The main issue is that wgsl does not have a custom stride attribute. For example
/// this `sm::Array` needs to have a stride of 16 to follow the uniform layout rules
struct A { a: sm::Array<f32x1> }
/// but we can't enforce that, so a `#[gpu_repr(uniform)] on shame's side is not easily possible.

/// To solve this problem and to further couple `TypeLayout` and `Repr`, we introduce new types and a trait.
struct GpuTypeLayout<T: TypeRepr>;
trait TypeRepr { } impl by structs Storage, Uniform, Packed
/// Which can be constructed like so
let storage = GpuTypeLayout::<Storage>::new(layoutable_type);
let packed  = GpuTypeLayout::<Packed>::new(layoutable_type);
/// However, the following is not possible
let uniform = GpuTypeLayout::<Uniform>::new(layoutable_type);
/// because one of the goals of `GpuTypeLayout` is to always be representable
/// as a wgsl type. To obtain a `GpuTypeLayout<Uniform>` we instead need to do
let uniform = GpuTypeLayout::<Uniform>::try_from(storage)?;
/// Which doesn't change the underlying `TypeLayout`, but now we can be sure
/// that it is valid to use in a uniform buffer. In other words, we are laying out
/// the `LayoutableType` according to storage layout rules and checking whether
/// it happens to also follow the uniform layout rules that way.
assert_eq!(storage.layout(), uniform.layout()); // .layout returns a `TypeLayout`

/// The main advantage of `GpuTypeLayout` compared to the current `TypeLayout` is that 
/// layout errors can be caught ahead of shader compile time and it also eliminates a ton of errors
/// that could occur due to having stronger guarantees.
/// In particular it will make a nice user interface for the storage/uniform buffer any api rework.
impl Any {
    pub fn storage_buffer_binding(ty: GpuTypeLayout<Storage>, ...) -> Any { ... }
}

/// One other important thing to mention is that `GpuLayout` has been adjusted as follows
trait GpuLayout: LayoutableType {
    // fn gpu_layout() -> TypeLayout; does not exist anymore, instead we have
    type GpuRepr: TypeReprStorageOrPacked;
    fn cpu_type_name_and_layout() -> ...;
}
/// And to actually obtain the `GpuTypeLayout` or `TypeLayout` we offer
pub fn gpu_type_layout<T: GpuLayout + ?Sized>() -> GpuTypeLayout<T::GpuRepr> {
    GpuTypeLayout::new(T::layoutable_type())
}
pub fn gpu_layout<T: GpuLayout + ?Sized>() -> TypeLayout { gpu_type_layout::<T>().layout() }

File structure

Now a bit about what each of the new files contains. We'll start with the type_layout/layoutable directory:

• mod.rs contains the definition of LayoutableType and it's nested enum variants such as Vector, SizedStruct and others. Display impls can also be found here
• align_size.rs contains all layout calculations necessary to define a TypeLayout for a LayoutableType. Once the storage/uniform buffer any api has been reworked slightly, this will be the only file containing layout calculations in all of shame.
• builder.rs contains an infallible builder implementation for SizedStruct and a fallible LayoutableType::struct_from_parts
• ir_compat.rs contains TryFrom conversion code of LayoutableType to StoreType and the other way around. The StoreType to LayoutableType conversion will be removed upon the storage/uniform buffer any api rework - it's currently functioning as temporary glue.

Otherwise there is mostly type_layout/mod.rs, which contains the TypeLayout, GpuTypeLayout, Repr definitions and type_layout/construction.rs, which contains TypeLayout::new_layout_for and the GpuTypeLayout construction and conversion code showcased above.

Final remarks

GpuTypeLayout is currently hooked up for cpu and gpu layout comparison, however it is not yet hooked up for storage/uniform buffer creation, which currently uses it's own layout calculation and verification functions - those can fully be replaced with GpuTypeLayout, which will be done in a follow up storage/uniform buffer any api rework PR.

The PR is ready for review!