wip: GPU skin palette compute pass (WOW_GPU_SKIN=1), not yet wired into the app
Standalone module and its palette shader; no caller references it yet, so the default CPU skin path is unchanged.
This commit is contained in:
chairmanpow
@@ -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<GpuSkinRange>,
|
||||
}
|
||||
|
||||
/// `WOW_GPU_SKIN=1`, read once.
|
||||
pub(crate) fn gpu_skin_enabled() -> bool {
|
||||
static ON: std::sync::OnceLock<bool> = 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::<GpuSkinStaged>()
|
||||
.add_plugins(ExtractResourcePlugin::<GpuSkinStaged>::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::<SkinPaletteNode>(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<GpuSkinStaged>) {
|
||||
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<RenderDevice>,
|
||||
asset_server: Res<AssetServer>,
|
||||
mut cache: ResMut<PipelineCache>,
|
||||
existing: Option<Res<SkinPipelines>>,
|
||||
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<Shader> =
|
||||
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::<SkinPipelines>() else {
|
||||
return Ok(());
|
||||
};
|
||||
let Some(shared) = world.get_resource::<SharedLightBuffer>() else {
|
||||
return Ok(());
|
||||
};
|
||||
let Some(data) = world.get_resource::<GpuSkinStaged>() else {
|
||||
return Ok(());
|
||||
};
|
||||
if data.ranges.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
let device = world.resource::<RenderDevice>();
|
||||
let queue = world.resource::<RenderQueue>();
|
||||
let cache = world.resource::<PipelineCache>();
|
||||
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(())
|
||||
}
|
||||
}
|
||||
@@ -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<f32>,
|
||||
world_row1: vec4<f32>,
|
||||
world_row2: vec4<f32>,
|
||||
ibp0: vec4<f32>,
|
||||
ibp1: vec4<f32>,
|
||||
ibp2: vec4<f32>,
|
||||
ibp3: vec4<f32>,
|
||||
}
|
||||
|
||||
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<storage, read> uploads: array<SkinUpload>;
|
||||
// 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<storage, read_write> palette: array<vec4<f32>>;
|
||||
@group(0) @binding(2) var<uniform> params: SkinParams;
|
||||
|
||||
@compute @workgroup_size(64)
|
||||
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
||||
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<f32>(r0.x, r1.x, r2.x, r0.w);
|
||||
palette[base + 1u] = vec4<f32>(r0.y, r1.y, r2.y, r1.w);
|
||||
palette[base + 2u] = vec4<f32>(r0.z, r1.z, r2.z, r2.w);
|
||||
}
|
||||
Reference in New Issue
Block a user