export_executorch.py

# SPDX-FileCopyrightText: Copyright 2026 Arm Limited and/or its affiliates <open-source-office@arm.com>
# SPDX-License-Identifier: Apache-2.0

import os
from pathlib import Path
from typing import Sequence

import torch
from arm_backend_monkey_patch import apply_arm_backend_monkey_patch
from executorch.backends.arm.quantizer import (
    VgfQuantizer,
    get_symmetric_quantization_config,
)
from executorch.backends.arm.quantizer.arm_quantizer import (
    QuantizationSpec,
    annotate_input_qspec_map,
    annotate_output_qspec,
    is_annotated,
    mark_node_as_annotated,
)
from executorch.backends.arm.common.annotation_meta import ArmAnnotationInfo
from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner
from executorch.exir.backend.partitioner import Partitioner, PartitionResult
from executorch.exir.capture._config import EdgeCompileConfig
from executorch.exir.passes.quantize_io_pass import extract_io_quant_params
from executorch.exir.dialects._ops import ops as exir_ops
from torch.fx.passes.infra.pass_base import PassResult
from torchao.quantization.pt2e.observer import FixedQParamsObserver
from torchao.quantization.pt2e.quantizer.quantizer import QuantizationAnnotation

from executorch_grid_sampler_backend import (
    GridSamplerOnlyPartitioner,
    pad_rgb_grid_sampler_channels_pass,
)
from utils import _emit_export_reports, _enable_vgf_debug_logging


apply_arm_backend_monkey_patch()


def _is_q_op(node) -> bool:
    target = str(node.target)
    return (
        "quantized_decomposed.quantize_per_tensor" in target
        or "quantized_decomposed.quantize_per_channel" in target
    )


def _is_dq_op(node) -> bool:
    target = str(node.target)
    return (
        "quantized_decomposed.dequantize_per_tensor" in target
        or "quantized_decomposed.dequantize_per_channel" in target
    )


class VgfPartitionerWrapper(Partitioner):
    """
    Wrap VGF partitioning to keep only true graph-IO Q/DQ at top level.

    Why this exists:
    - In INT+FP mode, Arm's TOSA partitioner does not automatically detag boundary
      Q/DQ nodes. That can cause VGF to swallow graph-boundary Q/DQ, and then
      `extract_io_quant_params` cannot match/strip them later.
    - The custom grid-sampler backend also relies on keeping the dequant before
      `grid_sampler` and the quant after it visible at top level. That backend
      wants to delegate and replace the whole dequant -> grid_sampler -> quant
      sequence as one quantized grid-sampler operation.

    How this differs from regular `VgfPartitioner`:
    - Regular VGF partitioning keeps whatever tags the TOSA partitioner emits.
    - This wrapper post-processes those tags and detags only:
      1) quantize nodes that directly consume placeholders
      2) dequantize nodes that directly feed the graph output
      3) quantize/dequantize nodes adjacent to `grid_sampler`
    - Internal quant/dequant inside delegated partitions are intentionally kept.
    """

    def __init__(
        self,
        inner_partitioner,
        *,
        quantized_input_idxs: Sequence[int] = (),
        quantized_output_idxs: Sequence[int] = (),
    ) -> None:
        super().__init__()
        self.inner = inner_partitioner
        self._quantized_input_idxs = frozenset(quantized_input_idxs)
        self._quantized_output_idxs = frozenset(quantized_output_idxs)

    def ops_to_not_decompose(self, ep):
        return self.inner.ops_to_not_decompose(ep)

    def partition(self, exported_program) -> PartitionResult:
        result = self.inner.partition(exported_program)
        gm = result.tagged_exported_program.graph_module
        user_input_names = list(result.tagged_exported_program.graph_signature.user_inputs)
        quantized_input_names = {
            user_input_names[idx]
            for idx in self._quantized_input_idxs
            if idx < len(user_input_names)
        }
        output_node = next((node for node in gm.graph.nodes if node.op == "output"), None)
        output_parent_to_idx = {}
        if output_node is not None:
            output_parent_to_idx = {
                parent: idx for idx, parent in enumerate(output_node.all_input_nodes)
            }

        for node in gm.graph.nodes:
            tag = node.meta.get("delegation_tag")
            if tag is None:
                continue

            if _is_q_op(node):
                is_graph_input_q = any(
                    inp.op == "placeholder"
                    and inp.name in quantized_input_names
                    for inp in node.all_input_nodes
                )
                feeds_grid_sampler = any(
                    "aten.grid_sampler_2d.default" in str(inp.target)
                    for inp in node.all_input_nodes
                )
                if is_graph_input_q or feeds_grid_sampler:
                    del node.meta["delegation_tag"]
            elif _is_dq_op(node):
                is_graph_output_dq = any(
                    user.op == "output"
                    and output_parent_to_idx.get(node) in self._quantized_output_idxs
                    for user in node.users
                )
                feeds_grid_sampler = any(
                    "aten.grid_sampler_2d.default" in str(user.target)
                    for user in node.users
                )
                if is_graph_output_dq or feeds_grid_sampler:
                    del node.meta["delegation_tag"]
        active_tags = {
            node.meta["delegation_tag"]
            for node in gm.graph.nodes
            if "delegation_tag" in node.meta
        }
        result.partition_tags = {
            tag: spec
            for tag, spec in result.partition_tags.items()
            if tag in active_tags
        }
        return result


