executorch 1.6.0__cp315-cp315-win_amd64.whl

executorch/devtools/debug_format/et_schema.py

@@ -0,0 +1,432 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+"""
+Intermediate Representation of ExecuTorch Concepts in Developer Tools
+"""
+
+from __future__ import annotations
+
+import operator
+import warnings
+
+from collections import defaultdict
+
+from enum import Enum
+from types import NoneType
+from typing import Any, Dict, List, Optional, Set, Tuple
+
+import torch
+from executorch import exir
+from executorch.devtools.debug_format.base_schema import (
+    Node,
+    OperatorGraph,
+    OperatorNode,
+    ValueNode,
+)
+
+from torch._higher_order_ops.auto_functionalize import (
+    auto_functionalized,
+    auto_functionalized_v2,
+)
+from torch._subclasses import FakeTensor
+
+
+# Keywords used in debug_format Metadata
+class RESERVED_METADATA_ARG(Enum):
+    DEBUG_HANDLE = "debug_handle"
+    MODULE_STACK = "nn_module_stack"
+    SOURCE_FN_STACK = "source_fn_stack"
+    MODULE_TYPE = "module_type"
+    PROFILE_START_TIME = "profile_start_time"
+    PROFILE_END_TIME = "profile_end_time"
+    LOAD_START_TIME = "load_start_time"
+    LOAD_END_TIME = "load_end_time"
+    MEMORY_USAGE = "memory_usage"
+    DEBUG_ENTRY = "debug_entry"
+    STACK_TRACE = "stack_trace"
+    DEBUG_DATA = "debug_data"
+
+    METRICS_KEYWORD = "metrics"
+    PROFILE_SUMMARY_COLDSTART = "Coldstart"
+    PROFILE_SUMMARY_AVERAGE = "Average"
+    PROFILE_SUMMARY_P90 = "P90"
+    PROFILE_SUMMARY_P10 = "P10"
+    PROFILE_SUMMARY_MIN = "Min"
+    PROFILE_SUMMARY_MAX = "Max"
+
+    AGGREGATED_OP_TABLE = "Aggregated op stats"
+    RUN_SUMMARY_INDIVIDUAL_RUNS_TABLE = "Run summary individual stats"
+    RUN_SUMMARY_TABLE = "Aggregated run summary stats"
+    OP_INSTANCE_SUMMARY_TABLE = "Individual op stats"
+
+    TABLES_KEYWORD = "tables"
+    KV_KEYWORD = "kv"
+    MODEL_LOAD_TIME_KEY = "Model load time (ms)"
+
+
+# Representation of an FX GraphModule as an OperatorGraph
+class FXOperatorGraph(OperatorGraph):
+    @staticmethod
+    def _get_node_name(node: torch.fx.Node) -> str:
+        if node.target == operator.getitem:
+            # pyre-ignore[9]: Incompatible variable type
+            node = node.args[0]
+            assert isinstance(
+                node, torch.fx.Node
+            ), f"First argument of getitem must be a torch fx node. Got {node.args[0]}"
+
+        # Adding the "_" to the node name prevents TensorBoard from collapsing
+        # nodes with similar names that only differ by an integer at the end of
+        # their name.
+        return node.name + "_"
+
+    @staticmethod
+    def _get_op_name(node: torch.fx.Node) -> str:
+        # pyre-ignore[16]: has no attribute `__name__`.
+        return node.target.__name__
+
+    # Given a node and its metadata (if containing module stack), update the provided module mappings
+    @staticmethod
+    def _update_module_mapping(
+        node: Node,
+        module_mapping: Dict[Tuple[str, str], List[Node]],
+        metadata: Dict[str, Any],
+    ):
+        if (
+            source_fn_stack := metadata.get("source_fn_stack")
+        ) is not None and "nn_module_stack" in metadata:
+            # (module name, module type)
+            source_fn = source_fn_stack[-1]
+            module_type = (
+                source_fn[1] if isinstance(source_fn[1], str) else source_fn[1].__name__
+            )
+            module_mapping[(source_fn[0], module_type)].append(node)
+
+    @staticmethod
+    def _parse_args(  # noqa: C901
+        node: torch.fx.Node,
+        nodes: Dict[str, Node],
+        const_count: int,
+        module_mapping: Dict[Tuple[str, str], List[Node]],
+        enable_module_hierarchy: bool,
+    ) -> Tuple[List[Node], int]:
+        inputs = []
+        op = node.op
+        name = node.name
+        args = node.args
+        kwargs = node.kwargs
+        named_args = None
+        if node.op == "call_function" and hasattr(node.target, "_schema"):
+            # pyre-ignore
+            named_args = node.target._schema.arguments
+
+        if node.op == "call_function" and (
+            node.target == auto_functionalized or node.target == auto_functionalized_v2
+        ):
+            # for functioanlized HOPs, args for the corresponding functional op are stored in kwargs
+            args = tuple(kwargs.values())
+
+        for index, arg in enumerate(args):
+            if isinstance(arg, torch.fx.node.Node):
+                if arg.target == exir.memory.alloc:
+                    continue
+                arg_name = FXOperatorGraph._get_node_name(arg)
+            elif isinstance(arg, (int, float, torch.dtype, str)):
+                # e.g. The "0" from node.args of squeeze_copy (mm_default, 0)
+                if named_args and len(named_args) > index:
+                    arg_name = named_args[index].name + "_" + str(const_count)
+                else:
+                    arg_name = "CONST_" + str(const_count)
+                const_count += 1
+                const_node = ValueNode(arg_name, val=str(arg))
+                nodes[arg_name] = const_node
+                if enable_module_hierarchy:
+                    FXOperatorGraph._update_module_mapping(
+                        const_node, module_mapping, node.meta
+                    )
+            elif isinstance(arg, list):
+                arg_name: List[str] = []
+                for list_arg in arg:
+                    if isinstance(list_arg, (int, float)):
+                        # Consider the whole list of ints/floats as a single constant and
+                        # stringify that.
+                        if named_args and len(named_args) > index:
+                            arg_name = [named_args[index].name + "_" + str(const_count)]
+                        else:
+                            arg_name = ["CONST_" + str(const_count)]
+                        const_count += 1
+                        const_node = ValueNode(arg_name[0], val=arg)
+                        nodes[arg_name[0]] = const_node
+                        if enable_module_hierarchy:
+                            FXOperatorGraph._update_module_mapping(
+                                const_node, module_mapping, node.meta
+                            )
+                        break
+                    elif isinstance(list_arg, torch.fx.node.Node):
+                        arg_name += [FXOperatorGraph._get_node_name(list_arg)]
+                    elif list_arg is None:
+                        arg_name += ["CONST_NONE" + str(const_count)]
+                        const_count += 1
+                        const_node = ValueNode(arg_name[-1], val=str(arg))
+                        nodes[arg_name[-1]] = const_node
+                        if enable_module_hierarchy:
+                            FXOperatorGraph._update_module_mapping(
+                                const_node, module_mapping, node.meta
+                            )
+                    else:
+                        raise Exception(
+                            f"Unsupported argument encountered in list {arg}, {type(arg[0])}"
+                        )
+            elif isinstance(arg, NoneType):
+                continue
+            else:
+                raise Exception(
+                    f"Unsupported argument encountered {op}, {name}, {arg}, {type(arg)}"
+                )
+
+            if isinstance(arg_name, list):
+                for val in arg_name:
+                    inputs.append(nodes[val])
+            else:
+                inputs.append(nodes[arg_name])
+        for _, node in kwargs.items():
+            # We can ignore the out kwarg as that's mostly used to pass in the output tensor
+            # which has been memory planned. The same value is also returned by the operator
+            # which is then consumed by other nodes in the graph.
+            if (
+                isinstance(node, torch.fx.node.Node)
+                and node.target == exir.memory.alloc
+            ):
+                continue
+            else:
+                warnings.warn(
+                    f"Unsupported kwarg encountered: {name}, {kwargs}", stacklevel=1
+                )
+
+        return inputs, const_count
+
+    # Given an FX GraphModule, parse it into an OperatorGraph
+    @staticmethod
+    def gen_operator_graph(
+        model: torch.fx.GraphModule,
+        skip_stack_trace: Optional[bool] = False,
+        enable_module_hierarchy: bool = False,
+    ) -> FXOperatorGraph:
+        graph: torch.fx.Graph = model.graph
+
+        nodes = {}
+        input_nodes = {}
+        output_nodes = {}
+        out_variant_output_nodes = set()
+        module_mapping = defaultdict(list)
+
+        const_count = 0
+        for fx_node in graph.nodes:
+            if (
+                fx_node.target == exir.memory.alloc
+                or fx_node.target == operator.getitem
+            ):
+                continue
+            op = fx_node.op
+            name = FXOperatorGraph._get_node_name(fx_node)
+            target = fx_node.target
+            args = fx_node.args
+            kwargs = fx_node.kwargs
+            metadata = FXOperatorGraph._extract_metadata(fx_node.meta, skip_stack_trace)
+            output_shapes = FXOperatorGraph._extract_output_shapes(
+                fx_node.meta.get("val")
+            )
+            dtype = FXOperatorGraph._extract_output_dtype(fx_node.meta.get("val")) or ""
+
+            assert (
+                op != "call_module"
+            ), f"Module call not yet supported in edge model graph [.toEdge()]: {name}, {str(target)}"
+            assert (
+                op != "call_method"
+            ), f"Call Method not yet supported in edge model graph [.toEdge()]: {name}, {str(target)}"
+
+            # Input
+            if op == "placeholder":
+                node = ValueNode(
+                    name,
+                    output_shapes=output_shapes,
+                    metadata=metadata,
+                    dtype=str(dtype),
+                )  # val is default arg
+                input_nodes[name] = node
+            # Constants
+            elif op == "get_attr":
+                node = ValueNode(
+                    name,
+                    output_shapes=output_shapes,
+                    metadata=metadata,
+                    dtype=str(dtype),
+                )
+            # Output
+            elif op == "output":
+                assert len(args) == 1
+                # Args of op=='output' is a wrapped list of return nodes ([ret_1, ret_2, ...], )
+                in_nodes = [
+                    (
+                        nodes[FXOperatorGraph._get_node_name(ret)]
+                        if ret is not None
+                        else []
+                    )
+                    for ret in args[0]
+                ]
+                node = ValueNode(
+                    name,
+                    inputs=in_nodes,
+                    output_shapes=output_shapes,
+                    metadata=metadata,
+                    dtype=str(dtype),
+                )
+                output_nodes[name] = node
+            # Op Calls
+            elif op == "call_function":
+                inputs, const_count = FXOperatorGraph._parse_args(
+                    fx_node, nodes, const_count, module_mapping, enable_module_hierarchy
+                )
+                named_args = []
+                if fx_node.op == "call_function" and hasattr(fx_node.target, "_schema"):
+                    named_args = [arg.name for arg in fx_node.target._schema.arguments]
+                node = OperatorNode(
+                    name,
+                    inputs=inputs,
+                    output_shapes=output_shapes,
+                    metadata=metadata,
+                    op=FXOperatorGraph._get_op_name(fx_node),
+                    named_args=named_args,
+                )
+                if enable_module_hierarchy:
+                    FXOperatorGraph._update_module_mapping(
+                        node, module_mapping, fx_node.meta
+                    )
+
+                for kwarg_name, kwarg in kwargs.items():
+                    if (
+                        isinstance(kwarg, torch.fx.node.Node)
+                        and kwarg.target == exir.memory.alloc
+                        and kwarg_name == "out"
+                    ):
+                        nodes[FXOperatorGraph._get_node_name(kwarg)] = node
+                        out_variant_output_nodes.add(
+                            FXOperatorGraph._get_node_name(kwarg)
+                        )
+            else:
+                raise Exception(f"Unsupported op type encountered {op}, {name}")
+
+            nodes[name] = node
+        return FXOperatorGraph._compose_op_graph(
+            "base",
+            nodes,
+            input_nodes,
+            output_nodes,
+            out_variant_output_nodes,
+            module_mapping,
+        )
+
+    @staticmethod
+    def _compose_op_graph(
+        name: str,
+        nodes: Dict[str, Node],
+        input_nodes: Dict[
+            str, Node | OperatorGraph
+        ],  # Never OperatorGraph, annotated for Pyre
+        output_nodes: Dict[
+            str, Node | OperatorGraph
+        ],  # Never OperatorGraph, annotated for Pyre
+        out_variant_output_nodes: Set[str],
+        module_mapping: Dict[
+            Tuple[str, str], List[Any]
+        ],  # Any used here for Pyre, list of Nodes
+    ):
+        # Generate Module Graphs
+        module_graphs: List[OperatorGraph] = []
+        for module_key, module_nodes in module_mapping.items():
+            module_element = OperatorGraph(
+                graph_name=module_key[0],
+                elements=module_nodes,
+                metadata={"module_type": module_key[1]},
+            )
+            module_graphs.append(module_element)
+
+            # Remove module modes from main graph
+            for node in module_nodes:
+                nodes.pop(node.name)
+
+        main_nodes = [
+            node
+            for name, node in nodes.items()
+            if name not in input_nodes
+            and name not in output_nodes
+            and name not in out_variant_output_nodes
+        ]
+        main_graph = FXOperatorGraph(
+            graph_name="forward", elements=main_nodes + module_graphs
+        )
+        input_graph = FXOperatorGraph(
+            graph_name="inputs", elements=list(input_nodes.values())
+        )
+        output_graph = FXOperatorGraph(
+            graph_name="outputs", elements=list(output_nodes.values())
+        )
+
+        return FXOperatorGraph(
+            graph_name=name,
+            elements=[input_graph, main_graph, output_graph],
+        )
+
+    # Given a dict, extract only the utilized metadata
+    @staticmethod
+    def _extract_metadata(
+        metadata: Dict[str, Any], skip_stack_trace: Optional[bool] = False
+    ) -> Dict[str, Any]:
+        ret = {}
+        if RESERVED_METADATA_ARG.DEBUG_HANDLE.value in metadata:
+            ret[RESERVED_METADATA_ARG.DEBUG_HANDLE.value] = metadata[
+                RESERVED_METADATA_ARG.DEBUG_HANDLE.value
+            ]
+        if not skip_stack_trace and RESERVED_METADATA_ARG.STACK_TRACE.value in metadata:
+            ret[RESERVED_METADATA_ARG.STACK_TRACE.value] = metadata[
+                RESERVED_METADATA_ARG.STACK_TRACE.value
+            ]
+        if RESERVED_METADATA_ARG.MODULE_STACK.value in metadata:
+            ret[RESERVED_METADATA_ARG.MODULE_STACK.value] = metadata[
+                RESERVED_METADATA_ARG.MODULE_STACK.value
+            ]
+        return ret
+
+    @staticmethod
+    def _extract_output_shapes(val: Any) -> Optional[List[List[int]]]:
+        if isinstance(val, (FakeTensor, torch.Tensor)):
+            # If val is a single tensor
+            return [list(val.shape)]
+        elif isinstance(val, tuple) and all(
+            isinstance(tensor, (FakeTensor, torch.Tensor)) for tensor in val
+        ):
+            # If val is a tuple of tensors
+            shapes = [list(fake_tensor.shape) for fake_tensor in val]
+            return shapes
+        else:
+            return None
+
+    @staticmethod
+    def _extract_output_dtype(val: Any) -> Optional[List[torch.dtype]]:
+        if isinstance(val, (FakeTensor, torch.Tensor)):
+            # If val is a single tensor
+            return [val.dtype]
+        elif isinstance(val, tuple) and all(
+            isinstance(tensor, (FakeTensor, torch.Tensor)) for tensor in val
+        ):
+            # If val is a tuple of tensors
+            dtypes = [fake_tensor.dtype for fake_tensor in val]
+            return dtypes
+        else:
+            return None

