| 318 | |
| 319 | |
| 320 | class ToOutVarPass(PassBase): |
| 321 | def __init__(self, ignore_to_out_var_failure: bool = False) -> None: |
| 322 | self.ignore_to_out_var_failure = ignore_to_out_var_failure |
| 323 | |
| 324 | def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 |
| 325 | """ |
| 326 | Converts all of the functions to contain an out variant if it does not exist |
| 327 | """ |
| 328 | missing_out_vars: Set[str] = set() |
| 329 | |
| 330 | def get_submodule(node: torch.fx.Node) -> torch.fx.GraphModule: |
| 331 | assert node.op == "get_attr" |
| 332 | return getattr(graph_module, node.target) |
| 333 | |
| 334 | for node in graph_module.graph.nodes: |
| 335 | if node.op != "call_function": |
| 336 | continue |
| 337 | |
| 338 | target = node.target |
| 339 | if target == torch.ops.higher_order.cond: |
| 340 | self.call(get_submodule(node.args[1])) |
| 341 | self.call(get_submodule(node.args[2])) |
| 342 | continue |
| 343 | if target == torch.ops.higher_order.map_impl: |
| 344 | self.call(get_submodule(node.args[0])) |
| 345 | continue |
| 346 | elif target == torch.ops.higher_order.while_loop: |
| 347 | self.call(get_submodule(node.args[0])) |
| 348 | self.call(get_submodule(node.args[1])) |
| 349 | continue |
| 350 | elif target == torch.ops.higher_order.scan: |
| 351 | # scan(combine_fn, init, xs, additional_inputs) |
| 352 | # combine_fn is at args[0] |
| 353 | self.call(get_submodule(node.args[0])) |
| 354 | continue |
| 355 | elif getattr(target, "__module__", None) in ("builtins", "_operator"): |
| 356 | continue |
| 357 | elif target in to_out_var_skiplist: |
| 358 | continue |
| 359 | elif _get_overload_schema(target).kind() == SchemaKind.inplace: |
| 360 | continue |
| 361 | if not isinstance( |
| 362 | target, (torch._ops.OpOverload, EdgeOpOverload, BackendOpOverload) |
| 363 | ): |
| 364 | raise RuntimeError(f"Require an op overload for target: {target}") |
| 365 | |
| 366 | op_name = target._schema.name |
| 367 | overload_name = target._schema.overload_name |
| 368 | if is_out_variant(op_name, overload_name): |
| 369 | # TODO (zhxchen17) Remove this after functionalization is always on. |
| 370 | if "out" in node.kwargs and isinstance(node.kwargs["out"], fx.Node): |
| 371 | out = node.kwargs["out"] |
| 372 | if out.target is not memory.alloc and len(out.users) == 1: |
| 373 | with graph_module.graph.inserting_before(node): |
| 374 | alloc = make_alloc_node( |
| 375 | graph_module, |
| 376 | node.meta["val"], |
| 377 | node.meta["tensor_meta"], |
no outgoing calls