Struct Specialization 

pub struct Specialization { /* private fields */ }

One export and its exact configuration within a module.

Implementations

impl Specialization

pub fn new(symbol: impl Into<String>) -> Self

Select an export with no configuration overrides.

Examples found in repository

examples/compile_npu.rs (line 12)

fn main() -> Result<()> {
    let packets = std::env::args().nth(1).unwrap_or_else(|| "256".into());
    let artifact = Compiler::for_target(None, &Target::xdna())?
        .module(include_str!("../tests/kernels/copy.xdna.loom"))
        .compile(
            &Specialization::new("copy")
                .with_config("copy.packets", packets)
                .with_report(ReportMode::Summary),
        )?;
    println!("{}", artifact.path().display());
    Ok(())
}

examples/gemm_pipeline.rs (line 33)

fn gpu_kernel(
    runtime: &Runtime,
    source: &str,
    symbol: &str,
    config: &[(&str, String)],
    elements: usize,
    bindings: Vec<BindingContract>,
) -> Result<GpuKernel> {
    let compiler = hrx::loom::Compiler::for_target(None, runtime.gpu()?.target())?;
    let module = compiler.module(source);
    let mut spec = hrx::loom::Specialization::new(symbol);
    for (key, value) in config {
        spec.set_config(*key, value.clone());
    }
    let artifact = module.compile(&spec)?;
    let stream = hrx::gpu::Stream::open()?;
    let raw = unsafe { stream.load_artifact(&artifact) }?;
    let mut constants = hrx::gpu::Constants::new();
    match raw.info().constant_byte_length {
        8 => constants.push(elements as u32)?,
        12 => constants.push(elements as u64)?,
        _ => return Err(hrx::Error::Message("unexpected scalar ABI".into())),
    }
    constants.push(1f32)?;
    unsafe {
        runtime.load_gpu_kernel(
            artifact.path(),
            symbol,
            [elements.div_ceil(256) as u32, 1, 1],
            [256, 1, 1],
            KernelContract {
                bindings,
                constants: constants.as_bytes().to_vec(),
            },
        )
    }
}
fn main() -> Result<()> {
    let args: Vec<_> = std::env::args().collect();
    let args = if args.len() == 1 {
        vec![
            args[0].clone(),
            String::new(),
            "gemm".into(),
            "8".into(),
            "8".into(),
            "8".into(),
        ]
    } else {
        args
    };
    if args.len() != 6 {
        return Err(hrx::Error::Message(
            "usage: gemm_pipeline <gemm.loom> <symbol> M K N".into(),
        ));
    }
    let parse = |i: usize| {
        args[i]
            .parse::<usize>()
            .map_err(|e| hrx::Error::Message(e.to_string()))
    };
    let (m, k, n) = (parse(3)?, parse(4)?, parse(5)?);
    let extent = |a: usize, b: usize, width: usize| {
        a.checked_mul(b)
            .filter(|&count| count > 0 && count <= (1 << 30))
            .and_then(|count| count.checked_mul(width))
            .ok_or_else(|| hrx::Error::Message("invalid matrix extent".into()))
    };
    let (a_bytes, b_bytes, c_bytes) = (extent(m, k, 2)?, extent(k, n, 2)?, extent(m, n, 4)?);
    let runtime = Runtime::new()?;
    let device = runtime.npu(0)?;
    let (source, spec) = if args[1].is_empty() {
        (
            include_str!("../tests/kernels/gemm_bf16.xdna.loom").to_string(),
            hrx::loom::Specialization::new("gemm"),
        )
    } else {
        (
            std::fs::read_to_string(&args[1])?,
            hrx::loom::Specialization::new(&args[2])
                .with_config("gemm.m", m.to_string())
                .with_config("gemm.k", k.to_string())
                .with_config("gemm.n", n.to_string()),
        )
    };
    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
        .module(&source)
        .compile(&spec)?;
    let npu = unsafe {
        device.load_artifact(
            &artifact,
            1,
            KernelContract {
                bindings: vec![
                    binding(a_bytes, Access::Read),
                    binding(b_bytes, Access::Read),
                    binding(c_bytes, Access::Write),
                ],
                constants: vec![],
            },
        )
    }?;
    let prep = gpu_kernel(
        &runtime,
        include_str!("../tests/kernels/euler.loom"),
        "krea2_euler",
        &[
            ("krea2.euler.grid_x", (m * k).div_ceil(256).to_string()),
            ("krea2.euler.grid_y", "1".into()),
        ],
        m * k,
        vec![
            binding(a_bytes, Access::ReadWrite),
            binding(a_bytes, Access::Read),
        ],
    )?;
    let epilogue = gpu_kernel(
        &runtime,
        include_str!("../tests/kernels/add_f32.loom"),
        "add_f32",
        &[("add.grid", (m * n).div_ceil(256).to_string())],
        m * n,
        vec![binding(c_bytes, Access::ReadWrite)],
    )?;
    let allocate = |bytes| runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()));
    let weights = allocate(b_bytes)?;
    let ones = allocate(a_bytes)?;
    for buffer in [&weights, &ones] {
        for word in buffer.map_write()?.as_chunks_mut::<2>().0 {
            word.copy_from_slice(&0x3f80u16.to_le_bytes());
        }
    }
    let prepare = || -> Result<(ExecutableGraph, Buffer)> {
        let a = allocate(a_bytes)?;
        let c = allocate(c_bytes)?;
        let mut graph = runtime.graph();
        graph.fill(a.view(), 0)?;
        graph.gpu(&prep, &[a.view(), ones.view()])?;
        graph.npu(&npu, &[a.view(), weights.view(), c.view()])?;
        graph.gpu(&epilogue, &[c.view()])?;
        Ok((graph.prepare()?, c))
    };
    let (first, c0) = prepare()?;
    let (second, c1) = prepare()?;
    for _ in 0..5 {
        first.submit()?.wait()?;
        second.submit()?.wait()?;
    }
    let before = runtime.statistics();
    let mut latency = Vec::with_capacity(50);
    for _ in 0..50 {
        let start = Instant::now();
        first.submit()?.wait()?;
        latency.push(start.elapsed().as_secs_f64());
    }
    let start = Instant::now();
    for _ in 0..50 {
        let a = first.submit()?;
        let b = second.submit()?;
        a.wait()?;
        b.wait()?;
    }
    let requests_per_second = 100.0 / start.elapsed().as_secs_f64();
    for c in [&c0, &c1] {
        if c.map_read()?
            .as_chunks::<4>()
            .0
            .iter()
            .any(|v| f32::from_le_bytes(*v) != k as f32 + 1.0)
        {
            return Err(hrx::Error::Message("GEMM pipeline output mismatch".into()));
        }
    }
    let after = runtime.statistics();
    assert_eq!(before.allocations, after.allocations);
    assert_eq!(before.imports, after.imports);
    assert_eq!(after.copied_bytes, 0);
    latency.sort_by(f64::total_cmp);
    println!(
        "M={m} K={k} N={n} p50_ms={:.3} p95_ms={:.3} pipelined_requests_s={requests_per_second:.2}",
        percentile(&latency, 50) * 1e3,
        percentile(&latency, 95) * 1e3
    );
    println!("{}", serde_json::to_string(&after)?);
    Ok(())
}

examples/startup_bench.rs (line 50)

