fuse_pointwise

torch_tvarant.compiler.fuse_pointwise(gm) → GraphModule

Collapse connected pointwise subgraphs into one JIT kernel.

Source code in torch_tvarant/compiler.py

def fuse_pointwise(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
    """Collapse connected pointwise subgraphs into one JIT kernel."""
    graph = gm.graph
    fused_groups = 0
    for root in list(graph.nodes):
        if _pw_kind(root) is None:
            continue
        if any(_pw_kind(u) is not None for u in root.users if isinstance(u, torch.fx.Node)):
            continue

        fused = {root}
        changed = True
        while changed:
            changed = False
            for n in graph.nodes:
                if n in fused or _pw_kind(n) is None:
                    continue
                users = [u for u in n.users if isinstance(u, torch.fx.Node)]
                if users and all(u in fused for u in users):
                    fused.add(n)
                    changed = True
        if len(fused) < 2:
            continue

        order = [n for n in graph.nodes if n in fused]
        ssa: dict[torch.fx.Node, int] = {}
        inputs: list[torch.fx.Node] = []
        ops: list[int] = []
        a: list[int] = []
        b: list[int] = []
        input_ids: list[int] = []
        alphas: list[float] = []
        consts: list[float] = []

        def emit_load(n: torch.fx.Node) -> int:
            if n in ssa:
                return ssa[n]
            ops.append(LOAD)
            a.append(-1)
            b.append(-1)
            input_ids.append(len(inputs))
            alphas.append(1.0)
            consts.append(0.0)
            inputs.append(n)
            ssa[n] = len(ops) - 1
            return ssa[n]

        def rec_input(src: object) -> int:
            if isinstance(src, torch.fx.Node):
                if src in fused:
                    return ssa[src]
                return emit_load(src)
            ops.append(CONST)
            a.append(-1)
            b.append(-1)
            input_ids.append(-1)
            alphas.append(1.0)
            consts.append(float(src))  # type: ignore[arg-type]
            return len(ops) - 1

        ok = True
        for n in order:
            kind_op = _pw_kind(n)
            if kind_op is None:
                ok = False
                break
            kind, op = kind_op
            if kind == "unary":
                ia = rec_input(n.args[0])
                ops.append(op)
                a.append(ia)
                b.append(-1)
                input_ids.append(-1)
                alphas.append(1.0)
                consts.append(0.0)
            elif kind == "binary":
                ia = rec_input(n.args[0])
                ib = rec_input(n.args[1])
                alpha = 1.0
                if len(n.args) > 2:
                    alpha = float(n.args[2])  # type: ignore[arg-type]
                alpha = float(n.kwargs.get("alpha", alpha))
                ops.append(op)
                a.append(ia)
                b.append(ib)
                input_ids.append(-1)
                alphas.append(alpha)
                consts.append(0.0)
            elif kind == "scalar":
                ia = rec_input(n.args[0])
                ops.append(SCALE)
                a.append(ia)
                b.append(-1)
                input_ids.append(-1)
                alphas.append(float(n.args[1]))  # type: ignore[arg-type]
                consts.append(0.0)
            else:
                ok = False
                break
            ssa[n] = len(ops) - 1
        if not ok or not inputs:
            continue

        with graph.inserting_before(root):
            new = graph.call_function(
                torch.ops.tvarant.pointwise,
                args=(inputs, ops, a, b, input_ids, alphas, consts),
            )
        root.replace_all_uses_with(new)
        for n in reversed(order):
            if not n.users:
                graph.erase_node(n)
        fused_groups += 1

    graph.lint()
    gm.recompile()
    last_log["pointwise_groups"] = fused_groups
    return gm