https://github.com/charles-r-earp/krnl

Developed for autograph.

• Similar functionality to CUDA and OpenCL.
• Supports GPU's and other Vulkan 1.2 capable devices.
• MacOS / iOS supported via MoltenVK.
• Kernels are written inline, entirely in Rust.
  • Simple iterator patterns can be implemented without unsafe.
  • Supports inline SPIR-V assembly.
  • DebugPrintf integration, generates backtraces for panics.
• Buffers on the host can be accessed natively as Vecs and slices.

krnlc

Kernel compiler for krnl.

• Built on spirv-builder.
• Supports dependencies defined in Cargo.toml.
• Uses spirv-tools to validate and optimize.
• Compiles to "krnl-cache.rs", so the crate will build on stable Rust.

Example

use krnl::{
    macros::module,
    anyhow::Result,
    device::Device,
    buffer::{Buffer, Slice, SliceMut},
};

#[module]
mod kernels {
    #[cfg(not(target_arch = "spirv"))]
    use krnl::krnl_core;
    use krnl_core::macros::kernel;

    pub fn saxpy_impl(alpha: f32, x: f32, y: &mut f32) {
        *y += alpha * x;
    }

    // Item kernels for iterator patterns.
    #[kernel]
    pub fn saxpy(alpha: f32, #[item] x: f32, #[item] y: &mut f32) {
        saxpy_impl(alpha, x, y);
    }

    // General purpose kernels like CUDA / OpenCL.
    #[kernel]
    pub fn saxpy_global(alpha: f32, #[global] x: Slice<f32>, #[global] y: UnsafeSlice<f32>) {
        use krnl_core::buffer::UnsafeIndex;

        let global_id = kernel.global_id();
        if global_id < x.len().min(y.len()) {
            saxpy_impl(alpha, x[global_id], unsafe { y.unsafe_index_mut(global_id) });
        }
    }
}

fn saxpy(alpha: f32, x: Slice<f32>, mut y: SliceMut<f32>) -> Result<()> {
    if let Some((x, y)) = x.as_host_slice().zip(y.as_host_slice_mut()) {
        x.iter()
            .copied()
            .zip(y.iter_mut())
            .for_each(|(x, y)| kernels::saxpy_impl(alpha, x, y));
        return Ok(());
    }
    if true {
        kernels::saxpy::builder()?
            .build(y.device())?
            .dispatch(alpha, x, y)
    } else {
        // or
        kernels::saxpy_global::builder()?
            .build(y.device())?
            .with_global_threads(y.len() as u32)
            .dispatch(alpha, x, y)
    }
}

fn main() -> Result<()> {
    let x = vec![1f32];
    let alpha = 2f32;
    let y = vec![0f32];
    let device = Device::builder().build().ok().unwrap_or(Device::host());
    let x = Buffer::from(x).into_device(device.clone())?;
    let mut y = Buffer::from(y).into_device(device.clone())?;
    saxpy(alpha, x.as_slice(), y.as_slice_mut())?;
    let y = y.into_vec()?;
    println!("{y:?}");
    Ok(())
}

v0.1.0

krnl was developed to replace core functionality in autograph:

• Only targets Vulkan, more portable than Metal / DX12.
  • Metal is supported via MoltenVK.
• GPGPU kernels implemented inline in Rust:
  • Kernels can be defined in the same file, near where they are invoked.
  • Modules allow sharing code between host and device.
  • Kernel bindings are type safe, checked at compile time.
  • Simple iterator patterns can be implemented without unsafe.
  • Supports specialization constants provided at runtime.
  • DeviceInfo includes useful properties:
    • Max / default threads per group.
    • Max / min threads per subgroup.
  • With DebugPrintf, kernel panics produce errors on the host.
  • krnlc generates a device crate and invokes spirv-builder.
    • spirv-builder / spirv-tools are compiled once on install.
    • Significantly streamlines and accelerates workflow.
  • Kernels are compressed to reduce package and binary size.
• Device operations readily execute:
  • Block until kernels / transfers can queue.
  • An operation can be queued while another is executing.
  • Reduced latency, better repeatability, reliability, and performance.
• Device buffers can be copied by the host if host visible.
• Large buffer copies are streamed rather than allocating a large temporary:
  • Reuses a few small buffers for transfers.
  • Overlaps host and device copies.
  • Performance significantly closer to CUDA.
  • Also streams between devices.
• Device buffers can be i32::MAX bytes (~2 GB, up from 256 MB).
• Scalar / ScalarBufferBase replaces Float / FloatBuffer:
  • Streamlined conversions between buffers.
• Buffers can be sliced.
• Supports wasm (without device feature).

MSRV: 1.70.0