Module fusion_reduction 

Reduction-pattern fusion pass — fuses reduction-chained element-wise regions (LayerNorm, softmax, …) that the linear element-wise fuser in crate::optimizer::fusion cannot capture.

Why a separate pass

The element-wise chain fuser merges only linear producer → consumer chains in which every fused node is the sole consumer of its predecessor. Reduction patterns break that rule on purpose:

  LayerNorm:                       Softmax:

       x                               x
       │                               │
    ┌──┴── mean (reduction)        ┌───┴── max (reduction)
    │       │                      │        │
    ▼       ▼                      ▼        ▼
    sub ◄───┘  (broadcast)         sub ◄────┘  (broadcast)
    │                               │
    ▼                               ▼
    var (reduction)                exp
    │                               │
    ▼                               ▼
    normalize ◄─ (broadcast)       sum (reduction)
    │                               │
    ▼                               ▼
    scale_shift                    divide ◄─ (broadcast)

A reduction node produces a small statistic that is broadcast back to one or more element-wise consumers, so the reduction output fans out to several nodes (out-degree ≥ 2) and the region re-converges at a single sink. That diamond shape is precisely what the chain fuser rejects (it treats any fan-out as “parallel branches, do not fuse”). Fusing these regions into one kernel removes the intermediate broadcast round-trips to device memory — the dominant cost of an un-fused LayerNorm/softmax.

What is fused

A set of nodes is a reduction-fusion region rooted at a fusible kernel r when:

1. r is a fusible NodeKind::KernelLaunch.
2. The region R = {r} ∪ members forms a single-entry, single-exit (SESE) subgraph: r dominates every member, and a unique member sink is reachable from every member and is the only node in R with a successor outside R (post-dominates the region).
3. Every member is a fusible kernel with a launch configuration compatible with r (same total thread count).
4. The region is closed: no edge leaves a non-sink member to a node outside R. This is what makes the rewrite semantics-preserving — every intermediate buffer is consumed entirely inside the fused kernel and is never observed by the rest of the graph.
5. The region contains genuine broadcast fan-out: at least one member has out-degree ≥ 2 inside R (the reduction → broadcast shape). This is what distinguishes a reduction region from a plain linear chain (handled by the element-wise fuser) and forces region size ≥ 3.

When a region matches, the pass emits a ReductionFusionGroup describing the merge. As with crate::optimizer::fusion, the graph itself is not mutated — the descriptor is consumed downstream by the PTX codegen layer, and the rewrite helper materialises the fused ComputeGraph when a concrete rewritten DAG is required (e.g. for the CPU simulator).

Algorithm

1. Topological + dominance analysis.
2. For each fusible kernel r in topological order that is not yet claimed, grow the dominated, fusible, closed region below r and test the SESE + fan-out conditions.
3. Accept the largest such region; mark its members claimed so they are not re-used by a later (smaller) region.

Structs

ReductionFusionGroup

A reduction-broadcast region that the pass merges into a single kernel.

ReductionFusionPlan

The complete reduction-fusion plan produced by analyse.

Enums

ReductionPattern

The classified shape of a fused reduction region.

Functions

analyse

Runs the reduction-fusion analysis pass on graph.

rewrite

Rewrites graph by collapsing every region in plan into a single fused kernel node, producing a new ComputeGraph.