executorch/devtools/etdump/etdump_schema_flatcc.fbs

@@ -0,0 +1,181 @@
+include "scalar_type.fbs";
+
+namespace etdump;
+
+// Identifier of a valid executor schema.
+file_identifier "ED00";
+// Extension of written files.
+file_extension "etdp";
+
+table Null {}
+
+table Tensor {
+  scalar_type:executorch_flatbuffer.ScalarType;
+  sizes:[long];
+  strides:[long];
+  offset:long;
+}
+
+table Int {
+  int_val:long;
+}
+
+table Bool {
+  bool_val:bool;
+}
+
+table Float {
+  float_val:float;
+}
+
+table Double {
+  double_val:double;
+}
+
+table String {
+  string_val:string;
+}
+
+table TensorList {
+  tensors:[Tensor];
+}
+
+enum ValueType : byte { Null, Int, Bool, Float, Double, Tensor, TensorList, String,}
+
+// This table is only necessary because flatbuffer can't support a list of
+// union elements directly, they need to be wrapped in a table.
+table Value {
+    val:ValueType;
+    tensor:Tensor;
+    tensor_list:TensorList;
+    int_value:Int;
+    float_value:Float;
+    double_value:Double;
+    bool_value:Bool;
+    output:Bool;
+}
+
+// This table contains all the details that we store to debug an op executed in the
+// runtime. Each entry here reprsents an op executed in the runtime during inference.
+table DebugEvent {
+  // The chain to which this instruction belongs to. For now it will always be 0,
+  // as chains are not used, but left in place in case they are used in the future.
+  chain_index:ulong;
+
+  // Runtime instruction id to which this event corresponds to.
+  instruction_id:int = -1;
+
+  // List of debug entries corresponding to this debug event. These could be
+  // scalars, tensors etc.
+  // The order of these entries is expected to be the same as the parameters
+  // that are passed into a operator.
+  debug_entry:Value;
+
+  // Integer based delegate debug identifier.
+  delegate_debug_id_int:int = -1;
+
+  // String based delegate debug identifier.
+  delegate_debug_id_str:string;
+
+  // Name assigned to this debug event by the runtime. If it is an operator
+  // call this will just be the name of the operator that was executed.
+  name:string;
+}
+
+// All the details pertaining to an allocation done in the runtime. The main
+// details that we care about currently are the allocator id (which directly maps
+// to the allocators stored in the RunData.allocators vector) and the size of the
+// allocation done from it.
+table AllocationEvent {
+  // Allocator id of the allocator from which this memory was allocated.
+  allocator_id:int;
+
+  // Size of allocation in bytes.
+  allocation_size:ulong;
+}
+
+// This table contains all the details we need to represent a profiling event that
+// has occurred in the runtime. These could be an operator profiling event or something
+// more generic like the total time taken to execute an inference loop.
+table ProfileEvent {
+  // Name assigned to this profiling event by the runtime. If it is an operator
+  // call this will just be the name of the operator that was executed.
+  name:string;
+
+  // The chain to which this instruction belongs to. For now it will always be 0,
+  // as chains are not used, but left in place in case they are used in the future.
+  chain_index:int;
+
+  // Runtime instruction id to which this event corresponds to.
+  instruction_id:int = -1;
+
+  // If this is a delegate event then these will be the corresponding delegate
+  // debug identifiers of the event which occurred in the delegate backend. This
+  // should be the same delegate debug identifier that was generated AOT.
+
+  // Integer based delegate debug identifier.
+  delegate_debug_id_int:int = -1;
+
+  // String based delegate debug identifier.
+  delegate_debug_id_str:string;
+
+  // Delegate debug metadata bytes.
+  delegate_debug_metadata:[ubyte];
+
+  // Time at which this event started. Could be in units of time or CPU cycles.
+  start_time:ulong;
+
+  // Time at which this event ended. Could be in units of time or CPU cycles.
+  end_time:ulong;
+}
+
+// This table contains all the details we need to represent a profiling, allocation, or
+// debug event that has occurred in the runtime. These could be an operator profiling
+// event or something more generic like the total time taken to execute an inference loop.
+// Since we can only keep one vector open at a time when building the flatbuffer, using a
+// generic Event type (as opposed to separate arrays for each type) allows us to add
+// different events as they occur.
+//
+// Exactly one of these fields must be non-null, and the rest must be set to null.
+table Event {
+  profile_event: ProfileEvent;
+
+  allocation_event: AllocationEvent;
+
+  debug_event: DebugEvent;
+}
+
+// Representation of an ExecuTorch memory allocator that is used in the runtime.
+table Allocator {
+  // Name of the allocator.
+  name:string;
+}
+
+// This table contains the total set of debug and profiling data that was
+// collected for a single run. It's not always necessary that the profiling
+// data and debug data map 1:1.
+table RunData {
+  name: string;
+
+  // If bundled input was used to run this model on device then this
+  // entry will contain the index of the bundled input that was used
+  // to run this specific block.
+  bundled_input_index:int = -1;
+
+  // List of allocators on which profiling was enabled in the runtime.
+  allocators:[Allocator];
+
+  events: [Event];
+}
+
+table ETDump {
+  // Schema version.
+  version:uint;
+
+  // This is the base level abstraction. One entry here will correspond
+  // to one run.
+  run_data:[RunData];
+
+}
+
+root_type ETDump;

