Apply algo to nodes in a submodule of the graph module.
(
algo: Callable[..., list[int]],
parent_graph_module: torch.fx.GraphModule,
alignment: int,
submodule_node: torch.fx.Node,
graph_signature: Optional[ExportGraphSignature] = None,
alloc_graph_input: bool = False,
)
| 1261 | |
| 1262 | |
| 1263 | def _handle_submodule( |
| 1264 | algo: Callable[..., list[int]], |
| 1265 | parent_graph_module: torch.fx.GraphModule, |
| 1266 | alignment: int, |
| 1267 | submodule_node: torch.fx.Node, |
| 1268 | graph_signature: Optional[ExportGraphSignature] = None, |
| 1269 | alloc_graph_input: bool = False, |
| 1270 | ) -> list[int]: |
| 1271 | """Apply algo to nodes in a submodule of the graph module.""" |
| 1272 | assert submodule_node.op == "get_attr" |
| 1273 | submodule = getattr(parent_graph_module, submodule_node.target) |
| 1274 | |
| 1275 | logging.debug(f"Planning memory for submodule {submodule_node.name}...") |
| 1276 | bufsizes = apply_algo( |
| 1277 | algo, |
| 1278 | submodule, |
| 1279 | alignment, |
| 1280 | graph_signature, |
| 1281 | alloc_graph_input=alloc_graph_input, |
| 1282 | alloc_graph_output=True, |
| 1283 | ) |
| 1284 | submodule.meta.update({"non_const_buffer_sizes": bufsizes}) |
| 1285 | logging.debug(f"Buffer sizes for submodule {submodule_node.name}: {bufsizes}") |
| 1286 | return bufsizes |
| 1287 | |
| 1288 | |
| 1289 | def _apply_algo_to_submodules( |
no test coverage detected