Source code for embedl_deploy._internal.core.plan

# Copyright (C) 2026 Embedl AB

"""Implementation of ``TransformationPlan`` functionality.

The plan-based workflow lets users inspect and edit pattern matches before
applying them to the model.
"""

import copy
import enum
import logging
import warnings
from collections.abc import Sequence
from dataclasses import dataclass, field
from typing import Any, TypedDict

import torch
from torch import fx, nn
from torch.export import export
from torch.fx.passes.shape_prop import ShapeProp

from embedl_deploy._internal.core.backend import get_backend
from embedl_deploy._internal.core.modules import symbolic_trace
from embedl_deploy._internal.core.patterns.main import (
    Pattern,
    PatternMatch,
    Phase,
)
from embedl_deploy._internal.core.patterns.recompositions.registry import (
    RECOMPOSITION_PATTERNS,
)
from embedl_deploy._internal.core.tree.state import ReplaceSession
from embedl_deploy._internal.core.tree.utils import remove_orphaned_modules

_LOG = logging.getLogger(__name__)




@dataclass
class TransformationPlan:
    """Editable transformation plan.

    Returned by :func:`~embedl_deploy.get_transformation_plan`. The ``matches``
    dict maps ``input_node_name → pattern_class_name → PatternMatch``. Toggle
    ``match.apply = False`` to skip specific matches before calling
    :func:`~embedl_deploy.apply_transformation_plan`.
    """

    #: The graph (not yet modified by replacements).
    model: fx.GraphModule
    #: Nested dict of discovered matches, keyed by the last matched node's name
    #: and the pattern class name.
    matches: dict[str, dict[str, PatternMatch]] = field(default_factory=dict)



class TransformationReport(TypedDict):
    """Summary of an ``apply_transformation_plan`` run."""

    applied: list[str]
    skipped: list[str]
    applied_count: int
    skipped_count: int
    total_count: int




@dataclass
class TransformationResult:
    """Result of applying a transformation plan.

    Returned by :func:`~embedl_deploy.apply_transformation_plan`.
    """

    #: The transformed model with fused / quantized modules.
    model: fx.GraphModule
    #: Summary of what was applied and what was skipped.
    report: TransformationReport
    #: The actual
    #: :class:`~embedl_deploy._internal.core.patterns.main.PatternMatch`
    #: objects that were applied.
    matches: list[PatternMatch]



def _matches_to_dict(
    matches: list[PatternMatch],
) -> dict[str, dict[str, PatternMatch]]:
    """Convert a flat match list to the nested-dict representation."""
    result: dict[str, dict[str, PatternMatch]] = {}
    for pm in matches:
        node_name = pm.tree_match.get_tree_nodes()[-1].name
        result.setdefault(node_name, {})
        pattern_name = pm.pattern.__name__
        if pattern_name in result[node_name]:
            raise ValueError(
                f"Duplicate match for ({node_name!r}, {pattern_name!r}); "
                "the nested-dict representation cannot hold both"
            )
        result[node_name][pattern_name] = pm
    return result


def _dict_to_matches(
    match_dict: dict[str, dict[str, PatternMatch]],
) -> list[PatternMatch]:
    """Flatten the nested dict back to a list (preserves insertion order)."""
    return [pm for per_node in match_dict.values() for pm in per_node.values()]


def _build_report(
    enabled: list[PatternMatch],
    skipped: list[PatternMatch],
) -> TransformationReport:
    """Build the summary dict for a TransformationResult."""
    return {
        "applied": [repr(m) for m in enabled],
        "skipped": [repr(m) for m in skipped],
        "applied_count": len(enabled),
        "skipped_count": len(skipped),
        "total_count": len(enabled) + len(skipped),
    }