fn main() -> Result<()> {
    let compiler_only = match std::env::args().nth(1).as_deref() {
        None => false,
        Some("--compiler-only") => true,
        Some(_) => {
            return Err(hrx::Error::Message(
                "usage: startup_bench [--compiler-only]".into(),
            ));
        }
    };
    let initial = memory()?;
    let start = Instant::now();
    let stream = if compiler_only {
        None
    } else {
        Some(Stream::open()?)
    };
    let open_ms = stream
        .as_ref()
        .map(|_| start.elapsed().as_secs_f64() * 1000.0);
    let after_open = if compiler_only { None } else { Some(memory()?) };
    let start = Instant::now();
    let compiler = match &stream {
        Some(stream) => hrx::loom::Compiler::for_stream(None, stream)?,
        None => hrx::loom::Compiler::resolve(None)?,
    };
    let compiler_ms = start.elapsed().as_secs_f64() * 1000.0;
    let after_compiler = memory()?;
    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
    let mut spec = Specialization::new("krea2_euler");
    spec.set_config("krea2.euler.grid_x", "1");
    spec.set_config("krea2.euler.grid_y", "1");
    let start = Instant::now();
    let artifact = module.compile(&spec)?;
    let first_compile_ms = start.elapsed().as_secs_f64() * 1000.0;
    let after_compile = memory()?;
    let start = Instant::now();
    let cached = module.compile(&spec)?;
    let cache_hit_ms = start.elapsed().as_secs_f64() * 1000.0;
    assert_eq!(artifact.bytes(), cached.bytes());
    println!(
        "{}",
        json!({"compiler_only": compiler_only, "open_ms": open_ms, "compiler_ms": compiler_ms,
            "first_compile_ms": first_compile_ms, "cache_hit_ms": cache_hit_ms,
            "initial": initial, "after_open": after_open,
            "after_compiler": after_compiler, "after_compile": after_compile,
            "compiler_identity": artifact.compiler_identity()})
    );
    Ok(())
}

examples/shared_roundtrip.rs (line 31)

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args().collect();
    if args.len() > 2 {
        return Err(hrx::Error::Message(
            "usage: shared_roundtrip [bytes]".into(),
        ));
    }
    let bytes = args
        .get(1)
        .map(String::as_str)
        .unwrap_or("1048576")
        .parse::<usize>()
        .map_err(|e| hrx::Error::Message(e.to_string()))?;
    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
        return Err(hrx::Error::Message(
            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
        ));
    }
    let runtime = Runtime::new()?;
    let device = runtime.npu(0)?;
    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
        .module(include_str!("../tests/kernels/copy.xdna.loom"))
        .compile(
            &hrx::loom::Specialization::new("copy")
                .with_config("copy.packets", (bytes / 4096).to_string()),
        )?;
    let binding = |bytes, access| BindingContract {
        bytes,
        access,
        alignment: 4,
        layout: "u32 passthrough".into(),
    };
    let contract = KernelContract {
        bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
        constants: vec![],
    };
    let kernel = unsafe { device.load_artifact(&artifact, 1, contract) }?;
    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let result = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let mut graph = runtime.graph();
    graph.fill(a.view(), 0x5a)?;
    graph.npu(&kernel, &[a.view(), c.view()])?;
    graph.copy(result.view(), c.view())?;
    let graph = graph.prepare()?;
    for _ in 0..10 {
        graph.submit()?.wait()?;
        let mapped = result.map_read()?;
        if mapped.iter().any(|&b| b != 0x5a) {
            return Err(hrx::Error::Message("GPU/NPU roundtrip mismatch".into()));
        }
    }
    println!(
        "GPU -> NPU -> GPU: all {bytes} bytes correct, 10 executions, {} prepared regions",
        graph.regions()
    );
    Ok(())
}

examples/shared_bench.rs (line 30)

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args().collect();
    if args.len() > 2 {
        return Err(hrx::Error::Message("usage: shared_bench [bytes]".into()));
    }
    let bytes = args
        .get(1)
        .map(String::as_str)
        .unwrap_or("1048576")
        .parse::<usize>()
        .map_err(|e| hrx::Error::Message(e.to_string()))?;
    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
        return Err(hrx::Error::Message(
            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
        ));
    }
    let runtime = Runtime::new()?;
    let device = runtime.npu(0)?;
    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
        .module(include_str!("../tests/kernels/copy.xdna.loom"))
        .compile(
            &hrx::loom::Specialization::new("copy")
                .with_config("copy.packets", (bytes / 4096).to_string()),
        )?;
    let binding = |bytes, access| BindingContract {
        bytes,
        access,
        alignment: 4,
        layout: "passthrough".into(),
    };
    let kernel = unsafe {
        device.load_artifact(
            &artifact,
            1,
            KernelContract {
                bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
                constants: vec![],
            },
        )
    }?;
    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let mut graph = runtime.graph();
    graph.fill(a.view(), 0x35)?;
    graph.npu(&kernel, &[a.view(), c.view()])?;
    let graph = graph.prepare()?;
    for _ in 0..5 {
        graph.submit()?.wait()?;
    }
    let mut samples = Vec::new();
    for _ in 0..100 {
        let start = Instant::now();
        graph.submit()?.wait()?;
        samples.push(start.elapsed().as_secs_f64());
    }
    if c.map_read()?.iter().any(|&b| b != 0x35) {
        return Err(hrx::Error::Message("incorrect benchmark output".into()));
    }
    samples.sort_by(f64::total_cmp);
    println!(
        "bytes={bytes} p50_us={:.2} p95_us={:.2}",
        percentile(&samples, 50) * 1e6,
        percentile(&samples, 95) * 1e6
    );
    Ok(())
}

examples/mall_probe.rs (line 27)

fn main() -> Result<()> {
    let seconds: f64 = std::env::args()
        .nth(1)
        .unwrap_or_else(|| "2".into())
        .parse()
        .map_err(|e: std::num::ParseFloatError| hrx::Error::Message(e.to_string()))?;
    if !seconds.is_finite() || !(0.01..=3600.0).contains(&seconds) {
        return Err(hrx::Error::Message("seconds must be in 0.01..=3600".into()));
    }
    let device = Device::open(Engine::Xdna, 0)?;
    let compiler = Compiler::for_target(None, device.target())?;
    println!("working_set_MiB GB/s dispatches/s");
    for mib in [1, 4, 8, 16, 32, 64] {
        let bytes = mib * 1024 * 1024;
        let artifact = compiler
            .module(include_str!("../tests/kernels/copy.xdna.loom"))
            .compile(
                &Specialization::new("copy")
                    .with_config("copy.packets", (bytes / 4096).to_string()),
            )?;
        let input = device
            .fabric()
            .allocate(bytes, std::slice::from_ref(&device))?;
        let output = device
            .fabric()
            .allocate(bytes, std::slice::from_ref(&device))?;
        input.write(0, &vec![0x5a; bytes])?;
        let program = unsafe {
            device.prepare_xdna(
                &artifact,
                1,
                &[
                    XdnaBinding {
                        buffer: &input,
                        offset: 0,
                        length: bytes,
                    },
                    XdnaBinding {
                        buffer: &output,
                        offset: 0,
                        length: bytes,
                    },
                ],
            )
        }?;
        unsafe { program.dispatch() }?.wait()?;
        let mut check = vec![0; bytes];
        output.read(0, &mut check)?;
        assert!(check.iter().all(|&b| b == 0x5a));
        let start = Instant::now();
        let mut count = 0;
        while start.elapsed() < Duration::from_secs_f64(seconds) {
            unsafe { program.dispatch() }?.wait()?;
            count += 1;
        }
        let elapsed = start.elapsed().as_secs_f64();
        println!(
            "{} {:.3} {:.1}",
            2 * mib,
            2.0 * bytes as f64 * count as f64 / elapsed / 1e9,
            count as f64 / elapsed
        );
    }
    Ok(())
}

Additional examples can be found in:

• examples/patch_bench.rs
• examples/stream_bench.rs
• examples/gpu_npu_parallel.rs

pub fn with_config( self, key: impl Into<String>, value: impl Into<String>, ) -> Self

Set one fully-qualified Loom configuration value.

Examples found in repository

examples/compile_npu.rs (line 13)

