burn
burn copied to clipboard
CubeCL: Compute Language Extension in Rust for Multi-target GPU kernels
The current JIT representation is very practical for targeting multiple GPU languages; however the kernels are very hard to read and maintain as it is composed of many declarative macros and is too close to assembly (for instance, there's no support for complex one-liners that need operation priority).
Thus, we will be developing a new language using procedural macros, that feels like writing plain Rust but compiles to a GPU language (using a backend's compiler and runtime like we have for wgpu at the moment). We call it CubeCL (Cube Compute Language), as a reference to the topology of workgroups often present in GPU languages.
I will be working for several weeks on this, migrating (again) and optimizing our jit kernels along the way. When it's done it must also be well-documented because it shall become the standard way of writing custom kernels in user space.
Here's a sketch of the steps to achieve (despite lots of uncertainty still):
- [x] Static analysis of Cube basic code
- [x] Codegen of JIT representation
- [ ] Support of all operators
- [ ] Support of GPU-specific instructions
- [ ] Migrate & optimize kernels
- [ ] Write an official doc