soft3/glia/run/backend/wgpu/kernels/softmax.rs

//! Softmax — stabilized by subtracting max.

use crate::backend::wgpu::device::{compute_pipeline, storage_ro, storage_rw, uniform, Device};

const SHADER: &str = r#"
struct Params {
    batch: u32,
    d: u32,
    _pad0: u32,
    _pad1: u32,
}

@group(0) @binding(0) var<storage, read> x: array<f32>;
@group(0) @binding(1) var<storage, read_write> y: array<f32>;
@group(0) @binding(2) var<uniform> params: Params;

var<workgroup> shared_max: array<f32, 256>;
var<workgroup> shared_sum: array<f32, 256>;

@compute @workgroup_size(256)
fn main(
    @builtin(workgroup_id) wg: vec3<u32>,
    @builtin(local_invocation_id) lid: vec3<u32>,
) {
    let b = wg.x;
    let tid = lid.x;
    if (b >= params.batch) { return; }

    // Pass 1: find max
    var local_max: f32 = -3.4e38;
    var j: u32 = tid;
    while (j < params.d) {
        local_max = max(local_max, x[b * params.d + j]);
        j = j + 256u;
    }
    shared_max[tid] = local_max;
    workgroupBarrier();

    var stride: u32 = 128u;
    while (stride > 0u) {
        if (tid < stride) {
            shared_max[tid] = max(shared_max[tid], shared_max[tid + stride]);
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let max_val = shared_max[0];

    // Pass 2: exp and partial sum
    var local_sum: f32 = 0.0;
    j = tid;
    while (j < params.d) {
        let e = exp(x[b * params.d + j] - max_val);
        y[b * params.d + j] = e;
        local_sum = local_sum + e;
        j = j + 256u;
    }
    shared_sum[tid] = local_sum;
    workgroupBarrier();

    stride = 128u;
    while (stride > 0u) {
        if (tid < stride) {
            shared_sum[tid] = shared_sum[tid] + shared_sum[tid + stride];
        }
        workgroupBarrier();
        stride = stride / 2u;
    }
    let inv_sum = 1.0 / shared_sum[0];

    // Pass 3: normalize
    j = tid;
    while (j < params.d) {
        y[b * params.d + j] = y[b * params.d + j] * inv_sum;
        j = j + 256u;
    }
}
"#;

pub fn dispatch(device: &Device, x: &wgpu::Buffer, batch: u32, d: u32) -> wgpu::Buffer {
    #[repr(C)]
    #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
    struct Params {
        batch: u32,
        d: u32,
        _pad0: u32,
        _pad1: u32,
    }
    let (pipeline, layout) = compute_pipeline(
        &device.device,
        SHADER,
        &[storage_ro(0), storage_rw(1), uniform(2)],
    );
    let out = device.alloc_f32((batch * d) as usize);
    let params = Params {
        batch,
        d,
        _pad0: 0,
        _pad1: 0,
    };
    let params_buf = device.upload_uniform(bytemuck::bytes_of(&params));

    let bg = device.device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: None,
        layout: &layout,
        entries: &[
            wgpu::BindGroupEntry {
                binding: 0,
                resource: x.as_entire_binding(),
            },
            wgpu::BindGroupEntry {
                binding: 1,
                resource: out.as_entire_binding(),
            },
            wgpu::BindGroupEntry {
                binding: 2,
                resource: params_buf.as_entire_binding(),
            },
        ],
    });

    let mut enc = device
        .device
        .create_command_encoder(&wgpu::CommandEncoderDescriptor::default());
    {
        let mut pass = enc.begin_compute_pass(&wgpu::ComputePassDescriptor::default());
        pass.set_pipeline(&pipeline);
        pass.set_bind_group(0, &bg, &[]);
        pass.dispatch_workgroups(batch, 1, 1);
    }
    device.queue.submit(std::iter::once(enc.finish()));
    out
}