def _collect_grid_position_subgraph_nodes(graph_module) -> list[object]:
    # Exclude only the grid-position subgraph that is used exclusively by the
    # sampling-grid path. If a producer also feeds non-grid computation, keep it
    # annotated so shared activations do not become a mixed plain/DQ boundary
    # later during Arm lowering.
    seed_nodes = []
    for node in graph_module.graph.nodes:
        target = str(node.target)
        if node.op != "call_function" or (
            "grid_sampler.default" not in target
            and "grid_sampler_2d.default" not in target
        ):
            continue
        if len(node.args) < 2:
            continue
        seed_nodes.append(node.args[1])

    collected = {node for node in seed_nodes if hasattr(node, "op")}
    changed = True

    while changed:
        changed = False
        for node in reversed(tuple(graph_module.graph.nodes)):
            if node in collected:
                continue
            if not hasattr(node, "users"):
                continue

            users = [user for user in node.users if hasattr(user, "op")]
            if not users:
                continue

            if all(user in collected for user in users):
                collected.add(node)
                changed = True

    return list(collected)


def _feeds_grid_sampler_grid_input(node) -> bool:
    # The grid placeholder may reach grid_sampler through inserted Q/DQ nodes,
    # so a direct-user check is not enough here. We use this to keep sampling
    # grids out of the channels-last image rewrite even after quantization.
    visited = set()

    def visit(current) -> bool:
        if current in visited or not hasattr(current, "users"):
            return False
        visited.add(current)

        for user in current.users:
            if (
                user.op == "call_function"
                and user.target == exir_ops.edge.aten.grid_sampler_2d.default
                and len(user.args) > 1
                and user.args[1] is current
            ):
                return True

            if user.op != "call_function":
                continue

            if _is_q_op(user) or _is_dq_op(user):
                if visit(user):
                    return True

        return False

    return visit(node)


