tokitai-operator 0.1.0

//! `dev_perf_smoke` — a lightweight perf smoke test for 5 small ops.
//!
//! Runs AddOp on 1K i64 elements, MatmulOp on a 32x32 pair, ReshapeOp
//! on a 4x4 -> 16 row-major copy, ReluOp on 512 elements, and
//! SoftmaxOp on a 1x8 row. Each op is timed with `std::time::Instant`
//! and printed in the greppable form `op=<name> us=<microseconds>`.
//!
//! The output is informational only. The timings are NOT
//! cross-machine comparable; do not enable the
//! `production_speedup_claim_allowed` flag from this output.
//! For real benchmarks, see `benches/` (criterion groups) and
//! the `paper_benchmarks` example.
//!
//! Run with:
//!   cargo run --offline --example dev_perf_smoke

use std::time::Instant;

use tokitai_operator::domain::DomainId;
use tokitai_operator::facade::Tokitai;
use tokitai_operator::object::{Shape, Tensor};
use tokitai_operator::op::{AddOp, MatmulOp, ReluOp, ReshapeOp, SoftmaxOp};

fn int_tensor(shape: Vec<usize>, data: Vec<i64>) -> Tensor<i64> {
    Tensor::dense_cpu(DomainId::new("integer"), Shape::from(shape), data)
}

fn time_op<F: FnMut() -> tokitai_operator::Result<()>>(name: &str, mut f: F) {
    let t0 = Instant::now();
    if let Err(e) = f() {
        eprintln!("  op={name} FAILED: {e}");
        return;
    }
    let elapsed = t0.elapsed();
    let us = elapsed.as_micros();
    // Greppable one-line format.
    println!("  op={name} us={us}");
}

fn smoke_add(cpu: &Tokitai) -> tokitai_operator::Result<()> {
    let n = 1024usize;
    let a = int_tensor(vec![n], (0..n as i64).collect());
    let b = int_tensor(vec![n], (0..n as i64).map(|x| x * 2).collect());
    time_op("AddOp[1K i64]", || {
        let _g = cpu.graph(|g| {
            let a = g.input_tensor(a.clone())?;
            let b = g.input_tensor(b.clone())?;
            g.op(AddOp, &[a, b])
        })?;
        Ok(())
    });
    Ok(())
}

fn smoke_matmul(cpu: &Tokitai) -> tokitai_operator::Result<()> {
    let n = 32usize;
    let a_data: Vec<i64> = (0..(n * n) as i64).collect();
    let b_data: Vec<i64> = (0..(n * n) as i64).map(|x| x % 7).collect();
    let a = int_tensor(vec![n, n], a_data);
    let b = int_tensor(vec![n, n], b_data);
    time_op("MatmulOp[32x32 i64]", || {
        let _g = cpu.graph(|g| {
            let a = g.input_tensor(a.clone())?;
            let b = g.input_tensor(b.clone())?;
            g.op(MatmulOp, &[a, b])
        })?;
        Ok(())
    });
    Ok(())
}

fn smoke_reshape(cpu: &Tokitai) -> tokitai_operator::Result<()> {
    let data: Vec<i64> = (0..16i64).collect();
    let src = int_tensor(vec![4, 4], data);
    // ReshapeOp takes 2 inputs: the data and a 1-D target-shape tensor.
    // The target-shape tensor stores the new shape in its *meta* dims
    // (parse_target_shape reads `meta.shape.dims`); its data vec is
    // ignored at planning time.
    let target = int_tensor(vec![16], vec![]);
    time_op("ReshapeOp[4x4->16]", || {
        let _g = cpu.graph(|g| {
            let src = g.input_tensor(src.clone())?;
            let tgt = g.input_tensor(target.clone())?;
            g.op(ReshapeOp, &[src, tgt])
        })?;
        Ok(())
    });
    Ok(())
}

fn smoke_relu(cpu: &Tokitai) -> tokitai_operator::Result<()> {
    let n = 512usize;
    let data: Vec<i64> = (0..n as i64).map(|x| x - 256).collect();
    let src = int_tensor(vec![n], data);
    time_op("ReluOp[512 i64]", || {
        let _g = cpu.graph(|g| {
            let src = g.input_tensor(src.clone())?;
            g.op(ReluOp, &[src])
        })?;
        Ok(())
    });
    Ok(())
}

fn smoke_softmax(cpu: &Tokitai) -> tokitai_operator::Result<()> {
    let data: Vec<i64> = vec![100, 200, 300, 400, 500, 600, 700, 800];
    let src = int_tensor(vec![1, 8], data);
    time_op("SoftmaxOp[1x8 i64]", || {
        let _g = cpu.graph(|g| {
            let src = g.input_tensor(src.clone())?;
            g.op(SoftmaxOp, &[src])
        })?;
        Ok(())
    });
    Ok(())
}

fn main() -> tokitai_operator::Result<()> {
    println!("== dev_perf_smoke (informational; not a cross-machine benchmark) ==");
    let cpu = Tokitai::cpu_only();
    println!("-- timings --");
    smoke_add(&cpu)?;
    smoke_matmul(&cpu)?;
    smoke_reshape(&cpu)?;
    smoke_relu(&cpu)?;
    smoke_softmax(&cpu)?;
    println!("-- notes --");
    println!("  - this example does NOT enable production_speedup_claim_allowed.");
    println!("  - timings are single-shot; for stable numbers use `benches/`.");
    println!("== dev_perf_smoke ok ==");
    Ok(())
}