3838
3939from executorch .backends .native .preprocess import NativeDelegateInfo
4040from 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
4850from executorch .exir ._serialize ._ptn import write_ptn
4951from torch .utils import _pytree as pytree
5052
5153if TYPE_CHECKING :
5254 from torch .export import ExportedProgram
5355 from torch .fx import GraphModule
5456
57+ _STATE_NAMESPACE = "state"
58+
5559
5660class 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:
255259def _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