Home · ASPLOS 2023 · Machine Learning Systems

MiniMalloc: A Lightweight Memory Allocator for Hardware-Accelerated Machine Learning

Michael D. Moffitt

Conference
ASPLOS 2023
Published
2023
Pages
238-252
Citations
5

Before a machine-learning model runs on an accelerator, the compiler must give every intermediate buffer a fixed address in a small on-chip memory so that no two live buffers overlap — a rectangle-packing puzzle with an astronomical number of arrangements; MiniMalloc searches only a special subset called canonical solutions, which it proves is enough to find any answer, prunes partial packings early by reasoning about space, and discards arrangements that another arrangement dominates, running orders of magnitude faster than the previous best method in a much smaller library.

The argument in one glance

  1. Problem

    Every buffer needs a fixed address, and the arrangements explode

    A compiler for a machine-learning accelerator knows in advance every intermediate buffer a model will create, when each one is born and dies, and how large it is. It must assign each buffer an offset in a fixed-size on-chip memory so that buffers alive at the same time never overlap. This is static memory allocation. It is equivalent to packing rectangles into a strip, and the number of ways to arrange them grows explosively with the number of buffers.

  2. What people did

    Fast greedy heuristics, or exact solvers

    Greedy allocators place buffers one at a time by simple rules and are fast, but on hard instances they either fail or waste memory. Exact methods, such as integer programming or constraint solvers, are guaranteed to find a placement if one exists.

  3. Why it fell short

    Exact solvers explore far too many equivalent arrangements

    A general-purpose solver treats every distinct placement as a distinct candidate. But most placements are trivial variations of each other, for example the same packing with one buffer floated slightly higher. Exploring all of them makes exact solvers slow and erratic, and the solver libraries themselves are large dependencies inside a compiler.

  4. This paper's idea

    Search only canonical solutions, infer space early, and drop dominated arrangements

    MiniMalloc is a purpose-built depth-first search. It restricts the search to canonical solutions, a special non-redundant form of packing that the author shows corresponds to the members of an algebraic lattice, so that any feasible instance has a canonical solution and nothing is lost. A spatial inference technique uses that structure to detect early that a partial placement cannot succeed and to backtrack sooner. A dominance check discards partial placements that leave gaps another placement would fill better. The abstract reports orders of magnitude speedups over the previous state of the art on many benchmarks, and a substantially smaller library.

Why this problem is worth solving

An accelerator such as a tensor processing unit (TPU) or the machine-learning block in a phone has a small, fast on-chip memory. A model’s intermediate results, its

buffers A block of memory holding one intermediate result of the model, with a fixed size and a fixed lifetime known to the compiler.
, must live there while they are in use. Unlike a general-purpose program, a machine-learning model has a fixed computation graph, so the compiler knows every buffer before the model runs: its size, the step at which it is produced, and the step after which it is no longer needed. Because all of this is known ahead of time, the compiler can decide the address of every buffer once, at compile time. That is
static memory allocation Assigning a fixed memory offset to every buffer at compile time, so that buffers whose lifetimes overlap never share bytes.
.

The rules are simple. Two buffers may share the same bytes only if their lifetimes do not overlap. No buffer may extend past the memory capacity. Draw time along the horizontal axis and memory along the vertical axis, and each buffer becomes a rectangle whose width is fixed by its lifetime and whose vertical position is the compiler’s choice. Static memory allocation is the task of sliding those rectangles up and down so that none overlap and all stay below the capacity line.

The stakes are practical. If the allocator finds no valid packing, the compiler must spill buffers to slower memory or refuse to compile the model. If the allocator is slow, compilation is slow, and on a phone compilation happens on the device. The companion TelaMalloc paper from the same conference, explained elsewhere on this site, reports that an exact solver was 68 percent of compile time for one production model and that some instances took minutes or even hours. The problem is NP-hard, so no method is fast on every possible input, but production instances have structure, and MiniMalloc is built to exploit it.

