kaio 0.0.4

Rust-native GPU kernel authoring framework. Write GPU compute kernels in Rust, automatically lower to PTX. Cross-platform (Windows + Linux), type-safe, no CUDA C++ required.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68

# KAIO


**Rust-native GPU kernel authoring framework.**

KAIO (καίω — to kindle, to ignite) lets developers write GPU compute
kernels in Rust and lower them to PTX for execution on NVIDIA GPUs.
A Rust alternative to OpenAI's Triton — Windows + Linux, automatic PTX
generation, Rust type safety, no CUDA C++ toolchain required.

## Quick Start


```rust
use kaio::prelude::*;

#[gpu_kernel(block_size = 256)]

fn saxpy(x: &[f32], y: &mut [f32], alpha: f32, n: u32) {
    let idx = thread_idx_x() + block_idx_x() * block_dim_x();
    if idx < n {
        y[idx] = alpha * x[idx] + y[idx];
    }
}
```

The `#[gpu_kernel]` macro lowers `saxpy` to PTX and generates a typed
`saxpy::launch()` function. PTX is generated once at first call and
cached for the process lifetime.

## Features


| Feature | Syntax | Status |
|---------|--------|--------|
| Arithmetic | `+`, `-`, `*`, `/`, `%`, `+=`, `-=`, `*=`, `/=` | ✅ |
| Comparisons | `<`, `<=`, `>`, `>=`, `==`, `!=` | ✅ |
| Control flow | `if`/`else`, `for`, `while` | ✅ |
| Array access | `a[idx]` (global memory) | ✅ |
| Shared memory | `shared_mem![f32; 256]` | ✅ |
| Synchronization | `bar_sync()` | ✅ |
| Warp shuffle | `shfl_sync_down/up/bfly()` | ✅ |
| Reductions | `block_reduce_sum()`, `block_reduce_max()` | ✅ |
| Type casts | `x as f32` | ✅ |
| Math builtins | `sqrt`, `exp`, `log`, `tanh`, `abs`, `min`, `max` | ✅ |
| Thread indices | `thread_idx_x()`, `block_idx_x()`, `block_dim_x()` | ✅ |
| FMA | `fma(a, b, c)` | ✅ |
| 2D blocks | `block_size = (16, 16)`, `thread_idx_y()` | ✅ |
| Tiled matmul | `kaio_ops::matmul()` (31% of cuBLAS) | ✅ |

## Architecture


| Crate | Description |
|-------|-------------|
| `kaio` | Umbrella crate — re-exports everything via `prelude` |
| `kaio-macros` | `#[gpu_kernel]` proc macro |
| `kaio-core` | PTX IR, instruction emitters, zero external dependencies |
| `kaio-runtime` | CUDA driver wrapper via cudarc |
| `kaio-ops` | Pre-built GPU operations (matmul, more planned) |

## Status


**Phase 4 complete** — tiled matmul (31% of cuBLAS sgemm on RTX 4090),
`kaio-ops` crate, 2D thread blocks, FMA, PTX inspection tools.
This is early development software. API will change.

See the [repository](https://github.com/dmriding/kaio) for the full
roadmap, CHANGELOG, and development logs.

## License


Licensed under either of Apache-2.0 or MIT at your option.