class VgfQuantizerWrapper(VgfQuantizer):
    """
    VGF quantizer wrapper that customizes graph IO and `grid_sampler` quantization.

    It applies caller-provided qspecs to graph inputs/outputs positionally, and
    also rewrites `grid_sampler` annotations so callers can separately control
    the image input, grid input, and output qspecs at the grid-sampler boundary.
    That keeps the quantized image path aligned with the separately lowered
    grid-sampler backend while still allowing a different quantization scheme
    for the sampling grid itself.
    """

    def __init__(
        self,
        compile_spec,
        *,
        input_qspecs: Sequence[QuantizationSpec | None],
        output_qspecs: Sequence[QuantizationSpec | None],
        grid_sampler_image_input_qspec: QuantizationSpec | None = None,
        grid_sampler_grid_input_qspec: QuantizationSpec | None = None,
        grid_sampler_output_qspec: QuantizationSpec | None = None,
    ) -> None:
        super().__init__(compile_spec)
        self._input_qspecs = tuple(input_qspecs)
        self._output_qspecs = tuple(output_qspecs)
        self._grid_sampler_image_input_qspec = grid_sampler_image_input_qspec
        self._grid_sampler_grid_input_qspec = grid_sampler_grid_input_qspec
        self._grid_sampler_output_qspec = grid_sampler_output_qspec

    def _clear_quantization_annotation(self, node) -> None:
        if not hasattr(node, "meta"):
            return
        node.meta.pop("quantization_annotation", None)
        node.meta.pop("is_annotated", None)
        custom_meta = node.meta.get("custom")
        if isinstance(custom_meta, dict):
            custom_meta.pop(ArmAnnotationInfo.CUSTOM_META_KEY, None)
            if len(custom_meta) == 0:
                node.meta.pop("custom", None)

    def _exclude_grid_position_subgraphs_from_int8(self, model) -> None:
        for node in _collect_grid_position_subgraph_nodes(model):
            self._clear_quantization_annotation(node)

    def _annotate_io(self, model, quantization_config):
        # Overrides the base quantizer's IO annotation hook so callers can
        # explicitly choose the boundary qspec for each graph input/output.
        del quantization_config
        placeholder_nodes = [
            node
            for node in model.graph.nodes
            if node.op == "placeholder" and len(node.users) > 0
        ]
        for idx, node in enumerate(placeholder_nodes):
            if idx >= len(self._input_qspecs):
                continue
            if is_annotated(node):
                continue
            input_qspec = self._input_qspecs[idx]
            if input_qspec is None:
                continue
            annotate_output_qspec(node, input_qspec)
            mark_node_as_annotated(node)

        # FX/export graphs normally have one `output` node whose inputs are the
        # actual graph outputs, so iterate that node's parents rather than
        # expecting one node per returned value.
        output_node = next(
            (node for node in model.graph.nodes if node.op == "output"),
            None,
        )
        if output_node is None or is_annotated(output_node):
            return

        did_annotate_any_output = False
        for idx, parent in enumerate(output_node.all_input_nodes):
            if idx >= len(self._output_qspecs):
                continue
            output_qspec = self._output_qspecs[idx]
            if output_qspec is None:
                continue
            annotate_input_qspec_map(output_node, parent, output_qspec)
            did_annotate_any_output = True

        if did_annotate_any_output:
            mark_node_as_annotated(output_node)

    def _annotate_grid_sample(self, model) -> None:
        for node in model.graph.nodes:
            if node.op != "call_function" or node.target not in {
                torch.ops.aten.grid_sampler.default,
                torch.ops.aten.grid_sampler_2d.default,
            }:
                continue

            quant_annotation = node.meta.get("quantization_annotation")
            if quant_annotation is None:
                quant_annotation = QuantizationAnnotation()
            if quant_annotation.input_qspec_map is None:
                quant_annotation.input_qspec_map = {}

            image_input_node = node.args[0]
            grid_input_node = node.args[1]
            input_qspec_map = dict(quant_annotation.input_qspec_map)

            if self._grid_sampler_image_input_qspec is not None:
                input_qspec_map[image_input_node] = self._grid_sampler_image_input_qspec

            if self._grid_sampler_grid_input_qspec is not None:
                input_qspec_map[grid_input_node] = self._grid_sampler_grid_input_qspec
            else:
                input_qspec_map.pop(grid_input_node, None)

            quant_annotation.input_qspec_map = input_qspec_map
            quant_annotation.output_qspec = self._grid_sampler_output_qspec
            node.meta["quantization_annotation"] = quant_annotation

    def _annotate_for_static_quantization_config(self, model):
        # This override is reached via the base quantizer's annotate() method
        # during PT2E prepare. Reuse the base static-annotation flow, then patch
        # up grid_sampler annotations for the custom backend boundary behavior.
        model = super()._annotate_for_static_quantization_config(model)
        # Keep the grid-position subgraph out of the default PT2E activation
        # annotation. The grid path may use a different quantization policy from
        # the image path, and we also want to avoid leaving a stray top-level
        # `quantize_per_tensor` after lowering.
        self._exclude_grid_position_subgraphs_from_int8(model)
        if (
            self._grid_sampler_image_input_qspec is not None
            or self._grid_sampler_grid_input_qspec is not None
            or self._grid_sampler_output_qspec is not None
        ):
            self._annotate_grid_sample(model)
        return model


