Memory Planning

In-place operation is a memory optimization where the output tensor of an operator reuses the memory buffer of one of its input tensors, provided the input is no longer needed for subsequent computation. This technique directly reduces the total number of live buffers at any point in the graph execution.

• Key Mechanism: The compiler performs liveness analysis to determine when an input tensor's lifetime ends. If an output can be written into a dead input's buffer, the allocation is merged.
• Example: An element-wise operation like ReLU can typically be performed in-place, as the input values are not needed after the activation is applied.
• Constraint: Not all operators support in-place execution. It is unsafe if the input is needed later (e.g., for the backward pass during training) or if the operation is not mathematically reversible for the optimization.

Memory reuse, or buffer sharing, is a technique where distinct tensors with non-overlapping lifetimes are assigned to the same physical memory region. Unlike in-place ops, the tensors are not producer-consumer related; they are independent but never alive simultaneously.

• Process: The compiler constructs a lifetime interference graph where nodes are tensors and edges indicate overlapping lifetimes. Tensors without edges between them are candidates for sharing.
• Objective: Minimize the chromatic number of this graph, which corresponds to the peak memory footprint.
• Real-World Impact: This is fundamental for running large models. For instance, in a transformer layer, the memory for the attention scores in layer N can be reused for the feed-forward network intermediate activations in layer N+1, as their lifetimes are sequential.

Memory pool allocation is a runtime strategy that pre-allocates a large, contiguous block of memory (the pool) from which individual tensor buffers are sub-allocated and de-allocated. This avoids the overhead of frequent system calls (malloc/free) and reduces memory fragmentation.

• Static Pools: Sized based on compile-time memory planning. The entire working memory for the graph is allocated once.
• Dynamic Pools: Can grow or shrink, often used with dynamic shapes where tensor sizes are not known at compile time.
• Advantage: Besides performance, it enables efficient implementation of memory reuse. The allocator tracks free blocks within the pool and assigns them to tensors based on the compiler's sharing plan.

Liveness analysis is the foundational compiler dataflow analysis that determines the lifetime of each tensor—the range of operations from its definition (creation) to its last use. It provides the essential input for both in-place optimization and memory reuse.

• Definition: A tensor is live at a program point if its value may be used at some future point. Its lifetime ends after its last use.
• Algorithm: Typically performed by traversing the computational graph in reverse topological order, propagating sets of live variables.
• Output: A lifetime interval for each tensor, which can be visualized on a timeline. Overlapping intervals cannot share memory.

The core objective of memory planning is peak memory minimization, formally modeled as an interval graph coloring problem. The goal is to assign memory offsets to tensor lifetime intervals such that the maximum concurrently used memory is as small as possible.

• NP-Hard Problem: Finding the absolute optimal mapping is computationally intractable for large graphs. Compilers use efficient heuristics.
• Common Heuristics: First-fit and best-fit algorithms, often ordered by tensor size or lifetime length.
• Trade-off: Aggressive minimization can increase memory fragmentation within the pool or complicate data layout for efficient access. The planner must balance peak usage with access performance.

For NPUs with a software-managed cache hierarchy (scratchpad SRAM), memory planning extends to orchestrating data movement between slow global memory (DRAM) and fast on-chip memory. This is often called static scratchpad allocation.

• Strategy: The compiler statically schedules DMA transfers to prefetch needed tensors into SRAM before computation and write back results afterward.
• Double Buffering: A common technique to hide memory transfer latency by using two buffers: one for computation while the other is being filled/evicted.
• Co-optimization: Memory planning for scratchpads is co-optimized with operator scheduling and tiling decisions to maximize data reuse in the fastest memory level.

The memory planning in these tools is powered by specific compiler algorithms:

• Liveness Analysis: Determines the program points where a tensor is first defined ('born') and last used ('dies'), establishing its lifetime interval.
• Graph Coloring: A classic register allocation algorithm adapted for memory, where tensors are nodes in an interference graph; connected (interfering) nodes must be assigned different memory 'colors'.
• In-Place Optimization: Identifies operators where an output tensor can safely reuse the memory of an input tensor (e.g., element-wise operations), reducing allocations.
• Memory Pooling (Arenas): Pre-allocates large blocks of memory from which individual tensor buffers are sub-allocated, minimizing system call overhead and fragmentation.

A compiler optimization that merges multiple adjacent operators into a single, compound kernel. This directly reduces the number of intermediate tensors that must be allocated by the memory planner, lowering kernel launch overhead and minimizing data movement between global and local memory.

• Example: Fusing a Convolution, Batch Normalization, and ReLU activation into one kernel.
• Impact: Creates larger, more complex operators that the memory planner must allocate buffers for, often simplifying the overall buffer allocation problem.

The compiler analysis that determines the dimensions of all tensors in a computational graph at compile time. This is a prerequisite for memory planning, as the planner requires precise tensor sizes to calculate buffer offsets and total allocation sizes.

• Enables: Accurate calculation of peak memory usage and the creation of a static memory allocation plan.
• Contrast with Dynamic Shapes: Graphs with dynamic shapes require more sophisticated, runtime memory management strategies.

An optimization that changes the execution sequence of independent operators. This can dramatically alter the liveness intervals of tensors, creating new opportunities for memory reuse that the memory planner can exploit to reduce peak memory consumption.

• Goal: Schedule operations to maximize the overlap of tensor lifetimes, allowing their memory to be aliased.
• Example: Reordering two independent branches of a graph to ensure one branch's outputs are consumed before the other branch's intermediates are produced.

A memory optimization technique primarily used during training. It strategically trades computation for memory by discarding certain intermediate activations during the forward pass and recomputing them during the backward pass. This is a form of memory planning that explicitly manages the lifetime of large tensors.

• Use Case: Essential for training very large models that exceed GPU/accelerator memory.
• Compiler Role: The compiler must decide which activations to checkpoint and insert the recomputation operations automatically.

A specific memory planning strategy where the compiler identifies operators that can safely write their output tensor into the memory buffer of one of their input tensors. This aliasing destroys the input data but eliminates the need for a separate output allocation.

• Conditions: Requires proving the input is not needed by any other operation after the in-place operator.
• Common Targets: Element-wise operations (e.g., ReLU) and certain reshaping operations.

The broader discipline of optimizing data movement across different levels of an NPU's memory subsystem (e.g., DRAM, shared SRAM, register files). Memory planning focuses on allocation in a single address space, while hierarchy management focuses on placement and movement between spaces.

• Relationship: The memory planner's buffer layout decisions directly influence the efficiency of subsequent data movement passes.
• Goal: Keep frequently accessed data in faster, smaller memories close to the compute units.