fn main() -> Result<()> {
    let packets = std::env::args().nth(1).unwrap_or_else(|| "256".into());
    let artifact = Compiler::for_target(None, &Target::xdna())?
        .module(include_str!("../tests/kernels/copy.xdna.loom"))
        .compile(
            &Specialization::new("copy")
                .with_config("copy.packets", packets)
                .with_report(ReportMode::Summary),
        )?;
    println!("{}", artifact.path().display());
    Ok(())
}

examples/shared_roundtrip.rs (line 32)

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args().collect();
    if args.len() > 2 {
        return Err(hrx::Error::Message(
            "usage: shared_roundtrip [bytes]".into(),
        ));
    }
    let bytes = args
        .get(1)
        .map(String::as_str)
        .unwrap_or("1048576")
        .parse::<usize>()
        .map_err(|e| hrx::Error::Message(e.to_string()))?;
    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
        return Err(hrx::Error::Message(
            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
        ));
    }
    let runtime = Runtime::new()?;
    let device = runtime.npu(0)?;
    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
        .module(include_str!("../tests/kernels/copy.xdna.loom"))
        .compile(
            &hrx::loom::Specialization::new("copy")
                .with_config("copy.packets", (bytes / 4096).to_string()),
        )?;
    let binding = |bytes, access| BindingContract {
        bytes,
        access,
        alignment: 4,
        layout: "u32 passthrough".into(),
    };
    let contract = KernelContract {
        bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
        constants: vec![],
    };
    let kernel = unsafe { device.load_artifact(&artifact, 1, contract) }?;
    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let result = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let mut graph = runtime.graph();
    graph.fill(a.view(), 0x5a)?;
    graph.npu(&kernel, &[a.view(), c.view()])?;
    graph.copy(result.view(), c.view())?;
    let graph = graph.prepare()?;
    for _ in 0..10 {
        graph.submit()?.wait()?;
        let mapped = result.map_read()?;
        if mapped.iter().any(|&b| b != 0x5a) {
            return Err(hrx::Error::Message("GPU/NPU roundtrip mismatch".into()));
        }
    }
    println!(
        "GPU -> NPU -> GPU: all {bytes} bytes correct, 10 executions, {} prepared regions",
        graph.regions()
    );
    Ok(())
}

examples/shared_bench.rs (line 31)

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args().collect();
    if args.len() > 2 {
        return Err(hrx::Error::Message("usage: shared_bench [bytes]".into()));
    }
    let bytes = args
        .get(1)
        .map(String::as_str)
        .unwrap_or("1048576")
        .parse::<usize>()
        .map_err(|e| hrx::Error::Message(e.to_string()))?;
    if bytes == 0 || bytes % 4096 != 0 || bytes > 64 * 1024 * 1024 {
        return Err(hrx::Error::Message(
            "bytes must be a multiple of 4096 in 4096..=67108864".into(),
        ));
    }
    let runtime = Runtime::new()?;
    let device = runtime.npu(0)?;
    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
        .module(include_str!("../tests/kernels/copy.xdna.loom"))
        .compile(
            &hrx::loom::Specialization::new("copy")
                .with_config("copy.packets", (bytes / 4096).to_string()),
        )?;
    let binding = |bytes, access| BindingContract {
        bytes,
        access,
        alignment: 4,
        layout: "passthrough".into(),
    };
    let kernel = unsafe {
        device.load_artifact(
            &artifact,
            1,
            KernelContract {
                bindings: vec![binding(bytes, Access::Read), binding(bytes, Access::Write)],
                constants: vec![],
            },
        )
    }?;
    let a = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let c = runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()))?;
    let mut graph = runtime.graph();
    graph.fill(a.view(), 0x35)?;
    graph.npu(&kernel, &[a.view(), c.view()])?;
    let graph = graph.prepare()?;
    for _ in 0..5 {
        graph.submit()?.wait()?;
    }
    let mut samples = Vec::new();
    for _ in 0..100 {
        let start = Instant::now();
        graph.submit()?.wait()?;
        samples.push(start.elapsed().as_secs_f64());
    }
    if c.map_read()?.iter().any(|&b| b != 0x35) {
        return Err(hrx::Error::Message("incorrect benchmark output".into()));
    }
    samples.sort_by(f64::total_cmp);
    println!(
        "bytes={bytes} p50_us={:.2} p95_us={:.2}",
        percentile(&samples, 50) * 1e6,
        percentile(&samples, 95) * 1e6
    );
    Ok(())
}

examples/mall_probe.rs (line 28)

fn main() -> Result<()> {
    let seconds: f64 = std::env::args()
        .nth(1)
        .unwrap_or_else(|| "2".into())
        .parse()
        .map_err(|e: std::num::ParseFloatError| hrx::Error::Message(e.to_string()))?;
    if !seconds.is_finite() || !(0.01..=3600.0).contains(&seconds) {
        return Err(hrx::Error::Message("seconds must be in 0.01..=3600".into()));
    }
    let device = Device::open(Engine::Xdna, 0)?;
    let compiler = Compiler::for_target(None, device.target())?;
    println!("working_set_MiB GB/s dispatches/s");
    for mib in [1, 4, 8, 16, 32, 64] {
        let bytes = mib * 1024 * 1024;
        let artifact = compiler
            .module(include_str!("../tests/kernels/copy.xdna.loom"))
            .compile(
                &Specialization::new("copy")
                    .with_config("copy.packets", (bytes / 4096).to_string()),
            )?;
        let input = device
            .fabric()
            .allocate(bytes, std::slice::from_ref(&device))?;
        let output = device
            .fabric()
            .allocate(bytes, std::slice::from_ref(&device))?;
        input.write(0, &vec![0x5a; bytes])?;
        let program = unsafe {
            device.prepare_xdna(
                &artifact,
                1,
                &[
                    XdnaBinding {
                        buffer: &input,
                        offset: 0,
                        length: bytes,
                    },
                    XdnaBinding {
                        buffer: &output,
                        offset: 0,
                        length: bytes,
                    },
                ],
            )
        }?;
        unsafe { program.dispatch() }?.wait()?;
        let mut check = vec![0; bytes];
        output.read(0, &mut check)?;
        assert!(check.iter().all(|&b| b == 0x5a));
        let start = Instant::now();
        let mut count = 0;
        while start.elapsed() < Duration::from_secs_f64(seconds) {
            unsafe { program.dispatch() }?.wait()?;
            count += 1;
        }
        let elapsed = start.elapsed().as_secs_f64();
        println!(
            "{} {:.3} {:.1}",
            2 * mib,
            2.0 * bytes as f64 * count as f64 / elapsed / 1e9,
            count as f64 / elapsed
        );
    }
    Ok(())
}

