Module kernel_fusion 

Elementwise Kernel-Fusion Planner

A CPU-side, compiler-style planner that operates over an operation DAG and decides which (primarily elementwise) operations can be fused into a single GPU kernel. No GPU execution happens here: the planner reasons purely about the graph topology, fusion legality, and the memory-bandwidth implications of fusion via a bytes-moved cost model.

Operation graph

The graph is a DAG of FusionOps. Each op has a stable id (equal to its insertion index), an OpKind, a list of inputs (the op ids of its producers), and an output tensor described by output_shape and dtype_bytes. The byte size of an op’s output tensor is product(output_shape) * dtype_bytes (an empty shape denotes a scalar, i.e. product == 1).

An op with no inputs is a source (leaf): it represents data that is already resident in memory (a parameter, gradient, or the result of an upstream subgraph). FusionGraph::validate guarantees the graph is acyclic and that every input references a real op.

Fusion legality

Two or more ops may be fused into one kernel iff:

1. every op in the group is fusible (an elementwise OpKind, see OpKind::is_fusible); barrier ops (MatMul, Reduce, Transpose) can never join an elementwise group and always break chains;
2. the members form a connected producer -> consumer chain in the DAG (a fusible edge connects a fusible producer to a fusible consumer);
3. the producer’s output shape is broadcast-compatible with the consumer’s output shape (NumPy trailing-dimension rule, optional via FusionPlanner::with_broadcast); and
4. fusing must not create a cycle in the group-contracted dependency graph. Greedily merging fusible edges can otherwise sandwich a barrier group between two halves of an elementwise group, which would require the fused kernel to run both before and after the barrier. Such merges are rejected so the inter-group schedule stays a DAG.

When an intermediate that is internal to a group is also consumed by an op outside the group, the intermediate cannot be elided: it is marked as a group output and materialized to memory (never illegally dropped).

Group formation

Groups are formed by traversing ops in topological order and greedily merging fusible producer -> consumer edges, subject to the legality checks above (barriers stay as singletons; cycle-creating merges are skipped). The result is a partition where every op belongs to exactly one group (singletons allowed) and the contracted group graph is acyclic.

Memory-bandwidth cost model

• Unfused bytes moved: for every op, read each of its distinct input tensors once and write its output tensor once; summed over all ops.
• Fused bytes moved: for every group, read the group’s external input tensors once (distinct producers outside the group) and write the group’s external output tensors (members consumed outside the group or that are graph terminals). Intermediates that stay inside the group live in registers and are not counted.

Because a singleton group reproduces exactly the unfused contribution of its single op, fusion can only ever remove traffic: bytes_fused <= bytes_unfused and speedup_estimate = bytes_unfused / bytes_fused >= 1.0 (a bandwidth-bound proxy).

Structs

FusionGraph

An append-only builder for an operation DAG.

FusionGroup

A group of ops fused into a single kernel.

FusionOp

A single operation in the fusion graph.

FusionPlan

The output of FusionPlanner::plan: the fusion groups plus a memory-bandwidth cost summary.

FusionPlanner

Plans elementwise kernel fusion over an operation DAG.

Enums

OpKind

The kind of an operation in the fusion graph.