Task Graph

A node represents an individual unit of computation or work within the graph. Each node is a task, which can be a function, a kernel (e.g., a matrix multiplication), or a data transformation. In neural network execution, a node often corresponds to a layer or a fused group of operations. Nodes are characterized by their computational cost, resource requirements (e.g., memory, registers), and the data they produce or consume.

A directed edge between two nodes represents a data dependency or a control dependency. It defines a producer-consumer relationship, where the output of the source task is required as input by the destination task. The direction of the edge imposes a partial ordering, ensuring tasks execute only after their predecessors have completed. Edges may be annotated with the size and type of data being transferred, which is critical for scheduling and memory allocation.

The acyclic property is fundamental: there are no cycles or loops in the graph. This guarantees that the computation can complete, as there is no circular wait condition. The DAG structure enables:

• Static analysis for scheduling and resource allocation.
• Identification of the critical path (the longest path through the graph), which determines the minimum possible execution time.
• Safe parallelization, as dependencies are explicit and non-circular.

Granularity refers to the size or computational weight of individual tasks. It is a key trade-off in graph design:

• Fine-grained tasks are small, numerous, and offer high parallelism but incur significant scheduling and synchronization overhead.
• Coarse-grained tasks are larger, reduce overhead, but may limit parallelism and load balancing. Optimizers often perform kernel fusion—merging multiple fine-grained operations into a single coarse-grained node—to minimize data movement and launch overhead on accelerators.

Task graphs implement a dataflow programming model, where execution is triggered by the availability of data. The semantics govern how data is passed:

• Value-based dataflow: Edges transfer actual data tensors or values.
• Token-based dataflow: Edges transfer control tokens or futures, with data fetched lazily. This model decouples the specification of what to compute from when and where to compute it, enabling dynamic, runtime schedulers to map tasks to available processors.

While not part of the static graph definition, the execution model defines how the graph is realized at runtime. Key concepts include:

• Dynamic Schedulers: Systems that map ready tasks (those with all dependencies satisfied) to worker threads or cores.
• Work Stealing: A load-balancing technique where idle workers steal tasks from the queues of busy workers.
• Barrier Synchronization: Points, often implicit at graph completion, where all tasks must finish before proceeding. The graph provides the blueprint; the scheduler orchestrates the live execution on parallel hardware.

In distributed training, a task graph orchestrates the data-parallel and model-parallel stages. Key tasks include:

• Data Loading & Augmentation: Parallel tasks fetch and preprocess mini-batches.
• Forward Pass: Computations flow layer-by-layer; dependencies enforce correct tensor shapes.
• Loss Calculation & Backward Pass: The autograd system dynamically builds a computation graph where each tensor operation becomes a node. The backward pass traverses this graph in reverse to compute gradients.
• Gradient Synchronization: An All-Reduce operation (e.g., using NCCL) is a collective task with dependencies on all local gradient calculations.
• Optimizer Step: The weight update task depends on the synchronized gradients. Frameworks like PyTorch and TensorFlow implicitly or explicitly manage these graphs to schedule work across GPUs or NPUs.

For low-latency model serving, systems like TensorRT, Triton Inference Server, or custom NPU runtimes use task graphs to manage continuous batching.

• Request Ingestion: Incoming queries become individual sub-graphs.
• Graph Fusion & Kernel Launch: The scheduler analyzes independent operations (e.g., matrix multiplications, activations) across a batch, fusing them into optimized kernels. Dependencies ensure fused kernels execute in the correct sequence.
• Memory Lifecycle Management: Tasks for allocating input buffers, executing the graph, and deallocating output buffers are explicitly modeled to prevent memory leaks and overlap computation with data transfer. This graph-based approach allows the scheduler to maximize hardware occupancy by packing operations from multiple requests into a single, efficient execution stream.

AI compilers (e.g., Apache TVM, MLIR, vendor SDKs) use multi-level task graphs to transform a high-level model into optimized NPU code.

1. High-Level Graph: The initial graph from a framework (ONNX, PyTorch IR) with abstract operations.
2. Operator Fusion Graph: The compiler applies kernel fusion rules, merging nodes like convolution, batch norm, and ReLU into a single compound task to minimize intermediate memory writes.
3. Scheduled Graph: Tasks are assigned to specific hardware units (e.g., tensor cores, vector units). Dependencies model data movement between on-chip SRAM, DRAM, and compute units.
4. Instruction Graph: The final graph represents low-level, scheduled instructions (microtasks) for the NPU's execution units. This graph is serialized into an executable command stream. This process is the core of hardware-aware model optimization.

