Graph Partitioning

The primary objective is to expose parallelism by splitting the computational graph into independent or semi-independent subgraphs that can be executed concurrently. This is critical for utilizing:

• Multi-core CPUs and many-core GPUs/NPUs.
• Pipeline parallelism where different stages of a model run on different hardware.
• Data parallelism for batch processing across devices.

Without partitioning, the graph executes as a single, sequential unit, leaving vast amounts of parallel hardware idle.

Partitioning maps different subgraphs to the most suitable processing unit based on operator characteristics. The compiler performs a cost-model analysis to decide placement, such as:

• Running matrix multiplications (GEMM) on an NPU's tensor cores.
• Executing control-flow heavy or scalar operations on a CPU.
• Offloading vision-centric convolutions to a GPU.

This objective maximizes overall system efficiency by leveraging the strengths of each hardware component.

A key goal is to reduce the performance penalty of data transfer between devices (e.g., CPU→GPU, GPU→NPU). The partitioner aims to:

• Cut edges with minimal data flow: Place highly interdependent nodes on the same device.
• Increase subgraph granularity: Create larger partitions to amortize communication overhead.
• Overlap computation and communication: Schedule data transfers for tensors needed later in the pipeline.

Excessive communication can easily become the bottleneck, negating the benefits of parallel execution.

Partitioning is essential for fitting large models into limited device memory (e.g., NPU SRAM, GPU VRAM). Objectives include:

• Splitting monolithic graphs that exceed a single device's memory capacity.
• Enabling model parallelism for massive layers that don't fit on one chip.
• Facilitating activation offloading/recomputation by creating natural checkpoint boundaries.

The partitioner works in tandem with the compiler's memory planner to schedule swaps and overlaps.

By creating isolated subgraphs, the compiler can apply targeted, hardware-specific optimizations to each partition. This includes:

• Kernel fusion within a partition tailored for a specific accelerator's instruction set.
• Precision calibration (e.g., INT8 quantization) for partitions running on supported units.
• Custom memory layouts (e.g., NHWC vs. NCHW) optimal for the target hardware.

Partitioning transforms a one-size-fits-all graph into specialized execution blocks.

The partitioner aims to distribute work evenly across available devices to prevent stragglers and maximize throughput. This involves:

• Estimating the computational cost (in FLOPs or cycles) of each node.
• Assigning nodes to create partitions with roughly equal execution time.
• Considering dynamic factors like kernel launch latency and current device load in runtime systems.

Poor load balancing leaves expensive hardware resources underutilized.

The Kernighan-Lin algorithm is a heuristic, iterative improvement algorithm for bisecting a graph into two roughly equal partitions while minimizing the number of edges cut between them. It operates by greedily swapping pairs of vertices between partitions to reduce the cut size.

• Key Mechanism: Starts with an initial bisection, then iteratively identifies and swaps vertex pairs that yield the greatest net reduction in edge cuts.
• Limitation: Primarily designed for bisection (2-way partitions) and requires an initial partition, which can affect final quality.
• Use Case: Foundational algorithm in physical design automation (VLSI) and early compiler partitioning; often used as a refinement step within larger multi-level frameworks.

Multi-level partitioning is a scalable framework that coarsens a graph, partitions the smaller version, and then uncoarsens and refines the partition. The METIS library is its canonical implementation for producing high-quality, balanced partitions.

• Three-Phase Process:
  1. Coarsening: Repeatedly merges vertices/edges to create a smaller, approximate graph.
  2. Partitioning: Applies an algorithm (like Kernighan-Lin) on the small, coarsest graph.
  3. Uncoarsening/Refinement: Projects the partition back to the original graph, refining at each level.
• Advantage: Dramatically faster and often higher-quality than direct partitioning of large graphs.
• Use Case: The de facto standard for static partitioning of large-scale graphs in scientific computing and NPU compilation for distributing subgraphs across cores.

Streaming partitioning algorithms process a graph as a sequence of vertices (a stream) and make immediate, irrevocable assignment decisions, optimizing for low memory overhead and single-pass processing.

• Linear Deterministic Greedy (LDG): Assigns each incoming vertex to the partition that currently hosts the most of its already-seen neighbors, balancing load via a capacity constraint.
• Characteristics:
  • Low Memory: Only partition state and a subset of the graph need to be in memory.
  • Heuristic: Decisions are local and greedy, so solution quality is typically lower than offline methods.
• Use Case: Partitioning extremely large graphs that cannot fit in memory, or for dynamic graphs where the structure evolves.

