# Runtime Dispatch > NumKong for Rust `Capabilities` answers for the CPU and each GPU group through its associated functions, and `Capabilities::configure_thread` enables CPU-specific acceleration features such as Intel AMX. It must be called once per thread before using AMX operations. ```rust use numkong::{Capabilities, Capability}; let enabled = Capabilities::cpu_enabled(); enabled.configure_thread().unwrap(); println!("dispatching to {enabled}"); // like "serial,neon,neonhalf" if enabled.contains(Capability::SapphireAmx) { println!("AMX enabled"); } for ordinal in 0..Capabilities::cuda_count_devices().unwrap() { println!("CUDA device {ordinal} runs {}", Capabilities::cuda_enabled(ordinal).unwrap()); // like "cuda,ampere,hopper" } ``` Each group reports a `Capabilities` set of `Capability` bits along two independent axes, plus the set dispatch uses: - `cpu_detected()`, `cuda_detected(ordinal)` and their ROCm and Metal kin: what the device can execute, from CPUID or HWCAP on the CPU and from the runtime on a GPU - `cpu_compiled()`, `cuda_compiled()` and their kin: what this binary contains, from the probes at build time - `cpu_enabled()`, `cuda_enabled(ordinal)` and their kin: what dispatch runs within, i.e. both axes at once Reach for the `_enabled` ones unless you specifically mean one of the raw axes. `cpu_detected()` describes the machine and says nothing about whether a kernel was compiled in, so a build whose ISA probes failed still reports your CPU's full feature set while containing no SIMD kernels at all. The crate keeps no process state: every CPU kernel call passes `Capabilities::CPUS`, and the library clamps it to the CPU capabilities it settles as it loads. No verb allocates: a GPU verb takes any memory its device reaches, which the runtime checks per call. `Stream::new(capabilities, ordinal)` makes a stream on a device, `Stream::default(capabilities)` names the vendor's default, and `stream.synchronize()` waits for it. `UnifiedAllocator::new(&stream)` allocates memory the host and that stream's device both address, the only unified memory the crate hands out. `stream.scope(|scope| ...)` runs the safe `_on` verbs, each queueing on that stream with its capabilities, and joins it before anything they borrow is released. `Capabilities::stream_init_cuda(ordinal)` and `Capabilities::stream_free_cuda(stream)` stay as the unsafe raw handles, with `rocm_` and `metal_` twins. Call `configure_thread` at the start of every thread that will use AMX operations. In a thread-pool setting, each worker thread needs its own call. The function is idempotent and cheap to call more than once on the same thread. `use numkong::Capabilities;` `use numkong::Capability;` `use numkong::Scope;` `use numkong::Stream;` ## Capabilities A set of capabilities, printed as comma-separated names like `serial,haswell`. ```rust pub struct Capabilities(_); ``` Examples: ``` use numkong::{Capabilities, Capability}; let enabled = Capabilities::cpu_enabled(); enabled.configure_thread()?; println!("dispatching to {enabled}"); if enabled.contains(Capability::SapphireAmx) { println!("AMX is enabled"); } for ordinal in 0..Capabilities::cuda_count_devices()? { println!("CUDA device {ordinal} runs {}", Capabilities::cuda_enabled(ordinal)?); } ``` Implements: `Clone`, `Copy`, `Debug`, `Default`, `StructuralPartialEq`, `PartialEq`, `Eq`, `Hash`, `From<[Capability](#Capability)>`, `BitOr<[Capability](#Capability)>`, `BitOr`, `Display` ## ANY Every capability, C's `nk_cap_any_k`. ```rust pub const ANY: Self = _; ``` ## CPUS Every CPU capability, the bits below the first GPU vendor's, C's `nk_cap_cpus_k`. ```rust pub const CPUS: Self = _; ``` Calls without a mask of their own pass it, as dispatch clamps it to what this CPU runs. ## GPUS Every GPU capability, which the CPU never detects or enables, C's `nk_cap_gpus_k`. ```rust pub const GPUS: Self = _; ``` ## bits The raw `nk_capability_t` mask. ```rust pub const fn bits(self) -> u64 ``` ## contains Whether `capability` is in this set. ```rust pub const fn contains(self, capability: Capability) -> bool ``` ## without This set less `capability`, like `enabled.without(Capability::Skylake)`. ```rust pub const fn without(self, capability: Capability) -> Self ``` ## iter The capabilities in this set, in bit order. ```rust pub fn iter(self) -> impl Iterator ``` ## cpu_detected Capabilities this CPU supports, whether or not their kernels were compiled in. ```rust pub fn cpu_detected() -> Self ``` ## cpu_compiled CPU capabilities whose kernels were compiled in, whether or not this CPU supports them. ```rust pub fn cpu_compiled() -> Self ``` ## cpu_enabled The CPU capabilities every kernel call runs within: [`Capabilities::cpu_detected`](#cpu_detected) & [`Capabilities::cpu_compiled`](#cpu_compiled), always with [`Capability::Serial`](#Serial). ```rust pub fn cpu_enabled() -> Self ``` ## configure_thread Sets up the calling thread for the CPU kernels in this set, usually [[`Capabilities::cpu_enabled`](#cpu_enabled)]: AMX tile permission on x86 Linux, fused BF16 dots on Arm. ```rust pub fn configure_thread(self) -> Result<(), Error> ``` Call it once per thread before using those kernels; it is idempotent. ## cuda_count_devices How many CUDA devices this process sees, zero without the runtime or without its kernels. ```rust pub fn cuda_count_devices() -> Result ``` ## cuda_detected Capabilities CUDA device `ordinal` supports, by the runtime's own numbering, whether or not their kernels were compiled in. ```rust pub fn cuda_detected(ordinal: usize) -> Result ``` ## cuda_compiled CUDA capabilities whose kernels were compiled in. ```rust pub fn cuda_compiled() -> Self ``` ## cuda_enabled The mask CUDA device `ordinal` dispatches with: [`Capabilities::cuda_detected`](#cuda_detected) & [`Capabilities::cuda_compiled`](#cuda_compiled). ```rust pub fn cuda_enabled(ordinal: usize) -> Result ``` ## stream_init_cuda A new `cudaStream_t` on CUDA device `ordinal`, which names that device to every call it is passed to; fails with [`Status::MissingGpu`](https://www.unum.cloud/docs/numkong/rust/quickstart.md#MissingGpu) past the last device or where this build lacks CUDA. ```rust pub unsafe fn stream_init_cuda(ordinal: usize) -> Result<*mut c_void, Error> ``` Safety: The caller owns the stream and frees it once, through [`Capabilities::stream_free_cuda`](#stream_free_cuda). ## stream_free_cuda Frees a stream [`Capabilities::stream_init_cuda`](#stream_init_cuda) made. ```rust pub unsafe fn stream_free_cuda(stream: *mut c_void) -> Result<(), Error> ``` Safety: [`stream`](#stream) must come from [`Capabilities::stream_init_cuda`](#stream_init_cuda) and not be freed yet, and nothing may use it afterwards, so whatever was queued on it is synchronized first. ## rocm_count_devices [`Capabilities::cuda_count_devices`](#cuda_count_devices), for ROCm. ```rust pub fn rocm_count_devices() -> Result ``` ## rocm_detected [`Capabilities::cuda_detected`](#cuda_detected), for ROCm. ```rust pub fn rocm_detected(ordinal: usize) -> Result ``` ## rocm_compiled [`Capabilities::cuda_compiled`](#cuda_compiled), for ROCm. ```rust pub fn rocm_compiled() -> Self ``` ## rocm_enabled [`Capabilities::cuda_enabled`](#cuda_enabled), for ROCm. ```rust pub fn rocm_enabled(ordinal: usize) -> Result ``` ## stream_init_rocm [`Capabilities::stream_init_cuda`](#stream_init_cuda), for ROCm, making a `hipStream_t`. ```rust pub unsafe fn stream_init_rocm(ordinal: usize) -> Result<*mut c_void, Error> ``` Safety: The caller owns the stream and frees it once, through [`Capabilities::stream_free_rocm`](#stream_free_rocm). ## stream_free_rocm [`Capabilities::stream_free_cuda`](#stream_free_cuda), for ROCm. ```rust pub unsafe fn stream_free_rocm(stream: *mut c_void) -> Result<(), Error> ``` Safety: [`stream`](#stream) must come from [`Capabilities::stream_init_rocm`](#stream_init_rocm), on the same terms. ## metal_count_devices [`Capabilities::cuda_count_devices`](#cuda_count_devices), for Metal, whose devices count in system order. ```rust pub fn metal_count_devices() -> Result ``` ## metal_detected [`Capabilities::cuda_detected`](#cuda_detected), for Metal. ```rust pub fn metal_detected(ordinal: usize) -> Result ``` ## metal_compiled [`Capabilities::cuda_compiled`](#cuda_compiled), for Metal. ```rust pub fn metal_compiled() -> Self ``` ## metal_enabled [`Capabilities::cuda_enabled`](#cuda_enabled), for Metal. ```rust pub fn metal_enabled(ordinal: usize) -> Result ``` ## stream_init_metal [`Capabilities::stream_init_cuda`](#stream_init_cuda), for Metal, making an `id`. ```rust pub unsafe fn stream_init_metal(ordinal: usize) -> Result<*mut c_void, Error> ``` Safety: The caller owns the stream and frees it once, through [`Capabilities::stream_free_metal`](#stream_free_metal). ## stream_free_metal [`Capabilities::stream_free_cuda`](#stream_free_cuda), for Metal. ```rust pub unsafe fn stream_free_metal(stream: *mut c_void) -> Result<(), Error> ``` Safety: [`stream`](#stream) must come from [`Capabilities::stream_init_metal`](#stream_init_metal), on the same terms. ## synchronize Waits for everything queued on [`stream`](#stream), a stream of this mask's group, null for its default, and null on the CPU. ```rust pub unsafe fn synchronize(self, stream: *mut c_void) -> Result<(), Error> ``` Safety: [`stream`](#stream) must be null or a live `cudaStream_t`, `hipStream_t` or `id` of this mask's group. ## Capability One capability, numbered like the C `nk_cap__k` bits: each capability group in a contiguous run, ascending by dispatch preference, with the GPU vendors' groups above bit 47. ```rust pub enum Capability ``` - `Serial` = `{ _ }`: Serial fallback capability without SIMD. Always available. - `Neon` = `{ _ }`: ARM NEON. - `NeonHalf` = `{ _ }`: ARM NEON FP16. - `NeonBfDot` = `{ _ }`: ARM NEON BF16. - `NeonFhm` = `{ _ }`: ARM NEON FP16 FML. - `NeonSdot` = `{ _ }`: ARM NEON i8 dot. - `NeonFp8` = `{ _ }`: ARM NEON FP8. - `Sve` = `{ _ }`: ARM SVE. - `SveHalf` = `{ _ }`: ARM SVE FP16. - `SveSdot` = `{ _ }`: ARM SVE i8 dot. - `SveBfDot` = `{ _ }`: ARM SVE BF16. - `Sve2` = `{ _ }`: ARM SVE2. - `Sme` = `{ _ }`: ARM SME. - `SmeF64` = `{ _ }`: ARM SME F64. - `SmeBi32` = `{ _ }`: ARM SME BI32I32. - `Haswell` = `{ _ }`: Intel AVX2. - `Alder` = `{ _ }`: Intel AVX2+VNNI. - `Sierra` = `{ _ }`: Intel AVXVNNIINT8. - `Skylake` = `{ _ }`: Intel AVX-512. - `IceLake` = `{ _ }`: Intel AVX-512 VNNI. - `Genoa` = `{ _ }`: AMD AVX-512 BF16. - `Turin` = `{ _ }`: AMD Turin AVX-512 CD. - `Sapphire` = `{ _ }`: Intel AVX-512 FP16. - `Diamond` = `{ _ }`: Intel AVX10.2. - `SapphireAmx` = `{ _ }`: Intel Sapphire AMX. - `GraniteAmx` = `{ _ }`: Intel Granite AMX FP16. - `DiamondAmx` = `{ _ }`: Intel Diamond Rapids AMX. - `Rvv` = `{ _ }`: RISC-V Vector. - `RvvBf16` = `{ _ }`: RISC-V Zvfbfwma. - `RvvHalf` = `{ _ }`: RISC-V Zvfh. - `RvvBb` = `{ _ }`: RISC-V Zvbb. - `V128` = `{ _ }`: WASM SIMD128. - `V128Relaxed` = `{ _ }`: WASM Relaxed SIMD. - `PowerVsx` = `{ _ }`: Power VSX 128-bit SIMD. - `LoongsonAsx` = `{ _ }`: LoongArch LASX 256-bit SIMD. - `Cuda` = `{ _ }`: NVIDIA baseline, running on every CUDA device. - `Ampere` = `{ _ }`: NVIDIA SM 8.0. - `Ada` = `{ _ }`: NVIDIA SM 8.9. - `Hopper` = `{ _ }`: NVIDIA SM 9.x. - `Blackwell` = `{ _ }`: NVIDIA SM 10.x. - `BlackwellRtx` = `{ _ }`: NVIDIA SM 12.x. - `BlackwellUltra` = `{ _ }`: NVIDIA SM 10.3. - `Rocm` = `{ _ }`: AMD baseline, running on every ROCm device. - `Cdna3` = `{ _ }`: AMD gfx942. - `Cdna4` = `{ _ }`: AMD gfx950. - `Cdna5` = `{ _ }`: AMD gfx1250. - `Metal` = `{ _ }`: Apple baseline, running on every Metal device. - `Apple9` = `{ _ }`: Apple GPU family 9. - `Apple10` = `{ _ }`: Apple GPU family 10. Implements: `Clone`, `Copy`, `Debug`, `StructuralPartialEq`, `PartialEq`, `Eq`, `Hash`, `From<[Capability](#Capability)>`, `BitOr`, `BitOr<[Capability](#Capability)>`, `Display` ## Scope The span of a [`Stream::scope`](#scope) body, which every `_on` verb takes to queue on that stream. ```rust pub struct Scope<'scope, 'env: 'scope> ``` A verb borrows what it reads and writes for `'scope`, which ends only once the stream is joined. ## stream The stream every verb of this scope queues on, whose capabilities it dispatches with. ```rust pub fn stream(&self) -> &Stream ``` ## Stream A stream of one capability group's device, where verbs and [`UnifiedAllocator`](https://www.unum.cloud/docs/numkong/rust/external-memory.md#UnifiedAllocator) blocks queue work: a vendor's default when null, the CPU's always null. ```rust pub struct Stream ``` Examples: ``` use numkong::{Capabilities, Stream}; let cpu = Stream::default(Capabilities::cpu_enabled()); assert!(cpu.as_raw().is_null()); // A vendor with no devices fails with `Status::MissingGpu`. for ordinal in 0..Capabilities::cuda_count_devices()? { let stream = Stream::new(Capabilities::cuda_enabled(ordinal)?, ordinal)?; stream.synchronize()?; } ``` Implements: `Debug`, `Drop` ## default The default stream of the default device of the group [`capabilities`](https://www.unum.cloud/docs/numkong/rust/external-memory.md#capabilities) name. ```rust pub fn default(capabilities: Capabilities) -> Self ``` ## new A new stream on device `ordinal` of the vendor [`capabilities`](https://www.unum.cloud/docs/numkong/rust/external-memory.md#capabilities) name, a `cudaStream_t`, a `hipStream_t` or an `id`, freed on drop. ```rust pub fn new(capabilities: Capabilities, ordinal: usize) -> Result ``` The CPU gets [`Stream::default`](#default). Fails with [`Status::MissingGpu`](https://www.unum.cloud/docs/numkong/rust/quickstart.md#MissingGpu) past the last device or where this build lacks the vendor. ## from_raw Borrows a stream the caller made and keeps owning, which is never freed here. ```rust pub unsafe fn from_raw(capabilities: Capabilities, handle: *mut c_void) -> Self ``` Safety: `handle` must be null or a live stream of the vendor [`capabilities`](https://www.unum.cloud/docs/numkong/rust/external-memory.md#capabilities) name, a `cudaStream_t`, a `hipStream_t` or an `id`, null on the CPU, and outlive the result. ## capabilities The capabilities every consumer of this stream dispatches with. ```rust pub fn capabilities(&self) -> Capabilities ``` ## as_raw The vendor's handle, null for a default stream. ```rust pub fn as_raw(&self) -> *mut c_void ``` ## synchronize Waits for everything queued on this stream, after which whatever its verbs wrote is readable from the host. ```rust pub fn synchronize(&self) -> Result<(), Error> ``` The CPU has nothing to wait for. ## scope Runs `body`, whose verbs queue on this stream, then joins the stream, even when `body` fails or unwinds, and reports the first failure, the body's or the join's. ```rust pub fn scope<'env, Output>( &'env self, body: impl for<'scope> FnOnce(&'scope Scope<'scope, 'env>) -> Result, ) -> Result ``` A verb borrows what it reads and writes for the whole scope, so none of it can be touched before the join. Examples: ``` use numkong::{Capabilities, ShapesPackedMatrix, Stream}; let triangle = [[1.0_f32, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]]; let packs = ShapesPackedMatrix::new(&[&triangle[..]])?; let stream = Stream::default(Capabilities::cpu_enabled()); let mut distances = [1.0_f64; 1]; stream.scope(|scope| packs.rmsds_on(scope, &packs, &mut distances))?; assert!(distances[0] < 1e-6); ```