A production RAG pipeline is a classic task graph that coordinates heterogeneous subsystems.

• Query Understanding & Rewriting: An initial LLM call to clarify the user intent.
• Vector Search: A parallel task queries a vector database using the rewritten query's embedding. This task is independent of subsequent generation but must complete before it.
• Re-Ranking & Synthesis: A separate task may re-rank retrieved chunks using a cross-encoder. The top-k chunks become inputs.
• Contextual Generation: The final LLM call task has a strict dependency on the retrieved context. The graph ensures the context is fully prepared before generation begins.
• Citation & Logging: Post-generation tasks format the output and log telemetry. Modeling this as a graph allows for async execution of non-critical path tasks (like logging) and easy insertion of caching nodes.

In agentic architectures, a task graph dynamically expands to represent a plan.

• Goal Decomposition: A planning LLM breaks a high-level objective ("analyze Q3 sales report") into a graph of subtasks: [fetch_data → clean_data → run_analysis → generate_summary].
• Tool Calling: Each subtask may require executing a tool (API, database query, code execution). These tool calls become nodes with dependencies on their inputs.
• Conditional & Loop Execution: The graph can include control-flow nodes. For example, an IF node based on the result of run_analysis might add a new investigate_anomaly task branch.
• Error Handling & Retry: Failed tasks can trigger recursive error correction sub-graphs. The system monitors the graph's critical path for timeouts or failures. Frameworks like LangGraph explicitly use this paradigm to manage state and execution flow across multi-step agentic workflows.

In large-scale ML pipelines, tools like Apache Airflow, Kubeflow Pipelines, and Meta's FBLearner use task graphs to manage data workflows.

• Extract, Transform, Load (ETL): Each data source has a dedicated extraction task. Transformation tasks (normalization, feature encoding) depend on specific extractions completing.
• Parallel Feature Computation: Independent feature calculations (e.g., computing rolling averages vs. one-hot encoding) are modeled as parallel nodes, significantly reducing wall-clock time.
• Validation & Joining: A task validates data quality before a final joining task merges all feature sets into a training dataset.
• Triggering Downstream Jobs: The final node in the graph often triggers the model training task graph. This clear dependency model ensures data is fresh and valid before costly training cycles begin on GPU/TPU/NPU clusters.

The critical path is the longest sequence of dependent tasks in a task graph from the start node to the finish node. It determines the minimum possible execution time (makespan) for the entire parallel computation. Optimizing a task graph often involves identifying and shortening its critical path through techniques like task fusion or asynchronous execution of non-critical dependencies.

Work stealing is a dynamic, decentralized load-balancing scheduling algorithm used to execute task graphs efficiently. Idle processors (or threads) 'steal' tasks from the queues of busy processors. This is highly effective for irregular or imbalanced task graphs where task costs are unknown at compile time, as it adaptively redistributes work without centralized coordination. Frameworks like Intel's Threading Building Blocks (TBB) and C++17's parallel algorithms often use this strategy.

A data race is a concurrency bug that can occur during the parallel execution of a task graph. It happens when two or more tasks access the same memory location concurrently, at least one access is a write, and the accesses are not ordered by explicit dependencies or synchronization. Correct task graph construction must define data dependencies (edges) to prevent races, ensuring deterministic execution. Tools like ThreadSanitizer are used to detect such violations.

SIMT is the execution model used by modern GPUs and many NPUs. It extends SIMD by issuing a single instruction to a warp or wavefront of threads, each executing it on its own data. This model is crucial for understanding how fine-grained parallel tasks (like individual neurons in a layer) are mapped to hardware. Task graphs for these architectures must be compiled into kernels that efficiently utilize the SIMT paradigm, managing thread divergence and memory coalescing.

A barrier is a synchronization primitive that forces all tasks in a parallel group to reach a specific point before any can proceed. In a task graph, barriers are often implicit at the join of multiple dependent tasks. They are essential for correctness in phases of computation (e.g., between forward and backward passes in training) but can create performance bottlenecks if they cause processors to idle. Asynchronous execution and double-buffering are techniques to hide this latency.

Amdahl's Law is a fundamental formula that predicts the maximum theoretical speedup of a parallel program: Speedup = 1 / (S + P/N), where S is the serial fraction, P is the parallelizable fraction, and N is the number of processors. It highlights that the serial sections of a task graph (nodes on the critical path with no parallelism) ultimately limit performance gains. This law drives the compiler optimization goal of minimizing serialization and maximizing task parallelism within graphs.