Operator Dispatch

The core function of dispatch is to select the appropriate kernel—a hardware-specific implementation of an operation—at runtime. This decision is based on a dispatch key composed of:

• Device Type (e.g., CPU, GPU, NPU)
• Data Type (e.g., float32, int8, bfloat16)
• Layout (e.g., contiguous, channels-last, sparse)
• Operator Schema (e.g., Conv2d, MatMul)

The runtime matches this key against a registry of available kernels to execute the operation with maximum efficiency for the given context.

Dispatch provides a clean abstraction layer between the logical computational graph and physical hardware. This enables:

• Single graph definition to run on multiple backends (CPU, GPU, NPU).
• Vendor-specific kernels (e.g., cuDNN for NVIDIA, MPS for Apple Silicon) to be plugged in without altering model code.
• Graceful fallback to a default CPU implementation if no optimized kernel is available for a target device.

This decoupling is fundamental for frameworks like PyTorch and TensorFlow, ensuring model portability across diverse deployment environments.

The dispatch mechanism must inspect the runtime context of each operator invocation to make the correct kernel choice. Key inspected attributes include:

• Tensor shapes and strides: Determines memory access patterns and suitability for vectorized kernels.
• Device placement: Identifies which physical or logical device holds the input tensors.
• Autograd state: In training frameworks, determines if a kernel needs to record operations for gradient computation.
• Stream/queue: For GPUs and other accelerators, ensures kernel execution is ordered correctly to maintain dependencies.

This inspection adds minimal but essential overhead to each operation.

Operator dispatch architectures balance two competing goals:

• Performance: Minimizing dispatch overhead is critical for fine-grained operations. Techniques include virtual function tables (vtables), hash maps for kernel lookup, and caching recently used kernel pointers.
• Flexibility: Supporting a wide array of kernels, devices, and data types requires a extensible registration system. New kernels can be registered at runtime by hardware libraries.

Ahead-of-Time (AOT) compilation can resolve dispatch decisions at compile time for a known hardware target, eliminating runtime overhead entirely but sacrificing flexibility.

In optimized inference runtimes, dispatch interacts closely with graph-level optimizations:

• Graph lowering may resolve generic operators into device-specific ones before dispatch.
• Operator fusion can create new, compound operators that require a dedicated fused kernel, registered in the dispatch table.
• Static dispatch: For graphs with known input types and shapes, the compiler can statically bind kernels, bypassing dynamic lookup. This is a key optimization in runtimes like TensorRT and TVM.

The dispatch system must remain consistent across both eager execution (operation-by-operation) and compiled graph execution modes.

Modern dispatch systems must efficiently handle specialized data formats:

• Sparse tensors: Operations on sparsely encoded data (e.g., CSR format) require kernels that skip zero computations. The dispatch key must include the sparse encoding type.
• Quantized operators: For INT8 or lower-precision inference, dispatch selects kernels that use integer arithmetic pipelines (e.g., using VNNI instructions on Intel CPUs, DP4A on NVIDIA GPUs). This often involves a separate registration path for quantized operator schemas (e.g., QuantizedConv2d).

Efficient dispatch for these paths is essential for performance in on-device and recommendation system workloads.

Graph lowering is the compiler process of transforming a high-level, hardware-agnostic intermediate representation (IR) of a neural network into a lower-level, target-specific IR or machine code. This creates the precise set of operations that the operator dispatch system must map to hardware kernels. Lowering decisions directly influence dispatch complexity by determining the granularity and type of operators presented to the runtime.

A hardware delegate is a framework component that offloads the execution of a subgraph to a dedicated accelerator (e.g., GPU, NPU). It acts as a specialized dispatch endpoint. The main dispatch logic often identifies subgraphs suitable for delegation and hands off control, after which the delegate manages its own internal kernel selection and execution on the target hardware.

Kernel auto-tuning is the automated search for optimal implementation parameters (e.g., tile sizes, thread counts) for a computational kernel on specific hardware. This process often occurs offline or at first dispatch. The resulting optimal kernel variant is then cached and selected by the dispatch mechanism for subsequent executions with matching input characteristics, maximizing hardware utilization.

Just-in-time (JIT) compilation generates optimized machine code for operators during program execution, often using runtime shape or data type information. JIT complements operator dispatch; the dispatch system may call a JIT compiler to generate a specialized kernel on-the-fly if no pre-compiled version exists, caching it for future use. This enables optimizations impossible with static ahead-of-time (AOT) compilation.

A cost model is an analytical or heuristic function that estimates the computational cost (latency, power) of executing a specific operator implementation. Advanced dispatch systems use cost models to make intelligent decisions between multiple valid kernels for the same operation, choosing the one predicted to be fastest or most efficient for the given runtime context (e.g., input size, current system load).

An Intermediate Representation (IR) is the abstract, platform-independent data structure that represents a neural network's computational graph within a compiler framework. The structure and operator set defined in the IR are the primary interface for the dispatch system. The dispatch mechanism maps nodes in this IR to concrete hardware kernels, making the IR design foundational to dispatch logic.