Struct CudaExecutable 

pub struct CudaExecutable { /* private fields */ }

Implementations

impl CudaExecutable

pub fn compile(graph: Graph) -> Self

JIT compile, stream-mode execution. Default entry point.

Honors RLX_CUDA_COMPILE_MODE=aot and RLX_CUDA_EXEC_MODE=graph|multistream:N.

Examples found in repository

examples/bench_matmul.rs (line 33)

fn bench(m: usize, k: usize, n: usize, warmup: usize, iters: usize) {
    let mut g = Graph::new("mm");
    let x = g.input("x", Shape::new(&[m, k], DType::F32));
    let w = g.param("w", Shape::new(&[k, n], DType::F32));
    let y = g.matmul(x, w, Shape::new(&[m, n], DType::F32));
    g.set_outputs(vec![y]);

    let mut exe = CudaExecutable::compile(g);
    let wv: Vec<f32> = (0..k * n).map(|i| (i as f32) * 1e-3).collect();
    exe.set_param("w", &wv);
    let xv: Vec<f32> = (0..m * k).map(|i| (i as f32) * 1e-3).collect();

    for _ in 0..warmup {
        let _ = exe.run(&[("x", &xv)]);
    }

    let t0 = Instant::now();
    for _ in 0..iters {
        let _ = exe.run(&[("x", &xv)]);
    }
    let dt = t0.elapsed().as_secs_f64() / iters as f64;
    let flops = 2.0 * (m * k * n) as f64;
    let gflops = flops / dt / 1e9;
    println!(
        "  M={:>5} K={:>5} N={:>5}   {:>8.3} ms   {:>8.1} GFLOP/s",
        m,
        k,
        n,
        dt * 1e3,
        gflops
    );
}

pub fn eager(graph: Graph, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>>

One-shot eager run. Compiles, executes once with the given inputs, and drops the executable. No persistent state.

pub fn compile_with( graph: Graph, compile_mode: CompileMode, exec_mode: ExecMode, ) -> Self

Full constructor with explicit compile + exec modes.

pub fn arena_ptr(&self) -> *const u8

Host buffer base for reading outputs after Self::run_slots. Offsets in the returned slot pairs are byte offsets into this buffer.

pub fn output_slots(&self) -> &[(usize, usize)]

pub fn run_slots(&mut self, inputs: &[&[f32]]) -> &[(usize, usize)]

Fast path: positional inputs, D2H into Self::host_arena, no per-output Vec.

pub fn set_active_extent(&mut self, extent: Option<(usize, usize)>)

Hint the next run to process only the first actual rows along the bucket axis (out of upper, the compile extent). Honored when every step in the schedule passes Step::safe_for_active_extent. Bypasses captured CUDA Graph (recorded at full extent) when active. See PLAN L1.

pub fn output_dtypes(&self) -> Vec<DType>

Declared graph-output dtypes, in graph.outputs order. Used by the runtime wrapper’s run_typed to narrow f32 outputs back to the declared dtype on the way out.

pub fn set_param(&mut self, name: &str, data: &[f32])

Examples found in repository

examples/bench_matmul.rs (line 35)

fn bench(m: usize, k: usize, n: usize, warmup: usize, iters: usize) {
    let mut g = Graph::new("mm");
    let x = g.input("x", Shape::new(&[m, k], DType::F32));
    let w = g.param("w", Shape::new(&[k, n], DType::F32));
    let y = g.matmul(x, w, Shape::new(&[m, n], DType::F32));
    g.set_outputs(vec![y]);

    let mut exe = CudaExecutable::compile(g);
    let wv: Vec<f32> = (0..k * n).map(|i| (i as f32) * 1e-3).collect();
    exe.set_param("w", &wv);
    let xv: Vec<f32> = (0..m * k).map(|i| (i as f32) * 1e-3).collect();

    for _ in 0..warmup {
        let _ = exe.run(&[("x", &xv)]);
    }

    let t0 = Instant::now();
    for _ in 0..iters {
        let _ = exe.run(&[("x", &xv)]);
    }
    let dt = t0.elapsed().as_secs_f64() / iters as f64;
    let flops = 2.0 * (m * k * n) as f64;
    let gflops = flops / dt / 1e9;
    println!(
        "  M={:>5} K={:>5} N={:>5}   {:>8.3} ms   {:>8.1} GFLOP/s",
        m,
        k,
        n,
        dt * 1e3,
        gflops
    );
}

pub fn set_param_bytes(&mut self, name: &str, data: &[u8])

Upload packed U8/I8 GGUF weights into the param slot (byte offset).

pub fn set_param_half(&mut self, name: &str, dtype: HalfDtype, bits: &[u16])

Upload a param as packed half-precision bits (u16 per element). Caller passes the raw IEEE-754 binary16 (F16) or BFloat16 (Bf16) bit pattern; the backend stores it in the half-arena side-buffer and skips the f32 slot entirely. Use cases: 2× weight-memory savings for inference, plus Tensor Core matmul via cublasGemmEx when both A and B (or just B) are stored half-precision.

When the same name is also set_param’d as f32, the half-arena entry takes precedence in the matmul dispatch. Use only one of the two for any given param.

pub fn run(&mut self, inputs: &[(&str, &[f32])]) -> Vec<Vec<f32>>

Examples found in repository

examples/bench_matmul.rs (line 39)

fn bench(m: usize, k: usize, n: usize, warmup: usize, iters: usize) {
    let mut g = Graph::new("mm");
    let x = g.input("x", Shape::new(&[m, k], DType::F32));
    let w = g.param("w", Shape::new(&[k, n], DType::F32));
    let y = g.matmul(x, w, Shape::new(&[m, n], DType::F32));
    g.set_outputs(vec![y]);

    let mut exe = CudaExecutable::compile(g);
    let wv: Vec<f32> = (0..k * n).map(|i| (i as f32) * 1e-3).collect();
    exe.set_param("w", &wv);
    let xv: Vec<f32> = (0..m * k).map(|i| (i as f32) * 1e-3).collect();

    for _ in 0..warmup {
        let _ = exe.run(&[("x", &xv)]);
    }

    let t0 = Instant::now();
    for _ in 0..iters {
        let _ = exe.run(&[("x", &xv)]);
    }
    let dt = t0.elapsed().as_secs_f64() / iters as f64;
    let flops = 2.0 * (m * k * n) as f64;
    let gflops = flops / dt / 1e9;
    println!(
        "  M={:>5} K={:>5} N={:>5}   {:>8.3} ms   {:>8.1} GFLOP/s",
        m,
        k,
        n,
        dt * 1e3,
        gflops
    );
}

pub fn run_read_outputs( &mut self, inputs: &[(&str, &[f32])], read_indices: Option<&[usize]>, ) -> Vec<Vec<f32>>

Run and read back only selected outputs (+ GPU handle feed outputs).

pub fn bind_gpu_handle(&mut self, name: &str, data: &[f32]) -> bool

pub fn has_gpu_handle(&self, name: &str) -> bool

pub fn set_gpu_handle_feed(&mut self, handle_name: &str, output_index: usize)

pub fn read_gpu_handle(&self, name: &str) -> Option<Vec<f32>>