Graph Lowering

Graph lowering operates across multiple levels of abstraction, systematically decomposing high-level operations. It begins with framework-specific operators (e.g., a PyTorch nn.Conv2d) and progresses through several Intermediate Representation (IR) dialects—such as linear algebra, loop nests, and vector operations—before reaching target-specific machine instructions. Each level applies legalization rules to convert operations into a form supported by the next, lower level. This layered approach enables hardware-agnostic high-level optimizations before committing to vendor-specific details.

A core function of lowering is legalization, the process of replacing or decomposing operations unsupported by the target hardware with equivalent, legal sequences. For an NPU, this might involve:

• Replacing a high-level Softmax with a manual decomposition using Exp, ReduceSum, and Div.
• Converting a 5D tensor operation into a series of lower-dimensional ops if the hardware only supports up to 4D.
• Mapping floating-point operations to fixed-point or block floating-point units. The legalizer uses a pattern-matching system to apply these hardware-specific rewriting rules.

Effective lowering depends on static shape inference, a compile-time analysis that propagates tensor dimension information throughout the graph. By resolving shapes before execution, the compiler can:

• Perform accurate memory planning, pre-allocating buffers of known size.
• Enable optimizations like constant folding and dead code elimination.
• Generate efficient, fixed-loop kernels without dynamic control flow overhead. This analysis transforms a dynamically-shaped graph from training frameworks into a static, predictable dataflow graph suitable for NPU execution.

NPUs and many accelerators excel at data-parallel computation but have limited support for complex control flow like loops and conditionals. A key lowering step is control flow flattening or linearization. This transforms structured control constructs into a dataflow representation using predicated execution or explicit computation graphs with conditional data edges. For example, a tf.while_loop may be unrolled a static number of times or lowered to a state machine where iterations are controlled by explicit data dependencies, eliminating branch instructions.

Lowering is not a single pass but a pipeline interleaved with optimization passes. Key optimizations performed during or just before lowering include:

• Operator Fusion: Merging adjacent nodes (e.g., Conv + BiasAdd + ReLU) into a single compound kernel.
• Layout Transformation: Converting tensor memory layouts (e.g., NCHW to NHWC) to match hardware access patterns.
• Common Subexpression Elimination (CSE): Removing duplicate calculations introduced during legalization.
• Constant Folding: Pre-computing operations on static weights. These optimizations are applied at the most appropriate abstraction level, often immediately after a legalization step creates new optimization opportunities.

The final stage of lowering is code generation, where the lowest-level IR is converted into executable code for the target. This involves:

• Instruction Selection: Mapping IR ops to specific NPU instructions (e.g., a MAC operation to a systolic array instruction).
• Register Allocation: Assigning virtual tensors to physical hardware registers or memory buffers.
• Scheduling: Determining the temporal order of instruction issuance to hide memory latency. This stage heavily utilizes the Vendor SDK and Intrinsic Mapping, emitting code that may be C/C++, assembly, or a proprietary binary format (e.g., TensorRT engine, CoreML mlmodel).

The XLA (Accelerated Linear Algebra) compiler within TensorFlow performs graph lowering through a defined sequence of Intermediate Representations (IRs). The process begins with the HLO (High-Level Optimizer) IR, which represents the graph after framework-level optimizations. XLA then lowers HLO through multiple passes:

• Target-specific legalization converts unsupported HLO ops to supported ones.
• Buffer assignment allocates memory for all tensors.
• Lowering to LLO (Low-Level Optimizer) IR, which is closer to machine code.
• Final lowering to LLVM IR for CPU/GPU or proprietary TPU IR for Google's hardware. This multi-stage approach allows for both hardware-agnostic and target-specific optimizations.

PyTorch's TorchInductor compiler, used by torch.compile, lowers graphs via the Aten IR (PyTorch's core tensor operation library). The lowering path is determined by the target backend:

• For NVIDIA GPUs, it lowers Aten ops to Triton IR, a high-level Python-like language for GPU programming, which is then compiled to PTX.
• For CPUs, it lowers Aten to OpenMP-annotated C++ loops via the Cpp backend.
• The process involves operator decomposition, where complex Aten ops are broken down into simpler primitives, followed by fusion and scheduling on the target IR. This design prioritizes flexibility and performance on diverse hardware.

TVM (Tensor Virtual Machine) employs a clear separation between computational and scheduling concerns. Graph lowering flows from Relay, a high-level functional IR for neural networks, down to TIR (Tensor IR).

• The Relay → TIR lowering involves legalization (mapping Relay operators to TIR primitives) and memory planning.
• TIR is a low-level, explicit loop-nest IR that represents detailed computation and hardware-aware schedules (e.g., loop tiling, vectorization).
• Finally, TIR is lowered to target code via code generation for backends like LLVM, CUDA, Vulkan, or vendor-specific NPU SDKs. This modularity is key to TVM's support for a wide array of accelerators.

