Got bored and implemented the AlphaTensor matrix multiplication algorithms in Rust with SIMD https://github.com/drbh/simd-alphatensor-rs

291 points drbh_

AI is really reaching an amazing point and DeepMind's recent breakthrough work on matrix multiplication is (AFAIK) the first computer-generated algorithm with superior complexity to human methods ever!

https://www.deepmind.com/blog/discovering-novel-algorithms-with-alphatensor

Anyway, it felt like this should be available for Rustaceans and that we should improve on the naive implementation by using SIMD when possible, and voila 🪄 https://github.com/drbh/simd-alphatensor-rs

This is conceptually the fastest and least computationally intensive implementation of a 4x4 \* 4x4 ever. However, practically is very likely not the fastest or least complex implementation (maybe it is... see update) - and is just a way to access and explore this awesome work in a great language.

For any devs who find this interesting - please please please contribute and feel free to add tests/benchmarks or improvements in any way!!

🙏

UPDATE

Benchmarks were added from a 2.6 GHz 6 Core i7 machine. Please always run benchmarks on your own device.

The cases are limited but in general, the simd_alphatensor seems to be wayyyyy faster 🦾

Benchmark Average Runtime
ndarry_2x2x2 314.53 ns
simd_alphatensor_2x2x2 13.821 ns
ndarry_4x4x4 402.60 ns
simd_alphatensor_4x4x4 92.732 ns