rlx-bench 0.2.2

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! Canonical benchmark patterns.
//!
//! Add new patterns here by implementing [`crate::BenchmarkPattern`].
//! Existing patterns are minimal but correct — they exercise the full
//! compile + dispatch path so timings reflect real backend cost
//! (kernel launch, parameter binding, output read), not just kernel
//! body execution.

use crate::{BenchmarkPattern, Tier};
use rlx_ir::infer::GraphExt;
use rlx_ir::{DType, Graph, NodeId, Op, Shape};

/// **L1 — single matmul.** `out = x @ w` with x: `[m, k]`, w: `[k, n]`.
/// Plain BLAS / kernel dispatch — no fusion. Use to measure raw
/// matmul throughput per device.
pub struct MatmulPattern {
    pub m: usize,
    pub k: usize,
    pub n: usize,
}

impl BenchmarkPattern for MatmulPattern {
    fn name(&self) -> &str {
        "matmul"
    }
    fn tier(&self) -> Tier {
        Tier::L1
    }

    fn build_graph(&self) -> Graph {
        let f = DType::F32;
        let mut g = Graph::new("matmul_bench");
        let x = g.input("x", Shape::new(&[self.m, self.k], f));
        let w = g.input("w", Shape::new(&[self.k, self.n], f));
        let out = g.mm(x, w);
        g.set_outputs(vec![out]);
        g
    }

    fn input_data(&self) -> Vec<(String, Vec<f32>)> {
        vec![
            ("x".to_string(), vec![1.0; self.m * self.k]),
            ("w".to_string(), vec![1.0; self.k * self.n]),
        ]
    }
}

/// **L1 — single LayerNorm.** `out = layer_norm(x, gamma, beta)` with
/// `x: [rows, hidden]`. Hits the per-row reduction kernel.
pub struct LayerNormPattern {
    pub rows: usize,
    pub hidden: usize,
}

impl BenchmarkPattern for LayerNormPattern {
    fn name(&self) -> &str {
        "layer_norm"
    }
    fn tier(&self) -> Tier {
        Tier::L1
    }

    fn build_graph(&self) -> Graph {
        let f = DType::F32;
        let mut g = Graph::new("ln_bench");
        let x = g.input("x", Shape::new(&[self.rows, self.hidden], f));
        let gamma = g.input("gamma", Shape::new(&[self.hidden], f));
        let beta = g.input("beta", Shape::new(&[self.hidden], f));
        let out = g.ln(x, gamma, beta, 1e-5);
        g.set_outputs(vec![out]);
        g
    }

    fn input_data(&self) -> Vec<(String, Vec<f32>)> {
        vec![
            ("x".to_string(), vec![1.0; self.rows * self.hidden]),
            ("gamma".to_string(), vec![1.0; self.hidden]),
            ("beta".to_string(), vec![0.0; self.hidden]),
        ]
    }
}

/// **L2 — composite: matmul → bias → relu.** `out = relu(x @ w + b)`.
/// Tests fusion: with `FuseMatMulBiasAct` enabled, the three IR ops
/// collapse into one `Op::FusedMatMulBiasAct`. Without fusion, three
/// separate kernel dispatches. This pattern is the canonical FFN
/// hidden-layer shape.
pub struct MatmulBiasReluPattern {
    pub m: usize,
    pub k: usize,
    pub n: usize,
}

impl BenchmarkPattern for MatmulBiasReluPattern {
    fn name(&self) -> &str {
        "matmul_bias_relu"
    }
    fn tier(&self) -> Tier {
        Tier::L2
    }

    fn build_graph(&self) -> Graph {
        let f = DType::F32;
        let mut g = Graph::new("matmul_bias_relu_bench");
        let x = g.input("x", Shape::new(&[self.m, self.k], f));
        let w = g.input("w", Shape::new(&[self.k, self.n], f));
        let b = g.input("b", Shape::new(&[self.n], f));
        let mm = g.mm(x, w);
        let biased = g.add(mm, b);
        let out = g.relu(biased);
        g.set_outputs(vec![out]);
        g
    }

    fn input_data(&self) -> Vec<(String, Vec<f32>)> {
        // Mix of negatives so relu actually does work (not pure pass-through).
        let x: Vec<f32> = (0..self.m * self.k)
            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
            .collect();
        let w: Vec<f32> = vec![0.5; self.k * self.n];
        let b: Vec<f32> = vec![-0.1; self.n];
        vec![
            ("x".to_string(), x),
            ("w".to_string(), w),
            ("b".to_string(), b),
        ]
    }
}

/// **L1 — batched 1D FFT (`Op::Fft`).**
///
/// Input is a constant 2N real-block tensor `[batch, 2·n]` so timed
/// iterations measure dispatch + kernel only (no host input upload).
/// `n` must be a power of two ≥ 2 for GPU native paths.
pub struct FftPattern {
    pub batch: usize,
    pub n: usize,
    pub inverse: bool,
}

impl FftPattern {
    pub fn label(&self) -> String {
        format!(
            "fft_b{}_n{}{}",
            self.batch,
            self.n,
            if self.inverse { "_inv" } else { "" }
        )
    }

    /// Bytes read+written per iteration (src copy + dst write, 2N f32 rows).
    pub fn traffic_bytes(&self) -> u64 {
        (self.batch * self.n * 2 * 2) as u64 * 4
    }
}

impl BenchmarkPattern for FftPattern {
    fn name(&self) -> &str {
        "fft"
    }

    fn tier(&self) -> Tier {
        Tier::L1
    }

    fn build_graph(&self) -> Graph {
        let mut g = Graph::new("fft_bench");
        let x = fft_const_block(&mut g, self.batch, self.n);
        let y = g.fft(x, self.inverse);
        g.set_outputs(vec![y]);
        g
    }

    fn input_data(&self) -> Vec<(String, Vec<f32>)> {
        Vec::new()
    }
}

fn fft_const_block(g: &mut Graph, batch: usize, n: usize) -> NodeId {
    let len = batch * n * 2;
    let mut data: Vec<f32> = Vec::with_capacity(len);
    for i in 0..len {
        data.push((i as f32 * 0.013).sin());
    }
    let mut bytes = Vec::with_capacity(len * 4);
    for v in &data {
        bytes.extend_from_slice(&v.to_le_bytes());
    }
    g.add_node(
        Op::Constant { data: bytes },
        vec![],
        Shape::new(&[batch, n * 2], DType::F32),
    )
}