xref: /aosp_15_r20/external/executorch/devtools/etdump/schema_flatcc.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
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# pyre-strict
8"""
9This file is the python representation of the schema contained in
10executorch/devtools/etdump/etdump_schema.fbs. Any changes made to that
11flatbuffer schema should accordingly be reflected here also.
12"""
13
14from dataclasses import dataclass
15from enum import Enum
16from typing import List, Optional
17
18from executorch.exir.scalar_type import ScalarType
19
20
21@dataclass
22class Tensor:
23    scalar_type: ScalarType
24    sizes: List[int]
25    strides: List[int]
26    offset: Optional[int]
27
28
29@dataclass
30class TensorList:
31    tensors: List[Tensor]
32
33
34@dataclass
35class Null:
36    pass
37
38
39@dataclass
40class Int:
41    int_val: int
42
43
44@dataclass
45class Bool:
46    bool_val: bool
47
48
49@dataclass
50class Double:
51    double_val: float
52
53
54@dataclass
55class Float:
56    float_val: float
57
58
59@dataclass
60class String:
61    string_val: str
62
63
64@dataclass
65class ContainerMetadata:
66    encoded_inp_str: str
67    encoded_out_str: str
68
69
70@dataclass
71class ValueType(Enum):
72    NULL = "Null"
73    INT = "Int"
74    BOOL = "Bool"
75    FLOAT = "Float"
76    DOUBLE = "Double"
77    TENSOR = "Tensor"
78    TENSOR_LIST = "TensorList"
79    STRING = "String"
80
81
82@dataclass
83class Value:
84    val: str  # Member of ValueType
85    tensor: Optional[Tensor]
86    tensor_list: Optional[TensorList]
87    int_value: Optional[Int]
88    float_value: Optional[Float]
89    double_value: Optional[Double]
90    bool_value: Optional[Bool]
91    output: Optional[Bool]
92
93
94@dataclass
95class DebugEvent:
96    name: Optional[str]
97    chain_index: int
98    instruction_id: int
99    delegate_debug_id_int: Optional[int]
100    delegate_debug_id_str: Optional[str]
101    debug_entry: Value
102
103
104# Note the differing value style is a result of ETDump string
105class PROFILE_EVENT_ENUM(Enum):
106    RUN_MODEL = "Method::execute"
107    OPERATOR_CALL = "OPERATOR_CALL"
108    DELEGATE_CALL = "DELEGATE_CALL"
109    LOAD_MODEL = "Program::load_method"
110
111
112@dataclass
113class ProfileEvent:
114    name: Optional[str]
115    chain_index: int
116    instruction_id: int
117    delegate_debug_id_int: Optional[int]
118    delegate_debug_id_str: Optional[str]
119    delegate_debug_metadata: Optional[bytes]
120    start_time: int
121    end_time: int
122
123
124@dataclass
125class AllocationEvent:
126    allocator_id: int
127    allocation_size: int
128
129
130@dataclass
131class Allocator:
132    name: str
133
134
135# Must have one of profile_event, allocation_event, or debug_event
136@dataclass
137class Event:
138    profile_event: Optional[ProfileEvent]
139    allocation_event: Optional[AllocationEvent]
140    debug_event: Optional[DebugEvent]
141
142
143@dataclass
144class RunData:
145    name: str
146    bundled_input_index: Optional[int]
147    allocators: Optional[List[Allocator]]
148    events: Optional[List[Event]]
149
150
151@dataclass
152class ETDumpFlatCC:
153    version: int
154    run_data: List[RunData]
155