New Breakthrough in Peak Memory Scheduling: Exact Implementation and Provable Optimality

Title

New Breakthrough in Peak Memory Scheduling: Exact Implementation and Provable Optimality

Meta Description

This article analyzes the latest research on peak memory scheduling: achieving reliable peak memory reduction through exact modeling and statement permutation, providing provable optimality certificates, and exploring the theoretical boundaries of selective recomputation, offering new ideas for edge AI inference optimization.

Article Content

Running real-time inference on edge devices makes activation memory a scarce resource. Operator ordering, as a memory optimization technique that does not increase latency, is extremely valuable. However, the traditional approach—handing the optimized order back to the compiler to regenerate code—is unreliable. Based on the latest research, this article analyzes how to achieve reliable peak memory reduction through exact modeling and statement permutation, provides provable optimality certificates, and explores the theoretical boundaries of selective recomputation.

The Unreliability of Compiler Code Regeneration

The research team conducted systematic tests on PyTorch Inductor and found that the code regeneration approach has serious deviations. Among 54 CPU configurations, 22 configurations had actual peak memory deviating from the predicted value; 9 of them were even worse than the original program. On GPU, 19 out of 50 configurations had the same problem, and numerical results also changed. This shows that relying on compiler code regeneration to implement operator order optimization is highly risky.

The reason is that the compiler may introduce additional temporary buffers or change memory allocation strategies during code generation, causing actual memory usage to deviate from the theoretical model. This unpredictability makes it difficult for developers to trust automated memory optimization workflows.

Peak memory refers to the maximum memory occupied by all active tensors (activations) at any point during program execution. During inference, the output tensor of each operator can be released once it is consumed by subsequent operators; different operator orders change the active intervals of tensors, thereby affecting peak memory. Unlike total memory allocation, peak memory determines whether the entire computation graph can fit on the device and is a hard constraint for edge deployment.

Operator ordering refers to rearranging the execution order of operators with no dependencies in the computation graph without changing data dependencies. For example, convolutions in two branches can be swapped so that the intermediate result of one branch is released earlier. Ideally, reordering does not increase any computation, so it is regarded as a "free" memory optimization technique.

Exact Implementation: A Statement Permutation Approach

To avoid the uncertainty introduced by compiler regeneration, the researchers proposed a method that directly permutes statements in the generated program. They first profile the generated program, then precisely model memory usage based on the read/write effects of the statements themselves, and finally permute the statements.

Theoretical proofs show that as long as memory allocation does not depend on historical state, this method can achieve the modeled peak memory byte-for-byte, and the output is bitwise identical to the original program. In all CPU runs, this goal was achieved. For GPU, the researchers further characterized why the caching allocator breaks this condition, providing direction for future improvements.

Statement permutation means swapping statement positions at the level of the low-level program generated by the compiler (such as loop nests or linear statement sequences), rather than rearranging operators on the high-level computation graph. Since the low-level program already specifies the read/write sets of each statement, swapping statements only requires checking whether data dependencies are broken, so the lifetime of each tensor can be precisely controlled. Compared with graph-level reordering, statement permutation bypasses the compiler's re-scheduling and buffer allocation, avoiding extra temporary variables introduced by code regeneration.

Bitwise identical means that the optimized program produces exactly the same output bit pattern as the original program under the same input. This property is crucial for numerically sensitive edge inference, because if reordering changes the floating-point accumulation order, it may introduce small errors that are amplified in subsequent layers. Statement permutation only changes the execution order of statements without changing the evaluation order of any arithmetic expression, thus guaranteeing bitwise consistency.

Provable Optimality Certificates

Knowing when to stop optimizing is equally important. The researchers proposed forward and reverse "trap" lower bounds, each of which can be computed via an O(n²) minimum cut calculation. These lower bounds can provide optimality certificates for the found operator order.

Experimental results show that on 94% of held-out wrapper graphs and all 16 microcontroller-level models, the certificate can confirm the optimality of the order in less than one second. This means that most inference graphs are already in an optimal scheduling state and require no further optimization. Branched networks are the exception—in such networks, reordering is free and beneficial, for example achieving a 10.0% peak memory reduction on DeepLabV3. In compilation training, this method can save up to 7.6% of memory.