def get_transformation_plan(
    graph_module: fx.GraphModule,
    patterns: Sequence[type[Pattern]],
) -> TransformationPlan:
    """Find all non-overlapping tree-based pattern matches.

    Each pattern's
    :meth:`~embedl_deploy._internal.core.patterns.main.Pattern.match` must
    return :class:`~embedl_deploy._internal.core.patterns.main.PatternMatch`
    objects with a populated ``tree_match``. Overlapping matches are resolved
    by marking later overlaps as ``apply=False``.

    Returns a :class:`~embedl_deploy.TransformationPlan` that can be inspected
    and edited before calling :func:`~embedl_deploy.apply_transformation_plan`.

    :param graph_module:
        The traced graph module to analyze.
    :param patterns:
        Patterns to search for. Order matters: patterns are matched in
        sequence, and earlier matches claim nodes first. Supply longest (most
        specific) patterns first to ensure they take priority over shorter ones
        when sub-graphs overlap.
    :returns:
        A :class:`~embedl_deploy.TransformationPlan` containing `graph_module`
        and ``matches`` (a nested dict of discovered pattern matches).

    Example:

    .. code-block:: python

        from embedl_deploy import get_transformation_plan, prepare_graph
        from embedl_deploy.backend import get_backend

        graph_module = prepare_graph(model, args)
        plan = get_transformation_plan(
            graph_module, patterns=get_backend().fusion_patterns
        )
        for node, pats in plan.matches.items():
            for name, match in pats.items():
                print(f"{node}: {name} apply={match.apply}")

    """
    # Strip torch.export shape-guard nodes that ShapeProp cannot evaluate.
    guards = [
        n
        for n in graph_module.graph.nodes
        if n.op == "call_module" and n.name.startswith("_guards")
    ]
    for node in guards:
        node.replace_all_uses_with(next(iter(node.args)))
        graph_module.graph.erase_node(node)
    if guards:
        graph_module.recompile()

    pattern_matches: list[PatternMatch] = []
    for pattern in patterns:
        pattern_matches.extend(pattern.match(graph_module))

    consumed: dict[fx.Node, str] = {}
    for pm in pattern_matches:
        name = pm.pattern.__name__
        tree_nodes = pm.tree_match.get_tree_nodes()
        overlap = set(tree_nodes) & consumed.keys()
        if overlap:
            pm.apply = False
            blockers = {consumed[n] for n in overlap}
            _LOG.debug(
                "Skipping %s match at %s: %d node(s) already consumed by %s",
                name,
                tree_nodes[-1].name,
                len(overlap),
                ", ".join(sorted(blockers)),
            )
            continue
        pm.apply = True
        for node in tree_nodes:
            consumed[node] = name

    return TransformationPlan(
        model=graph_module,
        matches=_matches_to_dict(pattern_matches),
    )



def _fake_args(graph_module: fx.GraphModule) -> list[torch.Tensor]:
    """Build fake inputs from placeholder ``tensor_meta``.

    Pins tensors and the graph module to the same device so ``torch.export``'d
    graphs with device-dispatched ops (e.g. SDPA) don't crash on cross-device
    tensors.

    :returns:
        A list of tensors matching the graph's placeholders, or an
        empty list if any placeholder lacks shape metadata.
    """
    try:
        device = next(graph_module.parameters()).device
    except StopIteration:
        device = torch.device("cpu")

    if device.type != "cpu":
        graph_module.to(device)

    args: list[torch.Tensor] = []
    for n in graph_module.graph.nodes:
        if n.op != "placeholder":
            continue
        meta = n.meta.get("tensor_meta")
        if meta is None or not hasattr(meta, "shape"):
            return []
        dtype = getattr(meta, "dtype", torch.float32)
        if dtype.is_floating_point:
            args.append(torch.randn(meta.shape, dtype=dtype, device=device))
        else:
            args.append(torch.zeros(meta.shape, dtype=dtype, device=device))
    return args


def propagate_shapes(
    graph_module: fx.GraphModule,
    args: Sequence[torch.Tensor],
) -> None:
    """Run ``ShapeProp`` in eval mode.

    Temporarily sets the model to eval mode to prevent batch-norm running-stat
    updates, then restores the original training state.

    :param graph_module:
        The graph module whose shapes should be refreshed.
    :param args:
        Example inputs for shape propagation.
    """
    training = graph_module.training
    graph_module.eval()
    with torch.no_grad():
        ShapeProp(graph_module).propagate(*args)  # type: ignore[no-untyped-call]
    graph_module.train(training)




