Skip to content

Commit b29af96

Browse files
JakeStevensfacebook-github-bot
authored andcommitted
Namespace immutable constants per method in to_native (#23891)
Summary: Exported methods can reuse a constant name with different immutable tensors, causing `to_native` to reject the combined program. Namespace every immutable constant as `<method>/<key>`, including single-method programs, and update its graph references, including quantization metadata and nested graphs. Identical tensors still deduplicate when saving the `.ptn`. Keys mutated by any method use shared `state/<key>` names throughout the program, preserving state shared by readers and writers. Update non-persistent buffer bindings and mutation outputs consistently, and reserve the `state` method namespace. Reject differing values for shared mutable keys and collisions between final names. Authored with assistance from OpenAI Codex. Reviewed By: digantdesai Differential Revision: D125294428
1 parent d70c276 commit b29af96

4 files changed

Lines changed: 300 additions & 36 deletions

File tree

‎backends/native/README.md‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,11 @@ quantize/dequantize ops. TorchAO int4 group-wise weights of `linear` and
3333
stored packed with their scheme in tensor metadata. Each method must lower wholly
3434
to one delegate; `to_native` raises otherwise.
3535

36+
Immutable constant keys are namespaced as `<method>/<key>`, even for a single
37+
method. A key mutated by any method becomes `state/<key>` in every method, so
38+
methods share its runtime state. Method names may not be `state` or start with
39+
`state/`. `to_native` raises if namespaced keys collide.
40+
3641
A `.ptn` is an uncompressed zip with fixed entry names:
3742

3843
| Entry | Contents |

‎backends/native/serialization/graph_serialize.py‎

Lines changed: 22 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
from enum import IntEnum
3232
from typing import (
3333
Any,
34+
Callable,
3435
cast,
3536
get_args,
3637
get_origin,
@@ -1487,39 +1488,48 @@ def _iter_graph_args(value: object) -> Iterator[GraphArg]:
14871488
yield from _iter_graph_args(item)
14881489

14891490

1490-
def collect_data_keys(program: Program) -> set[str]:
1491-
"""Return every out-of-line data key the program references.
1491+
def map_data_keys(method: Method, fn: Callable[[str], str]) -> None:
1492+
"""Replace every out-of-line data key in ``method`` with ``fn(key)``, in place.
14921493
14931494
Covers constant references (NamedTensorRef.data_key) plus the quant keys
14941495
(AffineGroup scale and zero-point) carried on any TensorMeta.quant, whether
14951496
it is attached to a constant or to a graph value (intermediates and I/O),
14961497
recursing into HOP subgraphs wherever a schema field carries a ``GraphArg``.
14971498
PackedQuant carries no external keys.
14981499
"""
1499-
keys: set[str] = set()
15001500

1501-
def add_meta(meta: Optional[TensorMeta]) -> None:
1501+
def map_meta(meta: Optional[TensorMeta]) -> None:
15021502
if meta is None or meta.quant is None:
15031503
return
15041504
scheme = meta.quant.scheme
15051505
if isinstance(scheme, AffineGroup):
1506-
keys.add(scheme.scale_data_key)
1506+
scheme.scale_data_key = fn(scheme.scale_data_key)
15071507
if scheme.zero_point_data_key:
1508-
keys.add(scheme.zero_point_data_key)
1508+
scheme.zero_point_data_key = fn(scheme.zero_point_data_key)
15091509

15101510
def visit_graph(graph: Graph) -> None:
15111511
for value in graph.tensor_values or []:
1512-
add_meta(value.meta)
1512+
map_meta(value.meta)
15131513
for node in graph.nodes:
15141514
for graph_arg in _iter_graph_args(node):
15151515
visit_graph(graph_arg.graph)
15161516

1517-
for method in program.methods:
1518-
for constant in method.constants or []:
1519-
keys.add(constant.data_key)
1520-
add_meta(constant.meta)
1521-
visit_graph(method.graph)
1517+
for constant in method.constants or []:
1518+
constant.data_key = fn(constant.data_key)
1519+
map_meta(constant.meta)
1520+
visit_graph(method.graph)
15221521

1522+
1523+
def collect_data_keys(program: Program) -> set[str]:
1524+
"""Return every out-of-line data key the program references."""
1525+
keys: set[str] = set()
1526+
1527+
def add(key: str) -> str:
1528+
keys.add(key)
1529+
return key
1530+
1531+
for method in program.methods:
1532+
map_data_keys(method, add)
15231533
return keys
15241534

15251535

‎exir/native/native.py‎

Lines changed: 55 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -38,20 +38,24 @@
3838

3939
from executorch.backends.native.preprocess import NativeDelegateInfo
4040
from executorch.backends.native.serialization.graph_serialize import (
41+
_compile_to_bytes,
4142
_same_tensor,
4243
collect_data_keys,
4344
deserialize_program,
45+
map_data_keys,
4446
merge_programs,
4547
validate_program,
4648
)
47-
from executorch.backends.native.serialization.schema import Program
49+
from executorch.backends.native.serialization.schema import OutputKind, Program
4850
from executorch.exir._serialize._ptn import write_ptn
4951
from torch.utils import _pytree as pytree
5052

5153
if TYPE_CHECKING:
5254
from torch.export import ExportedProgram
5355
from torch.fx import GraphModule
5456

57+
_STATE_NAMESPACE = "state"
58+
5559

5660
class NativeProgramManager:
5761
"""One or more methods lowered for the native runtime, ready to save.
@@ -255,13 +259,7 @@ def _delegate_info(lowered: object, method_name: str) -> NativeDelegateInfo:
255259
def _merge_constants(
256260
per_method: dict[str, dict[str, torch.Tensor]],
257261
) -> dict[str, torch.Tensor]:
258-
"""Union per-method constants by fqn.
259-
260-
serialize_program treats an fqn as one buffer across every method, so hold to
261-
that here rather than letting a later method silently clobber an earlier one.
262-
_same_tensor short-circuits on identity and shared storage, which is the usual
263-
case for a weight referenced by several methods.
264-
"""
262+
"""Merge constants by their final data keys, rejecting differing values."""
265263
merged: dict[str, torch.Tensor] = {}
266264
source: dict[str, str] = {}
267265
for method_name, constants in per_method.items():
@@ -347,6 +345,11 @@ def _normalize_programs(
347345
)
348346
if not method_name:
349347
raise ValueError("to_native: method names must be non-empty strings.")
348+
if method_name.split("/")[0] == _STATE_NAMESPACE:
349+
raise ValueError(
350+
f"to_native: method name {method_name!r} is reserved; "
351+
f"{_STATE_NAMESPACE!r} namespaces mutable state keys."
352+
)
350353
if not isinstance(program, TorchExportedProgram):
351354
raise TypeError(
352355
f"to_native: method {method_name!r} must map to an "
@@ -364,6 +367,11 @@ def to_native(
364367
``programs`` is a single ``ExportedProgram`` (lowered as the sole ``forward``
365368
method) or a dict mapping method name to ``ExportedProgram``.
366369
370+
Immutable constant keys are namespaced as ``<method>/<key>``, including for a
371+
single method. Keys mutated by any method become ``state/<key>`` in every
372+
method, so state remains shared across methods. Equal immutable tensors
373+
deduplicate on save.
374+
367375
Nothing about lowering is configurable yet. The passes, the partitioner, and
368376
the edge compile config are all owned here, because correctness of the package
369377
depends on them. Exposing the compile config is worth reconsidering.
@@ -378,9 +386,10 @@ def to_native(
378386
Raises:
379387
TypeError: If ``programs`` is not an ``ExportedProgram`` or a method-name
380388
dictionary containing only ``ExportedProgram`` values.
381-
ValueError: If the method dictionary is empty, a method name is empty, a
382-
method does not fully delegate to the native backend, a constant
383-
differs between methods, or a data key the graph references has no
389+
ValueError: If the method dictionary is empty, a method name is empty or
390+
in the reserved ``state`` namespace, a method does not fully delegate
391+
to the native backend, a mutable constant differs between methods,
392+
namespaced keys conflict, or a data key the graph references has no
384393
backing tensor.
385394
"""
386395
method_programs = _normalize_programs(programs)
@@ -437,13 +446,44 @@ def to_native(
437446
method_blobs[method_name] = module.processed_bytes
438447
method_constants[method_name] = _delegate_info(module, method_name).constants
439448

440-
ptg_blob = merge_programs(method_blobs)
441-
constants = _merge_constants(method_constants)
449+
program = deserialize_program(merge_programs(method_blobs))
450+
state_keys = _mutable_data_keys(program) | {
451+
buffer.fqn
452+
for method in program.methods
453+
for buffer in method.mutable_buffers or []
454+
}
455+
state = {key: f"{_STATE_NAMESPACE}/{key}" for key in state_keys}
456+
data_keys = set(state.values())
457+
for method in program.methods:
458+
constants_for_method = method_constants[method.name]
459+
namespaced = {
460+
key: f"{method.name}/{key}"
461+
for key in constants_for_method
462+
if key not in state
463+
}
464+
conflicts = set(namespaced.values()) & data_keys
465+
if conflicts:
466+
raise ValueError(
467+
"to_native: namespaced constant keys already exist: "
468+
f"{sorted(conflicts)}"
469+
)
470+
data_keys.update(namespaced.values())
471+
mapping = state | namespaced
472+
map_data_keys(method, lambda key, mapping=mapping: mapping.get(key, key))
473+
for buffer in method.mutable_buffers or []:
474+
buffer.fqn = state[buffer.fqn]
475+
for spec in method.output_specs or []:
476+
if spec.kind == OutputKind.BUFFER_MUTATION and spec.target is not None:
477+
spec.target = state.get(spec.target, spec.target)
478+
method_constants[method.name] = {
479+
mapping.get(key, key): tensor
480+
for key, tensor in constants_for_method.items()
481+
}
442482

443-
program = deserialize_program(ptg_blob)
483+
constants = _merge_constants(method_constants)
444484
_validate_merged_program(program, constants)
445485

446-
return NativeProgramManager(ptg=ptg_blob, constants=constants)
486+
return NativeProgramManager(ptg=_compile_to_bytes(program), constants=constants)
447487

448488

449489
__all__ = ["NativeProgramManager", "to_native"]

0 commit comments

Comments
 (0)