Hardware-aware partitioning extends traditional graph partitioning by incorporating constraints and costs specific to the target accelerator architecture, such as NPU memory hierarchy, core topology, and data movement penalties.

• Key Objectives:
  • Minimize Off-Chip Communication: Place heavily communicating operators in the same on-chip memory domain.
  • Respect Memory Limits: Ensure subgraphs fit within a core's local SRAM or a processor's cache.
  • Align with Dataflow: Partition to maximize reuse in systolic arrays or spatial accelerators.
• Strategy: Often uses a cost model within a multi-level or iterative framework, where the edge weight represents communication cost (e.g., latency, energy) rather than a simple unit count.
• Use Case: Essential for NPU compilers (like TVM, XLA) to map neural network graphs efficiently onto heterogeneous cores and memory banks.

This strategy uses operator clustering as a precursor to partitioning, grouping tightly-coupled operators (often intended for kernel fusion) into a single cluster before assigning clusters to hardware units.

• Process:
  1. Cluster Formation: Use heuristics (e.g., fusibility rules, data dependency strength) to form candidate clusters.
  2. Cluster as Node: Treat each cluster as a single node in a higher-level abstraction graph.
  3. Partition the Cluster Graph: Apply standard partitioning algorithms to this coarser graph.
• Benefit: Reduces the partitioning problem's scale and directly aligns partitions with fusion opportunities, minimizing intermediate memory transfers.
• Use Case: Common in ML compilers where graph fusion and partitioning are tightly integrated passes; a cluster becomes a partition target for a specific NPU core or functional unit.

Dynamic partitioning adjusts graph partitions at runtime based on observed performance metrics, load imbalance, or changing graph structures, as opposed to static compile-time partitioning.

• Drivers for Adaptation:
  • Load Imbalance: Runtime monitoring shows some processors are idle while others are overloaded.
  • Data-Dependent Workloads: The computational cost of operators varies with input data.
  • System Dynamics: Changes in available hardware resources (e.g., thermal throttling).
• Mechanisms: Can range from work stealing (where idle processors take work from others' queues) to more formal repartitioning triggered at interval boundaries.
• Challenge: Overhead of monitoring and repartitioning must be less than the performance gain.
• Use Case: Runtime systems for irregular graph algorithms (e.g., graph analytics) and advanced NPU schedulers handling variable batch sizes or multi-tenancy.

Operator clustering is a compiler technique that groups sets of operators within a computational graph into distinct clusters. This is often the immediate precursor to partitioning. Clusters are formed based on criteria like:

• Data locality and producer-consumer relationships
• Hardware affinity (e.g., which NPU core type an operator targets)
• Memory bandwidth constraints
• The goal is to create logical units that can be treated as a single entity for subsequent graph partitioning, fusion, or scheduling passes.

Graph fusion is a critical optimization that often follows or works in tandem with partitioning. It merges multiple adjacent operators or nodes within a computational graph—or within a partition—into a single, compound kernel. Key benefits include:

• Reduced kernel launch overhead by executing fused operators in one launch
• Minimized intermediate memory accesses by keeping data in registers or cache
• Enables more efficient use of NPU compute units Effective partitioning creates subgraphs where fusion opportunities are maximized, as operators destined for the same hardware unit are grouped together.

Memory planning is a compiler optimization pass that allocates memory buffers for tensors in a computational graph. It is deeply interdependent with partitioning. After a graph is partitioned:

• The planner must allocate memory for inputs, outputs, and intermediates within each partition
• It aims to minimize peak memory usage per partition, which is critical for devices with constrained memory
• Techniques like buffer reuse and in-place optimization are applied at the partition level Effective partitioning reduces the scope of the memory planning problem, making it more tractable and efficient.

Hardware-aware model optimization involves adapting a neural network's architecture or execution plan to leverage specific NPU features. Graph partitioning is a core tactic within this discipline. It requires the compiler to have a detailed model of the target hardware, including:

• Heterogeneous cores (e.g., vector, matrix, scalar units)
• Memory hierarchy (bandwidth and latency between levels)
• Interconnect topology between processors Partitioning decisions are made to place subgraphs on the hardware units where they will execute most efficiently, balancing load and minimizing data movement.

Parallelism and scheduling strategies determine how and when computational workloads are executed across NPU cores. Graph partitioning provides the units of work for these strategies. After partitioning:

• Static or dynamic schedulers assign partitions to available hardware resources
• Dependencies between partitions define the execution DAG (Directed Acyclic Graph)
• The goal is to maximize hardware utilization by overlapping computation and communication Partitioning effectively exposes the parallelism inherent in the model, enabling efficient scheduling across multiple processors, devices, or memory hierarchies.