blob: 48628fffbf9cdfca8c2d891018119742818b72cd [file] [log] [blame] [edit]
////////////////////////////////////////////////////////////////////////////////
// Utilities
////////////////////////////////////////////////////////////////////////////////
var<private> rand_seed : vec2<f32>;
fn rand() -> f32 {
rand_seed.x = fract(cos(dot(rand_seed, vec2<f32>(23.14077926, 232.61690225))) * 136.8168);
rand_seed.y = fract(cos(dot(rand_seed, vec2<f32>(54.47856553, 345.84153136))) * 534.7645);
return rand_seed.y;
}
////////////////////////////////////////////////////////////////////////////////
// Vertex shader
////////////////////////////////////////////////////////////////////////////////
struct RenderParams {
modelViewProjectionMatrix : mat4x4<f32>,
right : vec3<f32>,
up : vec3<f32>,
};
@binding(0) @group(0) var<uniform> render_params : RenderParams;
struct VertexInput {
@location(0) position : vec3<f32>,
@location(1) color : vec4<f32>,
@location(2) quad_pos : vec2<f32>, // -1..+1
};
struct VertexOutput {
@builtin(position) position : vec4<f32>,
@location(0) color : vec4<f32>,
@location(1) quad_pos : vec2<f32>, // -1..+1
};
@vertex
fn vs_main(in : VertexInput) -> VertexOutput {
var quad_pos = mat2x3<f32>(render_params.right, render_params.up) * in.quad_pos;
var position = in.position + quad_pos * 0.01;
var out : VertexOutput;
out.position = render_params.modelViewProjectionMatrix * vec4<f32>(position, 1.0);
out.color = in.color;
out.quad_pos = in.quad_pos;
return out;
}
////////////////////////////////////////////////////////////////////////////////
// Fragment shader
////////////////////////////////////////////////////////////////////////////////
@fragment
fn fs_main(in : VertexOutput) -> @location(0) vec4<f32> {
var color = in.color;
// Apply a circular particle alpha mask
color.a = color.a * max(1.0 - length(in.quad_pos), 0.0);
return color;
}
////////////////////////////////////////////////////////////////////////////////
// Simulation Compute shader
////////////////////////////////////////////////////////////////////////////////
struct SimulationParams {
deltaTime : f32,
seed : vec4<f32>,
};
struct Particle {
position : vec3<f32>,
lifetime : f32,
color : vec4<f32>,
velocity : vec3<f32>,
};
struct Particles {
particles : array<Particle>,
};
@binding(0) @group(0) var<uniform> sim_params : SimulationParams;
@binding(1) @group(0) var<storage, read_write> data : Particles;
@binding(2) @group(0) var texture : texture_2d<f32>;
@compute @workgroup_size(64)
fn simulate(@builtin(global_invocation_id) GlobalInvocationID : vec3<u32>) {
rand_seed = (sim_params.seed.xy + vec2<f32>(GlobalInvocationID.xy)) * sim_params.seed.zw;
let idx = GlobalInvocationID.x;
var particle = data.particles[idx];
// Apply gravity
particle.velocity.z = particle.velocity.z - sim_params.deltaTime * 0.5;
// Basic velocity integration
particle.position = particle.position + sim_params.deltaTime * particle.velocity;
// Age each particle. Fade out before vanishing.
particle.lifetime = particle.lifetime - sim_params.deltaTime;
particle.color.a = smoothstep(0.0, 0.5, particle.lifetime);
// If the lifetime has gone negative, then the particle is dead and should be
// respawned.
if (particle.lifetime < 0.0) {
// Use the probability map to find where the particle should be spawned.
// Starting with the 1x1 mip level.
var coord = vec2<i32>(0, 0);
for (var level = textureNumLevels(texture) - 1; level > 0; level = level - 1) {
// Load the probability value from the mip-level
// Generate a random number and using the probabilty values, pick the
// next texel in the next largest mip level:
//
// 0.0 probabilites.r probabilites.g probabilites.b 1.0
// | | | | |
// | TOP-LEFT | TOP-RIGHT | BOTTOM-LEFT | BOTTOM_RIGHT |
//
let probabilites = textureLoad(texture, coord, level);
let value = vec4<f32>(rand());
let mask = (value >= vec4<f32>(0.0, probabilites.xyz)) & (value < probabilites);
coord = coord * 2;
coord.x = coord.x + select(0, 1, any(mask.yw)); // x y
coord.y = coord.y + select(0, 1, any(mask.zw)); // z w
}
let uv = vec2<f32>(coord) / vec2<f32>(textureDimensions(texture));
particle.position = vec3<f32>((uv - 0.5) * 3.0 * vec2<f32>(1.0, -1.0), 0.0);
particle.color = textureLoad(texture, coord, 0);
particle.velocity.x = (rand() - 0.5) * 0.1;
particle.velocity.y = (rand() - 0.5) * 0.1;
particle.velocity.z = rand() * 0.3;
particle.lifetime = 0.5 + rand() * 2.0;
}
// Store the new particle value
data.particles[idx] = particle;
}
struct UBO {
width : u32,
};
struct Buffer {
weights : array<f32>,
};
@binding(3) @group(0) var<uniform> ubo : UBO;
@binding(4) @group(0) var<storage, read> buf_in : Buffer;
@binding(5) @group(0) var<storage, read_write> buf_out : Buffer;
@binding(6) @group(0) var tex_in : texture_2d<f32>;
@binding(7) @group(0) var tex_out : texture_storage_2d<rgba8unorm, write>;
////////////////////////////////////////////////////////////////////////////////
// import_level
//
// Loads the alpha channel from a texel of the source image, and writes it to
// the buf_out.weights.
////////////////////////////////////////////////////////////////////////////////
@compute @workgroup_size(64)
fn import_level(@builtin(global_invocation_id) coord : vec3<u32>) {
_ = &buf_in;
let offset = coord.x + coord.y * ubo.width;
buf_out.weights[offset] = textureLoad(tex_in, vec2<i32>(coord.xy), 0).w;
}
////////////////////////////////////////////////////////////////////////////////
// export_level
//
// Loads 4 f32 weight values from buf_in.weights, and stores summed value into
// buf_out.weights, along with the calculated 'probabilty' vec4 values into the
// mip level of tex_out. See simulate() in particle.wgsl to understand the
// probability logic.
////////////////////////////////////////////////////////////////////////////////
@compute @workgroup_size(64)
fn export_level(@builtin(global_invocation_id) coord : vec3<u32>) {
if (all(coord.xy < vec2<u32>(textureDimensions(tex_out)))) {
let dst_offset = coord.x + coord.y * ubo.width;
let src_offset = coord.x*2u + coord.y*2u * ubo.width;
let a = buf_in.weights[src_offset + 0u];
let b = buf_in.weights[src_offset + 1u];
let c = buf_in.weights[src_offset + 0u + ubo.width];
let d = buf_in.weights[src_offset + 1u + ubo.width];
let sum = dot(vec4<f32>(a, b, c, d), vec4<f32>(1.0));
buf_out.weights[dst_offset] = sum / 4.0;
let probabilities = vec4<f32>(a, a+b, a+b+c, sum) / max(sum, 0.0001);
textureStore(tex_out, vec2<i32>(coord.xy), probabilities);
}
}