def _set_tensors_to_channels_last_pass(graph_module):
    # Convert the input image tensor to channels-last memory format.
    # This is so that the VGF backend won't need to insert a transpose and we can alias inputs as images in scenario runner.
    for node in graph_module.graph.nodes:
        if node.op == "placeholder":
            if node.meta["val"].ndim == 4:
                # The second input to grid_sampler is the sampling grid, not an image.
                # Leaving that tensor in its native layout preserves the logical
                # [N, H, W, 2] coordinate order when we later serialize scenario
                # tensor dims and .npy contents from FX metadata.
                feeds_grid_sampler_grid_input = _feeds_grid_sampler_grid_input(node)
                if feeds_grid_sampler_grid_input:
                    continue

                node.meta["val"] = node.meta["val"].contiguous(
                    memory_format=torch.channels_last
                )

                # If the input is immediately quantised, change that node's output to channels last too, as the input to the VGF backend will be this quantised tensor
                for user in node.users:
                    if (
                        user.op == "call_function"
                        and user.target
                        == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
                    ):
                        user.meta["val"] = user.meta["val"].contiguous(
                            memory_format=torch.channels_last
                        )

    # Same for the output image
    for node in graph_module.graph.nodes:
        if node.op == "output":
            for output in node.args[0]:
                if output.meta["val"].ndim != 4:
                    continue

                output.meta["val"] = output.meta["val"].contiguous(
                    memory_format=torch.channels_last
                )

                # If the preceding node is a dequantize, change its input to channels last too, as the output of the VGF backend will be this quantised tensor
                if (
                    output.op == "call_function"
                    and output.target
                    == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
                ):
                    input = output.args[0]
                    input.meta["val"] = input.meta["val"].contiguous(
                        memory_format=torch.channels_last
                    )

    # Same for the inputs and outputs of grid sampler nodes, to avoid transposes on the VGF side of the boundary and to
    # allow the grid sampler backend to delegate them and alias as images
    for node in graph_module.graph.nodes:
        if (
            node.op == "call_function"
            and node.target == exir_ops.edge.aten.grid_sampler_2d.default
        ):
            # Input image
            input_image = node.args[0]
            input_image.meta["val"] = input_image.meta["val"].contiguous(
                memory_format=torch.channels_last
            )

            # If input image is quantised, the input to that dequantize node needs to be marked as channels_last too, since the grid sampler backend will want to delegate the dequant -> grid_sampler -> quant sequence as one unit and alias the dequant input as the image
            if (
                input_image.op == "call_function"
                and input_image.target
                == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default
            ):
                input = input_image.args[0]
                input.meta["val"] = input.meta["val"].contiguous(
                    memory_format=torch.channels_last
                )

            # Output image
            node.meta["val"] = node.meta["val"].contiguous(
                memory_format=torch.channels_last
            )

            # If the output image is quantised, the output of the quantize node after grid sampler also needs to be marked as channels_last, since the grid sampler backend will want to delegate the grid_sampler -> quant sequence as one unit and alias the quant output as the image
            for user in node.users:
                if (
                    user.op == "call_function"
                    and user.target
                    == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default
                ):
                    user.meta["val"] = user.meta["val"].contiguous(
                        memory_format=torch.channels_last
                    )

    return PassResult(graph_module=graph_module, modified=True)