MLIR (Multi-Level Intermediate Representation) formalizes graph lowering as dialect lowering. A computational graph enters MLIR in a high-level dialect (e.g., tosa, mhlo).

• Lowering is a series of conversion passes that translate operations from one dialect to another, progressively moving from abstract to concrete representations (e.g., mhlo → linalg → affine → llvm).
• Each dialect owns a set of operations, types, and transformations. The pattern rewrite engine applies lowering patterns to convert ops.
• This allows hardware vendors to define their own dialects (e.g., npu for a specific NPU) and lowering paths, creating a customizable and interoperable compilation stack.

ONNX Runtime lowers ONNX model graphs for execution via a combination of graph transformations and Execution Provider (EP) interfaces.

• First, graph optimization passes (constant folding, node fusion) are applied on the ONNX graph.
• The optimized graph is then partitioned and assigned to hardware-specific EPs (e.g., CUDA EP, TensorRT EP, OpenVINO EP).
• Each EP is responsible for the final lowering of its subgraph. For example, the CUDA EP may convert ONNX nodes to CUDA kernel calls, while the TensorRT EP performs a more aggressive lowering and fusion to TensorRT's engine format. This architecture decouples high-level optimizations from vendor-specific lowering.

Vendor SDKs for NPUs (like Qualcomm's AI Engine Direct or NVIDIA's TensorRT) perform the final, most hardware-specific stage of graph lowering.

• They ingest a framework-exported graph (e.g., ONNX, TensorFlow Lite) and apply vendor-optimized graph passes.
• This includes kernel selection (choosing the most efficient pre-written kernel for each op), memory allocation aligned with the NPU's memory hierarchy, and scheduling onto the NPU's cores.
• The output is a serialized engine or binary containing instructions and data tightly coupled to the NPU's microarchitecture. This stage is where abstract tensor operations become concrete hardware instructions.

An Intermediate Representation is a compiler's internal data structure or code used to represent a program between its source form and its target machine code. It is the central data structure upon which graph lowering and all other transformations operate.

• Purpose: Enables machine-independent analysis and optimization.
• Levels: A compilation pipeline often uses multiple IRs, transitioning from high-level (e.g., framework graphs) to mid-level (e.g., operator lists) to low-level (e.g., LLVM IR, vendor-specific IR).
• Example: MLIR is a modern, extensible framework for defining domain-specific IRs, crucial for AI/ML compilation.

Graph canonicalization is a compiler transformation that rewrites a computational graph into a standard, simplified form. It is a critical preparatory step before lowering to eliminate syntactic variations and ensure predictable optimization.

• Function: Converts semantically equivalent operations into a single, canonical form (e.g., rewriting Sub(x, y) as Add(x, Neg(y)) if that's the canonical form for subtraction).
• Benefit: Simplifies pattern matching for subsequent passes like fusion, constant folding, and lowering itself.
• Process: Involves algebraic simplification, operator normalization, and dead argument elimination.

Operator legalization is a specific lowering step that converts high-level or framework-specific operators into a set of primitive operations supported by the target hardware or lower-level IR.

• Mechanism: A pattern-matching and rewriting process. For example, a Softmax operator may be legalized into a sequence of Exp, ReduceSum, and Div operations.
• Target: The output is a graph consisting only of operations in the target instruction set or a vendor's supported kernel library.
• Challenge: Must preserve numerical semantics while mapping to available hardware primitives.

MLIR is a flexible, extensible compiler infrastructure that provides a modular system of interoperable intermediate representations ('dialects'). It is fundamentally designed to facilitate sophisticated multi-level lowering for domains like machine learning.

• Core Concept: Different dialects represent different levels of abstraction (e.g., tensor dialect for high-level ops, linalg for loop-based computation, llvm for low-level machine ops).
• Lowering Path: Graph lowering in MLIR is the process of progressively translating operations from a higher-level dialect to a lower-level one through dialect conversion passes.
• Use Case: Enables co-design of hardware and compilers by allowing custom, hardware-specific dialects.

Instruction selection is the compiler backend phase that maps the low-level operations from a legalized IR to specific sequences of machine instructions available on the target hardware. It is the final step of lowering to machine code.

• Process: Uses pattern matching (often via tree parsing) to find the most efficient instruction sequence for each IR operation or subgraph.
• Considerations: Must account for instruction costs, available registers, and side effects.
• Example: Selecting between a fused multiply-add (FMA) instruction or separate MUL and ADD instructions based on the target NPU's capabilities.

Control flow flattening is a lowering transformation that converts structured control flow (like loops, conditionals) into a simpler, goto-based or dataflow representation. This is often necessary for hardware with limited or explicit control flow mechanisms.

• Purpose: Eases lowering to hardware that lacks complex control flow primitives, such as many NPUs which prefer static, data-parallel execution.
• Technique: Transforms for and if statements into a graph of basic blocks connected by explicit branch operations.
• Trade-off: Simplifies hardware mapping but can obscure high-level structure for later optimizations.