Memory planning for inference

Overview

An inference program produces intermediate tensors, and each one is needed only from the instruction that writes it to the last instruction that reads it. Memory planning gives every tensor an offset in a scratch buffer, so the buffer can be reused over time. The constraint shared by all variants is that two tensors whose live ranges overlap must not overlap in memory. The goal is a buffer as small as possible 2210.16556 §4.3.

In the one paper cited here, the planner runs after the bitwidth of each tensor has been chosen, because the tensor sizes depend on that choice 2210.16556 §4.3. The variants below are therefore the planner's own design choices, plus the greedy first-fit setting the paper compares against internally. The page reflects one source only.

Variants

Greedy first-fit placement

Tensors are placed one at a time, each at the first offset where it fits. The paper names this kind of greedy allocation as a cause of fragmentation, which can make a feasible layout look infeasible 2210.16556 §1. On SqueezeNet, the paper reports 1.44 MB peak RAM for first-fit placement and 1.16 MB for exact-cover placement, with the same bitwidth map [arxiv:2210.16556 §6.7]. The paper's text on first-fit is not in the source cards, so the rule above is the usual definition, not the paper's. Card: MinUn: Accurate ML Inference on Microcontrollers.

Exact-cover placement

Each tensor becomes a rectangle. Its width is its live range over the instruction index, and its height is its size in bytes. An exact cover of the canvas with minimum height M gives an offset for every tensor inside an M-byte buffer 2210.16556 §4.3. The authors solve the cover with Knuth's Algorithm X using dancing links, and visit tensors in decreasing order of rectangle area 2210.16556 §4.3. The paper says that with coarsening constant k = 1 the minimum height is found whenever a solution exists 2210.16556 §5. Algorithm X can take exponential time in the worst case. The authors cap the memory step at two hours per run, and the paper gives no worst-case runtime for the packing step 2210.16556 §5.

Coarsened exact cover

Sizes are rounded up to a multiple of k, which shrinks the canvas to M/k rows and so shrinks the search 2210.16556 §4.3. Rounding up can cost memory: the buffer may be larger than the k = 1 result, so the optimum can be lost 2210.16556 §4.3. The paper uses k = 64 for its RNNPool face models and k = 1 for the others 2210.16556 §5.

Lower bound as target height

The smallest buffer that could work is the largest total size of live tensors at any single instruction, measured in units of k 2210.16556 §4.3. The paper calls this a lower bound and does not say what to do when no exact cover exists at that height 2210.16556 §4.3. It is best used as the starting height for an exact-cover search, as in the sketch below.

When to use which

The source gives design trade-offs, not a decision rule. Each row is what the cited text supports.

Implementation sketch

The common core is three steps: compute sizes and live ranges, find a starting height from the lower bound, and try an exact cover at that height. The sketch follows the exact-cover variant. The loop that raises the height when no cover exists is our addition, not the paper's.

# Plan scratch offsets for one bitwidth map  [arxiv:2210.16556 §4.3]
# Input: tensors with size in bytes and live range [start, end]
# Output: offset per tensor, and the buffer size in bytes

def plan_memory(tensors, num_instructions, k):
    # Round each size up to a multiple of k  [§4.3]
    units = {t: ceil(t.size / k) for t in tensors}
    # Lower bound: the largest total of live units at one instruction  [§4.3]
    height = max(
        sum(units[t] for t in tensors if t.start <= i <= t.end)
        for i in range(num_instructions)
    )
    while True:
        # Exact cover of a canvas of this height, using Algorithm X with
        # dancing links and tensors in decreasing area order  [§4.3]
        offsets = exact_cover(tensors, units, height)
        if offsets is not None:
            # Scale offsets and height back from units to bytes  [§4.3]
            return {t: o * k for t, o in offsets.items()}, height * k
        # No cover at this height. The paper gives no procedure here, so
        # raising the height by one unit is our addition  [§4.3]
        height += 1

Contradictions

• Overall RAM reduction against 16-bit baselines. The paper's text says the posit configuration cuts peak RAM by 2.9x to 5.1x even against half precision or bfloat16 2210.16556 §6.4. The appendix table's RAM columns give smaller ratios against those baselines, roughly 1.3x to 4x 2210.16556 app. D. The source cards do not settle which is right. These figures cover the whole compiler, not the placement step alone.
• Face model peak RAM. The introduction gives about 600 KB falling to about 180 KB 2210.16556 §1. The appendix gives 614K to 192K for the posit Face-B model 2210.16556 app. D. Check which configuration a figure refers to before citing it.

Papers

• MinUn: Accurate ML Inference on Microcontrollers, arXiv:2210.16556: a compiler for ARM microcontrollers with a bitwidth search and an exact-cover scratch planner. It reports results on SqueezeNet and on small microcontroller models. The paper's code link is not verified.