🔮 :: Million Particles Dancing

📝 Rust Code

use bytemuck::{Pod, Zeroable};
use nannou::prelude::*;
use nannou::wgpu;

const NUM_PARTICLES: u32 = 1_000_000;
const WORKGROUP_SIZE: u32 = 256;

fn main() {
    nannou::app(model).update(update).run();
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct Particle {
    pos: [f32; 2],
    vel: [f32; 2],
    color: [f32; 4],
}

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct Uniforms {
    mouse: [f32; 2],
    time: f32,
    dt: f32,
    mouse_down: u32,
    num_particles: u32,
    _pad: [u32; 2],
}

struct Model {
    uniform_buf: wgpu::Buffer,
    compute_pipeline: wgpu::ComputePipeline,
    compute_bind_group: wgpu::BindGroup,
    render_pipeline: wgpu::RenderPipeline,
    render_bind_group: wgpu::BindGroup,
    mouse_down: bool,
}

fn model(app: &App) -> Model {
    let w_id = app
        .new_window()
        .size(1440, 1440)
        .title("Hybrid: nannou + wgpu")
        .view(view)
        .mouse_pressed(|_, m: &mut Model, _| m.mouse_down = true)
        .mouse_released(|_, m: &mut Model, _| m.mouse_down = false)
        .build()
        .unwrap();

    let window = app.window(w_id).unwrap();
    let device = window.device();

    let mut particles = vec![
        Particle { pos: [0.0; 2], vel: [0.0; 2], color: [0.0; 4] };
        NUM_PARTICLES as usize
    ];
    for p in particles.iter_mut() {
        p.pos = [random_range(-1.0, 1.0), random_range(-1.0, 1.0)];
        p.vel = [random_range(-0.0005, 0.0005), random_range(-0.0005, 0.0005)];
        p.color = [random_range(0.2, 0.8), random_range(0.4, 1.0), 1.0, 0.7];
    }

    let particle_buf = device.create_buffer_init(&wgpu::BufferInitDescriptor {
        label: Some("particles"),
        contents: bytemuck::cast_slice(&particles),
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::VERTEX,
    });

    let uniform_buf = device.create_buffer_init(&wgpu::BufferInitDescriptor {
        label: Some("uniforms"),
        contents: bytemuck::bytes_of(&Uniforms {
            mouse: [0.0; 2], time: 0.0, dt: 1.0 / 60.0,
            mouse_down: 0, num_particles: NUM_PARTICLES, _pad: [0; 2],
        }),
        usage: wgpu::BufferUsages::UNIFORM | wgpu::BufferUsages::COPY_DST,
    });

    let cs_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("compute"),
        source: wgpu::ShaderSource::Wgsl(COMPUTE_SHADER.into()),
    });

    let compute_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: None,
        entries: &[
            wgpu::BindGroupLayoutEntry {
                binding: 0,
                visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Storage { read_only: false },
                    has_dynamic_offset: false, min_binding_size: None,
                },
                count: None,
            },
            wgpu::BindGroupLayoutEntry {
                binding: 1,
                visibility: wgpu::ShaderStages::COMPUTE,
                ty: wgpu::BindingType::Buffer {
                    ty: wgpu::BufferBindingType::Uniform,
                    has_dynamic_offset: false, min_binding_size: None,
                },
                count: None,
            },
        ],
    });

    let compute_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None, layout: &compute_bgl,
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: particle_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: uniform_buf.as_entire_binding() },
        ],
    });

    let compute_pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: None,
        layout: Some(&device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: None, bind_group_layouts: &[&compute_bgl], push_constant_ranges: &[],
        })),
        module: &cs_module,
        entry_point: "main",
    });

    let vs_fs_module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("render"),
        source: wgpu::ShaderSource::Wgsl(RENDER_SHADER.into()),
    });

    let render_bgl = device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
        label: None,
        entries: &[wgpu::BindGroupLayoutEntry {
            binding: 0,
            visibility: wgpu::ShaderStages::VERTEX,
            ty: wgpu::BindingType::Buffer {
                ty: wgpu::BufferBindingType::Storage { read_only: true },
                has_dynamic_offset: false, min_binding_size: None,
            },
            count: None,
        }],
    });

    let render_bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None, layout: &render_bgl,
        entries: &[wgpu::BindGroupEntry {
            binding: 0, resource: particle_buf.as_entire_binding(),
        }],
    });

    let render_pipeline = device.create_render_pipeline(&wgpu::RenderPipelineDescriptor {
        label: Some("render_pipeline"),
        layout: Some(&device.create_pipeline_layout(&wgpu::PipelineLayoutDescriptor {
            label: None, bind_group_layouts: &[&render_bgl], push_constant_ranges: &[],
        })),
        vertex: wgpu::VertexState {
            module: &vs_fs_module, entry_point: "vs_main", buffers: &[],
        },
        fragment: Some(wgpu::FragmentState {
            module: &vs_fs_module, entry_point: "fs_main",
            targets: &[Some(wgpu::ColorTargetState {
                format: nannou::Frame::TEXTURE_FORMAT,
                blend: Some(wgpu::BlendState {
                    color: wgpu::BlendComponent {
                        src_factor: wgpu::BlendFactor::SrcAlpha,
                        dst_factor: wgpu::BlendFactor::OneMinusSrcAlpha,
                        operation: wgpu::BlendOperation::Add,
                    },
                    alpha: wgpu::BlendComponent::OVER,
                }),
                write_mask: wgpu::ColorWrites::ALL,
            })],
        }),
        primitive: wgpu::PrimitiveState {
            topology: wgpu::PrimitiveTopology::PointList,
            ..Default::default()
        },
        depth_stencil: None,
        multisample: wgpu::MultisampleState {
            count: 4,
            mask: !0,
            alpha_to_coverage_enabled: false,
        },
        multiview: None,
    });

    Model {
        uniform_buf,
        compute_pipeline, compute_bind_group,
        render_pipeline, render_bind_group,
        mouse_down: false,
    }
}