def apply_transformation_plan(
    plan: TransformationPlan,
) -> TransformationResult:
    """Apply the enabled matches from `plan`.

    Only matches with ``apply=True`` are applied via
    :meth:`~embedl_deploy._internal.core.patterns.main.Pattern.replace`. The
    plan's model is modified in place. After replacement, dead code and
    orphaned submodules are removed, the graph is linted and recompiled, and
    shape metadata is re-propagated when available. The model's training state
    is preserved.

    :param plan:
        The plan to apply (from
        :func:`~embedl_deploy.get_transformation_plan`).
    :returns:
        A :class:`~embedl_deploy.TransformationResult` containing ``model``
        (transformed), ``report`` (summary), and ``matches`` (applied matches).
    :raises ValueError:
        If any nodes are included in more than one enabled pattern.

    Example:

    .. code-block:: python

        result = apply_transformation_plan(plan)
        print(result.report)
        torch.onnx.export(result.model, x, "deployed.onnx")

    """
    graph_module = plan.model

    pattern_matches = _dict_to_matches(plan.matches)
    enabled = [pm for pm in pattern_matches if pm.apply]
    skipped = [pm for pm in pattern_matches if not pm.apply]

    consumed: set[fx.Node] = set()
    for pm in enabled:
        tree_nodes = set(pm.tree_match.get_tree_nodes())
        if inter := consumed.intersection(tree_nodes):
            msg = f"Nodes: {inter} included in more than one enabled pattern."
            raise ValueError(msg)
        consumed.update(tree_nodes)

    with ReplaceSession():
        for pm in enabled:
            pm.pattern.replace(pm)

    if enabled:
        graph_module.graph.eliminate_dead_code()
        remove_orphaned_modules(graph_module)
        # Graph.lint lacks type stubs in torch
        graph_module.graph.lint()  # type: ignore[no-untyped-call]
        graph_module.recompile()

    fake = _fake_args(graph_module)
    if fake:
        propagate_shapes(graph_module, fake)

    report = _build_report(enabled, skipped)

    return TransformationResult(
        model=graph_module,
        report=report,
        matches=enabled,
    )



def export_graph_module(
    model: nn.Module,
    args: tuple[Any, ...],
) -> fx.GraphModule:
    """Export `model` and return a plain ``GraphModule``.

    Calls :func:`torch.export.export`, then removes the monkey-patched
    ``train``/``eval`` overrides that ``export`` installs (they raise on call)
    so the returned graph module behaves like a normal
    :class:`~torch.nn.Module`. The model's training state is preserved.

    :param model:
        The module to export.
    :param args:
        Example inputs for :func:`torch.export.export`.
    :returns:
        A :class:`~torch.fx.GraphModule` with standard ``train``/``eval``
        and the same training state as `model`.
    """
    gm = export(model, args).module()
    del gm.train
    del gm.eval
    gm.train(model.training)
    return gm


def export_trace(
    model: nn.Module,
    args: tuple[Any, ...],
) -> fx.GraphModule:
    """Trace `model` via ``torch.export.export`` and recompose.

    Exports the model to an aten-level graph, then applies recomposition
    patterns to reconstruct module-level nodes (conv, linear, etc.).

    :param model:
        The module to trace.
    :param args:
        Example inputs for :func:`torch.export.export`.
    :returns:
        A :class:`~torch.fx.GraphModule` with recomposed module nodes.
    """
    gm = export_graph_module(model, args)
    plan = get_transformation_plan(gm, RECOMPOSITION_PATTERNS)
    return apply_transformation_plan(plan).model




class Trace(enum.Enum):
    """Tracing method used to produce the graph."""

    NONE = "none"
    SYMBOLIC = "symbolic"
    EXPORT = "export"