examples/gemm_pipeline.rs (line 103)

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args().collect();
    let args = if args.len() == 1 {
        vec![
            args[0].clone(),
            String::new(),
            "gemm".into(),
            "8".into(),
            "8".into(),
            "8".into(),
        ]
    } else {
        args
    };
    if args.len() != 6 {
        return Err(hrx::Error::Message(
            "usage: gemm_pipeline <gemm.loom> <symbol> M K N".into(),
        ));
    }
    let parse = |i: usize| {
        args[i]
            .parse::<usize>()
            .map_err(|e| hrx::Error::Message(e.to_string()))
    };
    let (m, k, n) = (parse(3)?, parse(4)?, parse(5)?);
    let extent = |a: usize, b: usize, width: usize| {
        a.checked_mul(b)
            .filter(|&count| count > 0 && count <= (1 << 30))
            .and_then(|count| count.checked_mul(width))
            .ok_or_else(|| hrx::Error::Message("invalid matrix extent".into()))
    };
    let (a_bytes, b_bytes, c_bytes) = (extent(m, k, 2)?, extent(k, n, 2)?, extent(m, n, 4)?);
    let runtime = Runtime::new()?;
    let device = runtime.npu(0)?;
    let (source, spec) = if args[1].is_empty() {
        (
            include_str!("../tests/kernels/gemm_bf16.xdna.loom").to_string(),
            hrx::loom::Specialization::new("gemm"),
        )
    } else {
        (
            std::fs::read_to_string(&args[1])?,
            hrx::loom::Specialization::new(&args[2])
                .with_config("gemm.m", m.to_string())
                .with_config("gemm.k", k.to_string())
                .with_config("gemm.n", n.to_string()),
        )
    };
    let artifact = hrx::loom::Compiler::for_target(None, device.target())?
        .module(&source)
        .compile(&spec)?;
    let npu = unsafe {
        device.load_artifact(
            &artifact,
            1,
            KernelContract {
                bindings: vec![
                    binding(a_bytes, Access::Read),
                    binding(b_bytes, Access::Read),
                    binding(c_bytes, Access::Write),
                ],
                constants: vec![],
            },
        )
    }?;
    let prep = gpu_kernel(
        &runtime,
        include_str!("../tests/kernels/euler.loom"),
        "krea2_euler",
        &[
            ("krea2.euler.grid_x", (m * k).div_ceil(256).to_string()),
            ("krea2.euler.grid_y", "1".into()),
        ],
        m * k,
        vec![
            binding(a_bytes, Access::ReadWrite),
            binding(a_bytes, Access::Read),
        ],
    )?;
    let epilogue = gpu_kernel(
        &runtime,
        include_str!("../tests/kernels/add_f32.loom"),
        "add_f32",
        &[("add.grid", (m * n).div_ceil(256).to_string())],
        m * n,
        vec![binding(c_bytes, Access::ReadWrite)],
    )?;
    let allocate = |bytes| runtime.allocate(bytes, MemoryPlacement::Shared(device.clone()));
    let weights = allocate(b_bytes)?;
    let ones = allocate(a_bytes)?;
    for buffer in [&weights, &ones] {
        for word in buffer.map_write()?.as_chunks_mut::<2>().0 {
            word.copy_from_slice(&0x3f80u16.to_le_bytes());
        }
    }
    let prepare = || -> Result<(ExecutableGraph, Buffer)> {
        let a = allocate(a_bytes)?;
        let c = allocate(c_bytes)?;
        let mut graph = runtime.graph();
        graph.fill(a.view(), 0)?;
        graph.gpu(&prep, &[a.view(), ones.view()])?;
        graph.npu(&npu, &[a.view(), weights.view(), c.view()])?;
        graph.gpu(&epilogue, &[c.view()])?;
        Ok((graph.prepare()?, c))
    };
    let (first, c0) = prepare()?;
    let (second, c1) = prepare()?;
    for _ in 0..5 {
        first.submit()?.wait()?;
        second.submit()?.wait()?;
    }
    let before = runtime.statistics();
    let mut latency = Vec::with_capacity(50);
    for _ in 0..50 {
        let start = Instant::now();
        first.submit()?.wait()?;
        latency.push(start.elapsed().as_secs_f64());
    }
    let start = Instant::now();
    for _ in 0..50 {
        let a = first.submit()?;
        let b = second.submit()?;
        a.wait()?;
        b.wait()?;
    }
    let requests_per_second = 100.0 / start.elapsed().as_secs_f64();
    for c in [&c0, &c1] {
        if c.map_read()?
            .as_chunks::<4>()
            .0
            .iter()
            .any(|v| f32::from_le_bytes(*v) != k as f32 + 1.0)
        {
            return Err(hrx::Error::Message("GEMM pipeline output mismatch".into()));
        }
    }
    let after = runtime.statistics();
    assert_eq!(before.allocations, after.allocations);
    assert_eq!(before.imports, after.imports);
    assert_eq!(after.copied_bytes, 0);
    latency.sort_by(f64::total_cmp);
    println!(
        "M={m} K={k} N={n} p50_ms={:.3} p95_ms={:.3} pipelined_requests_s={requests_per_second:.2}",
        percentile(&latency, 50) * 1e3,
        percentile(&latency, 95) * 1e3
    );
    println!("{}", serde_json::to_string(&after)?);
    Ok(())
}

examples/gpu_npu_parallel.rs (line 170)

fn main() -> Result<()> {
    let args: Vec<_> = std::env::args().skip(1).collect();
    if args.len() > 4 || args.iter().any(|arg| arg == "--help") {
        println!(
            "usage: gpu_npu_parallel [gpu-elements=16777216] [npu-matrices=1024] [samples=21] [trace.json]"
        );
        return if args.len() <= 4 {
            Ok(())
        } else {
            Err(Error::Message("too many arguments".into()))
        };
    }
    let number = |index: usize, default: usize| -> Result<usize> {
        args.get(index).map_or(Ok(default), |value| {
            value
                .parse()
                .map_err(|_| Error::Message(format!("invalid integer: {value}")))
        })
    };
    let elements = number(0, 1 << 24)?;
    let matrices = number(1, 1024)?;
    let samples = number(2, 21)?;
    if !(256..=1 << 26).contains(&elements)
        || !elements.is_multiple_of(256)
        || !(1..=16384).contains(&matrices)
        || !(3..=1000).contains(&samples)
    {
        return Err(Error::Message("gpu-elements must be a multiple of 256 in 256..=67108864; npu-matrices in 1..=16384; samples in 3..=1000".into()));
    }
    let trace_path = args.get(3).map(PathBuf::from);
    let runtime = Runtime::new()?;
    let gpu = runtime.gpu()?;
    let npu = runtime.npu(0)?;
    println!(
        "GPU: {}; NPU: {}",
        gpu.target().as_str(),
        npu.target().as_str()
    );
    let gpu_compiler = Compiler::for_target(None, gpu.target())?;
    let gpu_artifact = gpu_compiler
        .import_cxx(CxxSource::new(
            "affine.cpp",
            format!(
                r#"
#include <hip/hip_runtime.h>
__global__ [[loom::workgroup_size(256, 1, 1), loom::workgroup_count({}, 1, 1)]]
void affine(const int* input, int* output) {{
  unsigned i = blockIdx.x * 256u + threadIdx.x;
  output[i] = input[i] * 3 + 7;
}}
"#,
                elements / 256
            ),
        ))?
        .compile(&Specialization::new("affine"))?;
    // SAFETY: the fixed launch covers exactly elements i32 values in each
    // nonaliasing binding; bounded input values cannot overflow the arithmetic.
    let gpu_kernel = unsafe {
        runtime.load_gpu_artifact(
            &gpu_artifact,
            &[],
            contract(&[(elements * 4, Access::Read), (elements * 4, Access::Write)]),
        )
    }?;
    let npu_artifact = Compiler::for_target(None, npu.target())?
        .module(include_str!("kernels/batched_gemm.xdna.loom"))
        .compile(&Specialization::new("gemm").with_config("gemm.batch", matrices.to_string()))?;
    // SAFETY: the pipeline consumes matrices pairs of 8x8 BF16 inputs and
    // produces the same number of 8x8 FP32 outputs on one NPU column.
    let npu_kernel = unsafe {
        npu.load_artifact(
            &npu_artifact,
            1,
            contract(&[
                (matrices * 128, Access::Read),
                (matrices * 128, Access::Read),
                (matrices * 256, Access::Write),
            ]),
        )
    }?;
    let input = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
    let output = runtime.allocate(elements * 4, MemoryPlacement::GpuLocal)?;
    let host_input = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
    let host_output = runtime.allocate(elements * 4, MemoryPlacement::HostVisible)?;
    let mut upload = runtime.graph();
    upload.copy(input.view(), host_input.view())?;
    let upload = upload.prepare()?;
    let mut download = runtime.graph();
    download.copy(host_output.view(), output.view())?;
    let download = download.prepare()?;
    let a = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
    let b = runtime.allocate(matrices * 128, MemoryPlacement::NpuLocal(npu.clone()))?;
    let c = runtime.allocate(matrices * 256, MemoryPlacement::NpuLocal(npu))?;

    let mut gpu_only = runtime.graph();
    gpu_only.gpu(&gpu_kernel, &[input.view(), output.view()])?;
    let gpu_only = gpu_only.prepare()?;
    let mut npu_only = runtime.graph();
    npu_only.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
    let npu_only = npu_only.prepare()?;

    // Disjoint buffers imply no dependency edge. One submit admits both
    // branches to separate GPU/NPU workers; one wait joins their completions.
    let mut parallel = runtime.graph();
    parallel.gpu(&gpu_kernel, &[input.view(), output.view()])?;
    parallel.npu(&npu_kernel, &[a.view(), b.view(), c.view()])?;
    let parallel = parallel.prepare()?;
    for pass in 0..3 {
        initialize(&host_input, &a, &b, pass * 2)?;
        upload.submit()?.wait()?;
        gpu_only.submit()?.wait()?;
        npu_only.submit()?.wait()?;
        download.submit()?.wait()?;
        validate(&host_output, &c, pass * 2)?;
        // Change both inputs so validation cannot accept the sequential run's
        // old output if either parallel branch fails to write its result.
        initialize(&host_input, &a, &b, pass * 2 + 1)?;
        upload.submit()?.wait()?;
        parallel.submit()?.wait()?;
        download.submit()?.wait()?;
        validate(&host_output, &c, pass * 2 + 1)?;
    }

    let before = runtime.statistics();
    let mut sequential_ms = Vec::with_capacity(samples);
    let mut parallel_ms = Vec::with_capacity(samples);
    for sample in 0..samples {
        // Alternate order to reduce systematic clock/thermal ordering bias.
        for concurrent in if sample % 2 == 0 {
            [false, true]
        } else {
            [true, false]
        } {
            if concurrent {
                parallel_ms.push(completed_ms(&parallel)?);
            } else {
                let start = Instant::now();
                gpu_only.submit()?.wait()?;
                npu_only.submit()?.wait()?;
                sequential_ms.push(start.elapsed().as_secs_f64() * 1000.0);
            }
        }
    }
    let after = runtime.statistics();
    if before.allocations != after.allocations || before.imports != after.imports {
        return Err(Error::Message(
            "replay unexpectedly allocated or imported storage".into(),
        ));
    }
    download.submit()?.wait()?;
    validate(&host_output, &c, 5)?;
    let sequential = median(&mut sequential_ms);
    let concurrent = median(&mut parallel_ms);
    println!("{elements} GPU elements + {matrices} NPU matrices; {samples} paired samples");
    println!(
        "sequential median: {sequential:.3} ms; parallel median: {concurrent:.3} ms; speedup: {:.2}x",
        sequential / concurrent
    );
    println!(
        "Both CPU references match exactly, including changed-input replays. No native allocations/imports during timed replay."
    );

    // Capture separately so tracing does not affect the benchmark samples.
    runtime.start_trace(32)?;
    for _ in 0..4 {
        parallel.submit()?.wait()?;
    }
    let trace = runtime
        .finish_trace()
        .ok_or_else(|| Error::Message("missing execution trace".into()))?;
    let overlap: f64 = trace
        .events
        .iter()
        .filter(|e| e.lane == 1)
        .flat_map(|gpu| {
            trace.events.iter().filter(|e| e.lane == 3).map(move |npu| {
                ((gpu.start_us + gpu.duration_us).min(npu.start_us + npu.duration_us)
                    - gpu.start_us.max(npu.start_us))
                .max(0.0)
            })
        })
        .sum();
    println!(
        "GPU/NPU host-observed overlap across four replays: {overlap:.1} us (not device timestamps)"
    );
    if let Some(path) = trace_path {
        std::fs::write(&path, trace.to_json()?)?;
        println!("Trace: {}", path.display());
    }
    Ok(())
}

