Compiling a neural net to C for a speedup
Published on: 2025-06-16 16:22:47
Compiling a Neural Net to C for a 1,744× speedup
2025-05-27 · About 25 minutes long
tl;dr: I trained a neural network (NN), with logic gates in the place of activation functions, to learn a 3×3 kernel function for Conway’s Game of Life. I wanted to see if I could speed up inference by extracting the learned logic circuit from the NN. So, I wrote some code to extract and compile the extracted logic circuit to bit-parallel C (with some optimizations to remove gates that don’t contribute to the output). I benchmarked the original NN against the extracted 300-line single-threaded C program.; compiling the NN to C resulted in a 1,744× speedup! Crazy, right? Here’s the repo: ~354 lines of Python/JAX, ~331 lines of C, if you want to reproduce it and/or mess around.
The longer story
While plumbing the intertubes (as one does), I came across this fun publication by the Self Organising Systems group at Google, about Differentiable Logic Cellular Automata. This research caught my attention (I
... Read full article.