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..eb64d7048a2 --- /dev/null +++ b/backends/native/serialization/__init__.py @@ -0,0 +1,25 @@ +# 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_program, +) + +__all__ = [ + "serialize_graph", + "serialize_program", + "deserialize_graph", + "deserialize_program", + "validate_graph", + "validate_program", +] diff --git a/backends/native/serialization/graph_serialize.py b/backends/native/serialization/graph_serialize.py new file mode 100644 index 00000000000..dfdfe5a16d9 --- /dev/null +++ b/backends/native/serialization/graph_serialize.py @@ -0,0 +1,845 @@ +# 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 via the ExecuTorch NamedDataStore, keyed by fqn. + +Uses the "runtime flatc" pattern: the schema and the flatc binary are shipped as +package resources, and flatc is invoked to convert JSON <-> binary. +""" + +import importlib.resources +import json +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, + ConstantRef, + FloatArg, + FloatListArg, + Graph, + GraphArg, + IntArg, + IntListArg, + InputKind as SchemaInputKind, + Method, + MutableBufferSpec, + NamedArgument, + Node, + NoneArg, + OpKind, + OptionalTensorListArg, + Output, + OutputKind, + OutputSpec, + Program, + ScalarType, + ScalarTypeArg, + StringArg, + SymInt, + SymIntArg, + SymIntListArg, + TensorArg, + TensorListArg, + TensorMeta, + TensorValue, +) + +from executorch.exir._serialize._dataclass import _json_to_dataclass +from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile + +from torch.export.graph_signature import InputKind + +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 - mirrors 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 and 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 _sym(x: object) -> SymInt: + if isinstance(x, int): + return SymInt(as_int=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 silently freeze a dynamic shape to its example value.) + concrete = x.node.maybe_as_int() + if concrete is not None: + return SymInt(as_int=concrete) + return SymInt(as_symbol=str(x)) + return SymInt(as_symbol=str(x)) + + +def _tensor_meta(t: torch.Tensor) -> TensorMeta: + sizes = [_sym(s) for s in t.shape] + try: + strides = [_sym(s) for s in t.stride()] + except RuntimeError: + # Some tensors (e.g. sparse) do not expose strides; omit them. + strides = [] + return TensorMeta( + dtype=_scalar_type(t.dtype), + sizes=sizes, + strides=strides, + requires_grad=bool(getattr(t, "requires_grad", False)), + ) + + +# --------------------------------------------------------------------------- +# Argument dispatch. +# --------------------------------------------------------------------------- + + +def _to_list_arg( + items: list[object], schema_type_hint: str | None = None +) -> ArgumentValue: + if len(items) == 0: + # Empty lists are ambiguous (IntList vs TensorList vs FloatList etc.). + # If we have a schema type hint (e.g. "int[]", "Tensor[]"), use it to + # emit a typed empty list. Otherwise fail loud so the caller must handle + # the gap explicitly rather than silently coercing. + 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 "symint" in hint or "sym_int" in hint: + return SymIntListArg(values=[]) + if "int" in hint or hint == "": + # int[] is the most common default (e.g. reduction dims), but still + # require explicit handling for non-schema path: fail loud if no 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." + ) + if all(isinstance(x, torch.fx.Node) for x in items): + return TensorListArg(names=[cast(torch.fx.Node, x).name for x in items]) + if all(x is None or isinstance(x, torch.fx.Node) 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]) + # Symbolic/mixed int lists (e.g. dynamic view/expand sizes like [s0, -1]). + # Symbolic dims arrive either as scalar SymInts or as references to in-graph + # sym_size nodes; serialize each element as a SymInt (as_int for concrete dims, + # as_symbol for a symbol or referenced value name) so they survive the round + # trip instead of being dropped. + if all( + isinstance(x, (int, torch.SymInt, torch.fx.Node)) + and not isinstance(x, bool) + for x in items + ): + values: list[SymInt] = [] + for x in items: + if isinstance(x, torch.fx.Node): + values.append(SymInt(as_symbol=x.name)) + else: + values.append(_sym(x)) + return SymIntListArg(values=values) + # No known variant fits (e.g. a mixed float/None/other list). Fail loud rather + # than silently coercing to ints or dropping to an empty list -- an + # unrepresentable arg means a real schema gap to handle explicitly. + raise ValueError( + f"Cannot serialize list argument with element types " + f"{sorted({type(x).__name__ for x in items})}: {items!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): + # 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) + return GraphArg(name=v.name, graph=sub) + return TensorArg(name=v.name) + 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)) + if isinstance(v, torch.memory_format): + return StringArg(value=str(v)) + if isinstance(v, torch.SymInt): + return SymIntArg(value=_sym(v)) + if isinstance(v, (list, tuple)): + return _to_list_arg(list(v), schema_type_hint=schema_type_hint) + # device/layout are single-target/always-strided for this backend, so a string + # is enough; anything else is an unhandled arg type -- fail loud rather than + # emit a lossy repr. + if isinstance(v, (torch.device, torch.layout)): + return StringArg(value=str(v)) + raise ValueError(f"Cannot serialize argument of type {type(v).__name__}: {v!r}") + + +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 + # Pass schema type string (e.g. "int[]", "Tensor[]") for empty-list disambiguation. + 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, + "get_attr": OpKind.GET_ATTR, + }[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]: + """Map get_attr node name -> submodule GraphModule for HOP subgraphs.""" + 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 + + +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, dict[str, torch.Tensor]]: + """Build a Graph from an fx GraphModule. + + ``graph_signature`` is the ExportGraphSignature for a top-level method, or + ``None`` for a higher-order-op subgraph (which has no user-input/constant + signature -- its placeholders are the operands passed by the enclosing op). + Returns (graph, constant_data); constant_data is empty for subgraphs. + """ + nodes: list[Node] = [] + subgraph_map = _subgraph_map(graph_module) + + # fx value name -> its tensor (from meta['val']). Captured for every node so + # tensor_values includes inputs, outputs, mutable buffers, and all intermediates. + val_by_name: dict[str, torch.Tensor] = {} + + def _all_tensor_values() -> list[TensorValue]: + # Preserve topological order (insertion order of val_by_name). + return [ + TensorValue(name=name, meta=_tensor_meta(tensor)) + for name, tensor in val_by_name.items() + ] + + output_names: list[str] = [] + + for node in graph_module.graph.nodes: + # HOP subgraphs are inlined into the referencing arg as GraphArgs; drop the + # get_attr node that names the submodule so it doesn't dangle. + if node.op == "get_attr" and node.name in subgraph_map: + continue + + kind = _op_kind(node.op) + target = None + inputs: list[NamedArgument] = [] + outputs: list[Output] = [] + + if node.op == "output": + out_vals = node.args[0] if node.args else () + if isinstance(out_vals, (list, tuple)): + for v in out_vals: + inputs.append(NamedArgument(name=None, arg=_argument(v, subgraph_map))) + if isinstance(v, torch.fx.Node): + output_names.append(v.name) + else: + inputs.append( + NamedArgument(name=None, arg=_argument(out_vals, subgraph_map)) + ) + if isinstance(out_vals, torch.fx.Node): + output_names.append(out_vals.name) + elif node.op == "get_attr": + target = str(node.target) + outputs = [Output(name=node.name)] + else: + alias_of = None + if node.op == "call_function": + target = serialize_operator(node.target) + inputs = _named_arguments(node, subgraph_map) + alias_of = _output_alias_of(node) + outputs = [Output(name=node.name, alias_of=alias_of)] + + val = node.meta.get("val") + if isinstance(val, torch.Tensor): + val_by_name[node.name] = val + + nodes.append( + Node( + name=node.name, + op_kind=kind, + target=target, + inputs=inputs or None, + outputs=outputs or None, + ) + ) + + if graph_signature is None: + # Subgraph: placeholders are the operands; no signature-derived data. + user_inputs = [ + node.name for node in graph_module.graph.nodes if node.op == "placeholder" + ] + graph = Graph( + nodes=nodes, + inputs=user_inputs or None, + outputs=output_names or None, + tensor_values=_all_tensor_values() or None, + ) + return graph, {} + + user_inputs = list(getattr(graph_signature, "user_inputs", []) or []) + + # Classify each graph output: real result vs a writeback into a persistent + # buffer / user input. Keyed by the output value's SSA name. + buffers_to_mutate = getattr(graph_signature, "buffers_to_mutate", None) or {} + user_inputs_to_mutate = getattr(graph_signature, "user_inputs_to_mutate", None) or {} + output_specs: list[OutputSpec] = [] + for name in output_names: + if name in buffers_to_mutate: + output_specs.append( + OutputSpec( + name=name, + kind=OutputKind.BUFFER_MUTATION, + target=buffers_to_mutate[name], + ) + ) + elif name in user_inputs_to_mutate: + output_specs.append( + OutputSpec( + name=name, + kind=OutputKind.USER_INPUT_MUTATION, + target=user_inputs_to_mutate[name], + ) + ) + else: + output_specs.append(OutputSpec(name=name, kind=OutputKind.USER_OUTPUT)) + + # Buffers the graph mutates in-place (e.g. KV caches). Keyed by fqn so the + # matching ConstantRef can be flagged as stateful rather than read-only. + mutated_fqns = set(buffers_to_mutate.values()) + + _INPUT_KIND_MAP = { + InputKind.PARAMETER: SchemaInputKind.PARAMETER, + InputKind.BUFFER: SchemaInputKind.BUFFER, + InputKind.CONSTANT_TENSOR: SchemaInputKind.CONSTANT_TENSOR, + } + + constant_refs: list[ConstantRef] = [] + 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 + arg = ispec.arg + name = getattr(arg, "name", None) + target_fqn = ispec.target + if name is None or target_fqn is None: + continue + # Non-persistent buffer (e.g. a zero-initialized KV cache): graph state that + # is zero-initialized at load. torch.export captures its current value in + # `constants`, but it must NOT be shipped as data -- record it as a mutable + # buffer (shape/dtype live in tensor_values; fqn is the cross-method sharing + # identity). Checked before the data lookup precisely because its value is + # present in `constants`. + if ispec.kind == InputKind.BUFFER and getattr(ispec, "persistent", True) is False: + 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 + # Ship contiguous data, and compute meta from the same contiguous tensor so + # the serialized strides match the bytes (a non-contiguous constant view + # would otherwise record strides that disagree with the shipped layout). + tensor = tensor.contiguous() + constant_refs.append( + ConstantRef( + name=name, + fqn=target_fqn, + meta=_tensor_meta(tensor), + kind=_INPUT_KIND_MAP[ispec.kind], + mutated=target_fqn in mutated_fqns, + ) + ) + constant_data[target_fqn] = tensor + + graph = Graph( + nodes=nodes, + inputs=user_inputs or None, + outputs=output_names or None, + tensor_values=_all_tensor_values() or None, + constants=constant_refs or None, + output_specs=output_specs or None, + mutable_buffers=mutable_buffers or None, + ) + return graph, constant_data + + +# --------------------------------------------------------------------------- +# JSON <-> 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() + + +def _same_tensor(a: torch.Tensor, b: torch.Tensor) -> bool: + if a.dtype != b.dtype or a.shape != b.shape: + return False + # Compare on CPU to handle potential device mismatches and ensure a + # deterministic comparison. + ca = a.detach().cpu() + cb = b.detach().cpu() + # torch.equal treats NaN != NaN, so for floating point we compare + # element-wise with an explicit NaN-aware clause: two tensors with identical + # NaN payloads should be considered equal for dedup purposes. + if ca.dtype in (torch.float32, torch.float16, torch.float64, torch.bfloat16): + eq = (ca == cb) | (torch.isnan(ca) & torch.isnan(cb)) + return bool(eq.all().item()) + return bool(torch.equal(ca, cb)) + + +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 via the NamedDataStore. + + Cross-method sharing is by fqn: bundling methods here 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/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, cdata = _build_graph_body( + graph_module, graph_signature, state_dict, constants + ) + method_objs.append(Method(name=name, graph=graph)) + + 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 graph.mutable_buffers or []: + 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 validate_graph(graph: Graph, available_data_keys: set[str] | None = None) -> None: + """Assert the graph is self-contained: every value reference resolves and, + when ``available_data_keys`` is given, every constant's data is present. + + This enforces that a ``.ptn`` (this graph) plus its ``.ptd`` (the named data, + whose keys are ``available_data_keys``) carry everything needed to run -- + nothing extra is required from the ``.pte``. Raises ``ValueError`` on the + first inconsistency. + + Checks: every ``TensorArg`` / tensor-list element references a defined value + (a node output, a user input, a constant, or a mutable buffer); every output + ``alias_of`` resolves; every input/output/mutable-buffer has tensor metadata; + and every ``ConstantRef.fqn`` is available in the data keys. Mutable buffers + are not data-backed, so they are exempt from the data-keys check. + """ + defined: set[str] = {node.name for node in graph.nodes} + defined.update(graph.inputs or []) + defined.update(c.name for c in (graph.constants or [])) + defined.update(mb.name for mb in (graph.mutable_buffers or [])) + meta_names = {tv.name for tv in (graph.tensor_values or [])} + + def check_ref(name: str, ctx: str) -> None: + if name and name not in defined: + raise ValueError(f"{ctx}: unresolved value reference {name!r}") + + for node in graph.nodes: + for out in node.outputs or []: + if out.alias_of: + check_ref(out.alias_of, f"node {node.name!r} output alias") + for na in node.inputs or []: + value = na.arg.value + if isinstance(value, TensorArg): + check_ref(value.name, f"node {node.name!r}") + elif isinstance(value, TensorListArg): + for n in value.names: + check_ref(n, f"node {node.name!r}") + elif isinstance(value, OptionalTensorListArg): + for n, present in zip(value.names, value.has_value): + if present: + check_ref(n, f"node {node.name!r}") + elif isinstance(value, GraphArg): + # HOP subgraph: self-contained (its own placeholder namespace). + validate_graph(value.graph, available_data_keys) + + for name in graph.outputs or []: + check_ref(name, "output") + + 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") + + for mb in graph.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 graph.constants or []: + if c.fqn not in available_data_keys: + raise ValueError( + f"constant {c.name!r} (fqn {c.fqn!r}) has no data in the " + f"provided keys (.ptd/named-data)" + ) + + +def validate_program( + program: Program, available_data_keys: set[str] | None = None +) -> None: + """Validate every method's graph (see ``validate_graph``).""" + for method in program.methods: + validate_graph(method.graph, 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..a36a81c7ea9 --- /dev/null +++ b/backends/native/serialization/native_graph.fbs @@ -0,0 +1,334 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// +// FlatBuffer schema for the ExecuTorch native backend - THIS IS THE SOURCE OF TRUTH. +// +// This is a GENERIC representation of an fx graph: a topological list of nodes, +// where each node carries an op-kind, a target op-name string, and generic +// arguments encoded through a tagged Argument union. It is 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 via the ExecuTorch NamedDataStore and +// referenced by fully-qualified name (ConstantRef.fqn). +// +// 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 + +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, + GET_ATTR = 3, +} + +// 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 ConstantRef 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, +} + +// ============================================================================= +// Symbolic ints and tensor metadata +// ============================================================================= + +// A dimension / value that is either concrete (as_int, with as_symbol empty) or +// symbolic (as_symbol holds the symbol name, e.g. "s0"). +table SymInt { + as_int: long = 0; + as_symbol: string; +} + +table TensorMeta { + dtype: ScalarType; + sizes: [SymInt]; + strides: [SymInt]; + requires_grad: bool = false; +} + +// Side-table entry: maps an fx value name 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. +table TensorArg { + name: string (required); +} + +table NoneArg {} + +table IntArg { + value: long; +} + +table FloatArg { + value: double; +} + +table BoolArg { + value: bool; +} + +table StringArg { + value: string (required); +} + +table ScalarTypeArg { + value: ScalarType; +} + +table SymIntArg { + value: SymInt (required); +} + +table IntListArg { + values: [long] (required); +} + +// A list mixing concrete ints and symbolic dims, e.g. a dynamic view/expand size +// like [1, s0, -1, 64]. Each element is a SymInt (concrete via as_int, or symbolic +// via as_symbol). Used whenever an int-list arg contains at least one SymInt. +table SymIntListArg { + values: [SymInt] (required); +} + +table FloatListArg { + values: [double] (required); +} + +table BoolListArg { + values: [bool] (required); +} + +// A list of tensor references (e.g. cat's input list). +table TensorListArg { + names: [string] (required); +} + +// A list of optional tensor references (Tensor?[]). names[i] is only meaningful +// when has_value[i] is true; otherwise the element is 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 referenced submodule attr in the parent graph; `graph` is the +// branch/body itself, inlined (mirrors torch._export.serde's GraphArgument). +// This makes Graph recursive. +table GraphArg { + name: string (required); + graph: Graph (required); +} + +// BC: APPEND ONLY - new members go at the end. +union ArgumentValue { + TensorArg, + NoneArg, + IntArg, + FloatArg, + BoolArg, + StringArg, + ScalarTypeArg, + SymIntArg, + IntListArg, + FloatListArg, + BoolListArg, + TensorListArg, + OptionalTensorListArg, + SymIntListArg, + GraphArg, +} + +table Argument { + value: ArgumentValue; +} + +// A positional or keyword argument. `name` is the operator-schema parameter name +// (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_). This is the source of truth for + // mutation: the op-name '_' suffix is only a convention and cannot say *which* + // args are written. Sourced AOT from the op schema at serialization time. + 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. + name: string (required); + // If non-empty, this output shares storage with the input value named here + // (op-schema view/alias annotation, e.g. aten.view/slice/transpose whose + // return is Tensor(a)). Whether that sharing is a write is determined by the + // aliased input's NamedArgument.mutated, so it is not duplicated here. Empty + // means the output owns fresh storage. + alias_of: 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 via NamedDataStore, referenced by fqn. +// ============================================================================= + +table ConstantRef { + // Placeholder SSA name in the graph. + name: string (required); + // Fully-qualified name; the key into the NamedDataStore / NamedDataMap. + fqn: 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 Graph.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 Graph.outputs). + name: string (required); + kind: OutputKind; + // For BUFFER_MUTATION: fqn of the mutated buffer. For USER_INPUT_MUTATION: + // the mutated user input's name. Empty for USER_OUTPUT. + target: string; +} + +// A mutable buffer that is graph state but is NOT data-backed: a non-persistent +// buffer (e.g. a zero-initialized KV cache) that is not saved in the state dict. +// It is neither a user input nor a ConstantRef. Its shape/dtype live in +// tensor_values (keyed by `name`); no bytes are shipped (zero-initialized at +// load). Cross-method sharing is by `fqn` -- an ASSERTED, validated contract (the +// bundled methods must come from one model namespace; see serialize_program), +// not an inference from name/fqn equality alone. +table MutableBufferSpec { + // Placeholder SSA name in THIS graph. + name: string (required); + // Stable identity used to correlate the same buffer across methods. + fqn: string (required); +} + +// ============================================================================= +// Root type: Program (one or more named methods) +// ============================================================================= + +// A graph (a single fx graph in topological order). Used both for a method body +// and, recursively, for higher-order-op subgraphs (see GraphArg). +table Graph { + // All nodes, in topological order (fx graph iteration order). + nodes: [Node] (required); + // Names of graph-level input placeholders (user inputs), in order. + inputs: [string]; + // Names of graph-level outputs, in order. + outputs: [string]; + // Per-output classification (user output vs buffer/user-input mutation), + // parallel to `outputs` by name. + output_specs: [OutputSpec]; + // Side table: value name -> tensor metadata. + tensor_values: [TensorValue]; + // Constant/parameter/buffer references. These bind a placeholder SSA name in + // THIS graph to a data fqn, so they live per-graph even though the underlying + // data (NamedDataStore, keyed by fqn) is shared across methods. + constants: [ConstantRef]; + // Non-persistent mutable buffers (e.g. KV caches): graph state with no shipped + // data, zero-initialized at load. See MutableBufferSpec. + mutable_buffers: [MutableBufferSpec]; +} + +// A named method within a program (e.g. "forward", "encode", "decode"). +table Method { + name: string (required); + graph: Graph (required); +} + +// Root: a program bundling one or more methods. Methods share constant data +// (NamedDataStore keyed by fqn); each method's Graph carries its own ConstantRef +// 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..60dbd3be84e --- /dev/null +++ b/backends/native/serialization/schema.py @@ -0,0 +1,279 @@ +# 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. + +These classes are the AOT representation of the generic fx-graph flatbuffer. +They are serialized to JSON via ``executorch.exir._serialize._dataclass._DataclassEncoder`` +and compiled to a flatbuffer with ``flatc`` (see graph_serialize.py). + +Convention (matches executorch/exir/schema.py): + - Union-typed fields MUST be annotated as string literals (e.g. ``"ArgumentValue"``) + so ``_DataclassEncoder`` emits the ``_type`` discriminator that flatc's + JSON union format requires. The union member dataclass names MUST equal the + corresponding flatbuffer table names. + - All other fields use real type annotations. + - Fields that are ``required`` in the .fbs are non-optional here; everything else + is Optional with a default so the flatc --json round-trip (which omits unset + vectors) deserializes cleanly. +""" + +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 + GET_ATTR = 3 + + +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 + + +@dataclass +class SymInt: + as_int: int = 0 + as_symbol: Optional[str] = None + + +@dataclass +class TensorMeta: + dtype: ScalarType + sizes: List[SymInt] + strides: List[SymInt] + requires_grad: bool = False + + +@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 + + +@dataclass +class FloatArg: + value: float + + +@dataclass +class BoolArg: + value: bool + + +@dataclass +class StringArg: + value: str + + +@dataclass +class ScalarTypeArg: + value: ScalarType + + +@dataclass +class SymIntArg: + value: SymInt + + +@dataclass +class IntListArg: + values: List[int] + + +@dataclass +class SymIntListArg: + values: List[SymInt] + + +@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: + # Union types must be specified as strings so _DataclassEncoder can see them. + # ArgumentValue is defined below (after Graph) because one of its members, + # 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 + + +# A value produced by a node. `alias_of`, when set, is the SSA name of an input +# value this output shares storage with (op-schema view annotation). +@dataclass +class Output: + name: str + alias_of: Optional[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 ConstantRef: + name: str + fqn: 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 + + +# A non-data-backed mutable buffer (e.g. a zero-initialized KV cache): graph state +# that is neither a user input nor a ConstantRef. Shape/dtype live in +# tensor_values (keyed by `name`); cross-method sharing is by `fqn`. +@dataclass +class MutableBufferSpec: + name: str + fqn: str + + +@dataclass +class Graph: + nodes: List[Node] + inputs: Optional[List[str]] = None + outputs: Optional[List[str]] = None + tensor_values: Optional[List[TensorValue]] = field(default=None) + constants: Optional[List[ConstantRef]] = field(default=None) + output_specs: Optional[List[OutputSpec]] = field(default=None) + mutable_buffers: Optional[List[MutableBufferSpec]] = field(default=None) + + +# A subgraph passed to a higher-order op (torch.cond / while_loop / map). Inlined +# nested graph (mirrors torch._export.serde's GraphArgument); makes Graph +# recursive. Defined after Graph so the `graph` annotation is a real class +# reference (a string forward-ref would not deserialize as a nested dataclass). +@dataclass +class GraphArg: + name: str + graph: Graph + + +# BC: APPEND ONLY — keep in sync with the union in native_graph.fbs. +ArgumentValue = Union[ + TensorArg, + NoneArg, + IntArg, + FloatArg, + BoolArg, + StringArg, + ScalarTypeArg, + SymIntArg, + IntListArg, + FloatListArg, + BoolListArg, + TensorListArg, + OptionalTensorListArg, + SymIntListArg, + GraphArg, +] + + +@dataclass +class Method: + name: str + graph: Graph + + +@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..f12891d3ef4 --- /dev/null +++ b/backends/native/test/test_serialize.py @@ -0,0 +1,463 @@ +# 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 + +import torch +import torch.nn as nn + +from executorch.backends.native.serialization import ( + deserialize_graph, + deserialize_program, + serialize_graph, + serialize_program, + validate_graph, + validate_program, +) +from executorch.backends.native.serialization.graph_serialize import ( + _named_arguments, + _output_alias_of, + _to_arg_value, + serialize_operator, +) +from executorch.backends.native.serialization.schema import ( + GraphArg, + InputKind, + IntArg, + OpKind, + OutputKind, + ScalarType, + StringArg, + SymIntListArg, + TensorArg, + TensorListArg, +) +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", "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, + ) + return _Round(edge_ep, deserialize_graph(data), 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([s.as_int for s in meta.sizes], [2, 3]) + + 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.fqn for c in r.graph.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): + graph = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)).graph + weight = next(c for c in graph.constants if "weight" in c.fqn) + self.assertEqual(weight.kind, InputKind.PARAMETER) + self.assertFalse(weight.mutated) + + def test_mutated_buffer_flagged(self): + graph = _roundtrip(_Counter(), (torch.randn(1),)).graph + count = next(c for c in graph.constants if "count" in c.fqn) + 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): + graph = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)).graph + self.assertTrue(graph.output_specs) + self.assertTrue( + all(s.kind == OutputKind.USER_OUTPUT for s in graph.output_specs) + ) + + def test_buffer_mutation_classified(self): + graph = _roundtrip(_Counter(), (torch.randn(1),)).graph + muts = [s for s in graph.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 graph.output_specs if s.kind == OutputKind.USER_OUTPUT] + self.assertEqual(len(users), 1) + + +class ValidateGraphTest(unittest.TestCase): + def test_self_contained_passes(self): + r = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)) + # validate_graph returns None and raises on any inconsistency; a + # self-contained graph must validate cleanly. + self.assertIsNone(validate_graph(r.graph, set(r.constants.keys()))) + + def test_missing_constant_data_raises(self): + graph = _roundtrip(nn.Linear(4, 4), (torch.randn(1, 4),)).graph + with self.assertRaises(ValueError): + validate_graph(graph, available_data_keys=set()) + + +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] must serialize as a SymIntListArg that + # preserves the symbol, not collapse to an empty IntListArg. + graph = self._dynamic_view().graph + sym_lists = [ + na.arg.value + for n in graph.nodes + for na in (n.inputs or []) + if isinstance(na.arg.value, SymIntListArg) + ] + self.assertTrue(sym_lists, "expected a SymIntListArg for the dynamic view size") + self.assertTrue( + any(s.as_symbol for sl in sym_lists for s in sl.values), + f"expected a symbolic dim, got {sym_lists}", + ) + for sl in sym_lists: + self.assertTrue(sl.values) + + 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 symbolic. + graph = self._dynamic_view().graph + symbolic_dims = [ + s + for tv in (graph.tensor_values or []) + for s in tv.meta.sizes + if s.as_symbol + ] + self.assertTrue(symbolic_dims, "dynamic dim was frozen 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 strides + # 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 + + graph = _roundtrip(M(), (torch.randn(4, 8),)).graph + w = next(c for c in graph.constants if "w" in c.fqn) + self.assertEqual([s.as_int for s in w.meta.strides], [8, 1]) + + +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. + self.assertFalse(any(n.op_kind == OpKind.GET_ATTR for n in graph.nodes)) + + +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.fqn for c in program.methods[0].graph.constants} + fqns_b = {c.fqn for c in program.methods[1].graph.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 -- 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),)) + graph = r.graph + self.assertTrue(graph.mutable_buffers, "expected a non-persistent buffer") + mb = graph.mutable_buffers[0] + self.assertIn("cache", mb.fqn) + # It is NOT a data-backed constant and ships no bytes. + self.assertNotIn("cache", {c.fqn for c in (graph.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 (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_graph(r.graph, set(r.constants.keys()))) 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