Apply algo to map/cond/while/scan nodes in the graph module. This method will popuate graph_module.meta["non_const_buffer_sizes"] for all submodules and return a bufsizes list that is the maximum size of all buffers.
(
algo: Callable[..., list[int]],
graph_module: torch.fx.GraphModule,
alignment: int,
graph_signature: Optional[ExportGraphSignature] = None,
)
| 1287 | |
| 1288 | |
| 1289 | def _apply_algo_to_submodules( |
| 1290 | algo: Callable[..., list[int]], |
| 1291 | graph_module: torch.fx.GraphModule, |
| 1292 | alignment: int, |
| 1293 | graph_signature: Optional[ExportGraphSignature] = None, |
| 1294 | ) -> list[int]: |
| 1295 | """Apply algo to map/cond/while/scan nodes in the graph module. |
| 1296 | |
| 1297 | This method will popuate graph_module.meta["non_const_buffer_sizes"] for |
| 1298 | all submodules and return a bufsizes list that is the maximum size of all |
| 1299 | buffers. |
| 1300 | """ |
| 1301 | |
| 1302 | # Bufsizes for submodules. |
| 1303 | bufsizes: list[int] = [] |
| 1304 | |
| 1305 | def _handle( |
| 1306 | submodule_node: torch.fx.Node, |
| 1307 | alloc_graph_input: bool = False, |
| 1308 | ) -> None: |
| 1309 | current_bufsizes = _handle_submodule( |
| 1310 | algo, |
| 1311 | graph_module, |
| 1312 | alignment, |
| 1313 | submodule_node, |
| 1314 | graph_signature, |
| 1315 | alloc_graph_input=alloc_graph_input, |
| 1316 | ) |
| 1317 | nonlocal bufsizes |
| 1318 | _merge_bufsizes(bufsizes, current_bufsizes) |
| 1319 | |
| 1320 | for cond_node in get_cond_nodes(graph_module): |
| 1321 | _handle(cast(torch.fx.Node, cond_node.args[1])) |
| 1322 | _handle(cast(torch.fx.Node, cond_node.args[2])) |
| 1323 | |
| 1324 | for while_node in get_while_nodes(graph_module): |
| 1325 | _handle(cast(torch.fx.Node, while_node.args[0])) |
| 1326 | _handle(cast(torch.fx.Node, while_node.args[1])) |
| 1327 | |
| 1328 | for map_node in get_map_nodes(graph_module): |
| 1329 | _handle(cast(torch.fx.Node, map_node.args[0]), alloc_graph_input=True) |
| 1330 | |
| 1331 | for scan_node in get_scan_nodes(graph_module): |
| 1332 | _handle(cast(torch.fx.Node, scan_node.args[0]), alloc_graph_input=True) |
| 1333 | |
| 1334 | # TODO: We can handle delegates the same way as map/cond/while. |
| 1335 | # Maybe populate the graph_module.meta["non_const_buffer_sizes"] for delegates. |
| 1336 | |
| 1337 | return bufsizes |
| 1338 | |
| 1339 | |
| 1340 | _CPU_KEY: tuple[DeviceType, int] = (DeviceType.CPU, 0) |
no test coverage detected