def export_model(
    model: torch.nn.Module,
    model_inputs,
    *,
    artifacts_root: str | Path = Path("artifacts"),
    intermediate_artifacts_root: str | Path | None = None,
    quantized_input_qspecs: Sequence[QuantizationSpec | None] = (),
    quantized_output_qspecs: Sequence[QuantizationSpec | None] = (),
    quantize_grid_sample_output: bool = True,
    enable_pt2e_quantization: bool = True,
) -> tuple[object, tuple[torch.Tensor, ...], dict]:
    verbose = os.environ.get("EXPORT_EXECUTORCH_VERBOSE", "0") == "1"
    _enable_vgf_debug_logging(verbose)
    example_inputs = tuple(model_inputs)
    if not example_inputs:
        raise ValueError("export_model requires at least one model input tensor.")

    quantized_input_qspecs = tuple(quantized_input_qspecs)
    quantized_output_qspecs = tuple(quantized_output_qspecs)
    if enable_pt2e_quantization and len(quantized_input_qspecs) != len(example_inputs):
        raise ValueError(
            "quantized_input_qspecs must have one entry per model input when PT2E quantization is enabled, "
            f"got {len(quantized_input_qspecs)} qspecs for {len(example_inputs)} inputs"
        )

    artifacts_root = Path(artifacts_root)
    artifacts_root.mkdir(parents=True, exist_ok=True)
    if intermediate_artifacts_root is None:
        intermediate_artifacts_root = artifacts_root
    intermediate_artifacts_root = Path(intermediate_artifacts_root)
    intermediate_artifacts_root.mkdir(parents=True, exist_ok=True)

    # Need FP for calculations of grid sample positions and int16 for the quantisation to 16-bit SNORM.
    vgf_compile_spec = VgfCompileSpec(
        "TOSA-1.0+INT+FP+int16"
    ).dump_intermediate_artifacts_to(
        str(intermediate_artifacts_root / "vgf_intermediates")
    )

    exported_program = torch.export.export(model, example_inputs, strict=True)

    captured_graph = exported_program.module()
    output_node = next(
        node for node in captured_graph.graph.nodes if node.op == "output"
    )
    output_count = len(output_node.all_input_nodes)
    if enable_pt2e_quantization and len(quantized_output_qspecs) != output_count:
        raise ValueError(
            "quantized_output_qspecs must have one entry per model output when PT2E quantization is enabled, "
            f"got {len(quantized_output_qspecs)} qspecs for {output_count} outputs"
        )
    quantized_input_idxs = [
        idx for idx, qspec in enumerate(quantized_input_qspecs) if qspec is not None
    ]
    quantized_output_idxs = [
        idx for idx, qspec in enumerate(quantized_output_qspecs) if qspec is not None
    ]

    # Pad RGB inputs to grid sampler to RGBA so we can alias as 4-channel images
    pad_rgb_grid_sampler_channels_pass(captured_graph)

    # PT2E quantization with Arm VGF quantizer (8-bit symmetric)
    if enable_pt2e_quantization:
        grid_position_qspec = QuantizationSpec(
            dtype=torch.int16,
            observer_or_fake_quant_ctr=FixedQParamsObserver.with_args(
                scale=1.0 / 32767.0,
                zero_point=0,
                dtype=torch.int16,
                qscheme=torch.per_tensor_symmetric,
                quant_min=-32767,
                quant_max=32767,
            ),
            quant_min=-32767,
            quant_max=32767,
            qscheme=torch.per_tensor_symmetric,
            is_dynamic=False,
        )
        internal_image_qspec = QuantizationSpec(
            dtype=torch.int8,
            observer_or_fake_quant_ctr=FixedQParamsObserver.with_args(
                scale=1.0 / 127.0,
                zero_point=0,
                dtype=torch.qint8,
                qscheme=torch.per_tensor_symmetric,
                quant_min=-127,
                quant_max=127,
            ),
            quant_min=-127,
            quant_max=127,
            qscheme=torch.per_tensor_symmetric,
            is_dynamic=False,
        )

        quantizer = VgfQuantizerWrapper(
            vgf_compile_spec,
            input_qspecs=quantized_input_qspecs,
            output_qspecs=quantized_output_qspecs,
            grid_sampler_image_input_qspec=internal_image_qspec,
            grid_sampler_grid_input_qspec=grid_position_qspec,
            grid_sampler_output_qspec=(
                internal_image_qspec if quantize_grid_sample_output else None
            ),
        )
        symmetric_int8_config = get_symmetric_quantization_config(
            is_per_channel=True,
            is_qat=False,
            is_dynamic=False,
            act_qmin=-127,
            act_qmax=127,
            weight_qmin=-127,
            weight_qmax=127,
        )
        # `set_io(...)` enables the base IO-annotation path. This wrapper then
        # replaces that behavior in `_annotate_io(...)` using the caller-provided
        # per-input and per-output qspecs, but the base quantizer still expects a
        # real IO config here and may use it as a fallback/default.
        quantizer.set_global(symmetric_int8_config).set_io(symmetric_int8_config)
        quantized_graph = quantizer.quantize_with_submodules(
            captured_graph,
            calibration_samples=[example_inputs],
            is_qat=False,
        )
        exported = torch.export.export(quantized_graph, example_inputs)
    else:
        exported = torch.export.export(captured_graph, example_inputs, strict=True)

    from executorch.exir import to_edge_transform_and_lower

    vgf_partitioner = VgfPartitionerWrapper(
        VgfPartitioner(vgf_compile_spec),
        quantized_input_idxs=quantized_input_idxs,
        quantized_output_idxs=quantized_output_idxs,
    )
    grid_sampler_partitioner = GridSamplerOnlyPartitioner()
    edge_program = to_edge_transform_and_lower(
        exported,
        transform_passes=[
            _set_tensors_to_channels_last_pass,  # Convert image tensors to channels-last after any RGB->RGBA padding.
        ],
        partitioner=[vgf_partitioner],
        compile_config=EdgeCompileConfig(_check_ir_validity=False),
    )
    # IMPORTANT TO DO SEPARATE LOWERING HERE, as the multi-partiioner logic is busted with op decomposition!
    # It will decompose all the ops as our custom backend doesn't specify ops to not decompose, and then VGF backend doesn'#t support them (e.g. affine_grid -> linspace -> int64 ops)
    edge_program = edge_program.to_backend(grid_sampler_partitioner)
    if enable_pt2e_quantization:
        io_quant_params = extract_io_quant_params(
            edge_program,
            input_idxs=quantized_input_idxs,
            output_idxs=quantized_output_idxs,
        )
    else:
        io_quant_params = {"inputs": {}, "outputs": {}}
    _emit_export_reports(edge_program, verbose, io_quant_params)
    return edge_program, example_inputs, io_quant_params