Compute Graph

The compute graph is a Directed Acyclic Graph (DAG), meaning data flows in one direction along edges (arrows) and contains no cycles. This structure is fundamental because it:

• Enables automatic differentiation for backpropagation by providing a clear path for gradient flow.
• Allows for topological sorting, which determines a valid execution order for all operations.
• Prevents infinite loops during execution, ensuring deterministic forward and backward passes.
• Facilitates parallelism by identifying independent subgraphs that can be computed concurrently.

Each node in the graph represents a distinct mathematical operation (op). These are the fundamental units of computation. Common op types include:

• Element-wise ops: Addition, multiplication, activation functions (ReLU, Sigmoid).
• Linear algebra ops: Matrix multiplication (MatMul), convolution (Conv2D).
• Reduction ops: Sum, mean, max.
• Control flow ops: Conditional branches, loops (often represented as higher-level constructs).
• Specialized ops: Layer normalization, attention mechanisms. The op defines the computation performed and its required input/output tensor shapes.

Edges represent the multidimensional data arrays (tensors) that flow between operations. They are the inputs and outputs of nodes. Key aspects include:

• Tensor metadata (shape, data type) is propagated through the graph, enabling shape inference.
• Data dependencies are explicitly shown; an op executes only when all its input tensors are ready.
• Intermediate tensors can be large; a primary goal of graph optimization is to minimize their memory footprint through techniques like in-place operations and memory reuse.
• The flow of tensors defines the graph's dataflow semantics.

Compute graphs can be static (defined before execution) or dynamic (constructed during execution).

Static Graphs (e.g., TensorFlow 1.x, ONNX):

• The entire graph is defined and optimized ahead-of-time (AOT).
• Enables aggressive optimizations like operator fusion and constant folding.
• Less flexible; control flow must be represented as special graph nodes.

Dynamic Graphs (e.g., PyTorch eager mode):

• The graph is built on-the-fly as ops are executed.
• Offers greater flexibility and easier debugging.
• Optimization opportunities are more limited at runtime. Modern frameworks often use a hybrid approach: develop with dynamic graphs, then export to a static graph for production deployment.

The explicit graph representation enables a suite of crucial compiler-style optimizations performed before execution:

• Operator Fusion: Combines multiple sequential ops (e.g., Conv2D + BatchNorm + ReLU) into a single kernel to reduce memory accesses and launch overhead.
• Constant Folding: Pre-computes parts of the graph that only depend on constant values.
• Dead Code Elimination: Removes ops whose outputs are not used.
• Common Subexpression Elimination: Identifies and reuses identical computations.
• Layout Transformation: Changes the memory layout of tensors (e.g., NCHW to NHWC) to better match hardware. These transformations reduce inference latency and memory footprint, which is critical for on-device deployment.

The high-level compute graph must be lowered into hardware-specific instructions. This process involves:

• Graph Partitioning: Splitting the graph to run different subgraphs on optimal processors (e.g., CPU, GPU, NPU).
• Kernel Selection: Choosing the most efficient pre-written kernel or generating a new one for each op on the target hardware (e.g., using TensorRT or MLIR).
• Memory Scheduling: Allocating and scheduling memory buffers for tensors to minimize transfers and maximize cache locality.
• Scheduling: Determining the precise order of kernel launches to hide memory latency and maximize hardware utilization. The final output is an optimized, executable program for the target accelerator.

Modern frameworks operate in two primary modes. Eager execution (PyTorch default, TF eager) evaluates operations immediately, aiding debugging. Graph mode (TF's tf.function, PyTorch's torch.jit.trace/script) captures operations into a static compute graph for optimizations like constant folding and operator fusion. This graph is essential for export to formats like ONNX and for deployment optimizers.

For on-device deployment, lightweight runtimes consume pre-optimized compute graphs.

• TensorFlow Lite: Converts TensorFlow graphs into a flatbuffer format (.tflite). The TFLite interpreter executes the graph with optimized kernels for mobile CPUs, GPUs (via delegates), and NPUs.
• Core ML: Apple's framework. Models are converted to its .mlmodel format, where the graph is optimized for Apple Silicon (Neural Engine, GPU, CPU).
• PyTorch Mobile: Provides an optimized runtime for executing TorchScript graphs on Android and iOS, supporting lightweight just-in-time (JIT) optimizations.

JIT compilation (e.g., PyTorch's JIT, TensorFlow's XLA in JIT mode) occurs at runtime. The framework's graph executor analyzes the compute graph and generates optimized machine code tailored to the available hardware just before execution. This allows for optimizations based on actual input shapes and available system resources, though it introduces a one-time compilation overhead.

AOT compilation (e.g., TensorFlow's XLA AOT, TVM) happens before deployment. The compute graph is fully compiled into a standalone, optimized executable or library (e.g., a .so file or C++ code) for a specific target platform (e.g., ARM64). This eliminates runtime compilation overhead and is critical for predictable, low-latency inference on resource-constrained edge devices and microcontrollers.

A compiler optimization that merges multiple sequential operations (nodes) in a compute graph into a single, compound kernel. This reduces the overhead of launching multiple kernels and minimizes intermediate tensor writes to memory.

• Example: Fusing a Convolution, Batch Normalization, and ReLU activation into one operation.
• Benefit: Dramatically improves inference speed and reduces memory bandwidth pressure, which is critical for edge devices.

An optimization that evaluates subgraphs consisting of only constant values at compile time, replacing them with a single constant node in the final executable graph.

• Example: Calculating (weights * 2.0) + bias where all values are fixed, and storing just the result.
• Benefit: Eliminates redundant computation at runtime, simplifying the graph and reducing latency.

The process of fully compiling a model's compute graph into an optimized, platform-specific executable binary before deployment. This contrasts with Just-In-Time (JIT) compilation, which happens at runtime.

• Use Case: Essential for edge deployment where runtime resources are scarce and startup latency must be minimized.
• Tools: Executables are produced by compilers like TVM, Apache TVM, or hardware-specific SDKs.

The specialized process of translating a neural network's compute graph into a sequence of microcode instructions for a dedicated Neural Processing Unit (NPU). This involves graph lowering, operator mapping, and detailed memory scheduling.

• Complexity: NPUs have unique memory hierarchies and execution units, requiring a custom compiler (e.g., Qualcomm AI Engine Direct, MediaTek NeuroPilot).
• Goal: To maximize data reuse and parallel execution within the constraints of the NPU's on-chip memory and compute fabric.