• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7# flake8: noqa: F401
8import functools
9import inspect
10import os
11import random
12import unittest
13from typing import Callable, Dict, Optional, Tuple, Type
14
15import executorch.exir as exir
16
17import executorch.exir.control_flow as control_flow
18
19# @manual=//executorch/extension/pytree:pybindings
20import executorch.extension.pytree as pytree
21
22import torch
23
24from executorch.devtools.bundled_program.core import BundledProgram
25from executorch.devtools.bundled_program.serialize import (
26    serialize_from_bundled_program_to_flatbuffer,
27)
28
29from executorch.devtools.bundled_program.util.test_util import (
30    get_common_executorch_program,
31    SampleModel,
32)
33
34kernel_mode = None  # either aten mode or lean mode
35try:
36    from executorch.extension.pybindings.portable_lib import (
37        _load_bundled_program_from_buffer,
38        _load_for_executorch_from_buffer,
39        _load_for_executorch_from_bundled_program,
40    )
41
42    kernel_mode = "lean"
43except ImportError as e:
44    print(e)
45    pass
46
47try:
48    from executorch.extension.pybindings.aten_lib import (  # @manual=//executorch/extension/pybindings:aten_lib
49        _load_bundled_program_from_buffer,
50        _load_for_executorch_from_buffer,
51        _load_for_executorch_from_bundled_program,
52    )
53
54    assert kernel_mode is None
55    kernel_mode = "aten"
56except ImportError as e:
57    print(e)
58    pass
59
60assert kernel_mode is not None
61
62
63class BundledProgramE2ETest(unittest.TestCase):
64    def test_sample_model_e2e(self):
65        executorch_program, method_test_suites = get_common_executorch_program()
66        eager_model = SampleModel()
67
68        bundled_program = BundledProgram(executorch_program, method_test_suites)
69
70        bundled_program_buffer = serialize_from_bundled_program_to_flatbuffer(
71            bundled_program
72        )
73
74        executorch_bundled_program = _load_bundled_program_from_buffer(
75            bundled_program_buffer
76        )
77
78        executorch_module = _load_for_executorch_from_bundled_program(
79            executorch_bundled_program
80        )
81
82        for method_name in eager_model.method_names:
83            executorch_module.load_bundled_input(
84                executorch_bundled_program,
85                method_name,
86                0,
87            )
88            executorch_module.plan_execute(method_name)
89            executorch_module.verify_result_with_bundled_expected_output(
90                executorch_bundled_program,
91                method_name,
92                0,
93            )
94