pub fn set_config( &mut self, key: impl Into<String>, value: impl Into<String>, ) -> &mut Self

Set one fully-qualified Loom configuration value in place.

Examples found in repository

examples/gemm_pipeline.rs (line 35)

fn gpu_kernel(
    runtime: &Runtime,
    source: &str,
    symbol: &str,
    config: &[(&str, String)],
    elements: usize,
    bindings: Vec<BindingContract>,
) -> Result<GpuKernel> {
    let compiler = hrx::loom::Compiler::for_target(None, runtime.gpu()?.target())?;
    let module = compiler.module(source);
    let mut spec = hrx::loom::Specialization::new(symbol);
    for (key, value) in config {
        spec.set_config(*key, value.clone());
    }
    let artifact = module.compile(&spec)?;
    let stream = hrx::gpu::Stream::open()?;
    let raw = unsafe { stream.load_artifact(&artifact) }?;
    let mut constants = hrx::gpu::Constants::new();
    match raw.info().constant_byte_length {
        8 => constants.push(elements as u32)?,
        12 => constants.push(elements as u64)?,
        _ => return Err(hrx::Error::Message("unexpected scalar ABI".into())),
    }
    constants.push(1f32)?;
    unsafe {
        runtime.load_gpu_kernel(
            artifact.path(),
            symbol,
            [elements.div_ceil(256) as u32, 1, 1],
            [256, 1, 1],
            KernelContract {
                bindings,
                constants: constants.as_bytes().to_vec(),
            },
        )
    }
}

examples/startup_bench.rs (line 51)

fn main() -> Result<()> {
    let compiler_only = match std::env::args().nth(1).as_deref() {
        None => false,
        Some("--compiler-only") => true,
        Some(_) => {
            return Err(hrx::Error::Message(
                "usage: startup_bench [--compiler-only]".into(),
            ));
        }
    };
    let initial = memory()?;
    let start = Instant::now();
    let stream = if compiler_only {
        None
    } else {
        Some(Stream::open()?)
    };
    let open_ms = stream
        .as_ref()
        .map(|_| start.elapsed().as_secs_f64() * 1000.0);
    let after_open = if compiler_only { None } else { Some(memory()?) };
    let start = Instant::now();
    let compiler = match &stream {
        Some(stream) => hrx::loom::Compiler::for_stream(None, stream)?,
        None => hrx::loom::Compiler::resolve(None)?,
    };
    let compiler_ms = start.elapsed().as_secs_f64() * 1000.0;
    let after_compiler = memory()?;
    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
    let mut spec = Specialization::new("krea2_euler");
    spec.set_config("krea2.euler.grid_x", "1");
    spec.set_config("krea2.euler.grid_y", "1");
    let start = Instant::now();
    let artifact = module.compile(&spec)?;
    let first_compile_ms = start.elapsed().as_secs_f64() * 1000.0;
    let after_compile = memory()?;
    let start = Instant::now();
    let cached = module.compile(&spec)?;
    let cache_hit_ms = start.elapsed().as_secs_f64() * 1000.0;
    assert_eq!(artifact.bytes(), cached.bytes());
    println!(
        "{}",
        json!({"compiler_only": compiler_only, "open_ms": open_ms, "compiler_ms": compiler_ms,
            "first_compile_ms": first_compile_ms, "cache_hit_ms": cache_hit_ms,
            "initial": initial, "after_open": after_open,
            "after_compiler": after_compiler, "after_compile": after_compile,
            "compiler_identity": artifact.compiler_identity()})
    );
    Ok(())
}

examples/patch_bench.rs (line 76)

