CubeCL: GPU Kernels in Rust for CUDA, ROCm, and WGPU

Summary

Multi-platform high-performance compute language extension for Rust. With CubeCL, you can program your GPU using Rust, taking advantage of zero-cost abstractions to develop maintainable, flexible, and efficient compute kernels. CubeCL currently fully supports functions, generics, and structs, with partial support for traits, methods and type inference. As the project evolves, we anticipate even broader support for Rust language primitives, all while maintaining optimal performance. Example Simply annotate functions with the cube attribute to indicate that they should run on the GPU. use cubecl :: prelude :: * ; # [ cube ( launch_unchecked ) ] /// A [Line] represents a contiguous series of elements where SIMD operations may be available. /// The runtime will automatically use SIMD instructions when possible for improved performance. fn gelu_array < F : Float > ( input : & Array < Line < F > > , output : & mut Array < Line < F > > ) { if ABSOLUTE_POS < input . len ( ) { output [ ABSOLUTE_POS ] = gelu_scalar ( input [ ABSOLUTE_POS ] ) ; } } # [ cube ] fn gelu_scalar < F : Float > ( x : Line < F > ) -> Line < F > { // Execute the sqrt function at comptime. let sqrt2 = F :: new ( comptime ! ( 2.0f32 . sqrt ( ) ) ) ; let tmp = x / Line :: new ( sqrt2 ) ; x * ( Line :: erf ( tmp ) + 1.0 ) / 2.0 } You can then launch the kernel using the autogenerated gelu_array::launch_unchecked function. pub fn launch < R : Runtime > ( device : & R :: Device ) { let client = R :: client ( device ) ; let input = & [ - 1. , 0. , 1. , 5. ] ; let vectorization = 4 ; let output_handle = client . empty ( input . len ( ) * core :: mem :: size_of :: < f32 > ( ) ) ; let input_handle = client . create ( f32 :: as_bytes ( input ) ) ; unsafe { gelu_array :: launch_unchecked :: < f32 , R > ( & client , CubeCount :: Static ( 1 , 1 , 1 ) , CubeDim :: new ( input . len ( ) as u32 / vectorization , 1 , 1 ) , ArrayArg :: from_raw_parts :: < f32 > ( & input_handle , input . len ( ) , vectorization as u8 ) ,...