| # 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-strict |
| |
| import copy |
| import re |
| import reprlib |
| from dataclasses import fields |
| from enum import IntEnum |
| from typing import Any, List, Optional, TextIO |
| |
| import torch |
| from executorch.exir.error import ExportError, ExportErrorType, InternalError |
| |
| from executorch.exir.schema import ( |
| Bool, |
| BoolList, |
| DelegateCall, |
| Double, |
| DoubleList, |
| EValue, |
| Frame, |
| FrameList, |
| FreeCall, |
| Int, |
| IntList, |
| JumpFalseCall, |
| KernelCall, |
| MoveCall, |
| Null, |
| OptionalTensorList, |
| Program, |
| ScalarType, |
| String, |
| Tensor, |
| TensorList, |
| TensorShapeDynamism, |
| ) |
| |
| |
| def _scalar_type_str(scalar_type: ScalarType) -> str: |
| type2str = { |
| ScalarType.BYTE: "bt", |
| ScalarType.CHAR: "c", |
| ScalarType.SHORT: "s", |
| ScalarType.INT: "i", |
| ScalarType.LONG: "l", |
| ScalarType.HALF: "h", |
| ScalarType.FLOAT: "f", |
| ScalarType.DOUBLE: "d", |
| ScalarType.COMPLEX32: "c32", |
| ScalarType.COMPLEX64: "c64", |
| ScalarType.COMPLEX128: "c128", |
| ScalarType.BOOL: "b", |
| ScalarType.QINT8: "qi8", |
| ScalarType.QUINT8: "qui8", |
| ScalarType.QINT32: "qi32", |
| ScalarType.BFLOAT16: "bf16", |
| ScalarType.QUINT4x2: "qui4x2", |
| ScalarType.QUINT2x4: "qui2x4", |
| } |
| if not (ret := type2str.get(scalar_type, None)): |
| raise RuntimeError(f"Unrecognized scalar_type: {scalar_type}") |
| else: |
| return ret |
| |
| |
| def _is_dynamic_shape_tensor(tensor: Tensor) -> bool: |
| return tensor.shape_dynamism != TensorShapeDynamism.STATIC |
| |
| |
| def _format_evalue( # noqa: C901 |
| evalue: EValue, show_meminfo: bool, mark_dynamic_shape_tensor: bool |
| ) -> str: |
| evstr = "\033[34m" |
| if isinstance(evalue.val, Tensor): |
| tensor = evalue.val |
| if tensor.data_buffer_idx > 0: |
| assert not _is_dynamic_shape_tensor( |
| tensor |
| ), "A constant tensor can not be dynamic shape" |
| evstr += "CT" # constant tensor |
| assert tensor.allocation_info is None |
| else: |
| if mark_dynamic_shape_tensor: |
| if tensor.shape_dynamism == TensorShapeDynamism.DYNAMIC_BOUND: |
| evstr += "UB" # upper bound tensor will be shown as 'UBT' |
| elif tensor.shape_dynamism == TensorShapeDynamism.DYNAMIC_UNBOUND: |
| evstr += "DU" # dynamic unbound tensor will be shown as 'DUT' |
| evstr += "T" |
| if show_meminfo: |
| if tensor.allocation_info: |
| evstr += f"m{tensor.allocation_info.memory_id}.{tensor.allocation_info.memory_offset}" |
| else: |
| evstr += "m." |
| evstr += f"{tensor.sizes}{_scalar_type_str(tensor.scalar_type)}" |
| elif isinstance(evalue.val, TensorList): |
| evstr += "TL" |
| tensorlist = evalue.val |
| # pyre-ignore |
| evstr += str(tensorlist.items) |
| elif isinstance(evalue.val, OptionalTensorList): |
| evstr += "OTL" |
| optionaltensorlist = evalue.val |
| # pyre-ignore |
| evstr += str(optionaltensorlist.items) |
| elif isinstance(evalue.val, IntList): |
| evstr += "IL" |
| intlist = evalue.val |
| # pyre-ignore |
| evstr += str(intlist.items) |
| elif isinstance(evalue.val, DoubleList): |
| evstr += "DL" |
| doublelist = evalue.val |
| # pyre-ignore |
| evstr += str(doublelist.items) |
| elif isinstance(evalue.val, BoolList): |
| evstr += "BL" |
| boollist = evalue.val |
| # pyre-ignore |
| evstr += str(boollist.items) |
| elif isinstance(evalue.val, Int): |
| intval = evalue.val |
| evstr += f"I{intval.int_val}" |
| elif isinstance(evalue.val, Double): |
| doubleval = evalue.val |
| evstr += f"D{doubleval.double_val}" |
| elif isinstance(evalue.val, Bool): |
| boolval = evalue.val |
| evstr += f"B{int(boolval.bool_val)}" # print 0, 1 since it's shorter than false, true |
| elif isinstance(evalue.val, String): |
| stringval = evalue.val |
| evstr += f"S{stringval.string_val}" |
| elif isinstance(evalue.val, Null): |
| evstr += "N" # for null |
| else: |
| raise RuntimeError(f"Unrecognized type of evalue: {evalue}") |
| evstr += "\033[0m" |
| return evstr |
| |
| |
| def print_program( # noqa: C901 |
| program: Program, |
| show_meminfo: bool = True, |
| mark_dynamic_shape_tensor: bool = False, |
| out: Optional[TextIO] = None, |
| ) -> None: |
| """ |
| Dump the instruction list of a program in a more human readable fashion. |
| |
| The dump follows the following BNF syntax (I combime some regex syntax |
| so the grammar becomes shorter. The grammar is not strict but the main |
| purpose is to let people understand the dump): |
| ``` |
| PROGRAM: (INSTRUCTION)+ |
| INSTRUCTION: SEQUENCE_NO ':' (CALL_KERNEL | JUMP_FALSE) |
| JUMP_FALSE: 'JF' '(' EVALUE ')' '->' TARGET_SEQUENCE_NO |
| CALL_KERNEL: OVERLOADDED_OP_NAME ARGS |
| ARGS: EVALUE | ARGS ',' EVALUE |
| EVALUE: EVALUE_IDX ( TENSOR | INT | BOOL | ...) |
| INT: 'I' ACTUAL_INT_VALUE |
| BOOL: 'B' ZERO_OR_ONE |
| CONST_TENSOR_PREFIX: 'CT' |
| TENSOR: ('T' | CONST_TENSOR_PREFIX) (MEM_ALLOCATION_INFO)? TENSOR_SHAPE TENSOR_DTYPE |
| TENSOR_SHAPE: '[' dim0_size, dim1_size, ..., last_dim_size ']' |
| MEM_ALLOCATION_INFO: PLANNED_MEM_INFO | UNPLANNED_MEM_INFO |
| PLANNED_MEM_INFO: 'm' MEM_LAYER_ID '.' MEM_LAYER_OFFSET |
| UNPLANNED_MEM_INFO: 'm.' |
| ``` |
| |
| To make the dump easier to read, it's colored as follows: |
| 1. input/output EValues are marked as red |
| 2. EValue types (or more specifically tensor types with size and dtype) are marked as blue |
| """ |
| execution_plan = program.execution_plan[0] |
| operators = execution_plan.operators |
| delegates = execution_plan.delegates |
| chain = execution_plan.chains[0] |
| instructions = chain.instructions |
| inputs: List[int] = execution_plan.inputs |
| outputs: List[int] = execution_plan.outputs |
| values: List[EValue] = execution_plan.values |
| |
| def _format_arg(evalue_idx: int) -> str: |
| def _get_io_index(iolist: List[int], target_evalue_idx: int) -> int: |
| """ |
| The list is short enough so linear scan is proper. |
| """ |
| for io_idx, evalue_idx in enumerate(iolist): |
| if evalue_idx == target_evalue_idx: |
| return io_idx |
| return -1 |
| |
| argstr = str(evalue_idx) |
| if (input_idx := _get_io_index(inputs, evalue_idx)) >= 0: |
| argstr += f"\033[31mI{input_idx}\033[0m" |
| if (output_idx := _get_io_index(outputs, evalue_idx)) >= 0: |
| argstr += f"\033[31mO{output_idx}\033[0m" |
| |
| # EValue type |
| evalue = values[evalue_idx] |
| return argstr + _format_evalue(evalue, show_meminfo, mark_dynamic_shape_tensor) |
| |
| print( |
| f"The program contains the following {len(instructions)} instructions", file=out |
| ) |
| for idx, instr in enumerate(instructions): |
| print(f"{idx:3}: ", end="", file=out) |
| if isinstance(instr.instr_args, KernelCall): |
| kernel = instr.instr_args |
| op = operators[kernel.op_index] |
| args = kernel.args |
| |
| opname = f"{op.name}.{op.overload}" if op.overload else op.name |
| argstr = ",".join(map(_format_arg, args)) |
| print(f"{opname} {argstr}", file=out) |
| elif isinstance(instr.instr_args, DelegateCall): |
| delegate = instr.instr_args |
| backend = delegates[delegate.delegate_index] |
| args = delegate.args |
| backend_id = f"{backend.id}" |
| argstr = ",".join(map(_format_arg, args)) |
| print(f"{backend_id} {argstr}", file=out) |
| elif isinstance(instr.instr_args, JumpFalseCall): |
| jfcall = instr.instr_args |
| print( |
| f"JF ({_format_arg(jfcall.cond_value_index)}) -> {jfcall.destination_instruction}", |
| file=out, |
| ) |
| elif isinstance(instr.instr_args, MoveCall): |
| move_call = instr.instr_args |
| print( |
| f"MOVE {_format_arg(move_call.move_from)} -> {_format_arg(move_call.move_to)}", |
| file=out, |
| ) |
| elif isinstance(instr.instr_args, FreeCall): |
| print(f"FREE {_format_arg(instr.instr_args.value_index)}", file=out) |
| else: |
| raise InternalError(f"Unsupport instruction type {instr}") |
| |
| |
| # pyre-ignore |
| def pretty_print(obj: Any, indent: int = 0, out: Optional[TextIO] = None) -> None: |
| """ |
| Pretty prints the given object which is of the Program type and any of its |
| attribute’s types. |
| """ |
| if isinstance(obj, torch.fx.GraphModule): |
| raise ExportError( |
| ExportErrorType.INVALID_INPUT_TYPE, |
| "pretty_print() does not accept GraphModule as input.", |
| ) |
| |
| # Instruction types are IntEnum object |
| if isinstance(obj, IntEnum): |
| print(int(obj), end="", file=out) |
| return |
| |
| primitives = (int, str, bool, float, type(None)) |
| if isinstance(obj, primitives): |
| print(obj, end="", file=out) |
| return |
| |
| if isinstance(obj, bytes): |
| r = reprlib.Repr() |
| r.maxother = 1024 |
| print(r.repr(obj), end="", file=out) |
| return |
| |
| if isinstance(obj, list): |
| if len(obj) < 10 and all(isinstance(elem, int) for elem in obj): |
| print(obj, end="", file=out) |
| return |
| print("[", file=out) |
| for index, elem in enumerate(obj): |
| print(" " * (indent + 1), end="", file=out) |
| pretty_print(elem, indent + 1, out=out) |
| print(f"(index={index}),", file=out) |
| print(" " * indent + "]", end="", file=out) |
| return |
| |
| inline = all( |
| isinstance(getattr(obj, field.name), primitives) for field in fields(obj) |
| ) |
| end = "" if inline else "\n" |
| print(f"{type(obj).__name__}(", end=end, file=out) |
| for i, _field in enumerate(fields(obj)): |
| if not inline: |
| print(" " * (indent + 1), end="", file=out) |
| print(_field.name + "=", end="", file=out) |
| pretty_print(getattr(obj, _field.name), indent + 1, out=out) |
| if i < len(fields(obj)) - 1: |
| print(", ", end="", file=out) |
| print("", end=end, file=out) |
| if not inline: |
| print(" " * indent, end="", file=out) |
| print(")", end="" if indent else "\n", file=out) |
| |
| |
| def pretty_print_stacktraces(obj: FrameList) -> str: |
| """ |
| Pretty prints the traceback for one instruction |
| """ |
| pretty = "Traceback (most recent call last): \n" |
| for frame in obj.items: |
| pretty += f' File "{frame.filename}", ' |
| pretty += f"line {str(frame.lineno)}, in {frame.name}\n" |
| pretty += f"{frame.context} \n" |
| pretty += "\n" |
| return pretty |
| |
| |
| def add_cursor_to_graph(graph: torch.fx.Graph, finding_node: torch.fx.Node) -> str: |
| """ |
| Insert a cursor at the node location in the fx.Graph. |
| e.g: |
| # graph(): |
| # %x : [#users=1] = placeholder[target=x] |
| # %param : [#users=1] = get_attr[target=param] |
| # %add : [#users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {}) |
| # --> %linear : [#users=1] = call_module[target=linear](args = (%add,), kwargs = {}) |
| # %clamp : [#users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0}) |
| # return clamp |
| |
| This is mostly used for error reporting |
| """ |
| |
| new_graph = copy.deepcopy(graph) |
| |
| found_at = -1 |
| for ix, node in enumerate(graph.nodes): |
| if node == finding_node: |
| found_at = ix |
| |
| # This is heavily based on __str__ method of fx.Graph |
| def _format_graph(graph: torch.fx.Graph, offending_node_idx: int) -> str: |
| s = "graph():" |
| for ix, node in enumerate(graph.nodes): |
| node_str = node.format_node() |
| if node_str: |
| if ix != offending_node_idx: |
| s += "\n " + node_str |
| else: |
| s += "\n--> " + node_str |
| return s |
| |
| return _format_graph(new_graph, found_at) |
| |
| |
| def _stacktrace_to_framelist(stacktrace: str) -> FrameList: |
| """Creates a frame list from a stacktrace string.""" |
| pattern = r'File "(.*?)", line (\d+), in (.*?)\n' |
| matches = re.findall(pattern, stacktrace) |
| mapped_frame_list = [ |
| Frame( |
| filename=match[0], |
| lineno=int(match[1]), |
| name=match[2], |
| context=stacktrace.split("\n")[i * 2 + 1].strip(), |
| ) |
| for i, match in enumerate(matches) |
| ] |
| return FrameList(mapped_frame_list) |
| |
| |
| def inspect_node(graph: torch.fx.Graph, node: torch.fx.Node) -> str: |
| """ |
| Inspect a node by highlighting the node in the graph as well as the stacktrace. |
| |
| Args: |
| graph: The graph containing the node |
| node: The node to be inspected |
| |
| Return: A string. An example output is: |
| |
| _param_constant0 error_msg: Here is the failing node in the graph module: |
| graph(): |
| %arg0_1 : [num_users=1] = placeholder[target=arg0_1] |
| --> %_param_constant0 : [num_users=1] = get_attr[target=_param_constant0] |
| %_param_constant1 : [num_users=1] = get_attr[target=_param_constant1] |
| %aten_convolution_default : [num_users=2] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%arg0_1, %_param_constant0, %_param_constant1, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {}) |
| %_param_constant2 : [num_users=1] = get_attr[target=_param_constant2] |
| %_param_constant3 : [num_users=1] = get_attr[target=_param_constant3] |
| %aten_convolution_default_1 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%aten_convolution_default, %_param_constant2, %_param_constant3, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {}) |
| %aten_add_tensor : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.add.Tensor](args = (%aten_convolution_default, %aten_convolution_default_1), kwargs = {}) |
| %_param_constant4 : [num_users=1] = get_attr[target=_param_constant4] |
| %_param_constant5 : [num_users=1] = get_attr[target=_param_constant5] |
| %aten_convolution_default_2 : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.convolution.default](args = (%aten_add_tensor, %_param_constant4, %_param_constant5, [1, 1], [0, 0], [1, 1], False, [0, 0], 1), kwargs = {}) |
| %aten_gelu_default : [num_users=1] = call_function[target=executorch.exir.dialects.edge._ops.aten.gelu.default](args = (%aten_convolution_default_2,), kwargs = {}) |
| return [aten_gelu_default] |
| This node _param_constant0 has metadata of: |
| The node stacktrace: |
| Traceback (most recent call last): |
| File "/tmp/ipykernel_1204253/3382880687.py", line 7, in forward |
| return self.test_model(x) |
| File "/mnt/xarfuse/uid-25337/7b86ad0c-seed-nspid4026532987_cgpid2707357-ns-4026532984/torch/nn/modules/module.py", line 1528, in _call_impl |
| return forward_call(*args, **kwargs) |
| File "/tmp/ipykernel_1204253/712280972.py", line 10, in forward |
| a = self.conv1(x) |
| |
| """ |
| graph_str_with_cursor = add_cursor_to_graph(graph, node) |
| error_msg = ( |
| f"Here is the node in the graph module:\n" |
| f"{graph_str_with_cursor}\n" |
| f"This node {node} has metadata of:\n" |
| ) |
| # Node spec error message |
| if hasattr(node.meta, "spec"): |
| error_msg += f"The node spec:\n{node.meta['spec']}\n" |
| |
| # Stacktrace error message |
| if "stack_trace" in node.meta: |
| framelist = _stacktrace_to_framelist(node.meta["stack_trace"]) |
| error_msg += f"The node stacktrace:\n{pretty_print_stacktraces(framelist)}\n" |
| return error_msg |
