Shape Inference

Shape inference is performed statically, before the model is executed. The compiler analyzes the graph's operations and propagates known input shapes through the network using each operator's shape propagation function. This allows for:

• Memory planning: Pre-allocating buffers of the correct size.
• Kernel selection: Choosing the optimal implementation based on input/output dimensions.
• Error detection: Catching dimension mismatches early (e.g., attempting to multiply a [2,3] matrix with a [4,5] matrix).

Each operation in the graph has a well-defined shape function. The inference engine applies these functions sequentially in a topologically sorted order. For example:

• Element-wise ops (ReLU, Tanh): Output shape equals input shape.
• Matrix Multiplication: For matmul(A, B) where A is [m, k] and B is [k, n], the output is [m, n].
• Convolution: Output spatial dimensions are derived from input size, kernel size, padding, and stride.
• Concatenation: Output shape depends on the axis and the sizes of all inputs along that axis.

Real-world models often have dynamic dimensions (e.g., variable batch size, sequence length). Shape inference must handle these using symbolic expressions. A dimension may be represented not as a fixed integer but as a variable (e.g., batch, seq_len) or an expression (e.g., (input_height - kernel_size + 2*padding) / stride + 1). The system tracks these relationships, allowing optimizations that are parametric in the symbolic variables.

Accurate shape information unlocks several critical graph-level optimizations:

• Operator Fusion: Knowing tensor sizes is essential to evaluate the memory and compute trade-off of fusing ops.
• Constant Folding: If an op's inputs are fully known constants, its output can be precomputed, but only if the output shape is known.
• Static Memory Planning: The peak memory footprint can be calculated, and a single arena allocator can be set up with optimal buffer sizes and reuse strategies.
• Kernel Auto-Tuning: The search space for optimal kernel parameters is defined by the problem size (e.g., M, N, K for GEMM).

Shape inference is tightly coupled with a compiler's Intermediate Representation (IR). The IR nodes are annotated with shape attributes (e.g., tensor<2x?xf32>). Dedicated shape inference passes run on the IR, updating these annotations. In frameworks like MLIR, shape inference may be modeled as a type inference problem, where the 'type' of a value includes its shape.

Shape inference can fail or be incomplete. Common scenarios include:

• Truly dynamic ops: Operations like NonMaxSuppression have output sizes dependent on runtime data values.
• Control flow: Branches in if statements or loops may produce different shapes.
• Unknown initial shapes: If a graph input has an unspecified shape (?), downstream shapes may remain unknown. In such cases, the framework may defer to runtime shape calculation, allocate dynamic buffers, or require user-provided hints to proceed with compilation.

A compile-time optimization that pre-allocates and reuses memory buffers for tensors by analyzing their lifetimes across the computational graph. It relies on the results of shape inference to know the exact size of each tensor buffer. This eliminates the overhead of dynamic memory allocation during inference and minimizes the peak memory footprint, which is critical for deployment on memory-constrained edge devices.

• Key Input: A shape-inferred graph where every tensor's dimensions are known.
• Mechanism: Constructs a memory allocation plan by identifying when tensors are produced and consumed, allowing their memory to be safely reused (in-place or via a memory arena).
• Benefit: Predictable, low-latency inference startup and execution.

A compiler optimization that combines multiple sequential operations (nodes) in a computational graph into a single, more efficient fused kernel. Shape inference is a prerequisite, as the compiler must know the intermediate tensor shapes to validate the legality of the fusion and to generate the correct fused kernel code.

• Common Fusions: Convolution followed by BatchNorm and ReLU activation (Conv-BN-ReLU).
• Primary Benefit: Reduces memory bandwidth pressure by keeping intermediate results in fast registers or caches instead of writing them to and reading them from main memory.
• Secondary Benefit: Eliminates the overhead of launching multiple small, separate kernels.

A compile-time optimization that evaluates and replaces subgraphs comprised entirely of compile-time constants with a single constant node containing the precomputed result. Shape inference is often required to execute these constant subgraphs symbolically or to validate the shape of the resulting constant tensor.

• Example: A graph segment that calculates a fixed-size embedding lookup table from constants can be folded into a single constant tensor.
• Impact: Eliminates runtime computation for static portions of the model, simplifying the graph and reducing execution time.
• Relation to Shape Inference: The folding process itself may need to infer shapes to execute operations on constant tensors.

The process of transforming a high-level, hardware-agnostic Intermediate Representation (IR) of a neural network into a lower-level, target-specific IR or machine code. Shape inference is a critical early phase in this pipeline. The known tensor shapes enable lower-level optimizations and guide the selection of efficient hardware-specific kernels (e.g., choosing a 3x3 depthwise convolution kernel for a mobile NPU).

• Pipeline Stage: Typically occurs after high-level graph optimizations (fusion, folding) and before final code generation.
• Dependency: Lowering passes (like converting a MatMul to a series of hardware-specific GEMM calls) require precise knowledge of input and output shapes.

Involves transforming the in-memory arrangement (layout) of tensor data to better match the access patterns of computational kernels and the underlying hardware. Shape inference provides the tensor dimensions, which are essential for planning these transformations.

• Common Transformations: Converting between NHWC (batch, height, width, channels) and NCHW layouts to optimize for different hardware (e.g., NVIDIA GPUs vs. Intel CPUs).
• Benefit: Improves cache locality and enables efficient vectorized/SIMD memory loads, significantly speeding up operations like convolution.
• Challenge: May require inserting explicit layout transpose operations, the cost of which must be weighed against the compute savings.

A compiler pass that identifies and removes operations in a computational graph whose outputs do not contribute to the final model output. Accurate shape inference helps the compiler perform more aggressive dead code analysis by understanding data dependencies through the graph.

• Example: Removing debug-only operations like Print or unused branches of conditional logic.
• Impact: Reduces both computation and memory usage for the live graph.
• Connection to Shape Inference: Understanding how data flows (and where it stops) is easier when the shapes of all intermediate values are known, making dead nodes more identifiable.