Darkweb007/cuda-kernels

1"""The fusion pass.2 3Rule: two adjacent elementwise nodes A -> B (A's output feeds B) may be4fused into the same group iff:5  1. Both A and B are elementwise ops (see ops.is_elementwise).6  2. A's output has exactly one consumer in the whole graph (namely B).7 8Condition (2) is the standard real-world fusion constraint: if A's output9were used by two different downstream nodes, fusing A into just one of them10would mean recomputing A's work for the other consumer (or keeping a11separate materialized copy anyway), which is a genuine cost-model decision,12not something to do silently. So by default we keep A un-fused into B in13that case, and A ends up as its own single-node group (or fused with ITS14single-consumer... no, it has two, so it stays separate).15 16Implementation: union-find over node output names. Group membership is a17maximal weakly-connected component under the "fusable edge" relation above.18Groups are then topologically ordered (in original node order, which is19already topological since the Graph is built via sequential .add() calls).20"""21 22from __future__ import annotations23 24from dataclasses import dataclass, field25from typing import Dict, List26 27from .graph import Graph, Node28from .ops import is_elementwise29 30 31class _UnionFind:32    def __init__(self):33        self.parent: Dict[str, str] = {}34 35    def find(self, x: str) -> str:36        self.parent.setdefault(x, x)37        while self.parent[x] != x:38            self.parent[x] = self.parent[self.parent[x]]39            x = self.parent[x]40        return x41 42    def union(self, a: str, b: str) -> None:43        ra, rb = self.find(a), self.find(b)44        if ra != rb:45            self.parent[ra] = rb46 47 48@dataclass49class FusionGroup:50    nodes: List[Node] = field(default_factory=list)51 52    @property53    def inputs(self) -> List[str]:54        """External inputs to this group: names read but not produced within it."""55        produced = {n.output for n in self.nodes}56        seen = []57        for n in self.nodes:58            for i in n.inputs:59                if i not in produced and i not in seen:60                    seen.append(i)61        return seen62 63    @property64    def output(self) -> str:65        """The single external-facing output of this group. Assumes groups66        formed by `fuse()` have exactly one node whose output escapes the67        group (true by construction: we only merge single-consumer edges)."""68        internal = {i for n in self.nodes for i in n.inputs}69        externally_visible = [n.output for n in self.nodes if n.output not in internal]70        if len(externally_visible) != 1:71            raise ValueError(72                f"FusionGroup expected exactly one external output, got {externally_visible}"73            )74        return externally_visible[0]75 76    def __repr__(self) -> str:77        body = "\n    ".join(repr(n) for n in self.nodes)78        return f"FusionGroup(\n    {body}\n  )"79 80 81def fuse(graph: Graph) -> List[FusionGroup]:82    consumer_count: Dict[str, int] = {}83    for node in graph:84        for i in node.inputs:85            consumer_count[i] = consumer_count.get(i, 0) + 186 87    uf = _UnionFind()88    for node in graph:89        uf.find(node.output)  # ensure registered90 91    for node in graph:92        if not is_elementwise(node.op):93            continue94        for input_name in node.inputs:95            producer = graph.producer_of(input_name)96            if producer is None:97                continue  # graph input, nothing to fuse with98            if not is_elementwise(producer.op):99                continue100            if consumer_count.get(producer.output, 0) != 1:101                continue  # producer feeds >1 consumer, don't fuse102            uf.union(producer.output, node.output)103 104    groups_by_root: Dict[str, FusionGroup] = {}105    order: List[str] = []106    for node in graph:  # graph.nodes is already topological order107        root = uf.find(node.output)108        if root not in groups_by_root:109            groups_by_root[root] = FusionGroup()110            order.append(root)111        groups_by_root[root].nodes.append(node)112 113    return [groups_by_root[r] for r in order]114