Save a lowered program whose named data is file-backed (pytorch#23693)

## Problem

A backend can give ExecuTorch its named data as `FileBackedData` instead
of `bytes`. The data then waits in a file on disk until the `.pte` or
`.ptd` is written, instead of sitting in memory a second time. The CUDA
backend stores its weights this way, and `NamedDataStore`,
`to_executorch()` and `write_tensor_data_to_file()` all handle it.

Saving such a lowered program with `exir.save()` fails, though:

```
TypeError: Object of type FileBackedData is not JSON serializable
```

The serializer writes the lowered module's named data as JSON, and its
encoder only knows `bytes`.

## Change

The encoder for named data now reads a file-backed buffer into `bytes`
when it reaches it, then encodes it like any other buffer. The lowered
module keeps its file-backed buffers, and the saved file is the same as
if the backend had given `bytes`. Loading did not need a change: it
already decodes every buffer back to `bytes`.

What this costs: like any buffer in this JSON format, the data is held
in memory while it is encoded, about four times the size of all named
data, the same as for `bytes`. The named data of a lowered module is
written once for each call to it, so a store shared by several lowered
modules is written several times, and each one loads back as its own
copy. That was already so for `bytes`, and changing it needs a format
change.

## Test plan

A new test in `exir/tests/test_serde.py` lowers a small graph with the
named-data test backend, which gives it two buffers, and makes one of
them a `FileBackedData` that owns its file, as the CUDA backend's are.
It checks that the saved program is the same as when both buffers are
`bytes`, that the module and its file still hold the data after saving,
and that the program loads back with both buffers in order. It fails
without this change with the error above, and passes with it.

Ran `exir/tests/test_serde.py` on Linux aarch64 with Python 3.12, using
this branch's `exir` package and schema over an installed nightly: 16
passed with the change, and 15 passed on `main` (all tests except the
new one). The new test also fails on versions of the change that close
the file after reading it, convert only the first buffer, write the
converted copy back into the module, or use another base64 alphabet.

Author: shoumikhin

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,

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",

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):