What people did before, and where it broke

  • Greedy heuristics. Sort the buffers by some rule, place each at the lowest offset where it fits, never reconsider. Where it broke. Fast, but with no backtracking a bad early choice cannot be undone, so on tight instances the heuristic either exceeds capacity or fails outright.
  • Exact solvers: integer linear programming (ILP) and constraint programming. Encode the packing as constraints and hand it to a general solver, which explores placements with backtracking and proves infeasibility when there is no answer. Where it broke. The solver does not know that most placements are redundant. If one buffer can sit at offset 100 or offset 104 with nothing changing, the solver may try both. Multiply such choices across hundreds of buffers and the search space becomes astronomically larger than it needs to be. Runtime is also unpredictable, and the solver library is a heavy dependency.
  • Hybrids that run a heuristic and check it with a solver. TelaMalloc couples the two, using the solver to detect dead ends early. Where it broke. It is far better than either alone, but it still carries a general-purpose solver and still searches the general space of placements.
  • The gap. Nobody had asked which placements actually need to be examined. If a small, well-defined subset of placements is guaranteed to contain a solution whenever one exists, a search restricted to that subset is both complete and dramatically smaller.

Where MiniMalloc sits

Compiling a model for an accelerator

  1. Model graph

    Operators and the buffers that flow between them.

  2. Scheduling

    Fixes the order of operators, which fixes the lifetime of every buffer.

  3. Static memory allocation new

    Assigns each buffer an offset so that live buffers never overlap and none exceeds capacity. MiniMalloc replaces the general solver here.

  4. Code generation

    Emits accelerator instructions with the chosen addresses baked in.

MiniMalloc takes a list of buffers, each with a start time, an end time and a size, plus a capacity, and returns an offset for each buffer or reports that no packing exists.

The idea

This section explains the three mechanisms named in the abstract and the project’s public description. The full paper, with its formal definitions and proofs, was not available; the descriptions below state what those sources say and explain the standard reasoning behind each, without inventing detail.

Canonical solutions. Take any valid packing and imagine letting every buffer fall downward until it rests on the capacity floor or on another buffer alive at the same time. The result is still valid: moving a buffer down never creates an overlap that was not already there. Packings with this settled, non-redundant structure are what the author calls

canonical solutions A valid packing in a special non-redundant form; the paper shows these form an algebraic lattice, and that any feasible instance has one.
. The key claim from the paper’s public description is that canonical solutions correspond to the members of an algebraic
lattice A partially ordered set in which any two elements have a unique least upper bound and greatest lower bound.
, and that restricting the search to them keeps the algorithm sound (it never returns an invalid packing) and complete (if any packing exists, it finds one). The benefit is that the search never explores two placements that differ only by floating a buffer upward into empty space, which removes an enormous number of equivalent branches.

Recursive depth-first search. MiniMalloc places buffers one at a time in a chosen order, trying candidate offsets for each. When a buffer cannot be placed, it backtracks to the previous decision and tries the next option. This is ordinary backtracking search; what makes it fast is the small space and the pruning. The open-source implementation documents several of the ordering and pruning rules it uses: an initial static ordering of buffers, a dynamic preference for buffers with smaller viable offsets, and a temporal decomposition that splits the timeline where it can be split.

Spatial inference. Before committing to a partial placement, the solver reasons about the space that the remaining buffers will need. If a lower bound on the height needed in some time section already exceeds the capacity, that partial placement cannot succeed, and the search backtracks immediately instead of discovering the failure many decisions later. The abstract says this technique exploits the canonical structure to backtrack more effectively than would otherwise be possible.

Dominance. Two partial placements may be different yet one is never better than the other. For example, a placement that leaves a gap into which an unplaced buffer would fit easily is dominated by the placement that fills the gap. The abstract describes a new mechanism for detecting such dominated partial solutions and eliminating them from consideration.

Tracing MiniMalloc on the project’s own example

Packing five buffers into 12 units of memory

This is the example shipped with the open-source code. Each buffer is written name: [start, end) x size. Offsets are the solver's actual output. The order of placement is illustrative; the paper's exact ordering rule was not available.

Step 1 of 6

The instance