def prepare_graph(
    model: nn.Module | fx.GraphModule,
    args: tuple[Any, ...],
    *,
    trace: Trace = Trace.EXPORT,
) -> fx.GraphModule:
    """Trace, deep-copy, and propagate shapes.

    Produces a :class:`~torch.fx.GraphModule` ready for pattern matching. The
    original model is deep-copied and never modified.

    :param model:
        The model to prepare. A :class:`~torch.fx.GraphModule` may be
        re-traced (e.g. to apply recomposition via ``Trace.EXPORT``) or
        used as-is with ``Trace.NONE``.
    :param args:
        Example inputs for tracing and shape propagation.
    :param trace:
        Tracing method. ``Trace.EXPORT`` (default) uses
        :func:`torch.export.export` followed by recomposition;
        ``Trace.SYMBOLIC`` uses :func:`~torch.fx.symbolic_trace`;
        ``Trace.NONE`` skips tracing (requires a
        :class:`~torch.fx.GraphModule`).
    :returns:
        A deep-copied :class:`~torch.fx.GraphModule` with shape metadata.
    """
    match trace:
        case Trace.NONE:
            if not isinstance(model, fx.GraphModule):
                msg = (
                    "Trace.NONE requires a GraphModule, "
                    f"got {type(model).__name__}"
                )
                raise TypeError(msg)
            graph_module = model
        case Trace.SYMBOLIC:
            graph_module = symbolic_trace(model)
        case Trace.EXPORT:
            graph_module = export_trace(model, args)
        case _:  # pragma: no cover
            raise ValueError(f"unsupported trace method: {trace!r}")

    graph_module = copy.deepcopy(graph_module)
    propagate_shapes(graph_module, args)
    return graph_module





def transform(
    model: nn.Module | fx.GraphModule,
    args: tuple[Any, ...],
    patterns: Sequence[type[Pattern]] | None = None,
    *,
    trace: Trace = Trace.EXPORT,
) -> TransformationResult:
    """Apply pattern transformations to `model` in one step.

    Conversion patterns are applied iteratively until no new matches are found,
    then fusion patterns are matched and applied in a single pass.
    Recomposition is handled automatically during tracing (see
    :func:`prepare_graph`). The original model is deep-copied and never
    modified.

    :param model:
        The model to transform. A :class:`~torch.fx.GraphModule` may be
        re-traced or used as-is depending on ``trace``.
    :param args:
        Example inputs for tracing.
    :param patterns:
        Patterns to match and apply. When ``None`` (the default), the
        active backend's ``conversion_patterns`` and ``fusion_patterns``
        are used. Order matters: patterns are matched in sequence, and
        earlier matches claim nodes first. Supply longest (most specific)
        patterns first to ensure they take priority over shorter ones
        when sub-graphs overlap.
    :param trace:
        Tracing method. ``Trace.EXPORT`` (default) uses
        :func:`torch.export.export` followed by recomposition;
        ``Trace.SYMBOLIC`` uses :func:`~torch.fx.symbolic_trace`;
        ``Trace.NONE`` skips tracing (requires a
        :class:`~torch.fx.GraphModule`).
    :returns:
        A :class:`~embedl_deploy.TransformationResult` containing ``model``
        (transformed), ``report`` (summary), and ``matches`` (applied matches).
        The report and matches reflect the **fusion pass only** --
        conversion matches become stale after fusion rewrites the graph,
        so only the final pass is reported. Use
        :func:`~embedl_deploy.get_transformation_plan` and
        :func:`~embedl_deploy.apply_transformation_plan` individually
        when per-phase reporting is needed.

    Example:

    .. code-block:: python

        from embedl_deploy import transform

        result = transform(model, args)
        deployable_model = result.model

    """
    if patterns is None:
        backend = get_backend()
        conversions = backend.conversion_patterns
        fusions = backend.fusion_patterns
        skipped: list[type[Pattern]] = []
    else:
        conversions = [p for p in patterns if p.phase == Phase.CONVERSION]
        fusions = [p for p in patterns if p.phase == Phase.FUSION]
        skipped = [
            p
            for p in patterns
            if p.phase not in (Phase.CONVERSION, Phase.FUSION)
        ]
    if skipped:
        names = ", ".join(p.__name__ for p in skipped)
        warnings.warn(
            f"transform() only applies conversion and fusion patterns; "
            f"skipping: {names}. Recomposition is handled automatically "
            f"by prepare_graph().",
            stacklevel=2,
        )

    graph_module = prepare_graph(model, args, trace=trace)

    if conversions:
        while True:
            conv_plan = get_transformation_plan(
                graph_module,
                conversions,
            )
            if not conv_plan.matches:
                break
            conv_result = apply_transformation_plan(conv_plan)
            graph_module = conv_result.model

    fuse_plan = get_transformation_plan(
        graph_module,
        fusions,
    )
    fuse_result = apply_transformation_plan(fuse_plan)

    return fuse_result