Minimum cut is a classic problem in graph theory: finding a set of edges in a weighted directed graph such that removing these edges disconnects the source from the sink and the sum of the weights of the removed edges is minimized. In this paper, the researchers model the peak memory lower bound as a cut problem on a graph, where each cut corresponds to a set of tensors that "must be active simultaneously," and the sum of their weights is the lower bound on peak memory under that schedule. Since the minimum cut can be solved in polynomial time, a provable lower bound can be computed for any schedule at the cost of O(n²).

Optimality certificate means that when the peak memory of a schedule equals this lower bound, it can be proven optimal and no further search is needed. This provides a reliable stopping criterion for the optimizer, avoiding wasted computational resources on graphs that are already optimal.

Theoretical Boundaries of Selective Recomputation

After reaching certified optimality, the researchers further explored selective recomputation with per-operator upper bounds. They proved that for any fixed bound, the problem is strongly NP-hard, even on graphs with treewidth 2; on series-parallel graphs, it exactly encodes the knapsack problem; and on bounded-treewidth graphs, it can be solved exactly.

This theoretical result reveals the computational complexity of recomputation optimization. In practical applications, for common tiny models where peak memory is determined by a single layer, selective recomputation is useless. However, for architectures that retain low-value activations for later reuse, it can bring benefits. On the public ReDimNet speaker network, although the operator order was already optimal, selective recomputation still reduced peak memory by 5.2-6.1% with less than 0.4% increase in MACs, and the output remained bitwise identical.

Selective recomputation means discarding some intermediate activations when memory is tight and recomputing them when needed later, trading computation for memory. Unlike recomputing the entire backward pass during training, selective recomputation sets a memory upper bound for each operator and only recomputes activations that exceed the bound.

Strongly NP-hard means that even if all numerical parameters are encoded in unary, no polynomial-time algorithm exists for the problem (unless P=NP). This conclusion shows that finding the optimal selective recomputation strategy for general computation graphs is computationally infeasible, and one must rely on approximation algorithms or design exact algorithms for specific graph structures (such as bounded-treewidth graphs). Treewidth measures how close a graph is to a tree; the smaller the treewidth, the smaller the state space of dynamic programming on the graph, and the easier the problem is to solve exactly.

Practical Implications and Future Directions

This research conveys several key messages. First, directly permuting statements in the generated program is more reliable than relying on compiler regeneration and can exactly achieve the predicted peak memory. Second, most inference graphs are already optimally scheduled, and blind optimization may yield little benefit. Finally, selective recomputation is valuable on specific architectures, but its theoretical complexity means that algorithms need to be designed for specific problems.

For edge AI developers, this means that when choosing memory optimization strategies, priority should be given to operator order adjustment, and optimality certificates should be used to avoid over-optimization. At the same time, for branched networks and specific architectures, reordering and selective recomputation can bring additional benefits. In the future, with a further understanding of GPU caching allocator behavior, exact implementation methods are expected to be applied on more hardware.

One-line Summary

Directly permuting statements in the generated program can precisely reduce peak memory and provide optimality certificates, while selective recomputation is strongly NP-hard.

Paragraph Summary

This article analyzes a study on memory optimization for edge inference. The study found that handing the optimized operator order back to the compiler for code regeneration is unreliable: on PyTorch Inductor, 22/54 CPU configurations had actual peak memory deviating from predictions, and numerical results also changed on GPU. To address this, the researchers proposed directly permuting statements in the generated program, precisely modeling memory based on read/write effects, achieving byte-for-byte predicted peak memory on CPU with bitwise identical output. They also proposed forward and reverse trap lower bounds as optimality certificates; 94% of held-out graphs and all 16 microcontroller models were proven optimal within 1 second. Furthermore, selective recomputation was proven strongly NP-hard, but on ReDimNet it achieved a 5.2-6.1% peak memory reduction with less than 0.4% increase in MACs.