Burn Deep Learning Framework 0.11.0 Released: Just-in-Time Automatic Kernel Fusion & Founding Announcement

130 points ksyiros

It's been approximately a year since I announced Burn, a novel Deep Learning framework in Rust. Since that initial announcement, the project's growth has exceeded my expectations. Many individuals from the Rust community have generously contributed to this project, enhancing and accelerating its development.

Kernel Fusion

Over my posts in the last year, I have outlined plans for automatic kernel fusion to enhance the performance of memory-bound operations. The first version of this feature is now unveiled with Burn v0.11.0. As envisaged, it is designed as a backend decorator, akin to how autodiff is already implemented in Burn. Backends wanting to support automatic kernel fusion should provide a list of optimizations derived from a stream of operations.

Currently, only our WGPU backend supports fusing element-wise operations, but the speedup observed is already quite impressive. For instance, you now have the capability to implement your own Gelu activation function with your own error function approximation, and the framework will compile the stream of operations into a single WGSL compute shader. The execution speed is more than ten times faster than without the new decorator (I will provide precise benchmarks in blog posts in the near future). The objective is to empower users with an easy-to-use tensor API while letting the framework optimize the execution, without requiring users to implement optimized kernels for each model and backend they want to target.
While there is ongoing work to enhance caching and introduce additional optimizations, we also aim to fully support Autotune, our solution for kernel selection, in a combination with our fusion strategy. Detailed explanations of this solution will also be provided in blog posts over the coming months. Once we are satisfied with the progress, we plan to implement more backends supporting our automatic kernel fusion strategy.

Full Release Note: https://github.com/tracel-ai/burn/releases/tag/v0.11.0

Tracel AI

In recent weeks, I also announced the founding of Tracel AI (https://tracel.ai/), a company focused on delivering efficient AI infrastructures. Currently, our primary focus is on growing Burn, with plans to develop complementary commercial products in the future.

Burn will remain open-source under the MIT-Apache dual license, as we believe it is the most effective way to realize our mission. So let's say that Deep Learning in Rust is here to stay and the project has gone way beyond the personal project it once was when I started it.