OmniMalloc: static memory allocation for neural networks

OmniMalloc is an open-source static memory allocation framework for neural networks. Ahead-of-time compilers for AI accelerators must place every tensor in scratchpad memory at compile time; doing this well is an NP-hard 2D packing problem that directly determines whether a network fits on-chip and how fast it runs.

OmniMalloc provides a collection of allocation algorithms behind a single interface, from fast heuristics to exact solver-based approaches, together with benchmarking and visualization tooling to compare them on real workloads.

I designed and open-sourced OmniMalloc while building the memory optimization passes of Axelera AI’s production compiler, where it is deployed today.

Code: github.com/fpedd/omnimalloc




Enjoy Reading This Article?

Here are some more articles you might like to read next: