|
27 | 27 | Program, |
28 | 28 | ) |
29 | 29 | 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 | +) |
31 | 36 | from executorch.exir.lowered_backend_module import ( |
32 | 37 | executorch_call_delegate, |
33 | 38 | get_lowered_submodules, |
@@ -60,7 +65,7 @@ def _test_equivalence(self, programs): |
60 | 65 | flow keeps one single-method delegate blob per method and would ship the |
61 | 66 | same constants in a .ptd. The graphs are compared modulo that wrapping, |
62 | 67 | 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. |
64 | 69 | """ |
65 | 70 | method_programs = ( |
66 | 71 | programs if isinstance(programs, dict) else {"forward": programs} |
@@ -113,9 +118,17 @@ def _test_equivalence(self, programs): |
113 | 118 | # Constants may be embedded in the .pte (pte_data) or shipped to a |
114 | 119 | # .ptd (external_data) depending on external_constants_tag; collect |
115 | 120 | # 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()) |
117 | 122 | 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 | + } |
119 | 132 |
|
120 | 133 | # Emission is outside the comparison, but keep proving the ExecuTorch path |
121 | 134 | # still completes end to end through the unchanged NamedDataStore branch. |
@@ -239,7 +252,7 @@ def forward(self, x): |
239 | 252 |
|
240 | 253 | ep = torch.export.export(ConstAdd(), (torch.randn(2, 4),)) |
241 | 254 | _, 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"}) |
243 | 256 |
|
244 | 257 |
|
245 | 258 | class ToNativeInputValidationTest(unittest.TestCase): |
@@ -393,6 +406,83 @@ def test_dropped_public_input_is_rejected(self): |
393 | 406 | self._check(self._graph_module(drop_input=True)) |
394 | 407 |
|
395 | 408 |
|
| 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 | + |
396 | 486 | class MutableBufferDedupTest(unittest.TestCase): |
397 | 487 | def test_identical_mutable_buffers_stay_separate(self): |
398 | 488 | """Equal bytes must not make two independently mutable buffers one. |
|
0 commit comments