fn update(app: &App, model: &mut Model, _update: Update) {
    let window = app.main_window();
    let device = window.device();
    let queue = window.queue();
    let rect = app.window_rect();

    let mouse = app.mouse.position();
    let uniforms = Uniforms {
        mouse: [mouse.x / (rect.w() * 0.5), mouse.y / (rect.h() * 0.5)],
        time: app.time,
        dt: 1.0 / 60.0,
        mouse_down: model.mouse_down as u32,
        num_particles: NUM_PARTICLES,
        _pad: [0; 2],
    };
    queue.write_buffer(&model.uniform_buf, 0, bytemuck::bytes_of(&uniforms));

    let mut encoder = device.create_command_encoder(&Default::default());
    {
        let mut pass = encoder.begin_compute_pass(&Default::default());
        pass.set_pipeline(&model.compute_pipeline);
        pass.set_bind_group(0, &model.compute_bind_group, &[]);
        pass.dispatch_workgroups((NUM_PARTICLES + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE, 1, 1);
    }
    queue.submit(Some(encoder.finish()));
}

fn view(_app: &App, model: &Model, frame: Frame) {
    let mut encoder = frame.command_encoder();

    {
        let mut rpass = encoder.begin_render_pass(&wgpu::RenderPassDescriptor {
            label: None,
            color_attachments: &[Some(wgpu::RenderPassColorAttachment {
                view: frame.texture_view(),
                resolve_target: None,
                ops: wgpu::Operations {
                    load: wgpu::LoadOp::Clear(wgpu::Color {
                        r: 0.02, g: 0.02, b: 0.05, a: 1.0,
                    }),
                    store: true,
                },
            })],
            depth_stencil_attachment: None,
        });

        rpass.set_pipeline(&model.render_pipeline);
        rpass.set_bind_group(0, &model.render_bind_group, &[]);
        rpass.draw(0..NUM_PARTICLES, 0..1);
    }
}

const COMPUTE_SHADER: &str = r#"
struct Particle {
    pos: vec2<f32>,
    vel: vec2<f32>,
    color: vec4<f32>,
};

struct Uniforms {
    mouse: vec2<f32>,
    time: f32,
    dt: f32,
    mouse_down: u32,
    num_particles: u32,
};

@group(0) @binding(0) var<storage, read_write> particles: array<Particle>;
@group(0) @binding(1) var<uniform> u: Uniforms;

fn hash(p: vec2<f32>) -> vec2<f32> {
    var q = vec2<f32>(
        dot(p, vec2<f32>(127.1, 311.7)),
        dot(p, vec2<f32>(269.5, 183.3))
    );
    return fract(sin(q) * 43758.5453) * 2.0 - 1.0;
}

fn noise2d(p: vec2<f32>) -> vec2<f32> {
    let sp = p * 2.0 + vec2<f32>(u.time * 0.05);
    return hash(floor(sp)) * 0.5 + hash(floor(sp) + vec2<f32>(1.0, 0.0)) * 0.25
         + hash(floor(sp) + vec2<f32>(0.0, 1.0)) * 0.25;
}

fn gaussian(dist: f32, sigma: f32) -> f32 {
    return exp(-0.5 * pow(dist / sigma, 2.0));
}

@compute @workgroup_size(256)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= u.num_particles) { return; }

    var p = particles[i];

    let nf = noise2d(p.pos) * 0.3;
    p.vel += nf * u.dt;

    let to_mouse = u.mouse - p.pos;
    let dist = max(length(to_mouse), 0.001);
    let g = gaussian(dist, 0.4);
    let dir = to_mouse / dist;

    if (u.mouse_down == 1u) {
        p.vel += dir * g * 0.8 * u.dt;
    } else {
        p.vel -= dir * g * 0.6 * u.dt;
    }

    let base_hue = fract(p.pos.x * 0.5 + u.time * 0.02);
    let hue = mix(base_hue, 0.0, smoothstep(0.0, 1.0, g));
    let c = 0.7 + g * 0.3;
    let rgb = clamp(
        abs(fract(vec3<f32>(hue, hue + 0.333, hue + 0.666)) * 6.0 - 3.0) - 1.0,
        vec3<f32>(0.0), vec3<f32>(1.0)
    ) * c;
    p.color = vec4<f32>(rgb, 0.6 + g * 0.3);

    p.vel *= 0.98;
    let spd = length(p.vel);
    if (spd > 0.15) { p.vel = (p.vel / spd) * 0.15; }

    p.pos += p.vel * u.dt * 20.0;

    if (p.pos.x < -1.0) { p.pos.x += 2.0; }
    if (p.pos.x >  1.0) { p.pos.x -= 2.0; }
    if (p.pos.y < -1.0) { p.pos.y += 2.0; }
    if (p.pos.y >  1.0) { p.pos.y -= 2.0; }

    particles[i] = p;
}
"#;

const RENDER_SHADER: &str = r#"
struct Particle {
    pos: vec2<f32>,
    vel: vec2<f32>,
    color: vec4<f32>,
};

@group(0) @binding(0) var<storage, read> particles: array<Particle>;

struct VsOut {
    @builtin(position) pos: vec4<f32>,
    @location(0) color: vec4<f32>,
};

@vertex
fn vs_main(@builtin(vertex_index) vi: u32) -> VsOut {
    let p = particles[vi];
    var out: VsOut;
    out.pos = vec4<f32>(p.pos, 0.0, 1.0);
    out.color = p.color;
    return out;
}

@fragment
fn fs_main(in: VsOut) -> @location(0) vec4<f32> {
    return in.color;
}
"#;