fn attention(tokens: usize, repack: bool) -> Case {
    let capacity = tokens + 32;
    let symbol = if repack {
        "krea2_attention_query32"
    } else {
        "krea2_attention_gqa_lds_f16_wmma"
    };
    let prefix = if repack {
        "krea2.attention_query32"
    } else {
        "krea2.attention_gqa_lds_f16_wmma"
    };
    let mut spec = Specialization::new(symbol).with_report(ReportMode::Summary);
    for (key, value) in [
        ("q_stride", 512),
        ("kv_stride", 128),
        ("out_stride", 512),
        ("tokens", tokens),
        ("token_capacity", capacity),
    ] {
        spec.set_config(format!("{prefix}.{key}"), value.to_string());
    }
    spec.set_config(format!("{prefix}.scale"), "0.08838834764831845");
    let query = |row: usize, head: usize| ((row + head * 3) % 9) as f32 * 0.0625 - 0.25;
    let key = |row: usize| (row % 7) as f32 * 0.0625 - 0.1875;
    let value = |row: usize, channel: usize| {
        (row % 7) as f32 * 0.125 - 0.375 + (channel % 8) as f32 * 0.0625
    };
    let q = halves((0..capacity * 512).map(|i| {
        if i / 512 < tokens {
            query(i / 512, (i % 512) / 128)
        } else {
            0.0
        }
    }));
    let k = halves((0..capacity * 128).map(|i| if i / 128 < tokens { key(i / 128) } else { 0.0 }));
    let v = halves((0..capacity * 128).map(|i| {
        let (row, channel) = if repack {
            (i % capacity, i / capacity)
        } else {
            (i / 128, i % 128)
        };
        if row < tokens {
            value(row, channel)
        } else {
            0.0
        }
    }));
    // Independent f64 softmax oracle. Inputs are constant across the 128 Q/K
    // channels, so their exact dot product is 128 times the scalar product.
    let mut expected = Vec::with_capacity(tokens * 512);
    for row in 0..tokens {
        for head in 0..4 {
            let mut numerator = 0.0;
            let mut denominator = 0.0;
            for token in 0..tokens {
                let score = (128.0 * f64::from(query(row, head)) * f64::from(key(token))
                    / 128f64.sqrt())
                .exp();
                numerator += score * f64::from(value(token, 0));
                denominator += score;
            }
            expected.extend(
                (0..128).map(|channel| numerator / denominator + (channel % 8) as f64 * 0.0625),
            );
        }
    }
    Case {
        name: format!("{symbol}-{tokens}"),
        source: if repack {
            include_str!("../native/qualification/attention_query32.loom")
        } else {
            include_str!("../native/qualification/attention_gqa_lds_f16_wmma.loom")
        },
        spec,
        grid: [(tokens / if repack { 32 } else { 16 }) as u32, 1, 1],
        block: [if repack { 256 } else { 128 }, 1, 1],
        indices: vec![tokens as u32],
        data: vec![q, k, v, vec![0; tokens * 512 * 2]],
        output: 3,
        expected,
        tolerance: 0.001,
    }
}
fn gemm(large: bool, integer: bool) -> Case {
    let (m, n, k) = if large {
        (256usize, 1024usize, 1024usize)
    } else {
        (256, 512, 256)
    };
    let symbol = if integer {
        "krea2_gemm_i8_256"
    } else {
        "h3_gemm_f16_fast_256b"
    };
    let prefix = if integer {
        "krea2.gemm_i8_256"
    } else {
        "h3.gemm_f16_fast_256b"
    };
    let mut spec = Specialization::new(symbol).with_report(ReportMode::Summary);
    for (name, value) in [
        ("k_size", k),
        ("n_size", n),
        ("k_stride", k),
        ("m_group", 1),
    ] {
        spec.set_config(format!("{prefix}.{name}"), value.to_string());
    }
    let a = |row: usize, col: usize| ((row * 3 + col * 5 + row * col) % 11) as i32 - 5;
    let w = |row: usize, col: usize| ((row * 7 + col * 3 + row * col) % 13) as i32 - 6;
    let bias = |col: usize| (col % 7) as f32 * 0.03125 - 0.09375;
    // Evaluate the full dot product once per distinct row-pattern pair. Pattern
    // repetition in the authored inputs is independent of the compiled program.
    let mut products = [[0f32; 13]; 11];
    for (row, values) in products.iter_mut().enumerate() {
        for (col, product) in values.iter_mut().enumerate() {
            let sum: i32 = (0..k).map(|i| a(row, i) * w(col, i)).sum();
            *product = sum as f32 / 64.0;
        }
    }
    let expected = (0..m * n)
        .map(|i| {
            let dot = products[(i / n) % 11][(i % n) % 13];
            float(half(if integer {
                bf16(dot)
            } else {
                dot + bias(i % n)
            }))
        })
        .collect();
    let data = if integer {
        vec![
            (0..m * k).map(|i| a(i / k, i % k) as i8 as u8).collect(),
            (0..n * k).map(|i| w(i / k, i % k) as i8 as u8).collect(),
            floats(vec![0.125; n]),
            floats(vec![0.125; m]),
            vec![0; m * n * 2],
        ]
    } else {
        vec![
            halves((0..m * k).map(|i| a(i / k, i % k) as f32 * 0.125)),
            halves((0..n * k).map(|i| w(i / k, i % k) as f32 * 0.125)),
            vec![0; m * n * 2],
            floats((0..n).map(bias)),
        ]
    };
    Case {
        name: format!("{symbol}-{m}x{n}x{k}"),
        source: if integer {
            include_str!("../native/qualification/gemm_i8_256.loom")
        } else {
            include_str!("../native/qualification/gemm_f16_fast_256b.loom")
        },
        spec,
        grid: [
            (n / if integer { 128 } else { 256 }) as u32,
            (m / if integer { 256 } else { 128 }) as u32,
            1,
        ],
        block: [256, 1, 1],
        indices: vec![m as u32],
        data,
        output: if integer { 4 } else { 2 },
        expected,
        // Permit subnormal residuals from the floating-point matrix accumulator.
        // Integer GEMM remains exact; the absolute floor is one fp16 minimum normal.
        tolerance: if integer { 0.0 } else { 2f64.powi(-14) },
    }
}
fn conv(height: usize, cin: usize, n: usize) -> Case {
    let symbol = "arcface_conv3x3_f16_wmma_bnprelu";
    let mut spec = Specialization::new(symbol).with_report(ReportMode::Details);
    for (key, value) in [
        ("height", height),
        ("width", height),
        ("stride", 1),
        ("cin_pad", cin),
        ("cin_stride", cin),
        ("k_size", 9 * cin),
        ("n_size", n),
    ] {
        spec.set_config(
            format!("arcface.conv3x3_f16_wmma_bnprelu.{key}"),
            value.to_string(),
        );
    }
    let m = height * height;
    // Constant power-of-two operands: exact mathematical convolution, including
    // zero padding at borders. Bias is zero and PReLU slope is one.
    let expected = (0..m)
        .flat_map(|p| {
            let y = p / height;
            let x = p % height;
            let rows = if y == 0 || y + 1 == height { 2 } else { 3 };
            let cols = if x == 0 || x + 1 == height { 2 } else { 3 };
            std::iter::repeat_n((rows * cols * cin) as f64 / 64.0, n)
        })
        .collect();
    Case {
        name: format!("arcface_conv-{height}x{height}x{cin}x{n}"),
        source: include_str!("../native/qualification/arcface_conv3x3_f16_wmma_bnprelu.loom"),
        spec,
        grid: [(n / 64) as u32, m.div_ceil(64) as u32, 1],
        block: [256, 1, 1],
        indices: vec![m as u32],
        data: vec![
            halves(std::iter::repeat_n(0.125, m * cin)),
            halves(std::iter::repeat_n(0.125, n * 9 * cin)),
            floats(std::iter::repeat_n(0.0, 9 * n)),
            vec![0; m * n * 2],
            floats(std::iter::repeat_n(1.0, n)),
        ],
        output: 3,
        expected,
        tolerance: 0.0,
    }
}

examples/stream_bench.rs (line 79)

