Kina.ai

Specification

Phase 1, Section 3: Runtime Implementation

Phase 1: Implementation — the shape-typed core compiled end-to-end. Section 3 of 7: Runtime Implementation.

1. What this section commits to

The Rust runtime — what executes the code the OCaml compiler produces. Specifically:

Section 2 covered the compiler — what produces the code. Section 3 covers the runtime — what runs it. Together they form the working Phase 1 stack; Section 4 adds the operation extensions; the rest handle quality, verification, and exit.

2. Architecture overview

The runtime has three jobs:

  1. Allocate and manage memory for tensors, on host and on GPU device(s).
  2. Dispatch kernels that the compiler can't inline into LLVM IR — primarily GPU operations, which call into vendor libraries (cuBLAS, cuDNN) or in-house naive CUDA kernels.
  3. Expose a stable C ABI that the compiler emits calls to and that external code (PyTorch via DLPack, eventually domain libraries) consumes.

The runtime is not a host language for the compiler. The compiler doesn't run on top of it. The compiler produces shared objects; the runtime provides the symbols those shared objects need at link time and the services they call at runtime.

2.1 Where the runtime sits in the pipeline

.kina source → [OCaml compiler] → .so file
                                       │
                                       │ at runtime, the .so is loaded by:
                                       ▼
                              ┌─────────────────┐
                              │  user program   │  (could be Python, OCaml, Rust)
                              └────────┬────────┘
                                       │ links / dlopens the .so
                                       │ calls runtime APIs through FFI
                                       ▼
                              ┌─────────────────┐
                              │  Rust runtime   │
                              │  (libkina.so) │
                              └────────┬────────┘
                                       │ dispatches to:
                                       ▼
                          cuBLAS / cuDNN / in-house kernels

Two shared objects participate: libkina.so (the runtime) and the user's compiled program.so (one per compiled program). The user's program calls into the runtime for any operation it can't inline.

2.2 What this means for FFI design

The runtime exposes a C ABI. Every public symbol is extern "C" with #[no_mangle]. Every function returns an error code (per §10). All public types are #[repr(C)]. No Rust-specific machinery (panics, Drop, lifetimes) crosses the FFI boundary — the C ABI is the contract.

This is the standard pattern for Rust libraries that interoperate with C/C++/other languages. It costs ergonomic flexibility on the runtime's internal API but is the only correct choice for a library that will be linked into programs in arbitrary host languages.

3. Project structure

The Rust crate lives at runtime/ under the monorepo root.

runtime/
├── Cargo.toml
├── build.rs                              — find CUDA toolkit, link cuBLAS/cuDNN, compile .cu kernels
├── cbindgen.toml                         — config for header generation
├── include/
│   └── kina_runtime.h                  — generated C header (committed for downstream use)
├── src/
│   ├── lib.rs                            — crate root; public re-exports
│   ├── error.rs                          — error type, FFI error codes
│   ├── device.rs                         — device descriptors, enumeration
│   ├── stream.rs                         — CUDA stream management
│   ├── memory/
│   │   ├── mod.rs                        — MemoryResource trait
│   │   ├── cuda.rs                       — CudaMemoryResource
│   │   ├── cuda_async.rs                 — CudaAsyncMemoryResource
│   │   ├── pool.rs                       — PoolMemoryResource<Inner>
│   │   └── logging.rs                    — LoggingMemoryResource<Inner>
│   ├── tensor/
│   │   ├── mod.rs                        — Tensor struct
│   │   ├── dtype.rs                      — DType enum + canonical FFI codes
│   │   ├── dlpack.rs                     — DLPack import/export
│   │   └── view.rs                       — non-owning Tensor views
│   ├── dispatch/
│   │   ├── mod.rs                        — kernel dispatch
│   │   ├── cublas.rs                     — cuBLAS wrappers
│   │   ├── cudnn.rs                      — cuDNN wrappers
│   │   └── inhouse.rs                    — in-house naive kernel calls
│   ├── kernels/                          — .cu source compiled at build time
│   │   ├── elementwise.cu
│   │   ├── reduce.cu
│   │   ├── softmax.cu
│   │   ├── gather.cu
│   │   └── repeat_interleave.cu
│   └── ffi/
│       ├── mod.rs                        — extern "C" entry points
│       ├── memory.rs                     — memory FFI
│       ├── tensor.rs                     — tensor FFI
│       └── dispatch.rs                   — kernel-dispatch FFI
├── tests/
│   ├── memory_resource.rs
│   ├── tensor_lifecycle.rs
│   ├── dlpack_roundtrip.rs               — the flagship interop demo
│   ├── cublas_correctness.rs
│   └── kernel_correctness.rs
└── benches/
    ├── allocator_benchmarks.rs
    └── dispatch_benchmarks.rs

