Skip to content

Commit 0bac518

Browse files
JakeStevensfacebook-github-bot
authored andcommitted
Namespace lifted tensor constants per method in to_native (#23891)
Summary: Edge lowering generates `_lifted_tensor_constant{i}` keys independently for each method. When methods use different scalar values under the same generated key, `to_native` rejects the combined program. Always namespace generated immutable lifted constants as `<method>/<key>`, including single-method programs. Update the constant dictionaries and graph references together through a shared `map_data_keys` traversal, including quantization metadata and nested graphs. Equal tensors under different keys still deduplicate when saving the `.ptn`. Mutable keys keep their existing names, parameter/buffer conflicts continue to raise, and explicit keys occupying a generated namespace are rejected. Authored with assistance from OpenAI Codex. Differential Revision: D125294428
1 parent 91958e6 commit 0bac518

4 files changed

Lines changed: 150 additions & 28 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: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,10 @@ needed in the future like `constant_methods` and ETRecord integration.
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
4747
the native backend, if a constant differs between methods, or if a data key the
48-
graph references has no backing tensor.
48+
graph references has no backing tensor. Generated immutable
49+
`_lifted_tensor_constant{i}` keys are always namespaced as `<method>/<name>`,
50+
including for a single method. Equal tensors under different keys are deduplicated
51+
on save. A namespaced key that conflicts with an existing constant key is rejected.
4952

5053
## Current limitations: one complete delegate per method
5154

‎exir/native/native.py‎

Lines changed: 29 additions & 8 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
)
@@ -379,8 +379,10 @@ def to_native(
379379
dictionary containing only ``ExportedProgram`` values.
380380
ValueError: If the method dictionary is empty, a method name is empty, a
381381
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.
382+
differs between methods, a namespaced key conflicts with an existing
383+
constant key, or a data key the graph references has no backing tensor.
384+
Generated immutable ``_lifted_tensor_constant{i}`` keys are always
385+
namespaced per method.
384386
"""
385387
method_programs = _normalize_programs(programs)
386388

@@ -437,13 +439,32 @@ def to_native(
437439
method_blobs[method_name] = module.processed_bytes
438440
method_constants[method_name] = _delegate_info(module, method_name).constants
439441

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

443-
program = deserialize_program(ptg_blob)
464+
constants = _merge_constants(method_constants)
444465
_validate_merged_program(program, constants)
445466

446-
return NativeProgramManager(ptg=ptg_blob, constants=constants)
467+
return NativeProgramManager(ptg=_compile_to_bytes(program), constants=constants)
447468

448469

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

‎exir/native/test/test_native.py‎

Lines changed: 95 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,12 @@
2727
Program,
2828
)
2929
from executorch.exir import to_edge_transform_and_lower
30-
from executorch.exir._serialize._ptn import ALIASES_ENTRY, PTG_ENTRY, SAFETENSORS_ENTRY
30+
from executorch.exir._serialize._ptn import (
31+
ALIASES_ENTRY,
32+
PTG_ENTRY,
33+
read_ptn,
34+
SAFETENSORS_ENTRY,
35+
)
3136
from executorch.exir.lowered_backend_module import (
3237
executorch_call_delegate,
3338
get_lowered_submodules,
@@ -60,7 +65,7 @@ def _test_equivalence(self, programs):
6065
flow keeps one single-method delegate blob per method and would ship the
6166
same constants in a .ptd. The graphs are compared modulo that wrapping,
6267
and the full key sets (safetensors owners + aliases) must match the
63-
delegates' named-data keys.
68+
delegates' named-data keys after namespacing lifted constants.
6469
"""
6570
method_programs = (
6671
programs if isinstance(programs, dict) else {"forward": programs}
@@ -113,9 +118,17 @@ def _test_equivalence(self, programs):
113118
# Constants may be embedded in the .pte (pte_data) or shipped to a
114119
# .ptd (external_data) depending on external_constants_tag; collect
115120
# both so the key set matches to_native's .ptn/.safetensors.
116-
et_keys |= set(nds.pte_data.keys())
121+
method_keys = set(nds.pte_data.keys())
117122
for key_to_entry in nds.external_data.values():
118-
et_keys |= set(key_to_entry.keys())
123+
method_keys |= set(key_to_entry.keys())
124+
et_keys |= {
125+
(
126+
f"{name}/{key}"
127+
if key.startswith("_lifted_tensor_constant")
128+
else key
129+
)
130+
for key in method_keys
131+
}
119132

120133
# Emission is outside the comparison, but keep proving the ExecuTorch path
121134
# still completes end to end through the unchanged NamedDataStore branch.
@@ -239,7 +252,7 @@ def forward(self, x):
239252

240253
ep = torch.export.export(ConstAdd(), (torch.randn(2, 4),))
241254
_, st_keys, _ = self._test_equivalence(ep)
242-
self.assertTrue(any("lifted" in k for k in st_keys))
255+
self.assertEqual(st_keys, {"forward/_lifted_tensor_constant0"})
243256

244257

245258
class ToNativeInputValidationTest(unittest.TestCase):
@@ -393,6 +406,83 @@ def test_dropped_public_input_is_rejected(self):
393406
self._check(self._graph_module(drop_input=True))
394407

395408

409+
class LiftedConstantNamespaceTest(unittest.TestCase):
410+
def _scalar_shift(self, offset: float):
411+
class Shift(nn.Module):
412+
def forward(self, x):
413+
return x + offset
414+
415+
return torch.export.export(Shift(), (torch.randn(2, 4),))
416+
417+
def _read(self, native) -> tuple[Program, dict[str, torch.Tensor]]:
418+
with tempfile.TemporaryDirectory() as td:
419+
path = os.path.join(td, "m.ptn")
420+
native.save(path)
421+
ptg, constants = read_ptn(path)
422+
return deserialize_program(ptg), constants
423+
424+
def _lifted_keys(self, program: Program, method_name: str) -> list[str]:
425+
(method,) = [m for m in program.methods if m.name == method_name]
426+
return [c.data_key for c in method.constants or [] if "lifted" in c.data_key]
427+
428+
def test_distinct_lifted_constants_are_namespaced(self):
429+
programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(2.0)}
430+
native = to_native(programs)
431+
self.assertEqual(native.methods, {"a", "b"})
432+
program, constants = self._read(native)
433+
(a_key,) = self._lifted_keys(program, "a")
434+
(b_key,) = self._lifted_keys(program, "b")
435+
self.assertEqual(a_key, "a/_lifted_tensor_constant0")
436+
self.assertEqual(b_key, "b/_lifted_tensor_constant0")
437+
torch.testing.assert_close(
438+
constants[a_key], torch.full_like(constants[a_key], 1.0), rtol=0, atol=0
439+
)
440+
torch.testing.assert_close(
441+
constants[b_key], torch.full_like(constants[b_key], 2.0), rtol=0, atol=0
442+
)
443+
444+
def test_identical_lifted_constants_dedup_on_save(self):
445+
programs = {"a": self._scalar_shift(1.0), "b": self._scalar_shift(1.0)}
446+
native = to_native(programs)
447+
with tempfile.TemporaryDirectory() as td:
448+
path = os.path.join(td, "m.ptn")
449+
native.save(path)
450+
ptg, constants = read_ptn(path)
451+
with zipfile.ZipFile(path) as pkg:
452+
owners = _safetensors_keys(pkg.read(SAFETENSORS_ENTRY))
453+
aliases = json.loads(pkg.read(ALIASES_ENTRY))
454+
455+
program = deserialize_program(ptg)
456+
(a_key,) = self._lifted_keys(program, "a")
457+
(b_key,) = self._lifted_keys(program, "b")
458+
self.assertEqual(a_key, "a/_lifted_tensor_constant0")
459+
self.assertEqual(b_key, "b/_lifted_tensor_constant0")
460+
self.assertEqual(owners, {a_key})
461+
self.assertEqual(aliases, {b_key: a_key})
462+
self.assertIs(constants[a_key], constants[b_key])
463+
464+
def test_parameter_fqn_collision_is_still_rejected(self):
465+
programs = {
466+
"a": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)),
467+
"b": torch.export.export(nn.Linear(4, 4), (torch.randn(1, 4),)),
468+
}
469+
with self.assertRaisesRegex(ValueError, "differs between methods"):
470+
to_native(programs)
471+
472+
def test_namespaced_key_collision_is_rejected(self):
473+
class ConstWithBuffer(nn.Module):
474+
def __init__(self):
475+
super().__init__()
476+
self.register_buffer("forward/_lifted_tensor_constant1", torch.zeros(4))
477+
478+
def forward(self, x):
479+
return x + self.get_buffer("forward/_lifted_tensor_constant1") + 1.0
480+
481+
ep = torch.export.export(ConstWithBuffer(), (torch.randn(4),))
482+
with self.assertRaisesRegex(ValueError, "namespaced lifted constant keys"):
483+
to_native(ep)
484+
485+
396486
class MutableBufferDedupTest(unittest.TestCase):
397487
def test_identical_mutable_buffers_stay_separate(self):
398488
"""Equal bytes must not make two independently mutable buffers one.

0 commit comments

Comments
 (0)