Skip to content

Commit f807a6e

Browse files
JakeStevensmeta-codesync[bot]
authored andcommitted
Namespace immutable constants per method in to_native (pytorch#23891)
Summary: Pull Request resolved: pytorch#23891 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 retain their original names throughout the program, preserving shared state for readers and writers. Reject differing values for shared mutable keys and collisions between final names. Authored with assistance from OpenAI Codex. Differential Revision: D125294428
1 parent 19f8894 commit f807a6e

4 files changed

Lines changed: 247 additions & 41 deletions

File tree

‎backends/native/serialization/graph_serialize.py‎

Lines changed: 22 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,6 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7-
# pyre-strict
8-
97
"""
108
Serialize a torch.fx graph into the native backend's generic flatbuffer format
119
(native_graph.fbs) and back.
@@ -31,6 +29,7 @@
3129
from enum import IntEnum
3230
from typing import (
3331
Any,
32+
Callable,
3433
cast,
3534
get_args,
3635
get_origin,
@@ -1487,39 +1486,48 @@ def _iter_graph_args(value: object) -> Iterator[GraphArg]:
14871486
yield from _iter_graph_args(item)
14881487

14891488

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

1501-
def add_meta(meta: Optional[TensorMeta]) -> None:
1499+
def map_meta(meta: Optional[TensorMeta]) -> None:
15021500
if meta is None or meta.quant is None:
15031501
return
15041502
scheme = meta.quant.scheme
15051503
if isinstance(scheme, AffineGroup):
1506-
keys.add(scheme.scale_data_key)
1504+
scheme.scale_data_key = fn(scheme.scale_data_key)
15071505
if scheme.zero_point_data_key:
1508-
keys.add(scheme.zero_point_data_key)
1506+
scheme.zero_point_data_key = fn(scheme.zero_point_data_key)
15091507

15101508
def visit_graph(graph: Graph) -> None:
15111509
for value in graph.tensor_values or []:
1512-
add_meta(value.meta)
1510+
map_meta(value.meta)
15131511
for node in graph.nodes:
15141512
for graph_arg in _iter_graph_args(node):
15151513
visit_graph(graph_arg.graph)
15161514

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)
1515+
for constant in method.constants or []:
1516+
constant.data_key = fn(constant.data_key)
1517+
map_meta(constant.meta)
1518+
visit_graph(method.graph)
1519+
15221520

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

15251533

‎exir/native/README.md‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -44,8 +44,14 @@ needed in the future like `constant_methods` and ETRecord integration.
4444

4545
Raises `TypeError` for the wrong input or method-value types. Raises `ValueError`
4646
for an empty method mapping or method name, if a method does not fully delegate to
47-
the native backend, if a constant differs between methods, or if a data key the
48-
graph references has no backing tensor.
47+
the native backend, if a mutable constant differs between methods, if namespaced
48+
keys conflict, or if a data key the graph references has no backing tensor.
49+
50+
All immutable constant keys are namespaced as `<method>/<key>`, including for a
51+
single method. Equal immutable tensors under different keys are deduplicated on
52+
save. A key mutated by any method keeps its original name in every method,
53+
including methods that only read it, preserving shared runtime state. Namespaced
54+
keys must not collide with each other or with preserved mutable keys.
4955

5056
## Current limitations: one complete delegate per method
5157

‎exir/native/native.py‎

Lines changed: 34 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,6 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7-
# pyre-strict
8-
97
"""
108
Direct native lowering.
119
@@ -37,9 +35,11 @@
3735

3836
from executorch.backends.native.preprocess import NativeDelegateInfo
3937
from executorch.backends.native.serialization.graph_serialize import (
38+
_compile_to_bytes,
4039
_same_tensor,
4140
collect_data_keys,
4241
deserialize_program,
42+
map_data_keys,
4343
merge_programs,
4444
validate_program,
4545
)
@@ -254,13 +254,7 @@ def _delegate_info(lowered: object, method_name: str) -> NativeDelegateInfo:
254254
def _merge_constants(
255255
per_method: dict[str, dict[str, torch.Tensor]],
256256
) -> dict[str, torch.Tensor]:
257-
"""Union per-method constants by fqn.
258-
259-
serialize_program treats an fqn as one buffer across every method, so hold to
260-
that here rather than letting a later method silently clobber an earlier one.
261-
_same_tensor short-circuits on identity and shared storage, which is the usual
262-
case for a weight referenced by several methods.
263-
"""
257+
"""Merge constants by their final data keys, rejecting differing values."""
264258
merged: dict[str, torch.Tensor] = {}
265259
source: dict[str, str] = {}
266260
for method_name, constants in per_method.items():
@@ -363,6 +357,10 @@ def to_native(
363357
``programs`` is a single ``ExportedProgram`` (lowered as the sole ``forward``
364358
method) or a dict mapping method name to ``ExportedProgram``.
365359
360+
Immutable constant keys are namespaced as ``<method>/<key>``, including for a
361+
single method. Keys mutated by any method keep their original names so state
362+
remains shared across methods. Equal immutable tensors deduplicate on save.
363+
366364
Nothing about lowering is configurable yet. The passes, the partitioner, and
367365
the edge compile config are all owned here, because correctness of the package
368366
depends on them. Exposing the compile config is worth reconsidering.
@@ -378,9 +376,9 @@ def to_native(
378376
TypeError: If ``programs`` is not an ``ExportedProgram`` or a method-name
379377
dictionary containing only ``ExportedProgram`` values.
380378
ValueError: If the method dictionary is empty, a method name is empty, a
381-
method does not fully delegate to the native backend, a constant
382-
differs between methods, or a data key the graph references has no
383-
backing tensor.
379+
method does not fully delegate to the native backend, a mutable
380+
constant differs between methods, namespaced keys conflict, or a data
381+
key the graph references has no backing tensor.
384382
"""
385383
method_programs = _normalize_programs(programs)
386384

@@ -437,13 +435,33 @@ def to_native(
437435
method_blobs[method_name] = module.processed_bytes
438436
method_constants[method_name] = _delegate_info(module, method_name).constants
439437

440-
ptg_blob = merge_programs(method_blobs)
441-
constants = _merge_constants(method_constants)
438+
program = deserialize_program(merge_programs(method_blobs))
439+
mutated_keys = _mutable_data_keys(program)
440+
data_keys = set(mutated_keys)
441+
for method in program.methods:
442+
constants_for_method = method_constants[method.name]
443+
mapping = {
444+
key: f"{method.name}/{key}"
445+
for key in constants_for_method
446+
if key not in mutated_keys
447+
}
448+
conflicts = set(mapping.values()) & data_keys
449+
if conflicts:
450+
raise ValueError(
451+
"to_native: namespaced constant keys already exist: "
452+
f"{sorted(conflicts)}"
453+
)
454+
data_keys.update(mapping.values())
455+
map_data_keys(method, lambda key, mapping=mapping: mapping.get(key, key))
456+
method_constants[method.name] = {
457+
mapping.get(key, key): tensor
458+
for key, tensor in constants_for_method.items()
459+
}
442460

443-
program = deserialize_program(ptg_blob)
461+
constants = _merge_constants(method_constants)
444462
_validate_merged_program(program, constants)
445463

446-
return NativeProgramManager(ptg=ptg_blob, constants=constants)
464+
return NativeProgramManager(ptg=_compile_to_bytes(program), constants=constants)
447465

448466

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

0 commit comments

Comments
 (0)