| 123 | |
| 124 | |
| 125 | class DebugPass(PassBase): |
| 126 | def __init__( |
| 127 | self, |
| 128 | msg: str = "", |
| 129 | enable_debug_pass: bool = True, |
| 130 | show_src: bool = False, |
| 131 | show_full_path: bool = False, |
| 132 | show_all_frames: bool = False, |
| 133 | path_filter: Optional[str] = None, |
| 134 | show_spec: bool = False, |
| 135 | log_filename: Optional[str] = None, |
| 136 | ) -> None: |
| 137 | """ |
| 138 | show_src: whether to show source code that generated each fx Node |
| 139 | show_full_path: whether to show the full path of source code or just the filename |
| 140 | show_all_frames: control for each node whether show only the last frame or all the frames. |
| 141 | path_filter: a regular expression to filter the path of the stackframes |
| 142 | log_filename: if provided, the output will also be written to this path. |
| 143 | Existing content in this file will be discarded. |
| 144 | """ |
| 145 | self.msg = msg |
| 146 | self.enable_debug_pass = enable_debug_pass |
| 147 | self.show_src = show_src |
| 148 | self.show_full_path = show_full_path |
| 149 | self.show_all_frames = show_all_frames |
| 150 | self.show_spec = show_spec |
| 151 | self.log_filename = log_filename |
| 152 | if path_filter: |
| 153 | self.path_filter_re = re.compile(path_filter) # pyre-ignore |
| 154 | else: |
| 155 | self.path_filter_re = None |
| 156 | |
| 157 | def call(self, graph_module: torch.fx.GraphModule) -> PassResult: |
| 158 | """ |
| 159 | Counts the number of operations and call_funciton operations. |
| 160 | """ |
| 161 | if not self.enable_debug_pass: |
| 162 | return PassResult(graph_module, False) |
| 163 | # it doesn't make sense to mute the DebugPass if user already |
| 164 | # specify self.enable_debug_pass to be true |
| 165 | with override_logger(filename=self.log_filename): |
| 166 | self.callWithLoggerEnabled(graph_module) |
| 167 | return PassResult(graph_module, True) |
| 168 | |
| 169 | def printFrames(self, node: fx.Node) -> None: |
| 170 | """ |
| 171 | The DebugPass maybe used for graph generated by both the old exir dispatch |
| 172 | tracer or the new pt2 tracer. |
| 173 | The former store 'stack_trace' field as a json string; |
| 174 | the latter store 'stack_trace' field as a free form string like: |
| 175 | ``` |
| 176 | File "/data/sandcastle/boxes/fbsource/buck-out/v2/gen/fbcode/20c706e99f51cf3a/executorch/test/end2end/__end2end__/end2end#link-tree/executorch/test/end2end/test_end2end.py", line 150, in forward |
| 177 | o = o * a |
| 178 | ``` |
| 179 | Make this method handle both format. In future, maybe we can drop the |
| 180 | support for old exir dispatch tracer. |
| 181 | """ |
| 182 | if ( |
no outgoing calls