fn main() -> Result<()> {
    let mut stream = Stream::open()?;
    let mut metrics = BTreeMap::new();
    let samples = std::env::var("HRX_BENCH_SAMPLES")
        .ok()
        .and_then(|value| value.parse::<usize>().ok())
        .filter(|&value| value > 0)
        .unwrap_or(9);
    let (ns, _) = measure(&mut stream, samples, 50_000, |stream| {
        for _ in 0..50_000 {
            let buffer = stream.scratch(1024 * 1024)?;
            black_box(&buffer);
            stream.recycle(buffer)?;
        }
        Ok(())
    })?;
    metrics.insert("scratch_reuse_ns", ns);
    let (ns, _) = measure(&mut stream, samples, 2_000, |stream| {
        for _ in 0..2_000 {
            black_box(stream.allocate(1024 * 1024)?);
        }
        Ok(())
    })?;
    metrics.insert("allocate_drop_ns", ns);

    let bytes = 256 * 1024 * 1024;
    let chunk = vec![0x3cu8; 4 * 1024 * 1024];
    let weights = stream.allocate(bytes)?;
    let (_, ns) = measure(&mut stream, samples, 1, |stream| {
        for offset in (0..bytes).step_by(chunk.len()) {
            stream.upload(weights.binding().slice(offset, chunk.len())?, &chunk)?;
        }
        Ok(())
    })?;
    metrics.insert(
        "queued_upload_gib_s",
        bytes as f64 / (1024f64.powi(3) * ns * 1e-9),
    );
    let mut check = [0u8; 16];
    stream.read_blocking(
        weights.try_slice(bytes - check.len(), check.len())?,
        &mut check,
    )?;
    assert_eq!(check, [0x3c; 16]);
    drop(weights);

    let compiler = hrx::loom::Compiler::for_stream(None, &stream)?;
    let module = compiler.module(include_str!("../tests/kernels/euler.loom"));
    let mut spec = hrx::loom::Specialization::new("krea2_euler");
    spec.set_config("krea2.euler.grid_x", "1");
    spec.set_config("krea2.euler.grid_y", "1");
    let artifact = module.compile(&spec)?;
    // This repository owns the source; use its exact dimensions and scalar layout.
    let kernel = unsafe { stream.load_artifact(&artifact)? };
    let sample = stream.allocate(512)?;
    let velocity = stream.allocate(512)?;
    let ones: Vec<u8> = (0..256).flat_map(|_| 0x3f80u16.to_le_bytes()).collect();
    stream.upload_blocking(sample.binding(), &ones)?;
    stream.upload_blocking(velocity.binding(), &ones)?;
    let mut constants = Constants::new();
    match kernel.info().constant_byte_length {
        8 => constants.push(256u32)?,
        12 => constants.push(256u64)?,
        n => return Err(hrx::Error::Message(format!("unexpected constant size {n}"))),
    }
    constants.push(0f32)?;
    let bindings = [sample.binding(), velocity.binding()];
    let (enqueue, complete) = measure(&mut stream, samples, 2048, |stream| {
        for _ in 0..2048 {
            // The source accesses exactly the 256 bf16 elements in each buffer.
            unsafe {
                stream.dispatch(&kernel, [1; 3], [256, 1, 1], &constants, &bindings)?;
            }
        }
        Ok(())
    })?;
    metrics.insert("dispatch_enqueue_ns", enqueue);
    metrics.insert("dispatch_complete_ns", complete);

    // Use disjoint allocations in both graphs. Even an idempotent store is a
    // write: concurrent Euler dispatches cannot share their in/out buffer.
    let mut graph_buffers = Vec::new();
    for _ in 0..32 {
        let sample = stream.allocate(512)?;
        let velocity = stream.allocate(512)?;
        stream.upload(sample.binding(), &ones)?;
        stream.upload(velocity.binding(), &ones)?;
        graph_buffers.push((sample, velocity));
    }
    stream.synchronize()?;
    let mut builder = stream.graph()?;
    let mut previous = None;
    for (sample, velocity) in &graph_buffers {
        let after = previous.as_slice();
        previous = Some(unsafe {
            builder.dispatch(
                after,
                &kernel,
                [1; 3],
                [256, 1, 1],
                &constants,
                &[sample.binding(), velocity.binding()],
            )?
        });
    }
    let chained = builder.finish()?;

    // The same disjoint workloads, this time allowing concurrent scheduling.
    let mut builder = stream.graph()?;
    for (sample, velocity) in &graph_buffers {
        unsafe {
            builder.dispatch(
                &[],
                &kernel,
                [1; 3],
                [256, 1, 1],
                &constants,
                &[sample.binding(), velocity.binding()],
            )?;
        }
    }
    let independent = builder.finish()?;
    let mut graphs = [chained, independent];
    let mut enqueue = [Vec::new(), Vec::new()];
    let mut complete = [Vec::new(), Vec::new()];
    for round in 0..samples + 3 {
        for arm in if round % 2 == 0 { [0, 1] } else { [1, 0] } {
            stream.synchronize()?;
            let start = Instant::now();
            for _ in 0..256 {
                stream.launch(&mut graphs[arm])?;
            }
            let submitted = start.elapsed();
            stream.synchronize()?;
            let elapsed = start.elapsed();
            if round >= 3 {
                enqueue[arm].push(submitted.as_secs_f64() * 1e9 / 256.0);
                complete[arm].push(elapsed.as_secs_f64() * 1e9 / (256.0 * 32.0));
            }
        }
    }
    for values in enqueue.iter_mut().chain(complete.iter_mut()) {
        values.sort_by(f64::total_cmp);
    }
    let ns = complete[0][samples / 2];
    let ns_free = complete[1][samples / 2];
    metrics.insert("graph_enqueue_ns_per_replay", enqueue[0][samples / 2]);
    metrics.insert("graph_complete_ns_per_kernel", ns);
    metrics.insert("graph_independent_ns_per_kernel", ns_free);
    metrics.insert("graph_scheduling_difference_ns_per_kernel", ns - ns_free);
    for (sample, _) in &graph_buffers {
        let mut output = vec![0; 512];
        stream.read_blocking(sample.binding(), &mut output)?;
        assert_eq!(output, ones);
    }
    let mut output = vec![0; 512];
    stream.read_blocking(sample.binding(), &mut output)?;
    assert_eq!(output, ones);
    println!("{}", serde_json::to_string(&metrics)?);
    Ok(())
}

pub fn set_symbol(&mut self, symbol: impl Into<String>) -> &mut Self

Select an export/root symbol in place.

pub fn replace_config(&mut self, config: BTreeMap<String, String>) -> &mut Self

Replace all fully-qualified Loom configuration values.

pub fn with_report(self, report: ReportMode) -> Self

Request or suppress the compiler resource report.

Examples found in repository

examples/compile_npu.rs (line 14)

fn main() -> Result<()> {
    let packets = std::env::args().nth(1).unwrap_or_else(|| "256".into());
    let artifact = Compiler::for_target(None, &Target::xdna())?
        .module(include_str!("../tests/kernels/copy.xdna.loom"))
        .compile(
            &Specialization::new("copy")
                .with_config("copy.packets", packets)
                .with_report(ReportMode::Summary),
        )?;
    println!("{}", artifact.path().display());
    Ok(())
}

examples/patch_bench.rs (line 68)

