(
self, arg, arg_type: Optional[torch._C.Argument] = None
)
| 696 | ) |
| 697 | |
| 698 | def serialize_input( |
| 699 | self, arg, arg_type: Optional[torch._C.Argument] = None |
| 700 | ) -> Argument: |
| 701 | import torch._inductor.ir as inductor_ir |
| 702 | |
| 703 | inductor_tensor_buffers = ( |
| 704 | inductor_ir.Buffer, |
| 705 | inductor_ir.ReinterpretView, |
| 706 | ) |
| 707 | |
| 708 | if isinstance(arg, torch.fx.Node): |
| 709 | if arg.op == "get_attr": |
| 710 | assert isinstance(arg.target, str) |
| 711 | attr = getattr(arg.graph.owning_module, arg.target) |
| 712 | |
| 713 | if isinstance(attr, torch.Tensor): |
| 714 | raise SerializeError( |
| 715 | "getattr nodes containing tensors should not appear in the graph" |
| 716 | ) |
| 717 | elif isinstance(attr, torch.fx.GraphModule): |
| 718 | with self.save_graph_state(): |
| 719 | graph = self.serialize_graph(attr) |
| 720 | return Argument.create( |
| 721 | as_graph=GraphArgument(name=arg.target, graph=graph) |
| 722 | ) |
| 723 | else: |
| 724 | raise SerializeError( |
| 725 | f"Unsupported getattr attribute {arg.target} with type: {type(attr)}" |
| 726 | ) |
| 727 | elif self.is_sym_int_arg(arg): |
| 728 | return Argument.create( |
| 729 | as_sym_int=SymIntArgument.create(as_name=arg.name) |
| 730 | ) |
| 731 | elif self.is_sym_bool_arg(arg): |
| 732 | return Argument.create( |
| 733 | as_sym_bool=SymBoolArgument.create(as_name=arg.name) |
| 734 | ) |
| 735 | else: |
| 736 | if isinstance(arg.meta["val"], ep.CustomObjArgument): |
| 737 | return Argument.create( |
| 738 | as_custom_obj=CustomObjArgument( |
| 739 | name=arg.name, class_fqn=arg.meta["val"].class_fqn |
| 740 | ) |
| 741 | ) |
| 742 | return Argument.create(as_tensor=TensorArgument(name=arg.name)) |
| 743 | elif isinstance(arg, inductor_tensor_buffers): |
| 744 | # Other branches are for arguments in fx node. |
| 745 | # This is a special branch for handling buffers (representing tensor arguments) |
| 746 | # for inductor's ExternalFallbackNode |
| 747 | # export_extern_kernel_node() is using this function to serialize arguments |
| 748 | arg_name = arg.get_name() |
| 749 | assert arg_name is not None, "Buffer must have valid name" |
| 750 | return Argument.create(as_tensor=TensorArgument(name=arg_name)) |
| 751 | elif isinstance(arg, torch.SymInt): |
| 752 | # This is a special branch for handling SymInt args in inductor's |
| 753 | # ExternalFallbackNode. |
| 754 | # For regular FX graph, SymInt arg should be a fx.Node with |
| 755 | # self.is_sym_int_arg(arg) being true |
no test coverage detected