rullama 0.2.0

Browser-resident Gemma 4 inference: pure Rust → WebAssembly + WebGPU. Loads Ollama's on-disk GGUF blobs and runs the forward pass on the local GPU via hand-written WGSL.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! M0 compute spike: dispatch a trivial WGSL kernel that doubles each f32, read the
//! result back. Validates wgpu device → pipeline → buffer → dispatch → readback on the
//! current platform (native or browser).
//!
//! Not part of the production inference path. Keep until M2 lands real kernels, then
//! retain as the simplest end-to-end smoke test.

use crate::backend::WgpuCtx;
use crate::error::{Result, RullamaError};

const SPIKE_WGSL: &str = r#"
@group(0) @binding(0) var<storage, read>       inp:  array<f32>;
@group(0) @binding(1) var<storage, read_write> outp: array<f32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
    let i = gid.x;
    if (i >= arrayLength(&inp)) { return; }
    outp[i] = inp[i] * 2.0;
}
"#;

/// Run the doubling kernel on `input` and return the output. Used in tests.
pub async fn compute_spike(input: &[f32]) -> Result<Vec<f32>> {
    let ctx = WgpuCtx::new().await?;
    run(&ctx, input).await
}

async fn run(ctx: &WgpuCtx, input: &[f32]) -> Result<Vec<f32>> {
    let device = &ctx.device;
    let queue = &ctx.queue;

    let bytes: u64 = (input.len() * core::mem::size_of::<f32>()) as u64;

    let inp_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("spike.in"),
        size: bytes,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST,
        mapped_at_creation: false,
    });
    queue.write_buffer(&inp_buf, 0, bytemuck::cast_slice(input));

    let out_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("spike.out"),
        size: bytes,
        usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC,
        mapped_at_creation: false,
    });

    let read_buf = device.create_buffer(&wgpu::BufferDescriptor {
        label: Some("spike.read"),
        size: bytes,
        usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ,
        mapped_at_creation: false,
    });

    let module = device.create_shader_module(wgpu::ShaderModuleDescriptor {
        label: Some("spike.wgsl"),
        source: wgpu::ShaderSource::Wgsl(SPIKE_WGSL.into()),
    });

    let pipeline = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor {
        label: Some("spike.pipeline"),
        layout: None,
        module: &module,
        entry_point: Some("main"),
        compilation_options: Default::default(),
        cache: None,
    });

    let bind_group = device.create_bind_group(&wgpu::BindGroupDescriptor {
        label: Some("spike.bg"),
        layout: &pipeline.get_bind_group_layout(0),
        entries: &[
            wgpu::BindGroupEntry { binding: 0, resource: inp_buf.as_entire_binding() },
            wgpu::BindGroupEntry { binding: 1, resource: out_buf.as_entire_binding() },
        ],
    });

    let mut encoder = device.create_command_encoder(&wgpu::CommandEncoderDescriptor {
        label: Some("spike.encoder"),
    });
    {
        let mut cpass = encoder.begin_compute_pass(&wgpu::ComputePassDescriptor {
            label: Some("spike.pass"),
            timestamp_writes: None,
        });
        cpass.set_pipeline(&pipeline);
        cpass.set_bind_group(0, &bind_group, &[]);
        let workgroups = input.len().div_ceil(64) as u32;
        cpass.dispatch_workgroups(workgroups, 1, 1);
    }
    encoder.copy_buffer_to_buffer(&out_buf, 0, &read_buf, 0, bytes);
    queue.submit(Some(encoder.finish()));

    let slice = read_buf.slice(..);
    let (sender, receiver) = futures_channel::oneshot::channel();
    slice.map_async(wgpu::MapMode::Read, move |r| {
        let _ = sender.send(r);
    });
    device
        .poll(wgpu::PollType::Wait { submission_index: None, timeout: None })
        .map_err(|e| RullamaError::Inference(format!("{e:?}")))?;
    receiver
        .await
        .map_err(|e| RullamaError::BufferMap(format!("{e}")))?
        .map_err(|e| RullamaError::BufferMap(format!("{e}")))?;

    let data = slice.get_mapped_range();
    let out: Vec<f32> = bytemuck::cast_slice(&data).to_vec();
    drop(data);
    read_buf.unmap();
    Ok(out)
}

#[cfg(all(test, not(target_arch = "wasm32")))]
mod tests {
    use super::*;

    #[test]
    fn doubles_a_buffer() {
        let _ = env_logger::builder().is_test(true).try_init();
        let input: Vec<f32> = (0..256).map(|i| i as f32).collect();
        let output = pollster::block_on(compute_spike(&input)).expect("spike");
        for (i, v) in output.iter().enumerate() {
            assert!((*v - input[i] * 2.0).abs() < 1e-6, "i={i} got={v}");
        }
    }
}