diff --git a/exir/serde/serialize.py b/exir/serde/serialize.py index 2df4ba79876..d9cd9a700a7 100644 --- a/exir/serde/serialize.py +++ b/exir/serde/serialize.py @@ -22,6 +22,7 @@ import torch import torch.export.exported_program as ep from executorch.exir import delegate +from executorch.exir._serialize._cord import FileBackedData from executorch.exir._serialize._named_data_store import NamedDataStoreOutput from executorch.exir.backend.backend_details import DelegateScratchSpec from executorch.exir.backend.compile_spec_schema import ( @@ -48,6 +49,14 @@ log: logging.Logger = logging.getLogger(__name__) +class _NamedDataEncoder(export_serialize.EnumEncoder): + # Read each file only when it is encoded, so its raw bytes are freed right after. + def default(self, obj): + if isinstance(obj, FileBackedData): + obj = obj.to_bytes() + return super().default(obj) + + class GraphModuleSerializer(export_serialize.GraphModuleSerializer): def __init__( self, @@ -291,7 +300,7 @@ def serialize_bytes(b: bytes) -> str: assert isinstance(serialized_artifact.exported_program, schema.ExportedProgram) serialized_processed_bytes = serialize_bytes(lowered_module.processed_bytes) - named_data_store = json.dumps(export_serialize._dataclass_to_dict(lowered_module.named_data_store_output),cls=export_serialize.EnumEncoder) if lowered_module.named_data_store_output else None + named_data_store = json.dumps(export_serialize._dataclass_to_dict(lowered_module.named_data_store_output),cls=_NamedDataEncoder) if lowered_module.named_data_store_output else None serialized_lowered_module = SerdeLoweredBackendModule( original_module=serialized_artifact.exported_program, diff --git a/exir/tests/targets.bzl b/exir/tests/targets.bzl index d8f21020d2a..ff161d755e6 100644 --- a/exir/tests/targets.bzl +++ b/exir/tests/targets.bzl @@ -80,6 +80,7 @@ def define_common_targets(is_fbcode = False): "//executorch/exir:lib", "//executorch/exir/backend:backend_api", "//executorch/exir/backend/test:backend_with_compiler_demo", + "//executorch/exir/backend/test:backend_with_named_data_map", "//executorch/exir/backend/test:op_partitioner_demo", "//executorch/exir/serde:serialize", "//executorch/backends/xnnpack/partition:xnnpack_partitioner", diff --git a/exir/tests/test_serde.py b/exir/tests/test_serde.py index 6696b95fd31..2a159c704f8 100644 --- a/exir/tests/test_serde.py +++ b/exir/tests/test_serde.py @@ -7,6 +7,7 @@ # pyre-strict import io +import os import tempfile import unittest from contextlib import contextmanager @@ -19,12 +20,17 @@ XnnpackFloatingPointPartitioner, ) from executorch.exir import to_edge +from executorch.exir._serialize._cord import FileBackedData from executorch.exir.backend.backend_api import CompileSpec, to_backend from executorch.exir.backend.test.backend_with_compiler_demo import ( BackendWithCompilerDemo, ) +from executorch.exir.backend.test.backend_with_named_data_map import ( + BackendWithNDMPartitioner, +) from executorch.exir.backend.test.op_partitioner_demo import AddMulPartitionerDemo +from executorch.exir.lowered_backend_module import get_lowered_submodules from executorch.exir.program._program import ( EdgeProgramManager, to_edge_transform_and_lower, @@ -246,6 +252,44 @@ def forward(self, a, x, b): edge_new = deserialize(serialize(edge.exported_program())) self.check_ep(edge.exported_program(), edge_new, inputs) + def test_delegate_with_file_backed_named_data(self) -> None: + class Model(nn.Module): + def forward(self, x): + return (x + x) - x + + edge = to_edge(export(Model(), (torch.randn(1, 2),), strict=True)).to_backend( + BackendWithNDMPartitioner() + ) + graph_module = edge.exported_program().graph_module + (_, lowered, _) = get_lowered_submodules(graph_module)[0] + output = lowered.named_data_store_output + assert output is not None + payload = bytes(range(256)) + output.buffers[1] = payload + other = output.buffers[0] + expected = serialize(edge.exported_program()).exported_program + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "data.bin") + with open(path, "wb") as f: + f.write(payload) + data = FileBackedData.move_from(path) + output.buffers[1] = data + + serialized = serialize(edge.exported_program()).exported_program + self.assertEqual(serialized, expected) + buffer = io.BytesIO() + exir.save(edge.exported_program(), buffer) + self.assertIs(lowered.named_data_store_output, output) + self.assertIs(output.buffers[1], data) + self.assertEqual(data.to_bytes(), payload) + + buffer.seek(0) + loaded = exir.load(buffer) + (_, loaded_lowered, _) = get_lowered_submodules(loaded.graph_module)[0] + loaded_output = loaded_lowered.named_data_store_output + assert loaded_output is not None + self.assertEqual(loaded_output.buffers, [other, payload]) + def test_delegate_xnnpack(self) -> None: class SimpleConv1DModel(nn.Module): def __init__(self):