Static Shape Inference

The core characteristic of static shape inference is that all tensor dimensions (e.g., batch size, sequence length, channel count) are resolved during compilation, not at runtime. This is achieved by propagating known constants and symbolic expressions through the graph.

• Inputs: Requires known or symbolically defined input shapes (e.g., [batch, 224, 224, 3]).
• Propagation: Applies operator-specific shape rules (e.g., a 2D convolution's output height is derived from input height, kernel size, padding, and stride).
• Result: Every tensor in the graph has a known, fixed shape before execution begins.

With known tensor shapes, the compiler can perform static memory planning (also called static memory allocation). This determines the exact memory footprint required for execution and allocates buffers ahead of time.

• Peak Memory Calculation: The compiler can calculate the maximum simultaneous memory usage, crucial for memory-constrained NPUs.
• Buffer Reuse: Identifies tensors with non-overlapping lifetimes, allowing their memory to be reused, dramatically reducing total allocation.
• No Runtime Allocator Overhead: Eliminates the latency and fragmentation associated with dynamic memory allocation (e.g., malloc/free).

Static shape knowledge is essential for advanced kernel fusion optimizations. The compiler can safely fuse multiple operators into a single compound kernel only if the dataflow and intermediate tensor shapes are fully known.

• Fusion Validation: Ensures fused kernels have consistent data layouts and access patterns.
• Intermediate Buffer Elimination: Allows the compiler to eliminate the memory allocation for the output of the first fused operator, passing data via registers or shared memory.
• Example: Fusing a Convolution, BatchNorm, and ReLU activation requires knowing the exact output shape of the convolution to inline the subsequent operations.

A major constraint of static shape inference is its inability to handle truly dynamic shapes determined at runtime. It operates under one of two regimes:

• Fully Static: All dimensions are numerical constants (e.g., [1, 128, 128, 64]). Common in embedded and edge AI deployments.
• Symbolically Static: Dimensions are defined by symbolic variables (e.g., [batch, seq_len, 768]). The shape is fixed in structure but parametric in size. The compiled kernel is specialized for this symbolic pattern.

Dynamic operations (e.g., non-padded sequences, dynamic slicing) often break static shape inference, forcing a fallback to a less-optimized execution path.

Static shape inference is distinct from dynamic shape inference, which occurs at runtime. The trade-offs are fundamental to system design.

Frameworks like TensorFlow (with tf.function) and PyTorch (with TorchScript/TorchDynamo) use static shape inference to compile graphs from dynamically-shaped Python code.

For NPU and GPU compilers, static shapes unlock hardware-specific optimizations that are impossible with dynamic shapes.

• Register Allocation: Compilers can statically assign tensor slices to hardware registers for the fastest possible access.
• Memory Layout Transformation: The compiler can confidently transform tensor layouts (e.g., from NHWC to NCHW) to match hardware-specific DMA (Direct Memory Access) patterns or SIMD (Single Instruction, Multiple Data) lane structures.
• Scheduling & Tiling: Loop tiling strategies and parallel workgroup schedules are calculated based on known tensor dimensions to maximize data locality and compute unit utilization.

This allows the generation of highly efficient, ahead-of-time (AOT) compiled binaries tailored to a specific workload shape.

A compiler optimization pass that allocates memory buffers for all tensors in a computational graph. Static shape inference is a prerequisite, as knowing the exact dimensions of each tensor at compile time allows the compiler to:

• Calculate the precise memory footprint of each buffer.
• Perform in-place optimization, where the memory of an output tensor can safely reuse the memory of an input tensor that is no longer needed.
• Minimize the peak memory usage of the entire graph execution, which is critical for memory-constrained NPUs. Without static shapes, memory planning must be conservative, leading to larger allocations and reduced efficiency.

A compiler optimization that evaluates expressions consisting entirely of compile-time constants and replaces them with their computed result. This optimization relies on static shape inference to propagate constant tensor values through the graph. For example, a reshape operation with a statically known input shape and constant target dimensions can be folded into a single constant tensor. This eliminates runtime computation and can sometimes enable further dead code elimination if the folded constant is never used.

A key optimization that merges multiple adjacent operators into a single, compound kernel. Static shape inference enables aggressive fusion by providing guarantees about tensor dimensions:

• It allows the compiler to statically allocate shared memory for intermediate results between fused ops.
• It verifies that data layouts between consecutive ops are compatible for fusion without runtime checks.
• It enables the generation of a single, optimized kernel with known loop bounds and memory access patterns, eliminating kernel launch overhead and unnecessary global memory writes/reads.

Low-level loop transformations that require precise knowledge of iteration spaces. Static shape inference provides the concrete bounds of tensor dimensions, allowing the compiler to:

• Apply loop tiling effectively: Partition loops into fixed-size blocks that fit into NPU cache hierarchies for optimal data reuse.
• Enable automatic vectorization: Transform scalar operations to use SIMD instructions when loop counts are known multiples of the vector width.
• Perform loop unrolling: Decide on unroll factors based on known, small iteration counts. These optimizations are far less effective or impossible with dynamic shapes, as they rely on compile-time constants for loop bounds and strides.

A compilation strategy that optimizes a graph for execution with quantized (e.g., INT8) data types. Static shape inference is critical in this process:

• It allows the compiler to insert fake quantization and dequantization nodes at precise points in the graph during the optimization phase.
• The compiler can statically determine the scaling factors and zero-points for tensors of known shape, enabling constant propagation of these parameters.
• It facilitates the fusion of quantization/dequantization ops with neighboring floating-point operations, a key step in generating efficient integer-only kernels for NPUs.

A compilation strategy where the entire program is compiled to a native binary before execution. Static shape inference is a foundational requirement for true AOT compilation of neural networks, as it enables:

• The generation of statically sized memory pools and data structures.
• The creation of fully unrolled and specialized kernels with no runtime shape dispatch logic.
• The elimination of all shape-related control flow and memory allocation overhead from the runtime, resulting in minimal latency and predictable performance—essential for deployment on edge NPUs.