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): ``` PROGRA
( # noqa: C901
program: Program,
show_meminfo: bool = True,
mark_dynamic_shape_tensor: bool = False,
out: Optional[TextIO] = None,
)
| 147 | |
| 148 | |
| 149 | def print_program( # noqa: C901 |
| 150 | program: Program, |
| 151 | show_meminfo: bool = True, |
| 152 | mark_dynamic_shape_tensor: bool = False, |
| 153 | out: Optional[TextIO] = None, |
| 154 | ) -> None: |
| 155 | """ |
| 156 | Dump the instruction list of a program in a more human readable fashion. |
| 157 | |
| 158 | The dump follows the following BNF syntax (I combime some regex syntax |
| 159 | so the grammar becomes shorter. The grammar is not strict but the main |
| 160 | purpose is to let people understand the dump): |
| 161 | ``` |
| 162 | PROGRAM: (INSTRUCTION)+ |
| 163 | INSTRUCTION: SEQUENCE_NO ':' (CALL_KERNEL | JUMP_FALSE) |
| 164 | JUMP_FALSE: 'JF' '(' EVALUE ')' '->' TARGET_SEQUENCE_NO |
| 165 | CALL_KERNEL: OVERLOADDED_OP_NAME ARGS |
| 166 | ARGS: EVALUE | ARGS ',' EVALUE |
| 167 | EVALUE: EVALUE_IDX ( TENSOR | INT | BOOL | ...) |
| 168 | INT: 'I' ACTUAL_INT_VALUE |
| 169 | BOOL: 'B' ZERO_OR_ONE |
| 170 | CONST_TENSOR_PREFIX: 'CT' |
| 171 | TENSOR: ('T' | CONST_TENSOR_PREFIX) (MEM_ALLOCATION_INFO)? TENSOR_SHAPE TENSOR_DTYPE |
| 172 | TENSOR_SHAPE: '[' dim0_size, dim1_size, ..., last_dim_size ']' |
| 173 | MEM_ALLOCATION_INFO: PLANNED_MEM_INFO | UNPLANNED_MEM_INFO |
| 174 | PLANNED_MEM_INFO: 'm' MEM_LAYER_ID '.' MEM_LAYER_OFFSET |
| 175 | UNPLANNED_MEM_INFO: 'm.' |
| 176 | ``` |
| 177 | |
| 178 | To make the dump easier to read, it's colored as follows: |
| 179 | 1. input/output EValues are marked as red |
| 180 | 2. EValue types (or more specifically tensor types with size and dtype) are marked as blue |
| 181 | """ |
| 182 | execution_plan = program.execution_plan[0] |
| 183 | operators = execution_plan.operators |
| 184 | delegates = execution_plan.delegates |
| 185 | chain = execution_plan.chains[0] |
| 186 | instructions = chain.instructions |
| 187 | inputs: List[int] = execution_plan.inputs |
| 188 | outputs: List[int] = execution_plan.outputs |
| 189 | values: List[EValue] = execution_plan.values |
| 190 | |
| 191 | def _format_arg(evalue_idx: int) -> str: |
| 192 | def _get_io_index(iolist: List[int], target_evalue_idx: int) -> int: |
| 193 | """ |
| 194 | The list is short enough so linear scan is proper. |
| 195 | """ |
| 196 | for io_idx, evalue_idx in enumerate(iolist): |
| 197 | if evalue_idx == target_evalue_idx: |
| 198 | return io_idx |
| 199 | return -1 |
| 200 | |
| 201 | argstr = str(evalue_idx) |
| 202 | if (input_idx := _get_io_index(inputs, evalue_idx)) >= 0: |
| 203 | argstr += f"\033[31mI{input_idx}\033[0m" |
| 204 | if (output_idx := _get_io_index(outputs, evalue_idx)) >= 0: |
| 205 | argstr += f"\033[31mO{output_idx}\033[0m" |
| 206 |