Operator Clustering

The primary goal is to identify groups of operators that can be legally and efficiently fused into a single, compound kernel. This targets hardware accelerators like NPUs and GPUs where kernel launch overhead is significant. Clustering groups operators that:

• Have compatible data types and tensor layouts.
• Can be executed within the same hardware execution unit (e.g., a tensor core).
• Benefit from reduced intermediate memory writes by keeping data in fast registers or shared memory.

For example, a sequence like Conv2D -> BiasAdd -> ReLU is a classic cluster candidate for fusion into a single kernel.

A core objective is to reduce costly data transfers between different levels of the memory hierarchy. By clustering operators, the compiler aims to:

• Keep intermediate tensors in high-bandwidth memory (e.g., SRAM, shared memory) instead of writing them back to slower global memory (DRAM).
• Enable in-place computation where possible, reusing memory buffers.
• Co-locate computations that access the same region of data, improving data locality and cache efficiency.

This directly addresses the memory bandwidth bottleneck that dominates many neural network workloads.

Clustering transforms a fine-grained graph into a coarser-grained task graph, which simplifies scheduling. This allows the runtime or compiler to:

• Treat a cluster as a single scheduling unit, reducing the complexity of dependency analysis.
• Map entire clusters to specific NPU cores or processing elements for better workload balance.
• Exploit pipeline parallelism between independent clusters.
• Apply dynamic batching more effectively at the cluster level rather than per-operator.

Clustering is used to partition the graph into hardware-agnostic and hardware-specific sections. The objective is to:

• Group operators that must be executed using vendor-specific SDK kernels or hardware intrinsics (e.g., a custom NPU activation function).
• Isolate these clusters so they can be lowered and optimized using the vendor's proprietary toolchain.
• Keep the rest of the graph in a portable intermediate representation (IR) for other optimizations.

This enables a hybrid compilation approach, leveraging both open-source and proprietary compiler stacks.

Clustering enables strategic use of different numerical precisions (e.g., FP32, FP16, INT8) within a computational graph. Objectives include:

• Grouping operators that can safely execute in lower precision (e.g., INT8 quantized convolutions) into distinct clusters.
• Inserting precision conversion operations (e.g., Quantize/Dequantize) at cluster boundaries.
• Allowing the compiler to apply precision-specific kernel optimizations and memory layouts per cluster.

This is critical for maximizing throughput and energy efficiency on NPUs with strong low-precision arithmetic support.

Clustering creates a manageable set of optimization targets for performance tuning. The compiler can:

• Profile execution at the cluster level to identify bottlenecks.
• Apply auto-tuning techniques (e.g., searching over tile sizes, unroll factors) specifically to the fused kernel generated for a cluster.
• Use profiling data to make better clustering decisions in subsequent compilations via Profile-Guided Optimization (PGO).
• Cache optimal kernel implementations or cluster configurations for specific hardware targets.

In TensorFlow via XLA, clustering is performed during the HLO (High-Level Optimizer) phase. The compiler uses a cost model to evaluate fusion profitability and forms clusters based on:

• Dataflow dependencies and producer-consumer relationships.
• Heuristic rules preventing cycles and respecting resource constraints (e.g., memory).
• A FusionMerger pass that greedily merges nodes to reduce kernel launch overhead and intermediate tensor materialization. Clusters are then lowered to LLVM IR for target-specific code generation.

TorchInductor, PyTorch's default compiler, performs clustering as part of its fusion strategy. It operates on a FX graph representation:

• The Fusion pass identifies fusion patterns (e.g., pointwise, reduction, matmul).
• It groups nodes into FusionGroup objects, which become single computational kernels.
• Clustering decisions are made via pattern matching on operator types and a heuristic cost model that estimates memory bandwidth savings versus increased register pressure.
• The fused kernel is then generated in Triton or C++/OpenMP for GPU/CPU backends.

TVM implements operator clustering in its Relay and Relax IRs through the FuseOps pass.

• The pass partitions the dataflow graph using a maximal fusion algorithm, subject to constraints like operator count and avoiding cycles.
• It integrates with TOPI (TVM Operator Inventory) to ensure each cluster maps to a valid, optimized kernel implementation.
• A key feature is target-aware clustering, where fusion rules differ for NPUs (e.g., favoring specific operator sequences like Conv+ReLU+Add) versus GPUs.
• The resulting clusters are then scheduled and lowered to TIR for final code generation.

MLIR provides a structured framework for clustering through its layered dialects.

• A high-level graph dialect (e.g., mhlo or tosa) is first canonicalized.
• Clustering is often expressed as a rewriting pattern that merges operations within a dialect or legalizes them to a fusion-friendly dialect like linalg.
• The affine and scf dialects then handle loop nesting for the fused cluster.
• Policies are defined via Pass Managers, allowing hardware-specific clustering strategies (e.g., for NPUs, clusters may be formed to maximize use of systolic arrays).