fn attention(tokens: usize, repack: bool) -> Case {
    let capacity = tokens + 32;
    let symbol = if repack {
        "krea2_attention_query32"
    } else {
        "krea2_attention_gqa_lds_f16_wmma"
    };
    let prefix = if repack {
        "krea2.attention_query32"
    } else {
        "krea2.attention_gqa_lds_f16_wmma"
    };
    let mut spec = Specialization::new(symbol).with_report(ReportMode::Summary);
    for (key, value) in [
        ("q_stride", 512),
        ("kv_stride", 128),
        ("out_stride", 512),
        ("tokens", tokens),
        ("token_capacity", capacity),
    ] {
        spec.set_config(format!("{prefix}.{key}"), value.to_string());
    }
    spec.set_config(format!("{prefix}.scale"), "0.08838834764831845");
    let query = |row: usize, head: usize| ((row + head * 3) % 9) as f32 * 0.0625 - 0.25;
    let key = |row: usize| (row % 7) as f32 * 0.0625 - 0.1875;
    let value = |row: usize, channel: usize| {
        (row % 7) as f32 * 0.125 - 0.375 + (channel % 8) as f32 * 0.0625
    };
    let q = halves((0..capacity * 512).map(|i| {
        if i / 512 < tokens {
            query(i / 512, (i % 512) / 128)
        } else {
            0.0
        }
    }));
    let k = halves((0..capacity * 128).map(|i| if i / 128 < tokens { key(i / 128) } else { 0.0 }));
    let v = halves((0..capacity * 128).map(|i| {
        let (row, channel) = if repack {
            (i % capacity, i / capacity)
        } else {
            (i / 128, i % 128)
        };
        if row < tokens {
            value(row, channel)
        } else {
            0.0
        }
    }));
    // Independent f64 softmax oracle. Inputs are constant across the 128 Q/K
    // channels, so their exact dot product is 128 times the scalar product.
    let mut expected = Vec::with_capacity(tokens * 512);
    for row in 0..tokens {
        for head in 0..4 {
            let mut numerator = 0.0;
            let mut denominator = 0.0;
            for token in 0..tokens {
                let score = (128.0 * f64::from(query(row, head)) * f64::from(key(token))
                    / 128f64.sqrt())
                .exp();
                numerator += score * f64::from(value(token, 0));
                denominator += score;
            }
            expected.extend(
                (0..128).map(|channel| numerator / denominator + (channel % 8) as f64 * 0.0625),
            );
        }
    }
    Case {
        name: format!("{symbol}-{tokens}"),
        source: if repack {
            include_str!("../native/qualification/attention_query32.loom")
        } else {
            include_str!("../native/qualification/attention_gqa_lds_f16_wmma.loom")
        },
        spec,
        grid: [(tokens / if repack { 32 } else { 16 }) as u32, 1, 1],
        block: [if repack { 256 } else { 128 }, 1, 1],
        indices: vec![tokens as u32],
        data: vec![q, k, v, vec![0; tokens * 512 * 2]],
        output: 3,
        expected,
        tolerance: 0.001,
    }
}
fn gemm(large: bool, integer: bool) -> Case {
    let (m, n, k) = if large {
        (256usize, 1024usize, 1024usize)
    } else {
        (256, 512, 256)
    };
    let symbol = if integer {
        "krea2_gemm_i8_256"
    } else {
        "h3_gemm_f16_fast_256b"
    };
    let prefix = if integer {
        "krea2.gemm_i8_256"
    } else {
        "h3.gemm_f16_fast_256b"
    };
    let mut spec = Specialization::new(symbol).with_report(ReportMode::Summary);
    for (name, value) in [
        ("k_size", k),
        ("n_size", n),
        ("k_stride", k),
        ("m_group", 1),
    ] {
        spec.set_config(format!("{prefix}.{name}"), value.to_string());
    }
    let a = |row: usize, col: usize| ((row * 3 + col * 5 + row * col) % 11) as i32 - 5;
    let w = |row: usize, col: usize| ((row * 7 + col * 3 + row * col) % 13) as i32 - 6;
    let bias = |col: usize| (col % 7) as f32 * 0.03125 - 0.09375;
    // Evaluate the full dot product once per distinct row-pattern pair. Pattern
    // repetition in the authored inputs is independent of the compiled program.
    let mut products = [[0f32; 13]; 11];
    for (row, values) in products.iter_mut().enumerate() {
        for (col, product) in values.iter_mut().enumerate() {
            let sum: i32 = (0..k).map(|i| a(row, i) * w(col, i)).sum();
            *product = sum as f32 / 64.0;
        }
    }
    let expected = (0..m * n)
        .map(|i| {
            let dot = products[(i / n) % 11][(i % n) % 13];
            float(half(if integer {
                bf16(dot)
            } else {
                dot + bias(i % n)
            }))
        })
        .collect();
    let data = if integer {
        vec![
            (0..m * k).map(|i| a(i / k, i % k) as i8 as u8).collect(),
            (0..n * k).map(|i| w(i / k, i % k) as i8 as u8).collect(),
            floats(vec![0.125; n]),
            floats(vec![0.125; m]),
            vec![0; m * n * 2],
        ]
    } else {
        vec![
            halves((0..m * k).map(|i| a(i / k, i % k) as f32 * 0.125)),
            halves((0..n * k).map(|i| w(i / k, i % k) as f32 * 0.125)),
            vec![0; m * n * 2],
            floats((0..n).map(bias)),
        ]
    };
    Case {
        name: format!("{symbol}-{m}x{n}x{k}"),
        source: if integer {
            include_str!("../native/qualification/gemm_i8_256.loom")
        } else {
            include_str!("../native/qualification/gemm_f16_fast_256b.loom")
        },
        spec,
        grid: [
            (n / if integer { 128 } else { 256 }) as u32,
            (m / if integer { 256 } else { 128 }) as u32,
            1,
        ],
        block: [256, 1, 1],
        indices: vec![m as u32],
        data,
        output: if integer { 4 } else { 2 },
        expected,
        // Permit subnormal residuals from the floating-point matrix accumulator.
        // Integer GEMM remains exact; the absolute floor is one fp16 minimum normal.
        tolerance: if integer { 0.0 } else { 2f64.powi(-14) },
    }
}
fn conv(height: usize, cin: usize, n: usize) -> Case {
    let symbol = "arcface_conv3x3_f16_wmma_bnprelu";
    let mut spec = Specialization::new(symbol).with_report(ReportMode::Details);
    for (key, value) in [
        ("height", height),
        ("width", height),
        ("stride", 1),
        ("cin_pad", cin),
        ("cin_stride", cin),
        ("k_size", 9 * cin),
        ("n_size", n),
    ] {
        spec.set_config(
            format!("arcface.conv3x3_f16_wmma_bnprelu.{key}"),
            value.to_string(),
        );
    }
    let m = height * height;
    // Constant power-of-two operands: exact mathematical convolution, including
    // zero padding at borders. Bias is zero and PReLU slope is one.
    let expected = (0..m)
        .flat_map(|p| {
            let y = p / height;
            let x = p % height;
            let rows = if y == 0 || y + 1 == height { 2 } else { 3 };
            let cols = if x == 0 || x + 1 == height { 2 } else { 3 };
            std::iter::repeat_n((rows * cols * cin) as f64 / 64.0, n)
        })
        .collect();
    Case {
        name: format!("arcface_conv-{height}x{height}x{cin}x{n}"),
        source: include_str!("../native/qualification/arcface_conv3x3_f16_wmma_bnprelu.loom"),
        spec,
        grid: [(n / 64) as u32, m.div_ceil(64) as u32, 1],
        block: [256, 1, 1],
        indices: vec![m as u32],
        data: vec![
            halves(std::iter::repeat_n(0.125, m * cin)),
            halves(std::iter::repeat_n(0.125, n * 9 * cin)),
            floats(std::iter::repeat_n(0.0, 9 * n)),
            vec![0; m * n * 2],
            floats(std::iter::repeat_n(1.0, n)),
        ],
        output: 3,
        expected,
        tolerance: 0.0,
    }
}

pub fn set_report(&mut self, report: ReportMode) -> &mut Self

Request or suppress the compiler resource report in place.

pub fn symbol(&self) -> &str

Export/root symbol selected by this specialization.

pub fn configuration(&self) -> &BTreeMap<String, String>

Fully-qualified compiler configuration.

pub fn report_mode(&self) -> ReportMode

Compiler evidence requested by this specialization.

Trait Implementations

impl Clone for Specialization

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for Specialization

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Default for Specialization

fn default() -> Self

Returns the “default value” for a type.