rlx-bench 0.2.13

PLAN L5 — uniform benchmark harness for RLX backends + patterns
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

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// Licensed under the GNU General Public License, version 3.

//! wgpu-only SMALL-n batched FFT: one-workgroup-per-row (current default) vs
//! multi-row on-chip packing (rows=floor(2048/n) per workgroup), a single
//! forward FFT, toggled in-process via RLX_FFT_MULTIROW with RLX_FFT_FAST=1.
//!
//! Small n (<=1024) is `single_inner_only`: the default launches one 256-thread
//! workgroup per row, leaving it underfilled for n<512. Multi-row packs several
//! rows per group so the 256 threads stay busy. Whether it wins depends on batch.
//!
//! ```sh
//! cargo run -p rlx-bench --release --example bench_fft_wgpu_multirow \
//!     --features gpu,native-gpu-fft
//! ```

use std::io::Write;

use rlx_driver::Device;
use rlx_ir::{DType, Graph, Op, Shape, Tick};
use rlx_runtime::{CompiledGraph, Session};

fn fft_graph(batch: usize, n: usize) -> Graph {
    let mut g = Graph::new("wgpu_multirow");
    let len = batch * n * 2;
    let mut bytes = Vec::with_capacity(len * 4);
    for i in 0..len {
        bytes.extend_from_slice(&((i as f32 * 0.013).sin()).to_le_bytes());
    }
    let x = g.add_node(
        Op::Constant { data: bytes },
        vec![],
        Shape::new(&[batch, n * 2], DType::F32),
    );
    let y = g.fft(x, false);
    g.set_outputs(vec![y]);
    g
}

fn set_var(key: &str, on: bool) {
    // SAFETY: single-threaded benchmark driver toggling a process-local gate.
    unsafe {
        std::env::set_var(key, if on { "1" } else { "0" });
    }
}

fn main() {
    // Keep the on-chip fast path on; A/B only the multi-row packing.
    set_var("RLX_FFT_FAST", true);
    let empty: &[(&str, &[f32])] = &[];
    let median = |c: &mut CompiledGraph| -> u64 {
        for _ in 0..5 {
            let _ = c.run(empty);
        }
        let mut s = Vec::with_capacity(25);
        for _ in 0..25 {
            let t0 = Tick::now();
            let _ = c.run(empty);
            s.push(Tick::now().elapsed_ns(t0));
        }
        s.sort_unstable();
        s[s.len() / 2]
    };

    println!("wgpu small-n FFT: 1-row-per-wg (old) vs multi-row packing (new), single forward FFT");
    println!(
        "  {:>5} {:>6}  {:>11} {:>11}  {:>8}",
        "n", "batch", "old µs", "new µs", "new/old"
    );
    let _ = std::io::stdout().flush();

    for &n in &[8usize, 32, 128, 256, 512, 1024] {
        for &batch in &[16usize, 64, 256, 512, 1024] {
            let mut c = Session::new(Device::Gpu).compile(fft_graph(batch, n));
            set_var("RLX_FFT_MULTIROW", false);
            let old = median(&mut c);
            set_var("RLX_FFT_MULTIROW", true);
            let new = median(&mut c);
            let per = |t: u64| t as f64 / 1000.0;
            println!(
                "  {n:>5} {batch:>6}  {:>9.1}µs {:>9.1}µs  {:>7.2}×",
                per(old),
                per(new),
                old as f64 / new.max(1) as f64,
            );
            let _ = std::io::stdout().flush();
        }
        println!();
    }
}