From 5409a0791ce8aa9649faf837948ba137ef5d9f02 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Mon, 27 Jul 2026 09:39:43 -0700 Subject: [PATCH] Add generic fx-graph serializer (NPTG) (#21032) Summary: Serialize a torch.fx graph to a standalone flatbuffer (native_graph.fbs, id "NFXG") and back: a topological node list with a tagged Argument union, tensor metadata side table, and by-reference constants (ConstantRef.fqn) whose bytes are shipped separately. Covers input/output classification, symbolic dims in list args, and a validate_graph() completeness check. Constants are never inlined. Differential Revision: D111753288 --- backends/native/BUCK | 31 + backends/native/serialization/__init__.py | 27 + .../native/serialization/graph_serialize.py | 1339 +++++++++++++++++ .../native/serialization/native_graph.fbs | 430 ++++++ backends/native/serialization/schema.py | 343 +++++ backends/native/test/BUCK | 15 + backends/native/test/__init__.py | 0 backends/native/test/test_serialize.py | 1022 +++++++++++++ pytest.ini | 1 + 9 files changed, 3208 insertions(+) create mode 100644 backends/native/BUCK create mode 100644 backends/native/serialization/__init__.py create mode 100644 backends/native/serialization/graph_serialize.py create mode 100644 backends/native/serialization/native_graph.fbs create mode 100644 backends/native/serialization/schema.py create mode 100644 backends/native/test/BUCK create mode 100644 backends/native/test/__init__.py create mode 100644 backends/native/test/test_serialize.py diff --git a/backends/native/BUCK b/backends/native/BUCK new file mode 100644 index 00000000000..98f1c3d5f54 --- /dev/null +++ b/backends/native/BUCK @@ -0,0 +1,31 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +# Export the generic fx-graph schema so it can be shipped as a package resource. +runtime.export_file( + name = "native_graph.fbs", + src = "serialization/native_graph.fbs", + visibility = ["//executorch/backends/native/..."], +) + +fbcode_target( + _kind = runtime.python_library, + name = "lib", + typing = True, + srcs = [ + "serialization/__init__.py", + "serialization/graph_serialize.py", + "serialization/schema.py", + ], + resources = { + ":native_graph.fbs": "serialization/native_graph.fbs", + }, + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/exir/_serialize:lib", + ], +) diff --git a/backends/native/serialization/__init__.py b/backends/native/serialization/__init__.py new file mode 100644 index 00000000000..049fa14e097 --- /dev/null +++ b/backends/native/serialization/__init__.py @@ -0,0 +1,27 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +from executorch.backends.native.serialization.graph_serialize import ( + deserialize_graph, + deserialize_program, + serialize_graph, + serialize_program, + validate_graph, + validate_method, + validate_program, +) + +__all__ = [ + "serialize_graph", + "serialize_program", + "deserialize_graph", + "deserialize_program", + "validate_graph", + "validate_method", + "validate_program", +] diff --git a/backends/native/serialization/graph_serialize.py b/backends/native/serialization/graph_serialize.py new file mode 100644 index 00000000000..7fa930d1299 --- /dev/null +++ b/backends/native/serialization/graph_serialize.py @@ -0,0 +1,1339 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +""" +Serialize a torch.fx graph into the native backend's generic flatbuffer format +(native_graph.fbs) and back. + +The format is a topological list of fx nodes: each node carries an op-kind, a +target op-name string, and generic arguments encoded through a tagged Argument +union. Values are referenced by fx SSA name; tensor metadata lives in a side +table. Constant tensor data is not embedded here; it is returned separately so the +caller can ship it separately in an external constant file, keyed by data_key. + +Uses the "runtime flatc" pattern: the schema and the flatc binary are shipped as +package resources, and flatc is invoked to convert between JSON and binary. +""" + +import importlib.resources +import json +import operator +import os +import tempfile +from dataclasses import fields, is_dataclass +from enum import IntEnum +from typing import Any, cast, get_args, get_origin, get_type_hints, Union + +import torch + +from executorch.backends.native.serialization.schema import ( + Argument, + ArgumentValue, + BoolArg, + BoolListArg, + Dim, + FloatArg, + FloatListArg, + Graph, + GraphArg, + InputKind as SchemaInputKind, + IntArg, + IntListArg, + Method, + MutableBufferSpec, + NamedArgument, + NamedTensorRef, + Node, + NoneArg, + OpKind, + OptionalTensorListArg, + Output, + OutputKind, + OutputSpec, + OutputValueKind, + Program, + ScalarType, + ScalarTypeArg, + StringArg, + TensorArg, + TensorListArg, + TensorMeta, + TensorValue, +) + +from executorch.exir._serialize._dataclass import _json_to_dataclass +from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile +from executorch.exir.tensor import dim_order_from_stride, stride_from_dim_order + +from torch.export.graph_signature import InputKind, TensorArgument +from torch.fx.experimental.symbolic_shapes import statically_known_true + +SCHEMA_VERSION = "1" +_SCHEMA_RESOURCE = "native_graph.fbs" +_FILE_STEM = "native_graph" + +_DTYPE_TO_SCALAR_TYPE: dict[torch.dtype, ScalarType] = { + torch.uint8: ScalarType.BYTE, + torch.int8: ScalarType.CHAR, + torch.int16: ScalarType.SHORT, + torch.int32: ScalarType.INT, + torch.int64: ScalarType.LONG, + torch.float16: ScalarType.HALF, + torch.float32: ScalarType.FLOAT, + torch.float64: ScalarType.DOUBLE, + torch.bool: ScalarType.BOOL, + torch.bfloat16: ScalarType.BFLOAT16, +} +# Optional dtypes not present in every torch build. +for _name, _st in ( + ("uint16", ScalarType.UINT16), + ("uint32", ScalarType.UINT32), + ("uint64", ScalarType.UINT64), +): + _dt = getattr(torch, _name, None) + if _dt is not None: + _DTYPE_TO_SCALAR_TYPE[_dt] = _st + + +# --------------------------------------------------------------------------- +# Operator name (target) serialization, mirroring torch._export.serde. +# --------------------------------------------------------------------------- + + +def _resolve_op_overload(target: object) -> "torch._ops.OpOverload | None": + """Return the aten OpOverload for an fx call_function target, or None. + + Edge-dialect ops (EdgeOpOverload) wrap the underlying aten OpOverload in _op, so + prefer that. A plain OpOverload also exposes _op, but there it is the C++ builtin + (empty __name__); unwrapping it loses the op name, so only unwrap when _op is + itself an OpOverload, otherwise use the target directly. Non-OpOverload callables + (sym builtins, operator.*, higher-order ops) return None. + """ + inner = getattr(target, "_op", None) + if isinstance(inner, torch._ops.OpOverload): + return inner + if isinstance(target, torch._ops.OpOverload): + return target + return None + + +def serialize_operator(target: object) -> str: + if isinstance(target, str): + return target + op = _resolve_op_overload(target) + if op is not None: + module = op.__module__.replace("torch._ops", "torch.ops") + return f"{module}.{op.__name__}" + # Fallback for non-OpOverload callables (e.g. operator.getitem, sym ops). + module = getattr(target, "__module__", "") or "" + name = getattr(target, "__name__", "") or str(target) + module = module.replace("torch._ops", "torch.ops") + return f"{module}.{name}" if module else name + + +# --------------------------------------------------------------------------- +# Metadata helpers. +# --------------------------------------------------------------------------- + + +def _scalar_type(dtype: torch.dtype) -> ScalarType: + st = _DTYPE_TO_SCALAR_TYPE.get(dtype) + if st is None: + raise ValueError(f"Unsupported dtype for native serialization: {dtype}") + return st + + +def _dim(x: object) -> Dim: + if isinstance(x, int): + return Dim(min=x, max=x) + if isinstance(x, torch.SymInt): + # maybe_as_int() returns an int only when the SymInt is genuinely constant; it + # does not specialize a symbolic dim (int(x) would guard to the hint and freeze + # a dynamic shape to its example value). + concrete = x.node.maybe_as_int() + if concrete is not None: + return Dim(min=concrete, max=concrete) + # Symbolic dim: record its value range from the ShapeEnv rather than a symbol + # name (the runtime uses concrete shapes; symbols would need a C++ symbolic + # engine). Unbounded above is max = -1. + lower, upper = 0, -1 + try: + vr = x.node.shape_env.bound_sympy(x.node.expr) + if vr.lower.is_finite: + lower = int(vr.lower) + if vr.upper.is_finite: + upper = int(vr.upper) + except AttributeError: + # No shape_env/expr/bound available; leave unbounded. + pass + except (ValueError, TypeError, RuntimeError): + # Sympy bound computation failed; leave unbounded. + pass + return Dim(min=lower, max=upper) + return Dim(min=int(x), max=int(x)) + + +def _dim_order(t: torch.Tensor) -> list[int]: + """Return axes sorted outer-to-inner by stride for a dim-order expressible tensor. + + Only dim_order is serialized, not raw strides, so the strides must be recoverable + from sizes and dim_order. Raise if they are not (for example sliced, as_strided, + broadcast, or overlapping layouts) instead of silently losing the layout. Symbolic + strides are supported when they match the strides implied by the (possibly + symbolic) sizes. + """ + ndim = t.dim() + if ndim == 0: + return [] + strides = tuple(t.stride()) + sizes = list(t.shape) + # dim_order_from_stride handles symbolic strides and rejects stride-0 layouts. + dim_order = [int(d) for d in dim_order_from_stride(strides)] + expected = stride_from_dim_order(sizes, dim_order) + for i in range(ndim): + # A size-1 dim only ever indexes 0, so its stride is arbitrary and need not + # match the reconstruction (e.g. channels-last with a size-1 channel). + if statically_known_true(sizes[i] == 1): + continue + # statically_known_true proves equality without adding shape guards. + if not statically_known_true(strides[i] == expected[i]): + raise ValueError( + f"tensor is not dim-order expressible: strides {strides} with sizes " + f"{tuple(sizes)} imply dim_order {tuple(dim_order)}, which reconstructs " + f"strides {tuple(expected)}. Only dim_order is serialized, so this " + f"layout (e.g. sliced/as_strided or overlapping) would be lost." + ) + return dim_order + + +def _tensor_meta(t: torch.Tensor) -> TensorMeta: + return TensorMeta( + dtype=_scalar_type(t.dtype), + sizes=[_dim(s) for s in t.shape], + dim_order=_dim_order(t), + ) + + +# --------------------------------------------------------------------------- +# Argument dispatch. +# --------------------------------------------------------------------------- + + +def _int_val_node(x: object) -> bool: + """True if x is an fx node whose produced value is an int / SymInt (e.g. a + sym_size or shape-arithmetic node). Such a node referenced as an arg is a + dynamic int (an IntArg carrying a ref), not a tensor.""" + if not isinstance(x, torch.fx.Node): + return False + val = x.meta.get("val") + return isinstance(val, (int, torch.SymInt)) and not isinstance(val, bool) + + +def _is_tensor_node(x: object) -> bool: + return isinstance(x, torch.fx.Node) and not _int_val_node(x) + + +def _is_optional_tensor_node(x: object) -> bool: + return x is None or _is_tensor_node(x) + + +def _is_dynamic_int_element(x: object) -> bool: + return ( + (isinstance(x, int) and not isinstance(x, bool)) + or isinstance(x, torch.SymInt) + or _int_val_node(x) + ) + + +def _check_no_subgraph_in_list( + items: list[object], subgraph_map: dict[str, torch.fx.GraphModule] | None +) -> None: + if subgraph_map is None: + return + for x in items: + if isinstance(x, torch.fx.Node) and x.name in subgraph_map: + raise ValueError( + f"higher-order-op subgraph {x.name!r} appears inside a list " + f"argument; only a direct subgraph arg can be inlined as a " + f"GraphArg. This graph shape is not supported." + ) + + +def _empty_typed_list_arg( + schema_type_hint: str | None, items: list[object] +) -> ArgumentValue: + hint = (schema_type_hint or "").lower() + if "tensor?" in hint: + return OptionalTensorListArg(names=[], has_value=[]) + if "tensor" in hint: + return TensorListArg(names=[]) + if "bool" in hint: + return BoolListArg(values=[]) + if "float" in hint or "double" in hint: + return FloatListArg(values=[]) + if "int" in hint or hint == "": + if schema_type_hint is not None: + return IntListArg(values=[]) + raise ValueError( + f"Cannot serialize empty list argument without explicit element type: " + f"empty list is ambiguous (e.g. Tensor[] vs int[]). " + f"Schema hint: {schema_type_hint!r}, items: {items!r}. " + f"If this is a schema default, ensure the op schema provides a typed " + f"default or handle this case in _named_arguments." + ) + + +def _dynamic_int_list_arg(items: list[object]) -> IntListArg: + values: list[int] = [] + refs: list[str] = [] + for x in items: + if isinstance(x, torch.fx.Node): + values.append(0) + refs.append(x.name) + elif isinstance(x, torch.SymInt): + concrete = x.node.maybe_as_int() + if concrete is None: + raise ValueError( + f"symbolic SymInt list element {x!r} is not backed by an fx " + f"node; dynamic int args must reference an in-graph value." + ) + values.append(concrete) + refs.append("") + else: + values.append(cast(int, x)) + refs.append("") + return IntListArg(values=values, refs=refs) + + +def _to_list_arg( + items: list[object], + subgraph_map: "dict[str, torch.fx.GraphModule] | None" = None, + schema_type_hint: str | None = None, +) -> ArgumentValue: + if len(items) == 0: + return _empty_typed_list_arg(schema_type_hint, items) + _check_no_subgraph_in_list(items, subgraph_map) + if all(_is_tensor_node(x) for x in items): + return TensorListArg(names=[cast(torch.fx.Node, x).name for x in items]) + if all(_is_optional_tensor_node(x) for x in items): + return OptionalTensorListArg( + names=[cast(torch.fx.Node, x).name if x is not None else "" for x in items], + has_value=[x is not None for x in items], + ) + if all(isinstance(x, bool) for x in items): + return BoolListArg(values=[cast(bool, x) for x in items]) + if all(isinstance(x, int) and not isinstance(x, bool) for x in items): + return IntListArg(values=[cast(int, x) for x in items]) + if all(isinstance(x, (int, float)) and not isinstance(x, bool) for x in items): + return FloatListArg(values=[float(cast(float, x)) for x in items]) + if all(_is_dynamic_int_element(x) for x in items): + return _dynamic_int_list_arg(items) + raise ValueError( + f"Cannot serialize list argument with element types " + f"{sorted({type(x).__name__ for x in items})}: {items!r}" + ) + + +def _node_to_arg_value( + v: torch.fx.Node, + subgraph_map: "dict[str, torch.fx.GraphModule] | None", +) -> ArgumentValue: + # A get_attr node that resolves to a submodule GraphModule is a + # higher-order-op subgraph (torch.cond branch, map body, ...). Inline it + # as a GraphArg rather than referencing it like a tensor value. + if subgraph_map is not None and v.name in subgraph_map: + sub = _build_graph_body(subgraph_map[v.name], None, {}, None)[0] + return GraphArg(name=v.name, graph=sub) + # An int-producing node (sym_size / arith on shapes) referenced as an arg is a + # dynamic int, not a tensor; serialize it as an IntArg carrying a ref. + if _int_val_node(v): + return IntArg(value=0, ref=v.name) + # A bool/float-producing node (e.g. a shape comparison used as a cond predicate) + # is a dynamic scalar referenced by SSA name, not a tensor. + val = v.meta.get("val") + if isinstance(val, (bool, torch.SymBool)): + return BoolArg(value=False, ref=v.name) + if isinstance(val, (float, torch.SymFloat)): + return FloatArg(value=0.0, ref=v.name) + return TensorArg(name=v.name) + + +def _scalar_to_arg_value(v: object) -> ArgumentValue: + if isinstance(v, bool): + return BoolArg(value=v) + if isinstance(v, int): + return IntArg(value=v) + if isinstance(v, float): + return FloatArg(value=v) + if isinstance(v, str): + return StringArg(value=v) + if isinstance(v, torch.dtype): + return ScalarTypeArg(value=_scalar_type(v)) + # device/layout/memory_format are single-target and always-strided for this + # backend, so a string is enough. + if isinstance(v, (torch.memory_format, torch.device, torch.layout)): + return StringArg(value=str(v)) + if isinstance(v, torch.SymInt): + # A raw SymInt arg is only representable if it is a constant; a genuinely + # symbolic one must reference an in-graph value (an fx node), handled by + # _node_to_arg_value. + concrete = v.node.maybe_as_int() + if concrete is not None: + return IntArg(value=concrete) + raise ValueError( + f"symbolic SymInt scalar arg {v!r} is not backed by an fx node; " + f"dynamic int args must reference an in-graph value." + ) + # Anything else is an unhandled arg type, so fail loud rather than emit a lossy + # repr. + raise ValueError(f"Cannot serialize argument of type {type(v).__name__}: {v!r}") + + +def _to_arg_value( + v: object, + subgraph_map: "dict[str, torch.fx.GraphModule] | None" = None, + schema_type_hint: str | None = None, +) -> ArgumentValue: + if v is None: + return NoneArg() + if isinstance(v, torch.fx.Node): + return _node_to_arg_value(v, subgraph_map) + if isinstance(v, (list, tuple)): + return _to_list_arg(list(v), subgraph_map, schema_type_hint=schema_type_hint) + return _scalar_to_arg_value(v) + + +def _argument( + v: object, + subgraph_map: "dict[str, torch.fx.GraphModule] | None" = None, + schema_type_hint: str | None = None, +) -> Argument: + return Argument( + value=_to_arg_value(v, subgraph_map, schema_type_hint=schema_type_hint) + ) + + +def _named_arguments( + node: torch.fx.Node, + subgraph_map: "dict[str, torch.fx.GraphModule] | None" = None, +) -> list[NamedArgument]: + op = _resolve_op_overload(node.target) + if op is None: + # No schema available (sym ops, operator.*, higher-order ops): serialize + # args as given, resolving any subgraph references to inline GraphArgs. + result = [ + NamedArgument(name=None, arg=_argument(a, subgraph_map)) for a in node.args + ] + result += [ + NamedArgument(name=k, arg=_argument(v, subgraph_map)) + for k, v in node.kwargs.items() + ] + return result + + # Materialize every schema argument, filling defaults for anything the call + # omitted, so the serialized node fully specifies the op invocation without the + # consumer needing to know the op's default values. + result = [] + for sarg, value, present in _bound_schema_args(node, op): + if not present: + # Required arg not provided; leave it out (invalid graph otherwise). + continue + # A written arg is one the op mutates in-place (schema Tensor(a!)); the op + # schema is the source of truth (see comment on NamedArgument.mutated). + mutated = sarg.alias_info is not None and sarg.alias_info.is_write + type_hint = str(getattr(sarg, "type", "") or "") + result.append( + NamedArgument( + name=sarg.name, + arg=_argument(value, subgraph_map, schema_type_hint=type_hint), + mutated=mutated, + ) + ) + return result + + +def _bound_schema_args( + node: torch.fx.Node, op: "torch._ops.OpOverload" +) -> list[tuple[Any, object, bool]]: + """Bind each schema argument to the value the call supplies. + + Returns one (schema_arg, value, present) triple per ``op._schema.arguments`` + position, in schema order. ``present`` is False for a required arg the call + omitted (an invalid graph); ``value`` is then meaningless. + """ + bound: list[tuple[Any, object, bool]] = [] + for i, sarg in enumerate(op._schema.arguments): + if i < len(node.args): + bound.append((sarg, node.args[i], True)) + elif sarg.name in node.kwargs: + bound.append((sarg, node.kwargs[sarg.name], True)) + elif sarg.has_default_value(): + bound.append((sarg, sarg.default_value, True)) + else: + bound.append((sarg, None, False)) + return bound + + +def _output_alias_of(node: torch.fx.Node) -> str | None: + """SSA name of the input value this node's output shares storage with, or None. + + A view/alias op's return is annotated ``Tensor(a)`` (read-only view) or + ``Tensor(a!)`` (in-place); the shared alias symbol matches one of the op's + input args. We report the first return that aliases an in-graph input value, + which covers the common single-return view and in-place ops. The write-ness of + the sharing is carried by the aliased input's NamedArgument.mutated, so it is + not repeated here. + """ + op = _resolve_op_overload(node.target) + if op is None: + return None + bound = _bound_schema_args(node, op) + arg_sets = [ + (set(sarg.alias_info.before_set) if sarg.alias_info is not None else set()) + for sarg, _, _ in bound + ] + for ret in op._schema.returns: + if ret.alias_info is None: + continue + ret_set = set(ret.alias_info.before_set) + if not ret_set: + continue + for (_, value, present), aset in zip(bound, arg_sets): + if present and (aset & ret_set) and isinstance(value, torch.fx.Node): + return value.name + return None + + +# --------------------------------------------------------------------------- +# Graph construction. +# --------------------------------------------------------------------------- + + +def _op_kind(op: str) -> OpKind: + try: + return { + "call_function": OpKind.CALL_FUNCTION, + "placeholder": OpKind.PLACEHOLDER, + "output": OpKind.OUTPUT, + }[op] + except KeyError: + raise ValueError(f"Unsupported fx op {op!r}") from None + + +def _subgraph_map( + graph_module: torch.fx.GraphModule, +) -> dict[str, torch.fx.GraphModule]: + out: dict[str, torch.fx.GraphModule] = {} + for node in graph_module.graph.nodes: + if node.op == "get_attr": + attr = getattr(graph_module, str(node.target), None) + if isinstance(attr, torch.fx.GraphModule): + out[node.name] = attr + return out + + +_INPUT_KIND_MAP: dict[InputKind, SchemaInputKind] = { + InputKind.PARAMETER: SchemaInputKind.PARAMETER, + InputKind.BUFFER: SchemaInputKind.BUFFER, + InputKind.CONSTANT_TENSOR: SchemaInputKind.CONSTANT_TENSOR, +} + + +def _is_non_persistent_buffer(ispec: object) -> bool: + return getattr(ispec, "kind", None) == InputKind.BUFFER and not getattr( + ispec, "persistent", True + ) + + +def _tensor_values_from( + val_by_name: dict[str, torch.Tensor], + names: list[str], + exclude: set[str] | None = None, +) -> list[TensorValue]: + seen: set[str] = set() + out: list[TensorValue] = [] + for name in names: + if name in seen: + continue + if exclude is not None and name in exclude: + continue + tensor = val_by_name.get(name) + if tensor is None: + continue + seen.add(name) + out.append(TensorValue(name=name, meta=_tensor_meta(tensor))) + return out + + +def _is_getitem(fx_node: torch.fx.Node) -> bool: + return fx_node.op == "call_function" and fx_node.target is operator.getitem + + +def _getitem_users( + graph_module: torch.fx.GraphModule, +) -> dict[str, dict[int, torch.fx.Node]]: + """Map each tuple-producing node name to {index: getitem_node} for its users.""" + users: dict[str, dict[int, torch.fx.Node]] = {} + for fx_node in graph_module.graph.nodes: + if not _is_getitem(fx_node): + continue + producer, idx = fx_node.args[0], fx_node.args[1] + if isinstance(producer, torch.fx.Node) and isinstance(idx, int): + users.setdefault(producer.name, {})[idx] = fx_node + return users + + +def _is_tensor_list_return(ret_schema: Any) -> bool: + rt = getattr(ret_schema, "real_type", None) or getattr(ret_schema, "type", None) + return isinstance(rt, torch.ListType) and isinstance( + rt.getElementType(), torch.TensorType + ) + + +def _scalar_output_kind(val: object) -> "OutputValueKind | None": + if isinstance(val, (bool, torch.SymBool)): + return OutputValueKind.BOOL + if isinstance(val, (int, torch.SymInt)): + return OutputValueKind.INT + if isinstance(val, (float, torch.SymFloat)): + return OutputValueKind.FLOAT + return None + + +def _element_names( + fx_node: torch.fx.Node, + getitem_users: dict[str, dict[int, torch.fx.Node]], + count: int, +) -> list[str]: + """SSA names of a multi/list result's elements: the folded getitem user's name, + or a synthetic name for a result no getitem extracts.""" + idx_map = getitem_users.get(fx_node.name, {}) + return [ + idx_map[i].name if i in idx_map else f"{fx_node.name}_out{i}" + for i in range(count) + ] + + +def _single_value_output( + fx_node: torch.fx.Node, val: object, val_by_name: dict[str, torch.Tensor] +) -> Output: + if isinstance(val, torch.Tensor): + val_by_name[fx_node.name] = val + alias_of = _output_alias_of(fx_node) if fx_node.op == "call_function" else None + return Output(name=fx_node.name, alias_of=alias_of) + kind = _scalar_output_kind(val) + if kind is not None: + return Output(name=fx_node.name, kind=kind) + return Output(name=fx_node.name) + + +def _tuple_element_outputs( + fx_node: torch.fx.Node, + meta_val: object, + getitem_users: dict[str, dict[int, torch.fx.Node]], + val_by_name: dict[str, torch.Tensor], +) -> list[Output]: + """One output per element of a tuple return (schema tuple or schema-less HOP). + + Nested tensor-list elements (a `Tensor[]` inside a tuple) are rejected; scalar + elements carry an INT/BOOL/FLOAT kind. + """ + idx_map = getitem_users.get(fx_node.name, {}) + outs: list[Output] = [] + for idx, m in enumerate(cast("list[object]", meta_val)): + gi = idx_map.get(idx) + name = gi.name if gi is not None else f"{fx_node.name}_out{idx}" + if isinstance(m, (list, tuple)): + raise ValueError( + f"node {fx_node.name!r} return {idx} is a nested tensor-list; nested " + f"list returns inside a tuple are not supported." + ) + if isinstance(m, torch.Tensor): + val_by_name[name] = m + outs.append(Output(name=name)) + continue + kind = _scalar_output_kind(m) + if kind is None: + raise ValueError( + f"node {fx_node.name!r} return {idx} has unsupported type " + f"{type(m).__name__}" + ) + outs.append(Output(name=name, kind=kind)) + return outs + + +def _node_outputs( + fx_node: torch.fx.Node, + getitem_users: dict[str, dict[int, torch.fx.Node]], + val_by_name: dict[str, torch.Tensor], +) -> list[Output]: + """Outputs for a producer node, recording each result's tensor metadata. + + Schema-return driven (like torch._export.serde): a tuple of returns (e.g. topk, + max.dim) becomes one output per return; a single ``Tensor[]`` return (e.g. split, + unbind) becomes one TENSOR_LIST output over the element names; scalar returns + carry an INT/BOOL/FLOAT kind. getitem users are folded in so their SSA names name + the results. Schema-less tuple returns (HOPs like torch.cond) fold by meta['val']. + """ + meta_val = fx_node.meta.get("val") + op = _resolve_op_overload(fx_node.target) if fx_node.op == "call_function" else None + if op is not None: + returns = op._schema.returns + if len(returns) == 0: + return [] + if len(returns) == 1 and _is_tensor_list_return(returns[0]): + elems = list(meta_val) if isinstance(meta_val, (list, tuple)) else [] + names = _element_names(fx_node, getitem_users, len(elems)) + for nm, m in zip(names, elems): + if isinstance(m, torch.Tensor): + val_by_name[nm] = m + return [Output(name="", kind=OutputValueKind.TENSOR_LIST, names=names)] + if len(returns) > 1: + return _tuple_element_outputs(fx_node, meta_val, getitem_users, val_by_name) + return [_single_value_output(fx_node, meta_val, val_by_name)] + if fx_node.op == "call_function" and isinstance(meta_val, (tuple, list)): + return _tuple_element_outputs(fx_node, meta_val, getitem_users, val_by_name) + return [_single_value_output(fx_node, meta_val, val_by_name)] + + +def _output_node_inputs( + fx_node: torch.fx.Node, + subgraph_map: dict[str, torch.fx.GraphModule], +) -> tuple[list[NamedArgument], list[str]]: + out_vals = fx_node.args[0] if fx_node.args else () + out_list = list(out_vals) if isinstance(out_vals, (list, tuple)) else [out_vals] + inputs: list[NamedArgument] = [] + output_names: list[str] = [] + for v in out_list: + inputs.append(NamedArgument(name=None, arg=_argument(v, subgraph_map))) + # This node's arguments are the full ordered result list (tensor refs, scalar + # refs, and literals). Graph.outputs lists only the tensor values, which carry + # metadata and can be mutation targets. + if isinstance(v, torch.fx.Node) and isinstance(v.meta.get("val"), torch.Tensor): + output_names.append(v.name) + return inputs, output_names + + +def _collect_fx_nodes( + graph_module: torch.fx.GraphModule, + subgraph_map: dict[str, torch.fx.GraphModule], +) -> tuple[list[Node], dict[str, torch.Tensor], list[str]]: + nodes: list[Node] = [] + val_by_name: dict[str, torch.Tensor] = {} + output_names: list[str] = [] + getitem_users = _getitem_users(graph_module) + + for fx_node in graph_module.graph.nodes: + if fx_node.op == "get_attr": + if fx_node.name in subgraph_map: + continue + raise ValueError( + f"get_attr node {fx_node.name!r} (target {fx_node.target!r}) cannot be " + f"serialized: the graph is not lifted. Serialize a lifted " + f"torch.export ExportedProgram (params/buffers/constants as inputs) " + f"instead of an unlifted module." + ) + + # getitem users of a tuple-producing node are folded into that node's + # outputs (see _node_outputs), so they are not emitted as their own nodes. + if _is_getitem(fx_node): + producer = fx_node.args[0] + if isinstance(producer, torch.fx.Node) and isinstance( + producer.meta.get("val"), (tuple, list) + ): + continue + + kind = _op_kind(fx_node.op) + target = None + inputs: list[NamedArgument] = [] + outputs: list[Output] = [] + + if fx_node.op == "output": + inputs, node_output_names = _output_node_inputs(fx_node, subgraph_map) + output_names.extend(node_output_names) + else: + if fx_node.op == "call_function": + target = serialize_operator(fx_node.target) + inputs = _named_arguments(fx_node, subgraph_map) + outputs = _node_outputs(fx_node, getitem_users, val_by_name) + + nodes.append( + Node( + name=fx_node.name, + op_kind=kind, + target=target, + inputs=inputs or None, + outputs=outputs or None, + ) + ) + + return nodes, val_by_name, output_names + + +def _validate_tensor_user_inputs(graph_signature: object) -> None: + for ispec in getattr(graph_signature, "input_specs", []) or []: + if ispec.kind == InputKind.USER_INPUT and not isinstance( + ispec.arg, TensorArgument + ): + raise ValueError( + f"unsupported non-tensor user input {ispec.arg!r}; " + f"only tensor user inputs are supported." + ) + + +def _build_output_specs( + output_names: list[str], + graph_signature: object, +) -> tuple[list[OutputSpec], set[str]]: + params_to_mutate = getattr(graph_signature, "parameters_to_mutate", None) or {} + if params_to_mutate: + raise ValueError( + f"parameter mutation is not supported: {sorted(params_to_mutate)}. " + f"Serialize a graph that does not mutate parameters." + ) + buffers_to_mutate = getattr(graph_signature, "buffers_to_mutate", None) or {} + user_inputs_to_mutate = ( + getattr(graph_signature, "user_inputs_to_mutate", None) or {} + ) + specs: list[OutputSpec] = [] + for name in output_names: + if name in buffers_to_mutate: + specs.append( + OutputSpec( + name=name, + kind=OutputKind.BUFFER_MUTATION, + target=buffers_to_mutate[name], + ) + ) + elif name in user_inputs_to_mutate: + specs.append( + OutputSpec( + name=name, + kind=OutputKind.USER_INPUT_MUTATION, + target=user_inputs_to_mutate[name], + ) + ) + else: + specs.append(OutputSpec(name=name, kind=OutputKind.USER_OUTPUT)) + mutated_fqns = set(buffers_to_mutate.values()) + return specs, mutated_fqns + + +def _extract_constants_and_mutable_buffers( + graph_signature: object, + state_dict: dict[str, object], + constants: dict[str, object] | None, + mutated_fqns: set[str], +) -> tuple[list[NamedTensorRef], dict[str, torch.Tensor], list[MutableBufferSpec]]: + constant_refs: list[NamedTensorRef] = [] + constant_data: dict[str, torch.Tensor] = {} + mutable_buffers: list[MutableBufferSpec] = [] + + for ispec in getattr(graph_signature, "input_specs", []) or []: + if ispec.kind not in _INPUT_KIND_MAP: + continue + name = getattr(getattr(ispec, "arg", None), "name", None) + target_fqn = getattr(ispec, "target", None) + if name is None or target_fqn is None: + continue + if _is_non_persistent_buffer(ispec): + mutable_buffers.append(MutableBufferSpec(name=name, fqn=target_fqn)) + continue + tensor = None + if target_fqn in state_dict: + tensor = state_dict[target_fqn] + elif constants is not None and target_fqn in constants: + tensor = constants[target_fqn] + if not isinstance(tensor, torch.Tensor): + continue + tensor = tensor.contiguous() + constant_refs.append( + NamedTensorRef( + name=name, + data_key=target_fqn, + meta=_tensor_meta(tensor), + kind=_INPUT_KIND_MAP[ispec.kind], + mutated=target_fqn in mutated_fqns, + ) + ) + constant_data[target_fqn] = tensor + + return constant_refs, constant_data, mutable_buffers + + +def _build_subgraph( + graph_module: torch.fx.GraphModule, + nodes: list[Node], + val_by_name: dict[str, torch.Tensor], + output_names: list[str], +) -> tuple[ + Graph, + list[NamedTensorRef], + list[OutputSpec], + list[MutableBufferSpec], + dict[str, torch.Tensor], +]: + user_inputs = [n.name for n in graph_module.graph.nodes if n.op == "placeholder"] + tensor_values = _tensor_values_from(val_by_name, list(val_by_name.keys())) + graph = Graph( + nodes=nodes, + inputs=user_inputs or None, + outputs=output_names or None, + tensor_values=tensor_values or None, + ) + return graph, [], [], [], {} + + +def _build_method_graph( + graph_module: torch.fx.GraphModule, + nodes: list[Node], + val_by_name: dict[str, torch.Tensor], + output_names: list[str], + graph_signature: object, + state_dict: dict[str, object], + constants: dict[str, object] | None, +) -> tuple[ + Graph, + list[NamedTensorRef], + list[OutputSpec], + list[MutableBufferSpec], + dict[str, torch.Tensor], +]: + _validate_tensor_user_inputs(graph_signature) + output_specs, mutated_fqns = _build_output_specs(output_names, graph_signature) + constant_refs, constant_data, mutable_buffers = ( + _extract_constants_and_mutable_buffers( + graph_signature, state_dict, constants, mutated_fqns + ) + ) + constant_names = {c.name for c in constant_refs} + tensor_values = _tensor_values_from( + val_by_name, + [n for n in val_by_name if n not in constant_names], + ) + user_inputs = list(getattr(graph_signature, "user_inputs", []) or []) + graph = Graph( + nodes=nodes, + inputs=user_inputs or None, + outputs=output_names or None, + tensor_values=tensor_values or None, + ) + return graph, constant_refs, output_specs, mutable_buffers, constant_data + + +def _build_graph_body( + graph_module: torch.fx.GraphModule, + graph_signature: object | None, + state_dict: dict[str, object], + constants: dict[str, object] | None, +) -> tuple[ + Graph, + list[NamedTensorRef], + list[OutputSpec], + list[MutableBufferSpec], + dict[str, torch.Tensor], +]: + subgraph_map = _subgraph_map(graph_module) + nodes, val_by_name, output_names = _collect_fx_nodes(graph_module, subgraph_map) + + if graph_signature is None: + return _build_subgraph(graph_module, nodes, val_by_name, output_names) + + return _build_method_graph( + graph_module, + nodes, + val_by_name, + output_names, + graph_signature, + state_dict, + constants, + ) + + +# --------------------------------------------------------------------------- +# JSON to/from flatbuffer via flatc. +# --------------------------------------------------------------------------- + + +def _encode(o: object) -> object: + """Recursively convert a dataclass tree to JSON-compatible primitives. + + Emits the ``_type`` discriminator BEFORE the union value (flatc + requires the union type field to precede the value) and omits None-valued + optional fields. + """ + if is_dataclass(o): + out: dict[str, object] = {} + hints = get_type_hints(type(o)) + for f in fields(o): + val = getattr(o, f.name) + if val is None: + continue + hint = hints[f.name] + if get_origin(hint) is Union and type(None) not in get_args(hint): + out[f"{f.name}_type"] = type(val).__name__ + out[f.name] = _encode(val) + else: + out[f.name] = _encode(val) + return out + if isinstance(o, IntEnum): + return int(o) + if isinstance(o, (list, tuple)): + return [_encode(x) for x in o] + return o + + +def _prepare_schema(out_dir: str) -> str: + data = importlib.resources.read_binary(__package__, _SCHEMA_RESOURCE) + schema_path = os.path.join(out_dir, _SCHEMA_RESOURCE) + with open(schema_path, "wb") as f: + f.write(data) + return schema_path + + +def _compile_to_bytes(root: object) -> bytes: + json_str = json.dumps(_encode(root)) + with tempfile.TemporaryDirectory() as td: + schema_path = _prepare_schema(td) + json_path = os.path.join(td, _FILE_STEM + ".json") + with open(json_path, "w") as f: + f.write(json_str) + _flatc_compile(td, schema_path, json_path) + bin_path = os.path.join(td, _FILE_STEM + ".bin") + with open(bin_path, "rb") as f: + return f.read() + + +_FLOAT_DTYPES: tuple[torch.dtype, ...] = ( + torch.float32, + torch.float16, + torch.float64, + torch.bfloat16, +) + + +def _same_storage(a: torch.Tensor, b: torch.Tensor) -> bool: + try: + return ( + statically_known_true(a.storage_offset() == b.storage_offset()) + and a.untyped_storage().data_ptr() == b.untyped_storage().data_ptr() + ) + except Exception: + return False + + +def _float_equal_with_nan(a: torch.Tensor, b: torch.Tensor) -> bool: + try: + if not (torch.isnan(a).any() or torch.isnan(b).any()): + return False + except RuntimeError: + return False + eq = (a == b) | (torch.isnan(a) & torch.isnan(b)) + return bool(eq.all().item()) + + +def _same_tensor(a: torch.Tensor, b: torch.Tensor) -> bool: + if a is b: + return True + if a.dtype != b.dtype or a.shape != b.shape: + return False + if _same_storage(a, b): + return True + + ca = a.detach() + cb = b.detach() + + if ca.device != cb.device: + if ca.device.type != "cpu": + ca = ca.cpu() + if cb.device.type != "cpu": + cb = cb.cpu() + + if torch.equal(ca, cb): + return True + + if ca.dtype in _FLOAT_DTYPES: + return _float_equal_with_nan(ca, cb) + return False + + +def serialize_program( + methods: dict[ + str, + tuple[ + torch.fx.GraphModule, object, dict[str, object], dict[str, object] | None + ], + ], +) -> tuple[bytes, dict[str, torch.Tensor]]: + """Serialize one or more named methods into a native flatbuffer Program. + + ``methods`` maps a method name (e.g. "forward") to a + ``(graph_module, graph_signature, state_dict, constants)`` tuple. Returns + (flatbuffer_bytes, constant_data), where constant_data maps a fully-qualified + name to the constant tensor, merged (deduped by fqn) across all methods. The + caller ships constant_data as the external constant file. + + Cross-method sharing is by fqn: bundling methods asserts they come from a single + model namespace, so an fqn is the same buffer/constant everywhere. That assertion + is validated; if two methods carry the same fqn with different data (constants) or + different shape/dtype (mutable buffers), this raises rather than silently aliasing + or clobbering. + """ + method_objs: list[Method] = [] + constant_data: dict[str, torch.Tensor] = {} + mutable_meta: dict[str, TensorMeta] = {} + for name, (graph_module, graph_signature, state_dict, constants) in methods.items(): + graph, constant_refs, output_specs, mutable_buffers, cdata = _build_graph_body( + graph_module, graph_signature, state_dict, constants + ) + method_objs.append( + Method( + name=name, + graph=graph, + constants=constant_refs or None, + output_specs=output_specs or None, + mutable_buffers=mutable_buffers or None, + ) + ) + + for fqn, tensor in cdata.items(): + prev = constant_data.get(fqn) + if prev is not None and not _same_tensor(prev, tensor): + raise ValueError( + f"method {name!r}: constant fqn {fqn!r} conflicts with different " + f"data already provided by another method. Methods bundled into " + f"one program must share the same data per fqn." + ) + constant_data[fqn] = tensor + + meta_by_name = {tv.name: tv.meta for tv in (graph.tensor_values or [])} + for mb in mutable_buffers: + meta = meta_by_name.get(mb.name) + if meta is None: + continue + prev_meta = mutable_meta.get(mb.fqn) + if prev_meta is not None and prev_meta != meta: + raise ValueError( + f"method {name!r}: mutable buffer fqn {mb.fqn!r} has a different " + f"shape/dtype than in another method; a shared buffer must match." + ) + mutable_meta[mb.fqn] = meta + + program = Program(version=SCHEMA_VERSION, methods=method_objs) + return _compile_to_bytes(program), constant_data + + +def serialize_graph( + graph_module: torch.fx.GraphModule, + graph_signature: object, + state_dict: dict[str, object], + constants: dict[str, object] | None = None, +) -> tuple[bytes, dict[str, torch.Tensor]]: + """Serialize a single fx graph as a one-method ("forward") Program. + + Convenience wrapper around ``serialize_program``. Returns (flatbuffer_bytes, + constant_data) as documented there. + """ + return serialize_program( + {"forward": (graph_module, graph_signature, state_dict, constants)} + ) + + +def deserialize_program(data: bytes) -> Program: + """Deserialize native flatbuffer bytes back into a Program dataclass.""" + with tempfile.TemporaryDirectory() as td: + schema_path = _prepare_schema(td) + bin_path = os.path.join(td, _FILE_STEM + ".bin") + with open(bin_path, "wb") as f: + f.write(data) + _flatc_decompile(td, schema_path, bin_path) + json_path = os.path.join(td, _FILE_STEM + ".json") + with open(json_path) as f: + obj = json.load(f) + return _json_to_dataclass(obj, Program) + + +def deserialize_graph(data: bytes) -> Graph: + """Deserialize and return the first method's Graph (single-method convenience).""" + program = deserialize_program(data) + if not program.methods: + raise ValueError("program has no methods") + return program.methods[0].graph + + +def _check_optional_tensor_list_refs( + value: OptionalTensorListArg, ctx: str, check_ref: Any +) -> None: + if len(value.names) != len(value.has_value): + raise ValueError( + f"{ctx}: OptionalTensorListArg names length {len(value.names)} does not " + f"match has_value length {len(value.has_value)}" + ) + for n, present in zip(value.names, value.has_value): + if present: + check_ref(n, ctx) + + +def _check_int_list_refs(value: IntListArg, ctx: str, check_ref: Any) -> None: + refs = value.refs + # refs is parallel to values; a set refs[i] references an in-graph int. + if refs is None: + return + if len(refs) != len(value.values): + raise ValueError( + f"{ctx}: IntListArg refs length {len(refs)} does not " + f"match values length {len(value.values)}" + ) + for r in refs: + if r: + check_ref(r, ctx) + + +def _check_arg_refs(value: ArgumentValue, ctx: str, check_ref: Any) -> None: + if isinstance(value, TensorArg): + check_ref(value.name, ctx) + elif isinstance(value, TensorListArg): + for n in value.names: + check_ref(n, ctx) + elif isinstance(value, OptionalTensorListArg): + _check_optional_tensor_list_refs(value, ctx, check_ref) + elif isinstance(value, IntArg): + # ref, when set, is an SSA reference to an in-graph scalar value. + if value.ref: + check_ref(value.ref, ctx) + elif isinstance(value, (BoolArg, FloatArg)): + if value.ref: + check_ref(value.ref, ctx) + elif isinstance(value, IntListArg): + _check_int_list_refs(value, ctx, check_ref) + elif isinstance(value, GraphArg): + # HOP subgraph: self-contained (its own placeholder namespace). + validate_graph(value.graph) + + +def _check_io_meta(graph: Graph, meta_names: set[str]) -> None: + for name in graph.inputs or []: + if name not in meta_names: + raise ValueError(f"input {name!r} missing tensor metadata") + for name in graph.outputs or []: + if name not in meta_names: + raise ValueError(f"output {name!r} missing tensor metadata") + + +def _define_node_outputs(node: Node, define: Any) -> None: + ctx = f"node {node.name!r}" + for out in node.outputs or []: + if out.kind == OutputValueKind.TENSOR_LIST: + for nm in out.names or []: + define(nm, ctx) + else: + define(out.name, ctx) + + +def _validate_node(node: Node, defined: set[str], define: Any, check_ref: Any) -> None: + ctx = f"node {node.name!r}" + for na in node.inputs or []: + _check_arg_refs(na.arg.value, ctx, check_ref) + for out in node.outputs or []: + if out.alias_of: + check_ref(out.alias_of, f"{ctx} output alias") + # A placeholder must restate a value already declared as an input or external + # binding; any other node defines a fresh value per output (a multi-output op + # defines one value per result). + if node.op_kind == OpKind.PLACEHOLDER: + if node.name not in defined: + raise ValueError( + f"{ctx}: placeholder is not a declared input or external binding" + ) + else: + _define_node_outputs(node, define) + + +def validate_graph(graph: Graph, defined_extra: set[str] | None = None) -> None: + """Assert a (pure) graph body is structurally self-contained: every value + reference resolves and every input/output has tensor metadata. Recurses into + every ``GraphArg`` subgraph. Raises ``ValueError`` on the first inconsistency. + + defined_extra supplies names defined outside the body (a method's constants and + mutable buffers, see validate_method); a HOP subgraph passes None since its + operands are its own placeholders. Method-level bindings (constant data + availability, mutable-buffer metadata) are checked in validate_method. + """ + meta_names = {tv.name for tv in (graph.tensor_values or [])} + defined: set[str] = set() + seen_nodes: set[str] = set() + + def define(name: str, ctx: str) -> None: + if name in defined: + raise ValueError(f"{ctx}: duplicate value definition {name!r}") + defined.add(name) + + for name in graph.inputs or []: + define(name, "input") + for name in sorted(defined_extra or ()): + define(name, "external binding") + + def check_ref(name: str, ctx: str) -> None: + if name and name not in defined: + raise ValueError(f"{ctx}: unresolved value reference {name!r}") + + # Define a node's values only after its inputs resolve, so forward references and + # cycles are rejected (nodes must be topologically ordered). Node names are unique + # across the graph. + for node in graph.nodes: + if node.name in seen_nodes: + raise ValueError(f"node {node.name!r}: duplicate node name") + seen_nodes.add(node.name) + _validate_node(node, defined, define, check_ref) + + for name in graph.outputs or []: + check_ref(name, "output") + + _check_io_meta(graph, meta_names) + + +def validate_method( + method: Method, available_data_keys: set[str] | None = None +) -> None: + """Validate a method: its graph body (see ``validate_graph``) plus its + constant/mutable-buffer bindings. + + This enforces that the method graph plus its external constant data (whose keys + are available_data_keys) carry everything needed to run. + Names bound by constants / mutable buffers are supplied to the body's reference + check; every mutable buffer must have tensor metadata; and, when + ``available_data_keys`` is given, every ``NamedTensorRef.data_key`` must be + present. Mutable buffers are not data-backed, so they are exempt from the + data-keys check. + """ + defined_extra: set[str] = {c.name for c in (method.constants or [])} + defined_extra.update(mb.name for mb in (method.mutable_buffers or [])) + validate_graph(method.graph, defined_extra=defined_extra) + + meta_names = {tv.name for tv in (method.graph.tensor_values or [])} + for mb in method.mutable_buffers or []: + if mb.name not in meta_names: + raise ValueError( + f"mutable buffer {mb.name!r} (fqn {mb.fqn!r}) missing tensor metadata" + ) + + if available_data_keys is not None: + for c in method.constants or []: + if c.data_key not in available_data_keys: + raise ValueError( + f"constant {c.name!r} (data_key {c.data_key!r}) has no data in the " + f"provided external constant keys" + ) + + +def validate_program( + program: Program, available_data_keys: set[str] | None = None +) -> None: + """Validate every method (see ``validate_method``).""" + for method in program.methods: + validate_method(method, available_data_keys) diff --git a/backends/native/serialization/native_graph.fbs b/backends/native/serialization/native_graph.fbs new file mode 100644 index 00000000000..ed77bfe28c6 --- /dev/null +++ b/backends/native/serialization/native_graph.fbs @@ -0,0 +1,430 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// +// FlatBuffer schema for the ExecuTorch native backend. This is the source of truth. +// +// A generic representation of an fx graph: a topological list of nodes, each with an +// op-kind, a target op-name string, and generic arguments encoded through a tagged +// Argument union. Modeled on torch._export.serde +// (caffe2/torch/_export/serde/schema.py). Nodes and tensors are referenced by name; +// tensor metadata lives in a side table. Constant tensor data is not stored here; it +// is shipped separately in an external constant file, keyed by NamedTensorRef.data_key. +// +// BACKWARD COMPATIBILITY RULES: +// - New fields in tables: APPEND ONLY (add at the end, with a default value) +// - New union members: APPEND ONLY (add at the end of the union) +// - New tables: Safe to add freely +// - New enum values: APPEND ONLY +// - NEVER remove, reorder, or change the type of existing fields/members +// +// NAME SCOPING. Which namespace each name/reference field belongs to: +// (1) SSA value names: scoped to a SINGLE Graph body. The method's top-level +// Graph and each HOP subgraph each have their OWN independent namespace; +// references resolve only within the enclosing Graph. Fields: Node.name, +// Output.name, Output.alias_of, TensorArg.name, TensorListArg.names, +// OptionalTensorListArg.names, TensorValue.name, Graph.inputs, Graph.outputs. +// (2) Method-level bindings into the METHOD's top-level Graph (never a subgraph): +// NamedTensorRef.name, MutableBufferSpec.name, OutputSpec.name. +// (3) Data keys / fqns, the program-wide external-constant / cross-method +// namespace: NamedTensorRef.data_key, AffineGroup.scale_data_key / +// zero_point_data_key, MutableBufferSpec.fqn. +// Not value references: NamedArgument.name (operator-schema parameter name), +// GraphArg.name (submodule attr label, debug only), Method.name (program-level). +// OutputSpec.target is dual-namespace by kind: an fqn (3) for BUFFER_MUTATION, +// a user-input SSA name (2) for USER_INPUT_MUTATION. + +namespace native_backend; + +// ============================================================================= +// Enums +// ============================================================================= + +// Mirrors ExecuTorch's ScalarType ordering (runtime/core/portable_type/scalar_type.h). +enum ScalarType : byte { + BYTE = 0, + CHAR = 1, + SHORT = 2, + INT = 3, + LONG = 4, + HALF = 5, + FLOAT = 6, + DOUBLE = 7, + BOOL = 11, + BFLOAT16 = 15, + UINT16 = 16, + UINT32 = 17, + UINT64 = 18, + // Values are pinned to ExecuTorch's ScalarType ids; gaps (complex/qint) are + // reserved and may be filled later at their canonical id (append-only). +} + +// fx node kind (node.op). +enum OpKind : byte { + CALL_FUNCTION = 0, + PLACEHOLDER = 1, + OUTPUT = 2, +} + +// Classifies a graph input placeholder. Mirrors torch.export InputKind for the +// subset the native backend serializes. USER_INPUT placeholders are also listed +// by name in Graph.inputs; PARAMETER/BUFFER/CONSTANT_TENSOR placeholders each +// additionally get a NamedTensorRef carrying this kind. +enum InputKind : byte { + USER_INPUT = 0, + PARAMETER = 1, + BUFFER = 2, + CONSTANT_TENSOR = 3, +} + +// Classifies a graph output. USER_OUTPUT is a real result; the mutation kinds +// mark an output value that writes back into a persistent buffer or a user input +// (its OutputSpec.target names the mutated buffer fqn / user-input name). A +// runtime returns USER_OUTPUTs and applies the mutations as state updates. +enum OutputKind : byte { + USER_OUTPUT = 0, + BUFFER_MUTATION = 1, + USER_INPUT_MUTATION = 2, +} + +// Discriminates what a Node's Output produces. TENSOR (default) is the legacy +// single-tensor shape (name + optional alias_of). TENSOR_LIST is a `Tensor[]` +// return whose element SSA names are in Output.names (name empty). INT/BOOL/FLOAT +// are scalar returns named by Output.name. Append-only. +enum OutputValueKind : byte { + TENSOR = 0, + TENSOR_LIST = 1, + INT = 2, + BOOL = 3, + FLOAT = 4, +} + +// ============================================================================= +// Quantization +// ============================================================================= +// A tensor's quant scheme rides on its TensorMeta, so it applies uniformly to graph +// I/O, intermediates (tensor_values), and constants (NamedTensorRef.meta). Absent +// means not quantized. dtype stays the storage dtype; the scheme says how to +// interpret it. The scheme set is append-only. + +// Affine group-wise quant along the last (reduction) axis: +// value = scale * ((qdata + quant_min) - zero_point) +// qdata is stored unsigned in [0, quant_max - quant_min], LSB-first when packed +// sub-byte (two int4 nibbles per byte), so quant_min is also the unsigned-to-logical +// storage offset (logical q = qdata + quant_min, e.g. -8 for int4). quant_min and +// quant_max give the logical range, hence the bitwidth and signedness. group_size is +// the block length along the last axis; 0 means per-channel (one group over the whole +// axis). Scales (and, if asymmetric, zero-points) are shipped out-of-line via the +// external constant file, keyed by data_key (empty zero_point_data_key means symmetric). +// Their shapes derive from the weight shape and group_size, but their dtypes do not, +// so they are recorded here. The dequant result dtype is the op output value's +// TensorMeta.dtype, not part of this scheme. The qdata itself (the enclosing +// TensorMeta) uses dtype=BYTE with logical sizes; its packed byte length follows from +// the logical shape and bitwidth, not prod(sizes). +table AffineGroup { + scale_data_key: string (required); + scale_dtype: ScalarType; // FLOAT, HALF, or BFLOAT16; not derivable + quant_min: int; + quant_max: int; + group_size: int; // 0 means per-channel (one group over the last axis) + zero_point_data_key: string; // empty means symmetric + zero_point_dtype: ScalarType; // meaningful only when zero_point_data_key is set +} + +// Opaque codec-defined packed layout (e.g. "gguf:q4k", "mxfp4", "nvfp4"). Block size, +// bit width, and scales are all implied by codec. The enclosing TensorMeta uses +// dtype=BYTE with logical sizes; the packed byte length is defined by codec and the +// logical shape and carried by the external constant blob (NamedTensorRef.data_key), +// never prod(sizes). +table PackedQuant { + codec: string (required); +} + +// Append-only; new schemes go at the end. +union QuantScheme { + AffineGroup, + PackedQuant, +} + +// Wrapper so the scheme union can be an optional field on TensorMeta (an optional +// union cannot be expressed directly, so it is nested in an optional table). +table QuantSpec { + scheme: QuantScheme; +} + +// ============================================================================= +// Tensor metadata +// ============================================================================= + +// A tensor dimension as a range. Static: min == max. Dynamic: min < max, or max < 0 +// for unbounded. No symbols; the runtime uses concrete shapes at invocation and +// plans memory from max. +table Dim { + min: long; + max: long = -1; +} + +// sizes is the logical shape and dtype the element type. When quant is set the tensor +// is stored packed: dtype is BYTE and the physical layout and byte length come from +// the quant scheme and its external constant blob, not prod(sizes). +table TensorMeta { + dtype: ScalarType; + sizes: [Dim]; + // Memory layout as a permutation of dim indices, outermost first (contiguous is + // [0, 1, ..., n-1]). Independent of the sizes. Empty means contiguous. + dim_order: [int]; + // Absent means the tensor is not quantized. + quant: QuantSpec; +} + +// Side-table entry: maps an fx value name (SSA, this graph) to its tensor metadata. +table TensorValue { + name: string (required); + meta: TensorMeta (required); +} + +// ============================================================================= +// Argument union: one variant per kind of value an fx arg/kwarg can hold. +// Each variant is a table so the union can be nested inside other tables. +// ============================================================================= + +// Reference to another value by its fx SSA name (in the same graph). +table TensorArg { + name: string (required); +} + +table NoneArg {} + +// A scalar int operand: the literal value, or (when ref is non-empty) a reference to +// an in-graph int value by SSA name (a sym_size or arith node), with value ignored. +table IntArg { + value: long; + ref: string; +} + +table FloatArg { + value: double; + // Non-empty ref references an in-graph float value by SSA name; value is ignored. + ref: string; +} + +table BoolArg { + value: bool; + // Non-empty ref references an in-graph bool value by SSA name; value is ignored. + ref: string; +} + +table StringArg { + value: string (required); +} + +table ScalarTypeArg { + value: ScalarType; +} + +// A list of ints. Element i is a literal (values[i]), or a reference to an in-graph +// int value by SSA name when the parallel refs array is present and refs[i] is +// non-empty (a sym_size or arith node; values[i] ignored). Empty refs means all +// literal; when present, len(refs) == len(values). Example: a dynamic view size +// [s0, -1] is values=[0, -1], refs=["", ""]. +table IntListArg { + values: [long] (required); + refs: [string]; +} + +table FloatListArg { + values: [double] (required); +} + +table BoolListArg { + values: [bool] (required); +} + +// A list of tensor references by SSA name (same graph; e.g. cat's input list). +table TensorListArg { + names: [string] (required); +} + +// A list of optional tensor references (Tensor?[]) by SSA name (same graph). +// names[i] is only meaningful when has_value[i] is true; otherwise None. +table OptionalTensorListArg { + names: [string] (required); + has_value: [bool] (required); +} + +// A subgraph passed to a higher-order op (torch.cond / torch.while_loop / map). +// name is the original submodule attr label (debug only, not a value reference); +// graph is the inlined branch/body, with its own SSA namespace, which makes Graph +// recursive (mirrors torch._export.serde's GraphArgument). +table GraphArg { + name: string (required); + graph: Graph (required); +} + +// Append-only; new members go at the end. +union ArgumentValue { + TensorArg, + NoneArg, + IntArg, + FloatArg, + BoolArg, + StringArg, + ScalarTypeArg, + IntListArg, + FloatListArg, + BoolListArg, + TensorListArg, + OptionalTensorListArg, + GraphArg, +} + +table Argument { + value: ArgumentValue; +} + +// A positional or keyword argument. `name` is the operator-schema parameter name +// (NOT a value reference; empty for positional-only where a name is unavailable). +table NamedArgument { + name: string; + arg: Argument (required); + // True if the op writes this input in-place (op-schema alias annotation + // Tensor(a!), e.g. self in scatter_add_). The op schema is the source of truth + // for mutation; the op-name '_' suffix is only a convention and cannot say which + // args are written. + mutated: bool = false; +} + +table KeyValue { + key: string (required); + value: string (required); +} + +// ============================================================================= +// Node +// ============================================================================= + +// A value produced by a node. +table Output { + // fx SSA name of the produced value (this graph). Empty for a TENSOR_LIST output + // (the element names live in `names`). + name: string (required); + // If non-empty, the SSA name (same graph) of an input this output shares storage + // with (op-schema view annotation, e.g. aten.view/slice/transpose returning + // Tensor(a)). Whether the sharing is a write is given by the aliased input's + // NamedArgument.mutated, so it is not duplicated here. Empty means fresh storage. + alias_of: string; + // What this output produces. TENSOR (default) uses `name` (+ optional alias_of). + // TENSOR_LIST uses `names`. INT/BOOL/FLOAT are scalar returns named by `name`. + kind: OutputValueKind; + // Element SSA names for a TENSOR_LIST (`Tensor[]`) return, in order. + names: [string]; +} + +table Node { + // fx node name (SSA name of this node's output). + name: string (required); + op_kind: OpKind; + // Fully-qualified op name, e.g. "torch.ops.aten.add.Tensor". Empty for + // placeholder/output nodes. + target: string; + inputs: [NamedArgument]; + // Output value(s) produced by this node (usually a single entry; multiple for + // tuple-returning ops). + outputs: [Output]; + metadata: [KeyValue]; +} + +// ============================================================================= +// Constants: data shipped separately in an external constant file, keyed by data_key. +// ============================================================================= + +table NamedTensorRef { + // SSA placeholder name in the method's graph (namespace 2). + name: string (required); + // Key into the external constant file (the tensor's fully-qualified name). Shared + // across methods: the same data_key means the same blob (see serialize_program). + data_key: string (required); + meta: TensorMeta (required); + // Whether this input is a parameter, buffer, or constant tensor. Never + // USER_INPUT (those are listed by name in Graph.inputs instead). + kind: InputKind; + // True for buffers the graph mutates in-place whose initial values are saved + // (persistent buffers). Their updated values are state that must persist across + // executions; unlike parameters and frozen constants they are not read-only. + // Non-persistent mutable buffers (e.g. zero-initialized KV caches) carry no + // data and are listed in Method.mutable_buffers instead. + mutated: bool = false; +} + +// Classifies a graph-level output value (by SSA name), so a runtime can tell a +// real result from a buffer/user-input mutation writeback. +table OutputSpec { + // SSA name of the output value (matches an entry in the method's Graph.outputs). + name: string (required); + kind: OutputKind; + // For BUFFER_MUTATION: the mutated buffer's fqn (data namespace 3). For + // USER_INPUT_MUTATION: the mutated user input's SSA name (method graph, + // namespace 2). Empty for USER_OUTPUT. + target: string; +} + +// Graph state that is not data-backed: a non-persistent buffer (e.g. a +// zero-initialized KV cache) not saved in the state dict. It is neither a user input +// nor a NamedTensorRef. Shape and dtype live in tensor_values (keyed by name); no +// bytes are shipped (zero-initialized at load). Cross-method sharing is by fqn, a +// validated contract (bundled methods must come from one model namespace; see +// serialize_program), not inferred from name equality alone. +table MutableBufferSpec { + // SSA placeholder name in the method's graph (namespace 2). + name: string (required); + // Stable cross-method identity (data namespace 3) correlating the same buffer. + fqn: string (required); +} + +// ============================================================================= +// Root type: Program (one or more named methods) +// ============================================================================= + +// A pure function body (fx graph in topological order). Used identically for a +// top-level method body and, recursively, for higher-order-op subgraphs (see +// GraphArg). Stateful method-level bindings (constants, output_specs, +// mutable_buffers) live on Method, not here. +table Graph { + // All nodes, in topological order (fx graph iteration order). + nodes: [Node] (required); + // Names of graph-level input placeholders (user inputs at the top level; + // operands for a subgraph), in order. + inputs: [string]; + // Names of graph-level outputs, in order. + outputs: [string]; + // Side table: value name to tensor metadata. + tensor_values: [TensorValue]; +} + +// A named method within a program (e.g. "forward", "encode", "decode"). Carries +// the pure graph body plus its stateful signature bindings. +table Method { + name: string (required); + graph: Graph (required); + // Constant/parameter/buffer references binding a method-graph placeholder to a + // data_key (external constant file, shared across methods). A subgraph has none; its + // params are lifted here and passed into the op as operands. + constants: [NamedTensorRef]; + // Per-output classification (user output vs buffer/user-input mutation), + // parallel to graph.outputs by name. + output_specs: [OutputSpec]; + // Non-persistent mutable buffers (e.g. KV caches): graph state with no shipped + // data, zero-initialized at load. See MutableBufferSpec. + mutable_buffers: [MutableBufferSpec]; +} + +// Root: a program bundling one or more methods. Methods share constant data +// (an external constant file keyed by data_key); each Method carries its own NamedTensorRef +// bindings. +table Program { + // Schema version for compatibility. + version: string; + methods: [Method] (required); +} + +root_type Program; + +file_identifier "NPTG"; diff --git a/backends/native/serialization/schema.py b/backends/native/serialization/schema.py new file mode 100644 index 00000000000..500fdc8f6a9 --- /dev/null +++ b/backends/native/serialization/schema.py @@ -0,0 +1,343 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +""" +Python dataclass mirror of native_graph.fbs, the source of truth for semantics. + +Serialized to JSON by executorch.exir._serialize._dataclass._DataclassEncoder and +compiled to a flatbuffer with flatc (see graph_serialize.py). + +Conventions (match executorch/exir/schema.py): + - Union fields are annotated as string literals (e.g. "ArgumentValue") so the + encoder emits the _type discriminator that flatc's JSON unions need. + Union member class names must equal the flatbuffer table names. + - required fbs fields are non-optional here; the rest are Optional with a default + so the flatc --json round-trip (which omits unset vectors) deserializes. + +See native_graph.fbs for field semantics and name scoping. +""" + +from dataclasses import dataclass, field +from enum import IntEnum +from typing import List, Optional, Union + + +class ScalarType(IntEnum): + BYTE = 0 + CHAR = 1 + SHORT = 2 + INT = 3 + LONG = 4 + HALF = 5 + FLOAT = 6 + DOUBLE = 7 + BOOL = 11 + BFLOAT16 = 15 + UINT16 = 16 + UINT32 = 17 + UINT64 = 18 + + +class OpKind(IntEnum): + CALL_FUNCTION = 0 + PLACEHOLDER = 1 + OUTPUT = 2 + + +class InputKind(IntEnum): + USER_INPUT = 0 + PARAMETER = 1 + BUFFER = 2 + CONSTANT_TENSOR = 3 + + +class OutputKind(IntEnum): + USER_OUTPUT = 0 + BUFFER_MUTATION = 1 + USER_INPUT_MUTATION = 2 + + +class OutputValueKind(IntEnum): + # Discriminates what a Node.Output produces. TENSOR (default) keeps the legacy + # single-tensor shape (name + optional alias_of). TENSOR_LIST is a `Tensor[]` + # return whose element SSA names are in `names` (name empty). INT/BOOL/FLOAT are + # scalar returns whose SSA name is `name`. + TENSOR = 0 + TENSOR_LIST = 1 + INT = 2 + BOOL = 3 + FLOAT = 4 + + +# Quantization. A tensor's quant scheme rides on its TensorMeta, so it applies to +# graph I/O, intermediates, and constants uniformly. Absent means not quantized; +# dtype stays the storage dtype and the scheme says how to interpret it. The scheme +# set is append-only and grows over time. + + +@dataclass +class AffineGroup: + # Affine group-wise quant along the last axis. See native_graph.fbs for the full + # dequant and storage contract (unsigned qdata, quant_min offset, out-of-line + # scales, packed byte length). + scale_data_key: str + scale_dtype: ScalarType + quant_min: int + quant_max: int + group_size: int = 0 + zero_point_data_key: Optional[str] = None + zero_point_dtype: ScalarType = ScalarType.INT + + +@dataclass +class PackedQuant: + # Opaque codec-defined packed layout (e.g. "gguf:q4k", "mxfp4", "nvfp4"). Block + # size, bit width, and scales are all implied by `codec`. See native_graph.fbs. + codec: str + + +# Append-only; keep in sync with the union in native_graph.fbs. +QuantScheme = Union[ + AffineGroup, + PackedQuant, +] + + +@dataclass +class QuantSpec: + # Wrapper so the scheme union can be optional on TensorMeta: the JSON codec only + # supports required string-annotated unions, so an optional union is nested in + # this table (cf. EValue/KernelTypes in exir). + scheme: "QuantScheme" + + +@dataclass +class Dim: + # A tensor dimension range. Static: min == max. Dynamic: min < max, or max < 0 + # for unbounded. No symbols; the runtime uses concrete shapes and plans from max. + min: int + max: int = -1 + + +@dataclass +class TensorMeta: + # sizes is the logical shape, dtype the element type. When quant is set the tensor + # is packed (dtype BYTE) and the physical layout and byte length come from the + # quant scheme and its external constant blob, not prod(sizes). + dtype: ScalarType + sizes: List[Dim] + # Memory layout as a permutation of dim indices, outermost first (contiguous is + # [0, 1, ..., n-1]). Independent of the sizes. Empty means contiguous. + dim_order: Optional[List[int]] = None + quant: Optional[QuantSpec] = None + + +@dataclass +class TensorValue: + name: str + meta: TensorMeta + + +# --------------------------------------------------------------------------- +# Argument union members. Class names must match the .fbs table names. +# --------------------------------------------------------------------------- + + +@dataclass +class TensorArg: + name: str + + +@dataclass +class NoneArg: + pass + + +@dataclass +class IntArg: + value: int + # Non-empty means a reference to an in-graph int value by SSA name (a sym_size or + # arith node), and value is ignored. Empty means the literal value. + ref: Optional[str] = None + + +@dataclass +class FloatArg: + value: float + # Non-empty ref references an in-graph float value by SSA name; value is ignored. + ref: Optional[str] = None + + +@dataclass +class BoolArg: + value: bool + # Non-empty ref references an in-graph bool value by SSA name; value is ignored. + ref: Optional[str] = None + + +@dataclass +class StringArg: + value: str + + +@dataclass +class ScalarTypeArg: + value: ScalarType + + +@dataclass +class IntListArg: + values: List[int] + # Parallel to values: a non-empty refs[i] makes element i a reference to an + # in-graph int value by SSA name (values[i] ignored). Empty means all literal. + refs: Optional[List[str]] = None + + +@dataclass +class FloatListArg: + values: List[float] + + +@dataclass +class BoolListArg: + values: List[bool] + + +@dataclass +class TensorListArg: + names: List[str] + + +@dataclass +class OptionalTensorListArg: + names: List[str] + has_value: List[bool] + + +@dataclass +class Argument: + # String annotation so the encoder emits the union discriminator. ArgumentValue is + # defined after Graph because GraphArg holds a nested Graph (recursive schema). + value: "ArgumentValue" + + +@dataclass +class NamedArgument: + arg: Argument + name: Optional[str] = None + # True if the op writes this input in-place (schema Tensor(a!)). + mutated: bool = False + + +@dataclass +class KeyValue: + key: str + value: str + + +# `alias_of`, when set, is the SSA name of an input value this output shares storage +# with (op-schema view annotation). `kind` discriminates the produced value: a +# TENSOR_LIST (`Tensor[]` return) carries its element SSA names in `names` with an +# empty `name`; INT/BOOL/FLOAT are scalar returns named by `name`. +@dataclass +class Output: + name: str + alias_of: Optional[str] = None + kind: OutputValueKind = OutputValueKind.TENSOR + names: Optional[List[str]] = None + + +@dataclass +class Node: + name: str + op_kind: OpKind + target: Optional[str] = None + inputs: Optional[List[NamedArgument]] = None + outputs: Optional[List[Output]] = None + metadata: Optional[List[KeyValue]] = None + + +@dataclass +class NamedTensorRef: + name: str + data_key: str + meta: TensorMeta + kind: InputKind = InputKind.CONSTANT_TENSOR + mutated: bool = False + + +@dataclass +class OutputSpec: + name: str + kind: OutputKind = OutputKind.USER_OUTPUT + target: Optional[str] = None + + +# Mutable buffer that is graph state but not data-backed (e.g. a zero-initialized KV +# cache). Shape and dtype live in tensor_values; cross-method sharing is by fqn. +@dataclass +class MutableBufferSpec: + name: str + fqn: str + + +@dataclass +class Graph: + # Pure function body, used for both a top-level method and a HOP subgraph. + # Method-level bindings (constants, output_specs, mutable_buffers) live on Method. + nodes: List[Node] + inputs: Optional[List[str]] = None + outputs: Optional[List[str]] = None + tensor_values: Optional[List[TensorValue]] = field(default=None) + + +# Subgraph passed to a higher-order op (torch.cond / while_loop / map), inlined as a +# nested Graph. Defined after Graph so `graph` is a real class reference; a string +# forward-ref would not deserialize as a nested dataclass. +@dataclass +class GraphArg: + name: str + graph: Graph + + +# Append-only; keep in sync with the union in native_graph.fbs. +ArgumentValue = Union[ + TensorArg, + NoneArg, + IntArg, + FloatArg, + BoolArg, + StringArg, + ScalarTypeArg, + IntListArg, + FloatListArg, + BoolListArg, + TensorListArg, + OptionalTensorListArg, + GraphArg, +] + + +@dataclass +class Method: + name: str + graph: Graph + # Constant/parameter/buffer references binding a method-graph placeholder to a + # data_key. A HOP subgraph has none; its params are lifted here as operands. + constants: Optional[List[NamedTensorRef]] = None + # Per-output classification (user output vs buffer/user-input mutation), parallel + # to graph.outputs by name. + output_specs: Optional[List[OutputSpec]] = None + # Non-persistent mutable buffers (e.g. KV caches): no shipped data, zero-init. + mutable_buffers: Optional[List[MutableBufferSpec]] = None + + +@dataclass +class Program: + methods: List[Method] + version: Optional[str] = None diff --git a/backends/native/test/BUCK b/backends/native/test/BUCK new file mode 100644 index 00000000000..c9339aeba68 --- /dev/null +++ b/backends/native/test/BUCK @@ -0,0 +1,15 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target( + _kind = runtime.python_test, + name = "test_serialize", + srcs = ["test_serialize.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/native:lib", + "//executorch/exir:lib", + ], +) diff --git a/backends/native/test/__init__.py b/backends/native/test/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/native/test/test_serialize.py b/backends/native/test/test_serialize.py new file mode 100644 index 00000000000..afa4e8b64c2 --- /dev/null +++ b/backends/native/test/test_serialize.py @@ -0,0 +1,1022 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from collections import namedtuple +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from executorch.backends.native.serialization import ( + deserialize_graph, + deserialize_program, + serialize_graph, + serialize_program, + validate_graph, + validate_method, + validate_program, +) +from executorch.backends.native.serialization.graph_serialize import ( + _build_output_specs, + _compile_to_bytes, + _dim_order, + _named_arguments, + _node_outputs, + _output_alias_of, + _to_arg_value, + serialize_operator, +) +from executorch.backends.native.serialization.schema import ( + AffineGroup, + Argument, + BoolArg, + BoolListArg, + Dim, + FloatArg, + FloatListArg, + Graph, + GraphArg, + InputKind, + IntArg, + IntListArg, + Method, + MutableBufferSpec, + NamedArgument, + Node, + NoneArg, + OpKind, + OptionalTensorListArg, + Output, + OutputKind, + OutputValueKind, + PackedQuant, + Program, + QuantSpec, + ScalarType, + ScalarTypeArg, + StringArg, + TensorArg, + TensorListArg, + TensorMeta, + TensorValue, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as edge_ops + + +class _Add(nn.Module): + def forward(self, x, y): + return x + y + + +class _Counter(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("count", torch.zeros(1)) + + def forward(self, x): + self.count.add_(1) + return x + self.count + + +class _KVCache(nn.Module): + # A non-persistent (not saved in state_dict) mutable buffer, like a KV cache. + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(4), persistent=False) + + def forward(self, x): + self.cache.add_(x) + return self.cache + 1.0 + + +_ADD_INPUTS = (torch.randn(2, 3), torch.randn(2, 3)) + +# Sentinel marking a positional slot that should be a fresh tensor placeholder. +_T = object() + +_Round = namedtuple("_Round", ["edge_ep", "method", "graph", "data", "constants"]) + + +def _roundtrip(model, example_inputs, dynamic_shapes=None) -> _Round: + ep = torch.export.export(model, example_inputs, dynamic_shapes=dynamic_shapes) + edge_ep = to_edge(ep).exported_program() + data, constants = serialize_graph( + edge_ep.graph_module, + edge_ep.graph_signature, + edge_ep.state_dict, + edge_ep.constants, + ) + method = deserialize_program(data).methods[0] + return _Round(edge_ep, method, method.graph, data, constants) + + +def _edge_method(model, example_inputs): + """(graph_module, signature, state_dict, constants) tuple for serialize_program.""" + edge_ep = to_edge(torch.export.export(model, example_inputs)).exported_program() + return ( + edge_ep.graph_module, + edge_ep.graph_signature, + edge_ep.state_dict, + edge_ep.constants, + ) + + +def _call_targets(graph): + return [n.target for n in graph.nodes if n.op_kind == OpKind.CALL_FUNCTION] + + +class SerializeRoundTripTest(unittest.TestCase): + def test_add_target_roundtrips(self): + graph = _roundtrip(_Add(), _ADD_INPUTS).graph + self.assertIn("torch.ops.aten.add.Tensor", _call_targets(graph)) + + def test_topological_order_preserved(self): + class Model(nn.Module): + def forward(self, x): + return torch.relu(x) + 1.0 + + r = _roundtrip(Model(), (torch.randn(2, 2),)) + expected = [n.name for n in r.edge_ep.graph_module.graph.nodes] + self.assertEqual(expected, [n.name for n in r.graph.nodes]) + + def test_file_identifier(self): + self.assertEqual(_roundtrip(_Add(), _ADD_INPUTS).data[4:8], b"NPTG") + + def test_placeholder_and_output_nodes(self): + graph = _roundtrip(_Add(), _ADD_INPUTS).graph + kinds = {n.op_kind for n in graph.nodes} + self.assertIn(OpKind.PLACEHOLDER, kinds) + self.assertIn(OpKind.OUTPUT, kinds) + self.assertIn(OpKind.CALL_FUNCTION, kinds) + self.assertTrue(graph.inputs) + self.assertTrue(graph.outputs) + + def test_input_tensor_metadata_recorded(self): + graph = _roundtrip(_Add(), _ADD_INPUTS).graph + by_name = {tv.name: tv for tv in graph.tensor_values or []} + meta = by_name[graph.inputs[0]].meta + self.assertEqual(meta.dtype, ScalarType.FLOAT) + self.assertEqual([d.max for d in meta.sizes], [2, 3]) + + def test_intermediate_metadata_emitted(self): + # Every tensor-valued node (not just graph I/O) records shape + dtype, so a + # consumer never re-infers intermediate metadata. + class Model(nn.Module): + def forward(self, x): + return torch.relu(x) + 1.0 + + graph = _roundtrip(Model(), (torch.randn(2, 2),)).graph + meta_names = {tv.name for tv in (graph.tensor_values or [])} + call_outputs = { + n.name for n in graph.nodes if n.op_kind == OpKind.CALL_FUNCTION + } + self.assertTrue(call_outputs) + # includes the intermediate relu output, not only the final add output + self.assertTrue(call_outputs <= meta_names) + + def test_scalar_and_tensor_list_args(self): + class CatModel(nn.Module): + def forward(self, x, y): + return torch.cat([x, y], dim=1) + + graph = _roundtrip(CatModel(), _ADD_INPUTS).graph + arg_types = {type(na.arg.value) for n in graph.nodes for na in (n.inputs or [])} + self.assertIn(TensorListArg, arg_types) + self.assertIn(IntArg, arg_types) + + def test_constants_referenced_by_fqn(self): + r = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)) + fqns = {c.data_key for c in r.method.constants} + self.assertTrue(any("weight" in f for f in fqns)) + self.assertTrue(any("bias" in f for f in fqns)) + # Raw data is returned separately, keyed by the same fqns. + self.assertEqual(set(r.constants.keys()), fqns) + + def test_tensor_args_reference_by_name(self): + graph = _roundtrip(_Add(), _ADD_INPUTS).graph + valid = {n.name for n in graph.nodes} + for node in graph.nodes: + for na in node.inputs or []: + if isinstance(na.arg.value, TensorArg): + self.assertIn(na.arg.value.name, valid) + + +class InputClassificationTest(unittest.TestCase): + def test_parameter_classified_and_not_mutated(self): + method = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)).method + weight = next(c for c in method.constants if "weight" in c.data_key) + self.assertEqual(weight.kind, InputKind.PARAMETER) + self.assertFalse(weight.mutated) + + def test_mutated_buffer_flagged(self): + method = _roundtrip(_Counter(), (torch.randn(1),)).method + count = next(c for c in method.constants if "count" in c.data_key) + self.assertEqual(count.kind, InputKind.BUFFER) + self.assertTrue(count.mutated) + + +class SerializeOperatorTest(unittest.TestCase): + def test_plain_op_overload(self): + # A plain aten OpOverload (e.g. an aliasing op) must serialize to its real + # name, not the bare "torch._ops.aten." from over-unwrapping its `_op`. + self.assertEqual( + serialize_operator(torch.ops.aten.unsqueeze.default), + "torch.ops.aten.unsqueeze.default", + ) + + def test_sym_size_overload(self): + self.assertEqual( + serialize_operator(torch.ops.aten.sym_size.int), + "torch.ops.aten.sym_size.int", + ) + + def test_edge_op_unwraps_to_aten(self): + self.assertEqual( + serialize_operator(edge_ops.edge.aten.mul.Tensor), + "torch.ops.aten.mul.Tensor", + ) + + +class OutputSpecTest(unittest.TestCase): + def test_user_output_classified(self): + method = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)).method + self.assertTrue(method.output_specs) + self.assertTrue( + all(s.kind == OutputKind.USER_OUTPUT for s in method.output_specs) + ) + + def test_buffer_mutation_classified(self): + method = _roundtrip(_Counter(), (torch.randn(1),)).method + muts = [s for s in method.output_specs if s.kind == OutputKind.BUFFER_MUTATION] + self.assertTrue(muts, "expected a BUFFER_MUTATION output spec") + self.assertTrue(any(s.target and "count" in s.target for s in muts)) + users = [s for s in method.output_specs if s.kind == OutputKind.USER_OUTPUT] + self.assertEqual(len(users), 1) + + def test_parameter_mutation_rejected(self): + sig = SimpleNamespace( + buffers_to_mutate={}, + user_inputs_to_mutate={}, + parameters_to_mutate={"w": "w"}, + ) + with self.assertRaisesRegex(ValueError, "parameter mutation is not supported"): + _build_output_specs([], sig) + + +class ValidateGraphTest(unittest.TestCase): + def test_self_contained_passes(self): + r = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)) + # validate_method returns None and raises on any inconsistency; a + # self-contained method must validate cleanly. + self.assertIsNone(validate_method(r.method, set(r.constants.keys()))) + + def test_missing_constant_data_raises(self): + r = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)) + with self.assertRaises(ValueError): + validate_method(r.method, available_data_keys=set()) + + @staticmethod + def _meta() -> TensorMeta: + return TensorMeta( + dtype=ScalarType.FLOAT, sizes=[Dim(min=1, max=1)], dim_order=[0] + ) + + @staticmethod + def _placeholder(name: str) -> Node: + return Node(name=name, op_kind=OpKind.PLACEHOLDER, outputs=[Output(name=name)]) + + @classmethod + def _call(cls, name: str, inputs=None) -> Node: + return Node( + name=name, + op_kind=OpKind.CALL_FUNCTION, + target="aten::relu", + inputs=inputs, + outputs=[Output(name=name)], + ) + + def test_placeholder_restating_input_passes(self): + g = Graph( + nodes=[ + self._placeholder("x"), + self._call( + "y", + [NamedArgument(name="self", arg=Argument(value=TensorArg("x")))], + ), + ], + inputs=["x"], + outputs=["y"], + tensor_values=[ + TensorValue(name="x", meta=self._meta()), + TensorValue(name="y", meta=self._meta()), + ], + ) + self.assertIsNone(validate_graph(g)) + + def test_duplicate_input_raises(self): + g = Graph(nodes=[], inputs=["x", "x"]) + with self.assertRaisesRegex(ValueError, "duplicate value definition 'x'"): + validate_graph(g) + + def test_node_redefining_input_raises(self): + g = Graph( + nodes=[self._call("x")], + inputs=["x"], + tensor_values=[TensorValue(name="x", meta=self._meta())], + ) + with self.assertRaisesRegex(ValueError, "duplicate value definition 'x'"): + validate_graph(g) + + def test_duplicate_node_raises(self): + g = Graph(nodes=[self._call("y"), self._call("y")]) + with self.assertRaisesRegex(ValueError, "duplicate node name"): + validate_graph(g) + + def test_undeclared_placeholder_raises(self): + g = Graph(nodes=[self._placeholder("z")]) + with self.assertRaisesRegex( + ValueError, "not a declared input or external binding" + ): + validate_graph(g) + + def test_duplicate_placeholder_raises(self): + g = Graph( + nodes=[self._placeholder("a"), self._placeholder("a")], + inputs=["a"], + tensor_values=[TensorValue(name="a", meta=self._meta())], + ) + with self.assertRaisesRegex(ValueError, "duplicate node name"): + validate_graph(g) + + def test_optional_tensor_list_length_mismatch_raises(self): + g = Graph( + nodes=[ + self._placeholder("a"), + self._call( + "y", + [ + NamedArgument( + name="indices", + arg=Argument( + value=OptionalTensorListArg( + names=["a", "b"], has_value=[True] + ) + ), + ) + ], + ), + ], + inputs=["a"], + tensor_values=[TensorValue(name="a", meta=self._meta())], + ) + with self.assertRaisesRegex(ValueError, "OptionalTensorListArg names length"): + validate_graph(g) + + def test_unresolved_reference_raises(self): + g = Graph( + nodes=[ + self._call( + "y", + [NamedArgument(arg=Argument(value=TensorArg("missing")))], + ) + ], + tensor_values=[TensorValue(name="y", meta=self._meta())], + ) + with self.assertRaisesRegex(ValueError, "unresolved value reference 'missing'"): + validate_graph(g) + + def test_int_list_refs_length_mismatch_raises(self): + g = Graph( + nodes=[ + self._placeholder("a"), + self._call( + "y", + [ + NamedArgument( + arg=Argument(value=IntListArg(values=[1, 2], refs=["a"])) + ) + ], + ), + ], + inputs=["a"], + tensor_values=[TensorValue(name="a", meta=self._meta())], + ) + with self.assertRaisesRegex(ValueError, "IntListArg refs length"): + validate_graph(g) + + def test_input_missing_metadata_raises(self): + g = Graph(nodes=[self._placeholder("x")], inputs=["x"]) + with self.assertRaisesRegex(ValueError, "input 'x' missing tensor metadata"): + validate_graph(g) + + def test_output_missing_metadata_raises(self): + g = Graph( + nodes=[ + self._placeholder("x"), + self._call( + "y", + [NamedArgument(arg=Argument(value=TensorArg("x")))], + ), + ], + inputs=["x"], + outputs=["y"], + tensor_values=[TensorValue(name="x", meta=self._meta())], + ) + with self.assertRaisesRegex(ValueError, "output 'y' missing tensor metadata"): + validate_graph(g) + + def test_mutable_buffer_missing_metadata_raises(self): + method = Method( + name="forward", + graph=Graph( + nodes=[self._placeholder("x")], + inputs=["x"], + tensor_values=[TensorValue(name="x", meta=self._meta())], + ), + mutable_buffers=[MutableBufferSpec(name="cache", fqn="m.cache")], + ) + with self.assertRaisesRegex( + ValueError, "mutable buffer 'cache' .* missing tensor metadata" + ): + validate_method(method) + + +class DynamicShapeTest(unittest.TestCase): + def _dynamic_view(self) -> _Round: + class M(nn.Module): + def forward(self, x): + return x.view(x.shape[0], -1) + 1.0 + + return _roundtrip( + M(), + (torch.randn(4, 8),), + dynamic_shapes={"x": {0: torch.export.Dim("b")}}, + ) + + def test_symbolic_size_list_roundtrips(self): + # A dynamic view size like [s0, -1] serializes as an IntListArg whose `refs` + # marks the symbolic element as a reference to an in-graph int value (the + # sym_size node), keeping the -1 as a literal, not collapsed or dropped. + graph = self._dynamic_view().graph + int_lists = [ + na.arg.value + for n in graph.nodes + for na in (n.inputs or []) + if isinstance(na.arg.value, IntListArg) + ] + ref_lists = [il for il in int_lists if il.refs and any(il.refs)] + self.assertTrue( + ref_lists, "expected an IntListArg with a ref for the dynamic view size" + ) + il = ref_lists[0] + # refs is parallel to values; the dynamic element has a non-empty ref, the + # -1 literal has an empty ref. + self.assertEqual(len(il.refs), len(il.values)) + self.assertTrue(any(r for r in il.refs)) + self.assertTrue(any(not r for r in il.refs)) + + def test_dynamic_dim_not_frozen_in_tensor_meta(self): + # int(sym) would specialize to the hint and freeze the dim; TensorMeta must + # keep it as a range (min != max), not a single concrete value. + graph = self._dynamic_view().graph + dynamic_dims = [ + d + for tv in (graph.tensor_values or []) + for d in tv.meta.sizes + if d.max != d.min + ] + self.assertTrue(dynamic_dims, "dynamic dim was frozen (min==max) in TensorMeta") + + +class ArgSerializationTest(unittest.TestCase): + def test_unrepresentable_scalar_arg_raises(self): + class Weird: + pass + + with self.assertRaises(ValueError): + _to_arg_value(Weird()) + + def test_device_and_layout_serialize_as_string(self): + self.assertIsInstance(_to_arg_value(torch.device("cpu")), StringArg) + self.assertIsInstance(_to_arg_value(torch.strided), StringArg) + self.assertIsInstance(_to_arg_value(torch.contiguous_format), StringArg) + + def test_default_args_are_materialized(self): + # aten.add.Tensor has a kwarg-only `alpha=1` default that x+y never passes; + # the serialized node must still carry it so the graph is self-describing. + graph = _roundtrip(_Add(), (torch.randn(3), torch.randn(3))).graph + add = next( + n for n in graph.nodes if n.target and n.target.endswith("add.Tensor") + ) + self.assertIn("alpha", {na.name for na in (add.inputs or [])}) + + def test_noncontiguous_constant_meta_matches_shipped_data(self): + # A non-contiguous constant is shipped contiguous, so its serialized layout + # (dim_order) must describe the contiguous layout, not the original view. + class M(nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.randn(8, 4).transpose(0, 1)) + + def forward(self, x): + return x + self.w + + method = _roundtrip(M(), (torch.randn(4, 8),)).method + w = next(c for c in method.constants if "w" in c.data_key) + self.assertEqual(w.meta.dim_order, [0, 1]) + + +class DimOrderTest(unittest.TestCase): + def test_contiguous(self): + self.assertEqual(_dim_order(torch.randn(2, 3, 4)), [0, 1, 2]) + + def test_permuted_contiguous(self): + t = torch.randn(2, 3, 4).permute(0, 2, 1) + self.assertEqual(_dim_order(t), [0, 2, 1]) + + def test_channels_last_with_size_one_channel(self): + # channels-last leaves the size-1 channel with an arbitrary stride, which + # must not be treated as a non-expressible layout. + t = torch.randn(2, 1, 3, 4).to(memory_format=torch.channels_last) + self.assertEqual(_dim_order(t), [0, 2, 1, 3]) + + def test_sliced_layout_raises(self): + t = torch.randn(4, 8)[:, :4] + with self.assertRaisesRegex(ValueError, "not dim-order expressible"): + _dim_order(t) + + +class MutationAndAliasTest(unittest.TestCase): + """Mutation and view aliasing are sourced from the op schema (Tensor(a!) / + Tensor(a)), not the op-name convention. These exercise the extraction helpers + directly on hand-built fx nodes (export/edge functionalize these ops away).""" + + @staticmethod + def _call(target, args): + g = torch.fx.Graph() + real = [g.placeholder(f"in{i}") if a is _T else a for i, a in enumerate(args)] + return g.call_function(target, tuple(real)) + + def test_inplace_op_marks_only_written_arg(self): + # add_(Tensor(a!) self, Tensor other, *, Scalar alpha=1): only self is + # mutated even though there are two tensor inputs. + node = self._call(torch.ops.aten.add_.Tensor, [_T, _T]) + mutated = {na.name: na.mutated for na in _named_arguments(node)} + self.assertTrue(mutated["self"]) + self.assertFalse(mutated["other"]) + + def test_inplace_op_output_aliases_written_input(self): + node = self._call(torch.ops.aten.add_.Tensor, [_T, _T]) + self.assertEqual(_output_alias_of(node), node.args[0].name) + + def test_view_op_output_aliases_input_without_mutation(self): + # view(Tensor(a) self, SymInt[] size) -> Tensor(a): read-only storage share. + node = self._call(torch.ops.aten.view.default, [_T, [6]]) + self.assertEqual(_output_alias_of(node), node.args[0].name) + self.assertFalse(any(na.mutated for na in _named_arguments(node))) + + def test_functional_op_has_no_alias_or_mutation(self): + node = self._call(torch.ops.aten.add.Tensor, [_T, _T]) + self.assertIsNone(_output_alias_of(node)) + self.assertFalse(any(na.mutated for na in _named_arguments(node))) + + def test_view_alias_survives_roundtrip(self): + g = torch.fx.Graph() + x = g.placeholder("x") + x.meta["val"] = torch.zeros(2, 3) + v = g.call_function(torch.ops.aten.view.default, (x, [6])) + v.meta["val"] = torch.zeros(6) + g.output((v,)) + gm = torch.fx.GraphModule(nn.Module(), g) + + data, _ = serialize_graph(gm, object(), {}, None) + graph = deserialize_graph(data) + view = next( + n for n in graph.nodes if n.target and n.target.endswith("view.default") + ) + self.assertEqual(view.outputs[0].alias_of, "x") + + +class SubgraphHOPTest(unittest.TestCase): + def _cond_program(self): + class CondModel(nn.Module): + def forward(self, pred, x): + def true_fn(x): + return x + 1.0 + + def false_fn(x): + return x - 1.0 + + return torch.cond(pred, true_fn, false_fn, (x,)) + + ep = torch.export.export(CondModel(), (torch.tensor(True), torch.randn(3))) + data, _ = serialize_graph( + ep.graph_module, ep.graph_signature, ep.state_dict, ep.constants + ) + return deserialize_graph(data) + + def test_cond_branches_serialized_as_inlined_subgraphs(self): + graph = self._cond_program() + graphargs = [ + na.arg.value + for n in graph.nodes + for na in (n.inputs or []) + if isinstance(na.arg.value, GraphArg) + ] + # torch.cond has a true and a false branch, each an inlined subgraph. + self.assertEqual(len(graphargs), 2) + self.assertTrue(all(ga.graph.nodes for ga in graphargs)) + + def test_cond_node_present_and_get_attr_dropped(self): + graph = self._cond_program() + self.assertTrue( + any(n.target and n.target.endswith("cond") for n in graph.nodes) + ) + # The get_attr nodes that named the branch submodules are inlined away: their + # names appear as GraphArg labels but never as nodes. + node_names = {n.name for n in graph.nodes} + ga_names = { + na.arg.value.name + for n in graph.nodes + for na in (n.inputs or []) + if isinstance(na.arg.value, GraphArg) + } + self.assertTrue(ga_names) + self.assertTrue(ga_names.isdisjoint(node_names)) + + def test_cond_graph_validates(self): + # Inlined subgraphs must themselves be self-contained (references resolve, + # I/O has tensor metadata); validate recurses into every GraphArg. + self.assertIsNone(validate_graph(self._cond_program())) + + def test_cond_symbool_predicate_is_bool_ref(self): + class M(nn.Module): + def forward(self, x): + def tf(t): + return t + 1 + + def ff(t): + return t - 1 + + return torch.cond(x.shape[0] > 4, tf, ff, (x,)) + + ep = torch.export.export( + M(), + (torch.randn(6),), + dynamic_shapes={"x": {0: torch.export.Dim("b", min=2, max=16)}}, + ) + data, _ = serialize_graph( + ep.graph_module, ep.graph_signature, ep.state_dict, ep.constants + ) + graph = deserialize_graph(data) + cond_node = next( + n for n in graph.nodes if n.target and n.target.endswith("cond") + ) + pred = cond_node.inputs[0].arg.value + self.assertIsInstance(pred, BoolArg) + self.assertTrue(pred.ref) + bool_out_names = { + o.name + for n in graph.nodes + for o in (n.outputs or []) + if o.kind == OutputValueKind.BOOL + } + self.assertIn(pred.ref, bool_out_names) + validate_graph(graph) + + def test_subgraph_in_list_arg_fails_loud(self): + # A HOP subgraph can only be inlined as a direct GraphArg. If a subgraph + # get_attr node shows up inside a list arg it must raise, not emit a dangling + # TensorListArg reference (the get_attr node is dropped from the graph). + g = torch.fx.Graph() + attr = g.get_attr("sub") + sg = torch.fx.Graph() + p = sg.placeholder("x") + sg.output((p,)) + sub_gm = torch.fx.GraphModule(nn.Module(), sg) + with self.assertRaises(ValueError): + _to_arg_value([attr], subgraph_map={attr.name: sub_gm}) + + +class MultiMethodTest(unittest.TestCase): + def test_program_bundles_named_methods(self): + methods = { + "add": _edge_method(_Add(), _ADD_INPUTS), + "linear": _edge_method(nn.Linear(4, 4), (torch.randn(1, 4),)), + } + data, constants = serialize_program(methods) + program = deserialize_program(data) + self.assertEqual({m.name for m in program.methods}, {"add", "linear"}) + validate_program(program, set(constants.keys())) + + def test_shared_constant_fqn_deduped_across_methods(self): + shared = nn.Linear(4, 4) + data, constants = serialize_program( + { + "a": _edge_method(shared, (torch.randn(1, 4),)), + "b": _edge_method(shared, (torch.randn(2, 4),)), + } + ) + program = deserialize_program(data) + # Both methods bind the weight/bias fqns; data is merged (deduped) by fqn. + fqns_a = {c.data_key for c in program.methods[0].constants} + fqns_b = {c.data_key for c in program.methods[1].constants} + self.assertEqual(fqns_a, fqns_b) + self.assertEqual(set(constants.keys()), fqns_a) + + def test_serialize_graph_is_single_forward_method(self): + data, _ = serialize_graph(*_edge_method(_Add(), _ADD_INPUTS)) + program = deserialize_program(data) + self.assertEqual([m.name for m in program.methods], ["forward"]) + + def test_conflicting_constant_fqn_across_methods_raises(self): + # Two independent Linear(4, 4) instances share fqns ("weight"/"bias") but + # hold different random data, so serialize_program must reject rather than + # silently clobber one method's data. + with self.assertRaises(ValueError): + serialize_program( + { + "a": _edge_method(nn.Linear(4, 4), (torch.randn(1, 4),)), + "b": _edge_method(nn.Linear(4, 4), (torch.randn(1, 4),)), + } + ) + + +class MutableBufferTest(unittest.TestCase): + def test_non_persistent_buffer_recorded_without_data(self): + r = _roundtrip(_KVCache(), (torch.randn(4),)) + method = r.method + self.assertTrue(method.mutable_buffers, "expected a non-persistent buffer") + mb = method.mutable_buffers[0] + self.assertIn("cache", mb.fqn) + # It is NOT a data-backed constant and ships no bytes. + self.assertNotIn("cache", {c.data_key for c in (method.constants or [])}) + self.assertNotIn(mb.fqn, r.constants) + # Shape/dtype are still available via the tensor_values side table. + self.assertIn(mb.name, {tv.name for tv in (method.graph.tensor_values or [])}) + + def test_mutable_buffer_graph_validates(self): + r = _roundtrip(_KVCache(), (torch.randn(4),)) + # Mutable buffers are exempt from the constant-data-keys check. + self.assertIsNone(validate_method(r.method, set(r.constants.keys()))) + + +class QuantSpecRoundTripTest(unittest.TestCase): + """The quant spec rides on TensorMeta. No producer sets it yet (the QDQ-fold + pass is future work), so these build the dataclasses directly and round-trip + them through flatc to lock the schema + optional-union wrapper.""" + + def _roundtrip_meta(self, quant) -> TensorMeta: + meta = TensorMeta( + dtype=ScalarType.CHAR, + sizes=[Dim(min=4, max=4)], + dim_order=[0], + quant=quant, + ) + graph = Graph( + nodes=[Node(name="w", op_kind=OpKind.PLACEHOLDER)], + tensor_values=[TensorValue(name="w", meta=meta)], + ) + program = Program(version="1", methods=[Method(name="forward", graph=graph)]) + out = deserialize_program(_compile_to_bytes(program)) + return out.methods[0].graph.tensor_values[0].meta + + def test_absent_quant_is_none(self): + self.assertIsNone(self._roundtrip_meta(None).quant) + + def test_affine_group_quant_roundtrips(self): + expected = AffineGroup( + scale_data_key="w.scale", + scale_dtype=ScalarType.HALF, + quant_min=-8, + quant_max=7, + group_size=32, + zero_point_data_key="w.zp", + zero_point_dtype=ScalarType.INT, + ) + scheme = self._roundtrip_meta(QuantSpec(scheme=expected)).quant.scheme + self.assertEqual(scheme, expected) + + def test_packed_quant_roundtrips(self): + # Weight-only formats (GGUF, MXFP4, NVFP4) carry only an opaque codec name. + scheme = self._roundtrip_meta( + QuantSpec(scheme=PackedQuant(codec="gguf:q4k")) + ).quant.scheme + self.assertIsInstance(scheme, PackedQuant) + self.assertEqual(scheme.codec, "gguf:q4k") + + +class NonTensorInputTest(unittest.TestCase): + def test_non_tensor_user_input_fails_loud(self): + # A scalar (non-tensor) user input cannot be represented, so serialization + # must raise rather than emit a graph with an untyped/valueless input. + class M(nn.Module): + def forward(self, x, n): + return x + n + + edge = to_edge(torch.export.export(M(), (torch.randn(3), 2))).exported_program() + with self.assertRaises(ValueError): + serialize_graph( + edge.graph_module, + edge.graph_signature, + edge.state_dict, + edge.constants, + ) + + +class MultiOutputTest(unittest.TestCase): + def test_topk_outputs_and_metadata(self): + class M(nn.Module): + def forward(self, x): + vals, idx = torch.topk(x, 2) + return vals + 1, idx + + r = _roundtrip(M(), (torch.randn(4),)) + graph = r.graph + producer = next( + n for n in graph.nodes if n.target and n.target.endswith("topk.default") + ) + # topk returns (values, indices): two distinctly named outputs. + self.assertEqual(len(producer.outputs), 2) + self.assertEqual(len({o.name for o in producer.outputs}), 2) + # Each result carries its own tensor metadata. + meta_names = {tv.name for tv in (graph.tensor_values or [])} + for out in producer.outputs: + self.assertIn(out.name, meta_names) + # getitem users are folded into the producer's outputs, not emitted as nodes. + self.assertFalse( + any(n.target and n.target.endswith("getitem") for n in graph.nodes) + ) + validate_method(r.method, set(r.constants.keys())) + + def test_split_tensor_list_output(self): + class M(nn.Module): + def forward(self, x): + a, b = torch.split(x, 2) + return a + 1, b + 1 + + r = _roundtrip(M(), (torch.randn(4),)) + graph = r.graph + producer = next(n for n in graph.nodes if n.target and "split" in n.target) + # A Tensor[] return is one TENSOR_LIST output holding the element names. + self.assertEqual(len(producer.outputs), 1) + out = producer.outputs[0] + self.assertEqual(out.kind, OutputValueKind.TENSOR_LIST) + self.assertEqual(len(out.names or []), 2) + meta_names = {tv.name for tv in (graph.tensor_values or [])} + for nm in out.names: + self.assertIn(nm, meta_names) + self.assertFalse( + any(n.target and n.target.endswith("getitem") for n in graph.nodes) + ) + validate_method(r.method, set(r.constants.keys())) + + def test_scalar_node_output_kind(self): + class M(nn.Module): + def forward(self, x): + n = x.shape[0] + return x + n + + r = _roundtrip( + M(), + (torch.randn(4, 3),), + dynamic_shapes={"x": {0: torch.export.Dim("b")}}, + ) + scalar_outs = [ + o + for n in r.graph.nodes + for o in (n.outputs or []) + if o.kind == OutputValueKind.INT + ] + self.assertGreaterEqual( + len(scalar_outs), 1, "expected an int-kind scalar output (sym_size)" + ) + + def test_zero_return_op_has_no_outputs(self): + g = torch.fx.Graph() + n = g.call_function(torch.ops.aten._assert_scalar.default, (True, "msg")) + n.meta["val"] = None + self.assertEqual(_node_outputs(n, {}, {}), []) + + +class DynamicScalarNodeArgTest(unittest.TestCase): + # A node producing a dynamic bool/float (e.g. a `x.shape[0] > 4` cond predicate) + # is referenced by SSA name, like a dynamic int, not serialized as a tensor. + def test_float_node_arg_uses_ref(self): + g = torch.fx.Graph() + n = g.placeholder("s") + n.meta["val"] = 1.5 + arg = _to_arg_value(n) + self.assertIsInstance(arg, FloatArg) + self.assertEqual(arg.ref, "s") + + def test_bool_node_arg_uses_ref(self): + g = torch.fx.Graph() + n = g.placeholder("b") + n.meta["val"] = True + arg = _to_arg_value(n) + self.assertIsInstance(arg, BoolArg) + self.assertEqual(arg.ref, "b") + + +class LiteralOutputTest(unittest.TestCase): + def test_literal_output_preserved(self): + class M(nn.Module): + def forward(self, x): + return x + 1, 7 + + r = _roundtrip(M(), (torch.randn(3),)) + out_node = next(n for n in r.graph.nodes if n.op_kind == OpKind.OUTPUT) + arg_vals = [na.arg.value for na in (out_node.inputs or [])] + # The output node's arguments are the full ordered result list. + self.assertEqual(len(arg_vals), 2) + self.assertIsInstance(arg_vals[0], TensorArg) + self.assertIsInstance(arg_vals[1], IntArg) + self.assertEqual(arg_vals[1].value, 7) + # The literal is not a tensor value, so it is absent from Graph.outputs. + self.assertEqual(len(r.graph.outputs or []), 1) + validate_method(r.method, set(r.constants.keys())) + + +class ArgValueDispatchTest(unittest.TestCase): + """Each ArgumentValue variant, exercised in isolation through _to_arg_value.""" + + @staticmethod + def _node(name: str, val: object) -> torch.fx.Node: + n = torch.fx.Graph().placeholder(name) + n.meta["val"] = val + return n + + def test_int_literal(self): + v = _to_arg_value(3) + self.assertIsInstance(v, IntArg) + self.assertEqual(v.value, 3) + self.assertIsNone(v.ref) + + def test_float_literal(self): + v = _to_arg_value(1.5) + self.assertIsInstance(v, FloatArg) + self.assertEqual(v.value, 1.5) + self.assertIsNone(v.ref) + + def test_bool_literal(self): + v = _to_arg_value(True) + self.assertIsInstance(v, BoolArg) + self.assertTrue(v.value) + self.assertIsNone(v.ref) + + def test_none(self): + self.assertIsInstance(_to_arg_value(None), NoneArg) + + def test_string(self): + v = _to_arg_value("nchw") + self.assertIsInstance(v, StringArg) + self.assertEqual(v.value, "nchw") + + def test_scalar_type(self): + v = _to_arg_value(torch.float16) + self.assertIsInstance(v, ScalarTypeArg) + self.assertEqual(v.value, ScalarType.HALF) + + def test_int_list_literal(self): + v = _to_arg_value([1, 2, 3]) + self.assertIsInstance(v, IntListArg) + self.assertEqual(v.values, [1, 2, 3]) + + def test_float_list_literal(self): + v = _to_arg_value([1.0, 2.0]) + self.assertIsInstance(v, FloatListArg) + self.assertEqual(v.values, [1.0, 2.0]) + + def test_bool_list_literal(self): + v = _to_arg_value([True, False]) + self.assertIsInstance(v, BoolListArg) + self.assertEqual(v.values, [True, False]) + + def test_int_node_uses_ref(self): + v = _to_arg_value(self._node("s", 5)) + self.assertIsInstance(v, IntArg) + self.assertEqual(v.ref, "s") + + def test_tensor_node_is_tensor_arg(self): + v = _to_arg_value(self._node("t", torch.zeros(2))) + self.assertIsInstance(v, TensorArg) + self.assertEqual(v.name, "t") + + def test_tensor_list(self): + v = _to_arg_value( + [self._node("a", torch.zeros(2)), self._node("b", torch.zeros(2))] + ) + self.assertIsInstance(v, TensorListArg) + self.assertEqual(v.names, ["a", "b"]) + + def test_optional_tensor_list(self): + v = _to_arg_value([self._node("a", torch.zeros(2)), None]) + self.assertIsInstance(v, OptionalTensorListArg) + self.assertEqual(v.names, ["a", ""]) + self.assertEqual(v.has_value, [True, False]) diff --git a/pytest.ini b/pytest.ini index 05aea9d4da6..949d918a963 100644 --- a/pytest.ini +++ b/pytest.ini @@ -76,6 +76,7 @@ testpaths = # backends backends/apple/coreml/test + backends/native/test backends/test/harness/tests backends/test/suite/tests backends/transforms