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 "
@@ -365,6 +368,11 @@ def to_native(
365368 ``programs`` is a single ``ExportedProgram`` (lowered as the sole ``forward``
366369 method) or a dict mapping method name to ``ExportedProgram``.
367370
371+ Immutable constant keys are namespaced as ``<method>/<key>``, including for a
372+ single method. Keys mutated by any method become ``state/<key>`` in every
373+ method, so state remains shared across methods. Equal immutable tensors
374+ deduplicate on save.
375+
368376 Nothing about lowering is configurable yet. The passes, the partitioner, and
369377 the edge compile config are all owned here, because correctness of the package
370378 depends on them. Exposing the compile config is worth reconsidering.
@@ -379,9 +387,10 @@ def to_native(
379387 Raises:
380388 TypeError: If ``programs`` is not an ``ExportedProgram`` or a method-name
381389 dictionary containing only ``ExportedProgram`` values.
382- ValueError: If the method dictionary is empty, a method name is empty, a
383- method does not fully delegate to the native backend, a constant
384- differs between methods, or a data key the graph references has no
390+ ValueError: If the method dictionary is empty, a method name is empty or
391+ in the reserved ``state`` namespace, a method does not fully delegate
392+ to the native backend, a mutable constant differs between methods,
393+ namespaced keys conflict, or a data key the graph references has no
385394 backing tensor.
386395 """
387396 method_programs = _normalize_programs (programs )
@@ -447,13 +456,44 @@ def to_native(
447456 method_blobs [method_name ] = module .processed_bytes
448457 method_constants [method_name ] = _delegate_info (module , method_name ).constants
449458
450- ptg_blob = merge_programs (method_blobs )
451- constants = _merge_constants (method_constants )
459+ program = deserialize_program (merge_programs (method_blobs ))
460+ state_keys = _mutable_data_keys (program ) | {
461+ buffer .fqn
462+ for method in program .methods
463+ for buffer in method .mutable_buffers or []
464+ }
465+ state = {key : f"{ _STATE_NAMESPACE } /{ key } " for key in state_keys }
466+ data_keys = set (state .values ())
467+ for method in program .methods :
468+ constants_for_method = method_constants [method .name ]
469+ namespaced = {
470+ key : f"{ method .name } /{ key } "
471+ for key in constants_for_method
472+ if key not in state
473+ }
474+ conflicts = set (namespaced .values ()) & data_keys
475+ if conflicts :
476+ raise ValueError (
477+ "to_native: namespaced constant keys already exist: "
478+ f"{ sorted (conflicts )} "
479+ )
480+ data_keys .update (namespaced .values ())
481+ mapping = state | namespaced
482+ map_data_keys (method , lambda key , mapping = mapping : mapping .get (key , key ))
483+ for buffer in method .mutable_buffers or []:
484+ buffer .fqn = state [buffer .fqn ]
485+ for spec in method .output_specs or []:
486+ if spec .kind == OutputKind .BUFFER_MUTATION and spec .target is not None :
487+ spec .target = state .get (spec .target , spec .target )
488+ method_constants [method .name ] = {
489+ mapping .get (key , key ): tensor
490+ for key , tensor in constants_for_method .items ()
491+ }
452492
453- program = deserialize_program ( ptg_blob )
493+ constants = _merge_constants ( method_constants )
454494 _validate_merged_program (program , constants )
455495
456- return NativeProgramManager (ptg = ptg_blob , constants = constants )
496+ return NativeProgramManager (ptg = _compile_to_bytes ( program ) , constants = constants )
457497
458498
459499__all__ = ["NativeProgramManager" , "to_native" ]
0 commit comments