CubeCL - kernels de GPU em Rust para CUDA, ROCm e WGPU

• CubeCL é uma extensão de linguagem multiplataforma de alto desempenho que permite escrever kernels de GPU em Rust
• Oferece suporte completo a funções, genéricos e structs, e suporte parcial a traits, métodos e inferência de tipos
• Suporta runtimes baseados em WGPU, CUDA e ROCm, e também está desenvolvendo um runtime JIT para CPU otimizado com instruções SIMD
• Fornece recursos de vetorização automática, otimização em tempo de compilação (Comptime) e autotuning para maximizar o desempenho do código
• Usa uma abordagem única que aproveita o sistema de macros procedurais do Rust para analisar e expandir código de kernel de GPU

---

Introdução ao CubeCL

• CubeCL é uma extensão da linguagem Rust para permitir escrever kernels de GPU diretamente em computação de alto desempenho
• Aproveita abstrações de custo zero para possibilitar o desenvolvimento de kernels eficiente e fácil de manter
• Atualmente oferece suporte completo a funções, genéricos e structs, e suporte parcial a traits, métodos e inferência de tipos

Exemplo simples de uso

• Usa o atributo #[cube] para indicar funções que serão executadas na GPU
• As funções gelu_array e gelu_scalar usam automaticamente instruções SIMD para melhorar o desempenho
• É possível executar kernels com a função launch_unchecked

Runtimes suportados

• WGPU: suporte multiplataforma a GPU via Vulkan, Metal, DirectX e WebGPU
• CUDA: suporte otimizado para GPUs NVIDIA
• ROCm/HIP: suporte a GPUs AMD (em desenvolvimento)
• Runtime JIT para CPU: em desenvolvimento, um runtime JIT para CPU otimizado com instruções SIMD baseado em Cranelift

Objetivos do CubeCL

• Permitir escrever kernels de computação de alto desempenho sem depender de hardware específico
• Melhorar a configurabilidade, reutilização, testabilidade e manutenibilidade do código com recursos de vetorização automática, comptime e ajuste automático
• Oferecer reutilização otimizada de buffers para alto throughput por meio de estratégias de gerenciamento de memória
• Desenvolver componentes de álgebra linear para construir um ecossistema de computação científica e de alto desempenho

Como o CubeCL funciona

• Usa uma abordagem única que aproveita o sistema de macros procedurais do Rust para analisar e expandir código de kernel de GPU
• Comptime: permite integrar facilmente otimizações em tempo de compilação sem transformar o código-fonte original
• Vetorização automática: vetoriza as entradas das funções CubeCL e determina o fator de vetorização de cada variável intermediária
• Integração com Rust: o código gerado continua sendo código Rust válido e pode ser distribuído em bundle sem dependência de um runtime específico

Design e recursos especiais

• O CubeCL é projetado em torno de cubos, e todas as APIs de computação são mapeadas ao hardware para usar representação 3D
• Vetorização automática: pode melhorar o desempenho usando instruções SIMD, e o fator de vetorização pode ser acessado diretamente dentro do kernel
• Comptime: modifica a IR do compilador em tempo de execução para oferecer otimização e flexibilidade
• Ajuste automático: executa pequenos benchmarks em tempo de execução para selecionar o kernel e a configuração ideais

Observações e histórico

• O CubeCL está atualmente em versão alfa, com partes ainda não refinadas
• Originalmente começou como o backend WebGPU do projeto Burn
• Depois foi expandido com a adição do runtime CUDA para uma IR genérica e um frontend em Rust
• Conceitos genéricos foram desenvolvidos para superar as diferenças entre CUDA e WebGPU