Five buffers must fit in a memory of capacity 12. Buffer b5 lives for the entire program. Buffers b3 and b4 cover the first and second halves. Buffers b1 and b2 are short and cover the first part of b3. Adding up sizes at any moment gives exactly 12, so there is no slack at all.

b1: [0,3) x 4 | b2: [3,9) x 4 | b3: [0,9) x 4 | b4: [9,21) x 4 | b5: [0,21) x 4 | capacity 12

Why it works

  • Redundant placements are the enemy, and canonical form removes them. Floating a buffer upward into empty space never helps, because every canonical packing is at least as tight. Because the search considers only settled packings, it explores each genuinely different arrangement once. Therefore the search space shrinks by a large factor with no loss of completeness.
  • The lattice structure makes early pruning possible. Because canonical solutions have a predictable structure, the solver can compute meaningful lower bounds on how much height each time section will need. When a bound exceeds capacity, the branch is dead, and the search knows it immediately.
  • Dominance removes near-duplicates that canonical form alone does not. Two canonical partial placements can still differ in ways that matter only in one direction. Discarding the dominated one avoids exploring a subtree whose every outcome is matched or beaten by the other.
  • A purpose-built search needs no general solver. MiniMalloc encodes its knowledge of the problem directly into the search rather than into constraints for a general engine. That is why the abstract can report both a large speedup and a substantial reduction in library size.

What it costs

  • It solves one problem. MiniMalloc is specialized to static allocation of buffers with fixed lifetimes and sizes. Constraints outside that model, such as alignment rules, bank conflicts or buffers that may be split, are not part of the abstract’s description and would need to be handled elsewhere or added to the search.
  • Worst-case time is still exponential. The problem is NP-hard. Canonical form and pruning shrink the search dramatically on real instances, but an adversarial instance can still force long backtracking. The open-source solver takes a timeout parameter for this reason.
  • Proofs carry the correctness burden. Completeness rests on the claim that every feasible instance has a canonical solution and that the pruning rules never discard the last one. Those proofs are in the paper and could not be checked here.
  • Baselines are not named in the abstract. The abstract reports orders of magnitude improvement over the previous state of the art on many benchmarks, without naming the baseline, the hardware or the benchmark suite. The open-source repository ships a set of eleven challenging instances, each with a capacity of 1,048,576 units.

The evaluation

Three ways to place buffers

PropertyGreedy heuristicGeneral exact solverMiniMalloc
Finds a packing whenever one existsNoYes
Search spaceNone, one pass
Speed on hard instancesFast, but may fail
Library sizeSmall
The MiniMalloc column reflects the abstract and the public repository. The paper's benchmark tables, baseline names and exact speedups were not available when this page was written.Underlined values explain themselves when clicked.

Check yourself

Check your understanding

  1. 1. Why can a compiler for a machine-learning accelerator decide buffer addresses statically, when a general-purpose program cannot?

  2. 2. What makes restricting the search to canonical solutions safe, rather than a shortcut that might miss the answer?

  3. 3. Predict what happens to the search if the spatial inference step is switched off but canonical-only search stays on.

  4. 4. In the five-buffer example, why does placing b3 at offset 8 instead of 4 get pruned?

Sources

The paper is “MiniMalloc: A Lightweight Memory Allocator for Hardware-Accelerated Machine Learning” by Michael D. Moffitt, Google, published at ASPLOS 2023, Volume 4, pages 238 to 252, DOI 10.1145/3623278.3624752. The code is open source at github.com/google/minimalloc.

This explainer is marked as a draft because the full text was not accessible. The paper is available only from the ACM Digital Library, which returned an anti-bot page, and no arXiv preprint, author-hosted copy or open repository version exists; the Google Research publication page links only to the DOI. Every claim attributed to the paper comes from its abstract or from the project’s own public description, which states that canonical solutions correspond to the members of an algebraic lattice and that restricting the search to them keeps the algorithm sound and complete. The list of search techniques comes from the documented options in the open-source solver. The worked example is the input and output file shipped with the repository. The figures on solver cost in production come from the TelaMalloc paper, explained separately on this site, and are attributed to it.

Related papers on Machine Learning Systems

See all 212 →