diff --git a/crates/benilla-world/src/gpu_skin.rs b/crates/benilla-world/src/gpu_skin.rs new file mode 100644 index 00000000..2df1f9fd --- /dev/null +++ b/crates/benilla-world/src/gpu_skin.rs @@ -0,0 +1,276 @@ +//! GPU skin palette (`WOW_GPU_SKIN=1`): the per-bone `world × inverse_bindpose` row arithmetic +//! and the row upload move to a compute pass, in place of the CPU's `write_rig_worlds` + +//! `upload_rig_palettes` write into the shared light buffer's palette region. The CPU still runs +//! `rig_worlds` (the sequential parent walk, billboard arms and anchor re-seats are main-world +//! work), stages the composed worlds + inverse bindposes, and keeps a small on-demand row source +//! for the mouseover picker. Default off: the CPU path is untouched and stays the source of truth. + +use bevy::asset::Handle; +use bevy::core_pipeline::core_3d::graph::{Core3d, Node3d}; +use bevy::math::Affine3A; +use bevy::prelude::*; +use bevy::render::extract_resource::{ExtractResource, ExtractResourcePlugin}; +use bevy::render::render_graph::{ + Node, NodeRunError, RenderGraphContext, RenderGraphExt, RenderLabel, +}; +use bevy::render::render_resource::binding_types::{ + storage_buffer_read_only_sized, storage_buffer_sized, uniform_buffer_sized, +}; +use bevy::render::render_resource::{ + BindingResource, BindGroupLayout, BindGroupLayoutDescriptor, BindGroupLayoutEntry, + BindGroupEntry, BufferBinding, BufferDescriptor, BufferUsages, CachedComputePipelineId, + ComputePassDescriptor, ComputePipelineDescriptor, PipelineCache, ShaderStages, +}; +use bevy::render::renderer::{RenderContext, RenderDevice, RenderQueue}; +use bevy::render::{Render, RenderApp, RenderSystems}; +use bevy::shader::Shader; + +use crate::lighting::SharedLightBuffer; +use crate::rig_palette::palette_region_offset; + +/// One dispatch's staged bones: the base bone index in the palette, the composed worlds (3 rows a +/// bone) and the inverse bindposes (4 rows a bone), both flattened `[f32; 4]`. +#[derive(Clone, Default)] +pub(crate) struct GpuSkinRange { + pub base_bone: u32, + pub palette_base_word: u32, + pub worlds: Vec<[f32; 4]>, + pub ibp: Vec<[f32; 4]>, +} + +/// Main world: the ranges staged since the last publish. `ExtractResource` clones it into the +/// render world each frame, where the compute node reads the clone; `Last` clears it after the +/// extract has run. `pub` because a `pub` system (`finalize_rig_worlds`) carries it as a param. +#[derive(Resource, Clone, Default, ExtractResource)] +pub struct GpuSkinStaged { + pub ranges: Vec, +} + +/// `WOW_GPU_SKIN=1`, read once. +pub(crate) fn gpu_skin_enabled() -> bool { + static ON: std::sync::OnceLock = std::sync::OnceLock::new(); + *ON.get_or_init(|| std::env::var_os("WOW_GPU_SKIN").is_some()) +} + +/// Flatten one rig's composed worlds + inverse bindposes into a [`GpuSkinRange`]. The upload rows +/// are `(x_axis[i], y_axis[i], z_axis[i], t[i])` per affine — the shader's weighted-row multiply +/// then lands on the CPU path's column-major rows byte-for-byte. +pub(crate) fn stage_range( + skin_base: u32, + worlds: &[GlobalTransform], + ibp: &[Mat4], + count: usize, +) -> GpuSkinRange { + let palette_base_word = (palette_region_offset() / 16) as u32 + 3 * skin_base; + let mut out = GpuSkinRange { + base_bone: skin_base, + palette_base_word, + worlds: Vec::with_capacity(count * 3), + ibp: Vec::with_capacity(count * 4), + }; + for b in 0..count.min(worlds.len()).min(ibp.len()) { + let w = worlds[b].affine(); + let wm = w.matrix3; + let wt = w.translation; + out.worlds.extend_from_slice(&[ + [wm.x_axis.x, wm.y_axis.x, wm.z_axis.x, wt.x], + [wm.x_axis.y, wm.y_axis.y, wm.z_axis.y, wt.y], + [wm.x_axis.z, wm.y_axis.z, wm.z_axis.z, wt.z], + ]); + let im = Affine3A::from_mat4(ibp[b]).matrix3; + let it = Affine3A::from_mat4(ibp[b]).translation; + out.ibp.extend_from_slice(&[ + [im.x_axis.x, im.y_axis.x, im.z_axis.x, it.x], + [im.x_axis.y, im.y_axis.y, im.z_axis.y, it.y], + [im.x_axis.z, im.y_axis.z, im.z_axis.z, it.z], + [0.0, 0.0, 0.0, 1.0], + ]); + } + out +} + +pub(crate) struct GpuSkinPlugin; + +impl Plugin for GpuSkinPlugin { + fn build(&self, app: &mut App) { + if !gpu_skin_enabled() { + return; + } + app.init_resource::() + .add_plugins(ExtractResourcePlugin::::default()) + .add_systems(Last, clear_gpu_skin_staged); + let render_app = app.sub_app_mut(RenderApp); + render_app + .add_systems(Render, prepare_skin_pipeline.in_set(RenderSystems::PrepareAssets)) + .add_render_graph_node::(Core3d, SkinPaletteLabel) + // Before the opaque pass, whose model shaders read the palette (the retained static + // pass shares the region; a one-frame lag on a just-swapped rig is invisible). + .add_render_graph_edges( + Core3d, + (Node3d::StartMainPass, SkinPaletteLabel, Node3d::MainOpaquePass), + ); + } +} + +/// The extract clones the staged ranges into the render world; clear them once it has run. +fn clear_gpu_skin_staged(mut staged: ResMut) { + staged.ranges.clear(); +} + +#[derive(Debug, Hash, PartialEq, Eq, Clone, RenderLabel)] +struct SkinPaletteLabel; + +/// The compute layout: uploads (read storage), the palette region (read-write storage), params. +fn layout_entries() -> [BindGroupLayoutEntry; 3] { + [ + storage_buffer_read_only_sized(false, None).build(0, ShaderStages::COMPUTE), + storage_buffer_sized(false, None).build(1, ShaderStages::COMPUTE), + uniform_buffer_sized(false, None).build(2, ShaderStages::COMPUTE), + ] +} + +#[derive(Resource)] +struct SkinPipelines { + /// The actual layout, built from the descriptor, for the bind group. + layout: BindGroupLayout, + /// The queued compute pipeline, ready once the cache reports it. + pipeline: CachedComputePipelineId, +} + +/// Queue the compute pipeline once; a per-frame rebuild would recompile it constantly +/// (`pipe_warm` logs each live compile). `existing` gates: build only while the resource is +/// absent, then the node reads it from then on. +fn prepare_skin_pipeline( + device: Res, + asset_server: Res, + mut cache: ResMut, + existing: Option>, + mut commands: Commands, +) { + if existing.is_some() { + return; + } + let entries = layout_entries(); + let layout = device.create_bind_group_layout("skin_palette_layout", &entries); + let descriptor = BindGroupLayoutDescriptor { + label: "skin_palette_layout".into(), + entries: entries.to_vec(), + }; + let shader: Handle = + asset_server.load("embedded://benilla_world/shaders/skin_palette.wgsl"); + let pipeline = cache.queue_compute_pipeline(ComputePipelineDescriptor { + label: Some("skin_palette".into()), + layout: vec![descriptor], + push_constant_ranges: Vec::new(), + shader, + shader_defs: Vec::new(), + entry_point: Some("main".into()), + zero_initialize_workgroup_memory: true, + }); + commands.insert_resource(SkinPipelines { layout, pipeline }); +} + +#[derive(Default)] +struct SkinPaletteNode; + +impl Node for SkinPaletteNode { + fn run<'w>( + &self, + _graph: &mut RenderGraphContext, + render_context: &mut RenderContext<'w>, + world: &'w World, + ) -> Result<(), NodeRunError> { + let Some(pipes) = world.get_resource::() else { + return Ok(()); + }; + let Some(shared) = world.get_resource::() else { + return Ok(()); + }; + let Some(data) = world.get_resource::() else { + return Ok(()); + }; + if data.ranges.is_empty() { + return Ok(()); + } + let device = world.resource::(); + let queue = world.resource::(); + let cache = world.resource::(); + let Some(pipeline) = cache.get_compute_pipeline(pipes.pipeline) else { + return Ok(()); + }; + + // One dispatch per staged range: each range is one rig's contiguous bones with its own + // palette base word, so a single `params` (base_bone 0, this range's word) covers it. + let encoder = render_context.command_encoder(); + for r in &data.ranges { + let bones = (r.worlds.len() / 3) as u32; + if bones == 0 { + continue; + } + let mut upload = Vec::with_capacity(bones as usize * 7 * 4); + for b in 0..bones as usize { + upload.extend_from_slice(&r.worlds[3 * b..3 * b + 3]); + upload.extend_from_slice(&r.ibp[4 * b..4 * b + 4]); + } + let staging = device.create_buffer(&BufferDescriptor { + label: Some("skin_palette_upload".into()), + size: (upload.len() * 4) as u64, + usage: BufferUsages::COPY_DST | BufferUsages::STORAGE, + mapped_at_creation: false, + }); + queue.write_buffer(&staging, 0, bytemuck::cast_slice(&upload)); + + let params = device.create_buffer(&BufferDescriptor { + label: Some("skin_palette_params".into()), + size: 16, + usage: BufferUsages::UNIFORM | BufferUsages::COPY_DST, + mapped_at_creation: false, + }); + queue.write_buffer( + ¶ms, + 0, + bytemuck::cast_slice(&[0u32, bones, r.palette_base_word, 0u32]), + ); + + let bind_group = device.create_bind_group( + "skin_palette_bind", + &pipes.layout, + &[ + BindGroupEntry { + binding: 0, + resource: BindingResource::Buffer(BufferBinding { + buffer: &staging, + offset: 0, + size: None, + }), + }, + BindGroupEntry { + binding: 1, + resource: BindingResource::Buffer(BufferBinding { + buffer: &shared.0, + offset: palette_region_offset(), + size: None, + }), + }, + BindGroupEntry { + binding: 2, + resource: BindingResource::Buffer(BufferBinding { + buffer: ¶ms, + offset: 0, + size: None, + }), + }, + ], + ); + + { + let mut pass = + encoder.begin_compute_pass(&ComputePassDescriptor::default()); + pass.set_pipeline(pipeline); + pass.set_bind_group(0, &bind_group, &[]); + pass.dispatch_workgroups((bones + 63) / 64, 1, 1); + } + } + Ok(()) + } +} \ No newline at end of file diff --git a/crates/benilla-world/src/shaders/skin_palette.wgsl b/crates/benilla-world/src/shaders/skin_palette.wgsl new file mode 100644 index 00000000..298bcc3e --- /dev/null +++ b/crates/benilla-world/src/shaders/skin_palette.wgsl @@ -0,0 +1,54 @@ +// GPU skin palette (`WOW_GPU_SKIN=1`): one invocation per bone writes the rig's three rows +// (`rig_from_joint × inverse_bindpose`, the CPU path's `write_rig_worlds` arithmetic) straight +// into the shared light buffer's palette region. Inputs are the per-bone world matrices the CPU +// already composed (`rig_worlds`) and the rig's inverse bindposes, staged each frame for the +// dirty ranges only. + +// One bone's upload: the world affine's three rows (columns of the matrix3, then the +// translation), then the inverse bindpose as a row-major 4x4. 7 vec4s a bone. +struct SkinUpload { + world_row0: vec4, + world_row1: vec4, + world_row2: vec4, + ibp0: vec4, + ibp1: vec4, + ibp2: vec4, + ibp3: vec4, +} + +struct SkinParams { + // The dispatch's base bone index and bone count, and the palette region's first row word. + base_bone: u32, + bone_count: u32, + palette_base_word: u32, + _pad: u32, +} + +@group(0) @binding(0) var uploads: array; +// The shared light buffer, bound as raw storage; `palette_base_word` is the palette region's +// offset in vec4 words. +@group(0) @binding(1) var palette: array>; +@group(0) @binding(2) var params: SkinParams; + +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= params.bone_count) { + return; + } + let u = uploads[params.base_bone + i]; + // world × ibp: the affine's three rows against the 4x4. + let w0 = u.world_row0; + let w1 = u.world_row1; + let w2 = u.world_row2; + let r0 = w0.x * u.ibp0 + w0.y * u.ibp1 + w0.z * u.ibp2 + w0.w * u.ibp3; + let r1 = w1.x * u.ibp0 + w1.y * u.ibp1 + w1.z * u.ibp2 + w1.w * u.ibp3; + let r2 = w2.x * u.ibp0 + w2.y * u.ibp1 + w2.z * u.ibp2 + w2.w * u.ibp3; + // The CPU writes `rows[r] = [m3.x_axis.x, m3.y_axis.x, m3.z_axis.x, t.x]`, i.e. the matrix's + // columns — the transpose of the above row-major multiply. Hand that transpose over here so + // the GPU rows byte-match the CPU path. + let base = params.palette_base_word + 3u * i; + palette[base] = vec4(r0.x, r1.x, r2.x, r0.w); + palette[base + 1u] = vec4(r0.y, r1.y, r2.y, r1.w); + palette[base + 2u] = vec4(r0.z, r1.z, r2.z, r2.w); +} \ No newline at end of file