|
37 | 37 |
|
38 | 38 | from executorch.backends.native.preprocess import NativeDelegateInfo |
39 | 39 | from executorch.backends.native.serialization.graph_serialize import ( |
| 40 | + _compile_to_bytes, |
40 | 41 | _same_tensor, |
41 | 42 | collect_data_keys, |
42 | 43 | deserialize_program, |
43 | 44 | merge_programs, |
| 45 | + rename_data_keys, |
44 | 46 | validate_program, |
45 | 47 | ) |
46 | 48 | from executorch.backends.native.serialization.schema import Program |
@@ -278,6 +280,46 @@ def _merge_constants( |
278 | 280 | return merged |
279 | 281 |
|
280 | 282 |
|
| 283 | +# Edge lowering generates these names independently for each method. |
| 284 | +_LIFTED_CONSTANT_PREFIX = "_lifted_tensor_constant" |
| 285 | + |
| 286 | + |
| 287 | +def _plan_lifted_constant_renames( |
| 288 | + per_method: dict[str, dict[str, torch.Tensor]], |
| 289 | + program: Program, |
| 290 | +) -> dict[str, dict[str, str]]: |
| 291 | + """Plan ``<method>/<key>`` names for conflicting immutable lifted constants. |
| 292 | +
|
| 293 | + Equal tensors stay shared; keys mutated by any method are never renamed. |
| 294 | + """ |
| 295 | + mutated_keys = _mutable_data_keys(program) |
| 296 | + taken = {key for constants in per_method.values() for key in constants} |
| 297 | + seen: dict[str, torch.Tensor] = {} |
| 298 | + renames: dict[str, dict[str, str]] = {} |
| 299 | + for method_name, constants in per_method.items(): |
| 300 | + mapping: dict[str, str] = {} |
| 301 | + for key, tensor in constants.items(): |
| 302 | + existing = seen.get(key) |
| 303 | + if existing is None: |
| 304 | + seen[key] = tensor |
| 305 | + continue |
| 306 | + if _same_tensor(existing, tensor): |
| 307 | + continue |
| 308 | + if not key.startswith(_LIFTED_CONSTANT_PREFIX) or key in mutated_keys: |
| 309 | + continue |
| 310 | + new_key = f"{method_name}/{key}" |
| 311 | + suffix = 1 |
| 312 | + while new_key in taken: |
| 313 | + suffix += 1 |
| 314 | + new_key = f"{method_name}/{key}#{suffix}" |
| 315 | + mapping[key] = new_key |
| 316 | + taken.add(new_key) |
| 317 | + seen[new_key] = tensor |
| 318 | + if mapping: |
| 319 | + renames[method_name] = mapping |
| 320 | + return renames |
| 321 | + |
| 322 | + |
281 | 323 | def _check_shared_mutable_buffers(program: Program) -> None: |
282 | 324 | """Hold the merged Program to serialize_program's cross-method guarantee. |
283 | 325 |
|
@@ -380,7 +422,8 @@ def to_native( |
380 | 422 | ValueError: If the method dictionary is empty, a method name is empty, a |
381 | 423 | method does not fully delegate to the native backend, a constant |
382 | 424 | differs between methods, or a data key the graph references has no |
383 | | - backing tensor. |
| 425 | + backing tensor. Conflicting generated ``_lifted_tensor_constant{i}`` |
| 426 | + keys are namespaced per method. |
384 | 427 | """ |
385 | 428 | method_programs = _normalize_programs(programs) |
386 | 429 |
|
@@ -435,12 +478,23 @@ def to_native( |
435 | 478 | method_name, |
436 | 479 | ) |
437 | 480 | method_blobs[method_name] = module.processed_bytes |
438 | | - method_constants[method_name] = _delegate_info(module, method_name).constants |
| 481 | + method_constants[method_name] = dict( |
| 482 | + _delegate_info(module, method_name).constants |
| 483 | + ) |
439 | 484 |
|
440 | 485 | ptg_blob = merge_programs(method_blobs) |
441 | | - constants = _merge_constants(method_constants) |
442 | | - |
443 | 486 | program = deserialize_program(ptg_blob) |
| 487 | + |
| 488 | + renames = _plan_lifted_constant_renames(method_constants, program) |
| 489 | + if renames: |
| 490 | + for method_name, mapping in renames.items(): |
| 491 | + rename_data_keys(program, method_name, mapping) |
| 492 | + constants_for_method = method_constants[method_name] |
| 493 | + for old_key, new_key in mapping.items(): |
| 494 | + constants_for_method[new_key] = constants_for_method.pop(old_key) |
| 495 | + ptg_blob = _compile_to_bytes(program) |
| 496 | + |
| 497 | + constants = _merge_constants(method_constants) |
444 | 498 | _validate_merged_program(program, constants) |
445 | 499 |
|
446 | 500 | return NativeProgramManager(ptg=ptg_blob, constants=constants) |
|
0 commit comments