Vendor NPU SDKs (e.g., NVIDIA TensorRT, Qualcomm SNPE, Google Edge TPU Compiler) implement aggressive, hardware-prescriptive clustering.

• TensorRT uses a layer and tensor fusion engine that profiles kernels to fuse adjacent operations, drastically reducing latency. It supports vertical (sequential) and horizontal (parallel) fusion.
• SNPE employs a partitioning step that clusters operators destined for the NPU versus CPU/DSP, followed by subgraph fusion within the NPU partition.
• These compilers use vendor-specific kernel libraries; clustering is often limited to pre-defined, highly optimized fusion patterns supported by the hardware.

Compiler implementations must balance multiple constraints:

• Memory Constraints: Clusters must fit within fast scratchpad memory (SRAM) of the NPU. Excessive fusion can lead to spilling to slower DRAM.
• Data Dependencies: Only forming clusters along producer-consumer edges to avoid cycles.
• Operator Compatibility: Fusing only compatible operations (e.g., element-wise with preceding convolution) that can be executed within a single kernel.
• Compilation Time vs. Performance: Exhaustive search for optimal clusters is NP-hard; compilers use greedy algorithms or pattern-based approaches.
• The goal is to maximize compute-to-memory-access ratio and minimize off-chip data movement.

Graph fusion is the compiler optimization that merges multiple adjacent operators into a single, compound kernel. This is the primary goal of operator clustering. The benefits are:

• Reduced kernel launch overhead by executing multiple operations in one GPU/NPU kernel call.
• Minimized intermediate memory accesses by keeping data in fast registers or cache between fused operations.
• Enables the use of vendor-optimized fused kernels (e.g., fused convolution-ReLU-BatchNorm).

While clustering identifies candidate groups, fusion performs the actual kernel generation.

Graph partitioning is the strategy of dividing a large computational graph into smaller, manageable subgraphs. It serves a broader purpose than clustering:

• Distributing workloads across multiple processors, devices (multi-NPU), or memory hierarchies.
• Isolation for heterogeneous execution, where different subgraphs run on CPUs, GPUs, or specialized NPU cores.
• Enabling pipelined execution across partitions.

Operator clustering can be seen as a fine-grained form of partitioning, where the partition boundary is defined by fusibility constraints and hardware unit affinity.

Operator reordering changes the execution sequence of independent operators in a graph. This optimization often precedes clustering to create better fusion opportunities.

Key principles:

• Improves data locality by bringing dependent operations closer together.
• Can reduce peak memory usage by allowing earlier release of large intermediate tensors.
• Must respect data dependencies; only commutative, independent operations can be reordered.

For example, reordering a sequence like (A -> B -> C) where A and C are independent of B might allow a new cluster (B -> A -> C) if A and C become adjacent and fusible.

This is a design-time strategy that adapts the neural network model itself to leverage specific NPU features, creating a graph that is inherently easier for the compiler to cluster and optimize.

Techniques include:

• Choosing operator types known to have efficient, fused implementations on the target NPU.
• Designing subgraph patterns that map cleanly to the NPU's vector/SIMD units or special function units (SFUs).
• Adjusting tensor layouts (e.g., NHWC vs. NCHW) to match the NPU's preferred memory access patterns.

This co-design approach between the model architect and the compiler stack ensures the computational graph presented for clustering is hardware-friendly from the start.

Kernel auto-tuning is the automated search for optimal implementation parameters for a given kernel or cluster. After the compiler decides which operators to fuse (clustering), auto-tuning decides how to best implement that fused cluster.

The search space includes:

• Tile sizes for blocking loops over data.
• Unroll factors for loops.
• Thread block dimensions and shared memory allocation strategies.
• Choices between different algorithm variants (e.g., Winograd vs. direct convolution).

This process uses empirical benchmarking on the target hardware to find the configuration that maximizes throughput or minimizes latency for the specific operator cluster.

MLIR is a compiler infrastructure that provides a modular system of interoperable intermediate representations ('dialects'). It is foundational for implementing modern operator clustering algorithms.

How it relates:

• Dialects for Abstraction: High-level ops (e.g., linalg.conv) can be gradually lowered through dialects, with clustering decisions made at the most appropriate abstraction level.
• Pattern Rewriting: Clustering is implemented as a set of rewrite patterns that match subgraphs and replace them with a fused op.
• Hardware-Specific Dialects: Vendor NPU SDKs often define their own MLIR dialects (e.g., tosa, tpu). Clustering rules are encoded as legalization patterns from a framework dialect (like TensorFlow or PyTorch) down to these hardware-specific dialects.

MLIR provides the structured framework to express and combine clustering with other graph optimizations.