Computational Graphs and DAGs

The fundamental unit of analysis in a deep learning compiler is not the individual instruction or the memory address, but the computational graph. While high-level frameworks often present an imperative interface to the developer, allowing line-by-line execution similar to standard Python, the compiler views the model as a static data flow network. This shift in perspective allows the compiler to reason about global data dependencies, memory lifespan, and operation granularity before a single kernel is launched on the hardware.

A computational graph is formally represented as a Directed Acyclic Graph (DAG). In this structure, nodes represent computational operations (such as matrix multiplication, convolution, or element-wise addition), and directed edges represent the flow of data (tensors) between these operations. This representation decouples the mathematical definition of the model from its execution strategy.

Consider a standard dense layer represented by the equation:

$$Y = \text{ReLU}(X W^T + b)$$

In a general-purpose programming language, this might be executed as a sequence of three distinct function calls. In a computational graph, it is a connected structure where the output of the matrix multiplication node flows directly into the addition node, and subsequently into the activation node. This structure explicitly encodes the dependency chain: the addition cannot proceed until the multiplication is complete, but it also implies that $X$ and $W$ are independent and could theoretically be loaded simultaneously.

Anatomy of the Graph

The graph consists of three primary components that the compiler must manage:

1. Nodes (Operators): These are the functional units. In high-level Intermediate Representations (IRs) like TVM Relay or XLA HLO, nodes carry semantic attributes. A Conv2D node, for example, stores static information such as kernel size, strides, padding, and dilation rates as attributes, rather than as input tensors.
2. Edges (Data Dependencies): Edges represent the tensors flowing through the system. Crucially, edges in a compiler graph carry type information, known as shape and data type (dtype). During the compilation process, "shape inference" passes propagate this information through the graph to determine memory allocation requirements without executing the code.
3. Source and Sink Nodes: The graph has explicit entry points (parameters or input data) and exit points (model predictions).

The acyclic nature of the DAG is strictly enforced in most tensor compilers. Cycles in data flow would imply a temporal paradox where an operation depends on its own future output. Recurrent Neural Networks (RNNs), which inherently contain loops, are typically handled by unrolling the loop into a long, acyclic sequence of operations for a fixed number of time steps, or by using specialized control flow operators that encapsulate the loop body as a sub-graph.

Visualizing Data Flow

To understand how a compiler sees a model, we can visualize a simplified subgraph of a residual block. This illustrates how data splits and merges, creating complex dependency patterns that the scheduler must resolve.

Representation of a residual block. The dashed line indicates the skip connection, creating a split in the data path that merges at the "Add" node. The compiler uses this topology to determine that the Input Tensor must be preserved in memory until the addition operation completes.

Data Flow vs. Control Flow

Standard compilers (like GCC or LLVM) excel at Control Flow Graphs (CFGs), where nodes represent basic blocks of instructions and edges represent jumps or branches. AI compilers, however, prioritize Data Flow Graphs (DFGs). In a DFG, the availability of data triggers execution, not a program counter.

This distinction complicates matters when models require conditional logic, such as dynamic execution in recursive networks or conditional generation in transformers. Deep learning compilers handle this through two primary methods:

1. Tracing/Unrolling: If the control flow depends on compile-time constants (e.g., a loop running exactly 4 times), the compiler unrolls the graph into a flat sequence. The if/else logic vanishes, replaced by the specific branch taken during the trace.
2. Functional Control Flow: For data-dependent control flow (e.g., logic that depends on a tensor value calculated at runtime), compilers introduce specialized nodes like If or While. Unlike a standard CPU branch, these nodes often take entire sub-graphs as arguments. An If node might accept a condition tensor, a "then" sub-graph, and an "else" sub-graph. This maintains the graph properties while supporting dynamic behavior.

Topological Ordering and Scheduling

A DAG defines dependencies but not a strict execution order. For the residual block shown above, the compiler knows conv1 must happen before bn1, but it needs to linearize this graph into a sequence of instructions for the hardware. This is achieved through topological sorting.

Topological sorting algorithms produce a linear ordering of vertices such that for every directed edge from $u$ to $v$, vertex $u$ comes before $v$ in the ordering. In complex graphs, multiple valid topological sorts exist. The compiler's scheduler chooses one that optimizes for specific metrics, such as:

• Peak Memory Usage: Executing branches in an order that allows intermediate tensors to be freed as early as possible.
• Cache Locality: Scheduling consumers immediately after producers to keep data in the L1 or L2 cache.

The following visualization demonstrates the concept of execution depth within a DAG. Nodes at the same depth depend only on nodes from previous depths, indicating operations that could theoretically be parallelized given sufficient hardware resources.

Operations grouped by their dependency depth. Steps with multiple blocks indicate potential for parallel execution or stream concurrency. The width of the bar at any step represents the "width" of the graph at that point.

Graph Granularity and Lowering

The complexity of the DAG depends heavily on the level of abstraction. At the highest level (Frontend IR), a node might represent a "Multi-Head Attention" block. This is a macroscopic view useful for graph-level optimizations like checking if the hardware has a specific accelerator for attention mechanisms.

As compilation proceeds, this graph is "lowered." The single Attention node explodes into dozens of smaller nodes: matrix multiplications, transposes, softmax functions, and element-wise divisions. The compiler continues to manage this as a DAG, but the graph size increases significantly. Eventually, these nodes are translated into loop nests and finally into hardware instructions.

Understanding the DAG is a prerequisite for the optimization techniques discussed in Chapter 2. Operator fusion, for instance, is simply the act of identifying specific sub-patterns in the DAG (like Conv2D followed by ReLU) and replacing them with a single "fused" node that performs both operations in one kernel launch. Without the global visibility provided by the graph structure, such optimizations would be impossible to apply safely.