A few notes on the layout:

3.1 Key dependencies

[dependencies]
thiserror = "1"                 # error type derivation
libc = "0.2"                    # C ABI types
cudarc = "0.10"                 # CUDA Driver API bindings (alternative: cust)
once_cell = "1"                 # lazy statics for global state (used sparingly)
parking_lot = "0.12"            # faster Mutex/RwLock than std

[build-dependencies]
cc = "1"                        # compile .cu kernels via build.rs
bindgen = "0.69"                # if we need to wrap CUDA headers directly
cbindgen = "0.26"               # generate the C header from Rust

Phase 1 commits to cudarc for CUDA Driver API access. Alternative cust exists; both are mature; cudarc has slightly better recent maintenance. Pinning to cudarc 0.10; upgrade in deliberate cycles (per Phase 0 §5.7's pinning policy applied to Rust dependencies).

4. Memory management

The largest single piece of runtime work, derived from the ecosystem-architecture forward track. Phase 1 commits to the trait + four implementations from that document.

4.1 The trait

pub trait MemoryResource: Send + Sync {
    /// Allocate `size` bytes on the appropriate device, ordered against `stream`.
    fn allocate(&self, size: usize, stream: &Stream) 
        -> Result<DevicePtr, AllocError>;

    /// Deallocate, ordered against `stream`. The size must match the original allocation.
    fn deallocate(&self, ptr: DevicePtr, size: usize, stream: &Stream) 
        -> Result<(), DeallocError>;

    /// Whether this resource supports stream-ordered allocation.
    fn supports_streams(&self) -> bool;

    /// Human-readable name for diagnostics.
    fn allocator_name(&self) -> &'static str;
}

The trait is Send + Sync because resources are shared across threads and across kernel launches on different streams. Stream-ordered semantics are explicit (the stream parameter); resources that don't support stream ordering implement them as fences before their underlying allocator calls.

4.2 The four initial implementations

CudaMemoryResource: thin wrapper around cudaMalloc / cudaFree. No pooling, no async. The simplest correct implementation; baseline for benchmarking; usable when other resources fail to initialize.

CudaAsyncMemoryResource: wraps cudaMallocAsync / cudaFreeAsync (CUDA 11.2+). Stream-ordered allocation with the driver's internal pool. Requires the device to have memory pool support; falls back to CudaMemoryResource if not.

PoolMemoryResource<Inner: MemoryResource>: pre-allocates a block from the inner resource (default 1 GiB) and sub-allocates from it via free-list management. Reduces fragmentation compared to per-call allocations; amortizes the cost of underlying calls. Configurable: pool size, alignment, growth policy.

LoggingMemoryResource<Inner: MemoryResource>: wraps any other resource and logs every alloc/free with size, stream, allocator, and stack trace (in debug builds). Used for diagnostics; not for production.

The composition pattern: a typical configuration is LoggingMemoryResource<PoolMemoryResource<CudaAsyncMemoryResource>>. Each layer adds behavior; the base layer hits the actual driver.

4.3 The current resource

A thread-local "current resource" pointer, with a global default. All allocations in the runtime — including those triggered by user code through the FFI — go through the current resource:

// Set the current resource for the current thread (e.g., during a benchmark).
pub fn with_resource<R: MemoryResource>(resource: &R, f: impl FnOnce()) { ... }

// The global default, used when no thread-local override is set.
pub fn set_default_resource(resource: Box<dyn MemoryResource>) { ... }

This lets callers configure the allocator without recompiling. It also lets diagnostics tools wrap allocations transparently (e.g., a memory-tracking dashboard sets the current resource to a tracking wrapper for the duration of a workload).

4.4 Acceptance criteria for memory management

Criterion Target
CudaMemoryResource correctness No leaks under stress test (1M alloc/free cycles)
CudaAsyncMemoryResource correctness Same; plus correct stream ordering verified
PoolMemoryResource performance At least 5× speedup over raw CudaMemoryResource on a many-small-alloc workload
LoggingMemoryResource overhead Less than 10% throughput cost in debug builds; near zero in release
Thread-local resource switching Working; benchmarked; documented
Send + Sync correctness Clippy + Miri verification on the test suite

5. The Tensor type

The single most ABI-critical type in the runtime.

5.1 Layout

#[repr(C)]
pub struct Tensor {
    // ---- DLPack-compatible prefix (must match DLTensor exactly) ----
    pub data: *mut u8,
    pub device: DLDevice,         // type + id
    pub ndim: i32,
    pub dtype: DLDataType,        // code + bits + lanes
    pub shape: *const i64,
    pub strides: *const i64,      // null if compact
    pub byte_offset: u64,

    // ---- Internal fields (after prefix) ----
    pub allocator: *const c_void, // type-erased MemoryResource pointer
    pub stream: StreamHandle,
    pub size_bytes: usize,        // for deallocation
    pub flags: TensorFlags,       // owned, view, etc.
    // Pad to alignment if needed.
}

The first seven fields exactly match DLTensor. The remaining fields are runtime-internal but follow the prefix in memory.

This means: a *const Tensor can be cast to *const DLTensor (zero cost) and read as DLPack. A *const DLTensor cannot be cast back without verification — the internal fields might be uninitialized — but it can be wrapped in a fresh Tensor with computed internal fields.

5.2 Lifecycle

Three creation paths:

  1. Owning allocation: Tensor::new(shape, dtype, device, stream) allocates through the current memory resource.
  2. Non-owning view: Tensor::view_of(parent, ...) produces a tensor sharing storage with another. Doesn't deallocate on drop.
  3. From DLPack: Tensor::from_dlpack(managed) adopts a tensor from another framework. Lifecycle managed by the deleter callback.

Three destruction paths:

  1. Owning Drop: deallocates through the recorded allocator.
  2. View Drop: no-op for storage; may release reference to parent.
  3. DLPack export: Tensor::to_dlpack() produces a DLManagedTensor whose deleter is our deallocation function. The Rust Tensor is consumed; ownership transfers to the consumer.

5.3 The DLPack interop module

A dedicated module (tensor/dlpack.rs) handles the DLPack C struct translations. Concretely:

#[repr(C)]
pub struct DLManagedTensor {
    pub dl_tensor: DLTensor,
    pub manager_ctx: *mut c_void,
    pub deleter: Option<unsafe extern "C" fn(*mut DLManagedTensor)>,
}

impl Tensor {
    /// Export this tensor as a DLPack DLManagedTensor.
    /// 
    /// Ownership transfers to the consumer; the consumer must call
    /// the deleter when finished.
    pub fn to_dlpack(self) -> Box<DLManagedTensor> { ... }

    /// Import a DLManagedTensor produced by another framework.
    /// 
    /// The resulting Tensor will call the original deleter when dropped.
    pub fn from_dlpack(managed: *mut DLManagedTensor) -> Result<Tensor, DLPackError> { ... }
}

The deleter contract: when the consumer finishes with the tensor, it calls the deleter with the managed pointer. The deleter restores the Rust ownership semantics (drops the original allocator handle, etc.). Getting this right is the most error-prone part of the runtime; comprehensive testing in tests/dlpack_roundtrip.rs is non-negotiable.

5.4 The interop demo (acceptance criterion)

The flagship validation for Phase 1's runtime work:

# Python script using PyTorch and our compiler
import torch
import kina  # the eventual Python binding; in Phase 1 may use ctypes directly

# 1. Create a tensor with our runtime via the compiled artifact
program = kina.load("matmul_program.so")
a = kina.zeros((128, 256), dtype="f32", device="cuda:0")
b = kina.zeros((256, 64), dtype="f32", device="cuda:0")

# 2. Run our compiled program: c = a @ b
c = program.matmul(a, b)

# 3. Hand c to PyTorch via DLPack (zero-copy)
c_torch = torch.from_dlpack(c.to_dlpack())

# 4. Run PyTorch operations on it
c_torch_relu = torch.nn.functional.relu(c_torch)

# 5. Hand back to our runtime
c_back = kina.from_dlpack(torch.utils.dlpack.to_dlpack(c_torch_relu))

# 6. Verify: c_back equals relu(a @ b) computed independently
expected = torch.relu(torch.zeros(128, 64))  # the expected result
assert torch.allclose(c_torch_relu.cpu(), expected)

This demo passing end-to-end is what "Phase 1 runtime DLPack interop works" means concretely.

6. Kernel dispatch

The runtime implements the FFI entry points the compiler emits calls to (per Section 2 §8).

6.1 The dispatch surface

Approximately 18 entry points, organized by operation class:

// Memory operations
RuntimeError runtime_alloc(Tensor* out, /* shape, dtype, device, stream */);
RuntimeError runtime_free(Tensor* tensor);

// Data movement
RuntimeError runtime_copy_to_device(Tensor* out, const Tensor* in, int device_id);
RuntimeError runtime_copy_to_host(Tensor* out, const Tensor* in);

// Matrix operations
RuntimeError runtime_gpu_matmul(Tensor* out, const Tensor* a, const Tensor* b, /* params */);

// Elementwise unary / binary / reductions
RuntimeError runtime_gpu_unary(Tensor* out, const Tensor* in, OpCode op);
RuntimeError runtime_gpu_binary(Tensor* out, const Tensor* a, const Tensor* b, OpCode op);
RuntimeError runtime_gpu_reduce(Tensor* out, const Tensor* in, int axis, OpCode op);

// Specialized
RuntimeError runtime_gpu_softmax(Tensor* out, const Tensor* in, int axis);
RuntimeError runtime_gpu_gather(Tensor* out, const Tensor* values, const Tensor* indices);
RuntimeError runtime_gpu_repeat_interleave(Tensor* out, const Tensor* in, int axis, int repeats);

// Casts and structural
RuntimeError runtime_gpu_cast(Tensor* out, const Tensor* in, DType target_dtype);
RuntimeError runtime_gpu_transpose(Tensor* out, const Tensor* in, const int* permutation);
RuntimeError runtime_gpu_reshape(Tensor* out, const Tensor* in, const int64_t* new_shape, int new_ndim);
RuntimeError runtime_gpu_broadcast(Tensor* out, const Tensor* in, const int64_t* target_shape);
RuntimeError runtime_gpu_slice(Tensor* out, const Tensor* in, int axis, int64_t start, int64_t stop);
RuntimeError runtime_gpu_concat(Tensor* out, const Tensor* a, const Tensor* b, int axis);

About 17 entry points. Each returns a RuntimeError code; output is via the out pointer; input tensors are const.

The discipline (per the bootstrap-and-self-modeling forward track §7.3): small, orthogonal, stateless. Each call is self-contained; no implicit state on the runtime side; Phase 2's effect handlers can intercept any of these.

6.2 Behind each entry point

Each FFI function dispatches to the appropriate backend:

The dispatch logic lives in dispatch/mod.rs with backend modules (cublas.rs, cudnn.rs, inhouse.rs).

6.3 In-house kernels

Compiled at runtime build time from the .cu files in src/kernels/. The build.rs script invokes nvcc to produce a static library linked into the runtime.

Each kernel is a straightforward CUDA implementation — no tile-IR magic, no autotuning, no specialization for hardware capabilities. Phase 1's goal is correctness and "within 3-5× of PyTorch performance"; aggressive optimization is Phase 5 work.

The list of in-house kernels:

Estimated total: maybe 40–60 kernel templates instantiated across dtypes; 1500–2500 lines of CUDA. Manageable.

7. CUDA toolchain integration

7.1 Linkage

The runtime links against:

build.rs finds the CUDA toolkit (via CUDA_PATH env var or standard paths), generates appropriate cargo:rustc-link-search and cargo:rustc-link-lib directives. CUDA 12.x is the pinned baseline (per Phase 0 §5.7 applied here).

Static vs. dynamic linkage: dynamic by default (smaller binary, system-installed CUDA). Static available via a Cargo feature flag for distribution scenarios.

7.2 Stream management

CUDA streams are fundamental to async dispatch. The runtime manages a pool of streams per device:

pub struct StreamPool {
    streams: Vec<Stream>,
    next: AtomicUsize,
}

impl StreamPool {
    pub fn checkout(&self) -> StreamHandle { ... }
    pub fn checkin(&self, handle: StreamHandle) { ... }
}

Each tensor carries a StreamHandle. Operations on the tensor are ordered against that stream. When two tensors with different streams are inputs to one operation, the runtime synchronizes (typically via cudaStreamWaitEvent) before launching.

This is the standard pattern; PyTorch and JAX both implement variations. Getting it right matters for both correctness (use-after-free across streams) and performance (avoiding spurious synchronization).

7.3 Device discovery

At runtime initialization:

pub fn enumerate_devices() -> Vec<DeviceInfo> { ... }

pub struct DeviceInfo {
    pub id: i32,
    pub name: String,
    pub compute_capability: (i32, i32),
    pub total_memory: usize,
    pub multi_processor_count: i32,
    // ... more from cudaDeviceProp
}

Used for capability checks (eventually; Phase 5's capability typing will lean on this) and for diagnostics. Phase 1 mostly just needs device count and basic properties.

7.4 Error translation

CUDA's cudaError_t and cuBLAS/cuDNN error codes get translated to the runtime's RuntimeError enum. The translation preserves enough information for diagnostics (the original CUDA error string) while flattening to a stable enum the FFI consumers can match on.

8. Error handling across the FFI boundary

A clean error model is essential for an FFI library. Phase 1 commits to error codes via return value, never panics across the FFI boundary.

8.1 The error type

Internal Rust code uses idiomatic Result<T, RuntimeError>:

#[derive(Debug, thiserror::Error)]
pub enum RuntimeError {
    #[error("out of memory: requested {requested} bytes")]
    OutOfMemory { requested: usize },

    #[error("CUDA error: {0}")]
    Cuda(String),

    #[error("invalid shape: {0}")]
    InvalidShape(String),

    #[error("DLPack error: {0}")]
    DLPack(String),

    #[error("kernel launch failed: {0}")]
    KernelLaunch(String),

    // ... more
}

8.2 The FFI surface

Across the C ABI, errors are an enum:

#[repr(C)]
pub enum FfiErrorCode {
    Success = 0,
    OutOfMemory = 1,
    CudaError = 2,
    InvalidShape = 3,
    DLPackError = 4,
    KernelLaunchFailure = 5,
    // ...
}

Plus a thread-local "last error" mechanism for retrieving the detailed message:

#[no_mangle]
pub extern "C" fn runtime_last_error_message() -> *const c_char { ... }

The caller checks the return code; if non-zero, calls runtime_last_error_message to get the human-readable detail. This is the libc/CUDA/Vulkan pattern; well-understood by anyone consuming a C library.

8.3 No panics across the boundary

Every public extern "C" function is wrapped:

#[no_mangle]
pub extern "C" fn runtime_gpu_matmul(...) -> FfiErrorCode {
    let result = std::panic::catch_unwind(|| {
        // The actual implementation
        do_matmul(...)
    });

    match result {
        Ok(Ok(())) => FfiErrorCode::Success,
        Ok(Err(e)) => { record_last_error(e); error_to_code(e) },
        Err(_) => { record_last_error("panic"); FfiErrorCode::InternalError },
    }
}

A panic in Rust crossing the FFI boundary is undefined behavior. catch_unwind is the only correct way to handle it. Every public entry point uses this wrapper; a procedural macro generates the boilerplate.

9. Implementation milestones (runtime-specific)

Section 1's three milestones, refined for the runtime:

9.1 Milestone 1 (months 1–2): CPU-only minimum

The runtime is a small static library at this point; only the host-side allocation paths are implemented. The compiler from Section 2 produces CPU code that runs without needing the runtime for the operations themselves; the runtime is needed for tensor lifecycle management around the compiled functions.

9.2 Milestone 2 (months 3–4): GPU support

This is the biggest milestone for the runtime. Most of the complexity lives here.

9.3 Milestone 3 (months 5–6): Polish and operation extensions

10. Acceptance criteria specific to the runtime

Criterion Target
All FFI entry points implemented and tested Each has at least one positive and one negative test
Tensor lifecycle correct under stress 1M alloc/dealloc cycles, no leaks (Valgrind/Miri)
DLPack roundtrip works with PyTorch The §5.4 demo passes in CI
Pool allocator performance ≥5× cuMalloc on a many-small-alloc workload
Stream ordering correct Operations on streams produce results in the right order; verified by stress test
No panics escape the FFI boundary Property tested with random invalid inputs
Send + Sync correctness Clippy + Miri clean
Error messages preserve information Every RuntimeError produces a unique, human-readable diagnostic
Header (kina_runtime.h) is up to date Generated by cbindgen in CI; mismatch fails the build
Runtime compiles cleanly on Linux + macOS Windows is later (Phase 5+); CI covers Linux + macOS

The runtime is under 8K lines of Rust as a rough budget. If it grows beyond, refactor before continuing — same discipline as the compiler's 15K budget (Section 2 §12.5).

11. Open issues introduced by this section

  1. Static vs. dynamic CUDA linkage. The default is dynamic; static is via Cargo feature. Production deployment scenarios may want static (fewer runtime dependencies); developer scenarios want dynamic (smaller binaries, faster builds). The default may need to flip later; flagged.

  2. Multi-GPU support in Phase 1. A single cuda:global(0) device is the Phase 1 target; multi-device sharding is Phase 3. But the runtime should not assume single-GPU — Device is parameterized by id from day one. The risk: Phase 1 testing only covers single-GPU, so multi-GPU bugs may lurk. Mitigation: at least one Phase 1 test exercises two devices to catch obvious issues.

  3. The kernel build pipeline. nvcc invocation from build.rs is platform-dependent and finicky. If it breaks on some configurations, contributors will hit it before contributing. The mitigation: clear error messages in build.rs, documented troubleshooting in the README.

  4. PyPI / cargo distribution story. Eventually the runtime needs to be installable as a normal Cargo dependency or, for Python users, as a PyPI wheel. Phase 1 is "git clone and build"; distribution is later (Phase 7+ likely). Worth noting that the runtime's design (single shared object, stable C ABI) keeps the path open.

  5. Error message localization. Phase 1's error messages are English-only. International deployment (if it ever matters) will want translation; the message-formatting layer should be designed for it. Cheap to do right now; expensive to retrofit.

  6. Memory manager for tabular data (Arrow). The trait allocates raw bytes (per the bootstrap track's framing). When typed dataframes land in Phase 5+, Arrow buffers (validity bitmaps, offset arrays, value buffers, dictionary indices) will be allocated through this trait. The trait is designed for it; flagged as an anticipated extension.

12. Section 4 preview

Section 4 covers the operation set extensions identified during Phase 0:

Each gets a typing rule (extending Phase 0's rules without contradicting them), a CPU lowering (through MLIR linalg), a GPU lowering (via the runtime FFI from this section), and tests.

After Section 4, the language has the complete operation set Phase 1 commits to. The remaining sections cover verification (5), testing and ABI stewardship (6), and exit criteria (7).

End of Phase 1, Section 3.