exir: add flatbuffer-to-program reader

This continues the work from #17333.

Change-Id: I35ac4cd5f6430ea89939453344c13e056b5c746c
Signed-off-by: Chizkiyahu Raful <chizkiyahu.raful@arm.com>

Author: chizkiyahu

exir/_serialize/_flatbuffer_program.py

@@ -8,12 +8,14 @@
 import enum
 import functools
 import importlib
+import pkgutil
 import tempfile
 
 from contextvars import ContextVar
 from dataclasses import fields, is_dataclass
 from functools import lru_cache
-from typing import Any, Dict, Optional
+from types import ModuleType
+from typing import Any, Dict, get_args, get_origin, get_type_hints, Optional, Union
 
 import flatbuffers  # pyre-ignore[21]
 from executorch.exir._serialize._flatbuffer import (
@@ -22,6 +24,7 @@
     _prepare_schema,
     _SchemaInfo,
 )
+from executorch.exir._serialize.generated import executorch_flatbuffer as _generated_fb
 from executorch.exir._serialize.generated.executorch_flatbuffer import (
     BackendDelegateInlineData as _BackendDelegateInlineData,
     Buffer as _Buffer,
@@ -33,6 +36,7 @@
 
 _T_CLASS_CACHE: Dict[type, type] = {}
 _FIELD_NAME_CACHE: Dict[type, tuple[tuple[str, str], ...]] = {}
+_TYPE_HINTS_CACHE: Dict[type, Dict[str, Any]] = {}
 _BUFFER_ALIGNMENT: ContextVar[int] = ContextVar("_BUFFER_ALIGNMENT", default=1)
 _DELEGATE_ALIGNMENT: ContextVar[int] = ContextVar("_DELEGATE_ALIGNMENT", default=1)
 
@@ -64,6 +68,15 @@ def _dataclass_field_map(dataclass_type: type) -> tuple[tuple[str, str], ...]:
     return mapping
 
 
+def _dataclass_type_hints(dataclass_type: type) -> Dict[str, Any]:
+    cached = _TYPE_HINTS_CACHE.get(dataclass_type)
+    if cached is not None:
+        return cached
+    type_hints = get_type_hints(dataclass_type)
+    _TYPE_HINTS_CACHE[dataclass_type] = type_hints
+    return type_hints
+
+
 def _create_aligned_byte_vector(builder: Any, data: bytes, alignment: int) -> int:
     if not _is_valid_alignment(alignment):
         raise ValueError(f"Bad alignment {alignment}")
@@ -194,6 +207,126 @@ def convert_program(val: Program) -> ProgramT:
     return _convert_dataclass(val)
 
 
+# The generated FlatBuffer Python modules import child tables/unions as modules
+# (for example, Program.ExecutionPlan becomes the ExecutionPlan module), but the
+# unpacking helpers later expect those globals to be the corresponding classes.
+# Rebind module globals like ExecutionPlan -> ExecutionPlan.ExecutionPlan so the
+# generated InitFromObj()/InitFromPackedBuf() code can instantiate nested types.
+def _patch_generated_module_aliases(module: ModuleType) -> None:
+    for name, maybe_module in vars(module).items():
+        if not isinstance(maybe_module, ModuleType):
+            continue
+        maybe_class = getattr(maybe_module, name, None)
+        if isinstance(maybe_class, type):
+            setattr(module, name, maybe_class)
+
+
+@lru_cache(maxsize=1)
+def _patch_generated_flatbuffer_aliases() -> None:
+    package_name = _generated_fb.__name__
+    for module_info in pkgutil.iter_modules(_generated_fb.__path__):
+        module = importlib.import_module(f"{package_name}.{module_info.name}")
+        _patch_generated_module_aliases(module)
+
+
+def _flatbuffer_dataclass_names(val: Any) -> tuple[str, Optional[str]]:
+    val_type_name = type(val).__name__
+    if val_type_name.endswith("T"):
+        return val_type_name, val_type_name[:-1]
+    return val_type_name, None
+
+
+def _matches_dataclass_union_type(
+    union_type: Any, val_type_name: str, val_dataclass_name: Optional[str]
+) -> bool:
+    if not is_dataclass(union_type):
+        return False
+    union_name = union_type.__name__
+    return union_name == val_type_name or (
+        val_dataclass_name is not None and union_name == val_dataclass_name
+    )
+
+
+def _matches_non_dataclass_union_type(union_type: Any, val: Any) -> bool:
+    if union_type is Any:
+        return True
+    if union_type is str and isinstance(val, (bytes, bytearray, memoryview)):
+        return True
+    union_origin = get_origin(union_type)
+    if union_origin is list and hasattr(val, "__iter__"):
+        return True
+    return isinstance(union_type, type) and isinstance(val, union_type)
+
+
+def _union_choice_from_value(union_types: tuple[Any, ...], val: Any) -> Any:
+    if val is None:
+        for union_type in union_types:
+            if union_type is type(None):
+                return union_type
+        return None
+
+    val_type_name, val_dataclass_name = _flatbuffer_dataclass_names(val)
+
+    for union_type in union_types:
+        if union_type is type(None):
+            continue
+        if _matches_dataclass_union_type(union_type, val_type_name, val_dataclass_name):
+            return union_type
+        if _matches_non_dataclass_union_type(union_type, val):
+            return union_type
+    return None
+
+
+def _convert_from_flatbuffer_value(val: Any, expected_type: Any) -> Any:
+    if val is None:
+        return None
+
+    origin = get_origin(expected_type)
+    if origin is list:
+        item_type = get_args(expected_type)[0]
+        return [_convert_from_flatbuffer_value(item, item_type) for item in val]
+
+    if origin is Union:
+        union_type = _union_choice_from_value(get_args(expected_type), val)
+        if union_type is None:
+            raise TypeError(
+                f"Could not match value type {type(val)} to {expected_type}"
+            )
+        if union_type is type(None):
+            return None
+        return _convert_from_flatbuffer_value(val, union_type)
+
+    if expected_type is bytes:
+        return _coerce_bytes(val)
+    if expected_type is str and isinstance(val, (bytes, bytearray, memoryview)):
+        return _coerce_bytes(val).decode("utf-8")
+    if is_dataclass(expected_type):
+        return _convert_from_flatbuffer_dataclass(val, expected_type)
+    if isinstance(expected_type, type) and issubclass(expected_type, enum.Enum):
+        if isinstance(val, expected_type):
+            return val
+        return expected_type(val)
+    if isinstance(expected_type, type):
+        return expected_type(val)
+    return val
+
+
+def _convert_from_flatbuffer_dataclass(val: Any, dataclass_type: type) -> Any:
+    result = {}
+    type_hints = _dataclass_type_hints(dataclass_type)
+    for src_name, dst_name in _dataclass_field_map(dataclass_type):
+        result[src_name] = _convert_from_flatbuffer_value(
+            getattr(val, dst_name), type_hints[src_name]
+        )
+    return dataclass_type(**result)
+
+
+def _flatbuffer_to_program(program_data: bytes) -> Program:
+    _patch_generated_flatbuffer_aliases()
+    program_t = ProgramT.InitFromPackedBuf(program_data)
+    return _convert_from_flatbuffer_dataclass(program_t, Program)
+
+
 @lru_cache(maxsize=1)
 def _get_schema_info(
     constant_tensor_alignment: Optional[int], delegate_alignment: Optional[int]
@@ -213,11 +346,7 @@ def _program_to_flatbuffer(
     constant_tensor_alignment: Optional[int] = None,
     delegate_alignment: Optional[int] = None,
 ) -> _FlatbufferResult:
-    """Converts a Program dataclass into binary flatbuffer data.
-
-    Unlike _program_json_to_flatbuffer(), this does not use JSON or invoke
-    flatc to build the binary.
-    """
+    """Converts a Program dataclass into binary flatbuffer data."""
     schema_info = _get_schema_info(constant_tensor_alignment, delegate_alignment)
     _set_pack_alignments(schema_info.tensor_alignment, schema_info.delegate_alignment)
     _install_fast_packers()

exir/_serialize/_program.py

@@ -17,11 +17,11 @@
 
 from executorch.exir._serialize._cord import Cord
 from executorch.exir._serialize._dataclass import _DataclassEncoder, _json_to_dataclass
-from executorch.exir._serialize._flatbuffer import (
-    _FlatbufferResult,
-    _program_flatbuffer_to_json,
+from executorch.exir._serialize._flatbuffer import _FlatbufferResult
+from executorch.exir._serialize._flatbuffer_program import (
+    _flatbuffer_to_program,
+    _program_to_flatbuffer,
 )
-from executorch.exir._serialize._flatbuffer_program import _program_to_flatbuffer
 from executorch.exir._serialize._named_data_store import (
     NamedDataStore,
     NamedDataStoreOutput,
@@ -757,9 +757,7 @@ def deserialize_pte_binary(program_data: bytes) -> PTEFile:
         segment_base_offset = eh.segment_base_offset
 
     # Parse the flatbuffer data.
-    program: Program = _json_to_program(
-        _program_flatbuffer_to_json(program_data[:program_size])
-    )
+    program: Program = _flatbuffer_to_program(program_data[:program_size])
 
     if segment_base_offset != 0:
         # Move segment data back into the Program.

exir/_serialize/test/test_flatbuffer_program.py

@@ -11,7 +11,10 @@
     _program_flatbuffer_to_json,
     _program_json_to_flatbuffer,
 )
-from executorch.exir._serialize._flatbuffer_program import _program_to_flatbuffer
+from executorch.exir._serialize._flatbuffer_program import (
+    _flatbuffer_to_program,
+    _program_to_flatbuffer,
+)
 from executorch.exir._serialize._program import _json_to_program, _program_to_json
 from executorch.exir.backend.compile_spec_schema import CompileSpec
 from executorch.exir.schema import (
@@ -172,6 +175,13 @@ def test_roundtrip_via_json(self) -> None:
         program2 = _json_to_program(_program_flatbuffer_to_json(result.data))
         self.assertEqual(program2, program)
 
+    def test_roundtrip_via_direct_python(self) -> None:
+        program = self._make_program()
+        result = _program_to_flatbuffer(
+            program, constant_tensor_alignment=32, delegate_alignment=64
+        )
+        self.assertEqual(_flatbuffer_to_program(result.data), program)
+
     def test_flatbuffer_paths_match(self) -> None:
         program = self._make_program()
         cases = [