CubeCL - kernels de GPU basados en Rust para CUDA, ROCm y WGPU

• CubeCL es una extensión de lenguaje multiplataforma y de alto rendimiento que permite escribir kernels de GPU en Rust
• Ofrece soporte completo para funciones, genéricos y structs, y soporte parcial para traits, métodos e inferencia de tipos
• Soporta runtimes basados en WGPU, CUDA y ROCm, y también está en desarrollo un runtime JIT para CPU optimizado con instrucciones SIMD
• Proporciona vectorización automática, optimización en tiempo de compilación (Comptime) y autotuning para maximizar el rendimiento del código
• Usa un enfoque único que aprovecha el sistema de macros procedimentales de Rust para analizar y expandir código de kernels de GPU

---

Introducción a CubeCL

• CubeCL es una extensión del lenguaje Rust para escribir directamente kernels de GPU orientados a cómputo de alto rendimiento
• Aprovecha abstracciones de costo cero para permitir un desarrollo de kernels eficiente y fácil de mantener
• Actualmente ofrece soporte completo para funciones, genéricos y structs, y soporte parcial para traits, métodos e inferencia de tipos

Ejemplo simple de uso

• Usa el atributo #[cube] para marcar funciones que se ejecutarán en la GPU
• Las funciones gelu_array y gelu_scalar usan automáticamente instrucciones SIMD para mejorar el rendimiento
• El kernel se puede ejecutar con la función launch_unchecked

Runtimes compatibles

• WGPU: soporte GPU multiplataforma mediante Vulkan, Metal, DirectX y WebGPU
• CUDA: soporte optimizado para GPU de NVIDIA
• ROCm/HIP: soporte para GPU de AMD (en desarrollo)
• Runtime JIT para CPU: se está desarrollando un runtime JIT para CPU, optimizado con instrucciones SIMD y basado en Cranelift

Objetivos de CubeCL

• Permitir escribir kernels de cómputo de alto rendimiento sin depender de un hardware específico
• Mejorar la componibilidad, reutilización, capacidad de prueba y mantenibilidad del código mediante vectorización automática, comptime y autoajuste
• Proporcionar reutilización optimizada de buffers para alto rendimiento mediante estrategias de gestión de memoria
• Desarrollar componentes de álgebra lineal para construir un ecosistema de cómputo científico y de alto rendimiento

Cómo funciona CubeCL

• Usa un enfoque único que aprovecha el sistema de macros procedimentales de Rust para analizar y expandir el código de kernels de GPU
• Comptime: permite integrar fácilmente optimizaciones en tiempo de compilación sin transformar el código fuente original
• Vectorización automática: vectoriza las entradas de las funciones de CubeCL para determinar los elementos de vectorización de cada variable intermedia
• Integración con Rust: el código generado sigue siendo código Rust válido y puede distribuirse en un bundle sin depender de un runtime específico

Diseño y funciones especiales

• CubeCL está diseñado alrededor del concepto de cubos, y todas las APIs de cómputo pueden mapearse al hardware para usar una representación 3D
• Vectorización automática: puede mejorar el rendimiento usando instrucciones SIMD, y se puede acceder directamente a los elementos de vectorización dentro del kernel
• Comptime: modifica el IR del compilador en tiempo de ejecución para ofrecer optimización y flexibilidad
• Autoajuste: ejecuta pequeños benchmarks en tiempo de ejecución para elegir el kernel y la configuración óptimos

Consideraciones e historia

• CubeCL se encuentra actualmente en versión alfa, por lo que aún tiene partes sin pulir
• Originalmente comenzó como el backend de WebGPU del proyecto Burn
• Luego se expandió con la adición del runtime de CUDA hacia un IR general y un frontend en Rust
• Se desarrollaron conceptos generales para superar las diferencias entre CUDA y WebGPU