executorch/devtools/etdump/scalar_type.fbs

@@ -0,0 +1,43 @@
+// Copyright (c) Meta Platforms, Inc. and affiliates.
+
+//
+// See README.md before modifying this file.
+//
+
+namespace executorch_flatbuffer;
+
+// The scalar data type.
+// Must match executorch/runtime/core/portable_type/tensor_impl.h
+enum ScalarType : byte {
+  BYTE = 0,
+  CHAR = 1,
+  SHORT = 2,
+  INT = 3,
+  LONG = 4,
+  HALF = 5,
+  FLOAT = 6,
+  DOUBLE = 7,
+  BOOL = 11,
+  QINT8 = 12,
+  QUINT8 = 13,
+  QINT32 = 14,
+  BFLOAT16 = 15,
+  QUINT4X2 = 16,
+  QUINT2X4 = 17,
+  BITS16 = 22,
+  FLOAT8E5M2 = 23,
+  FLOAT8E4M3FN = 24,
+  FLOAT8E5M2FNUZ = 25,
+  FLOAT8E4M3FNUZ = 26,
+  UINT16 = 27,
+  UINT32 = 28,
+  UINT64 = 29,
+  // Types currently not implemented.
+  // COMPLEXHALF = 8,
+  // COMPLEXFLOAT = 9,
+  // COMPLEXDOUBLE = 10,
+  // BITS1x8 = 18,
+  // BITS2x4 = 19,
+  // BITS4x2 = 20,
+  // BITS8 = 21,
+}