[fx] Add a faster method for inserting positional argument.

test/expect/TestFXAPIBackwardCompatibility.test_class_member_back_compat-fx_backcompat_class_members.expect

@@ -7,7 +7,7 @@ torch.fx.immutable_collections.immutable_dict ['clear', 'pop', 'popitem', 'updat
 torch.fx.immutable_collections.immutable_list ['append', 'clear', 'extend', 'insert', 'pop', 'remove']
 torch.fx.interpreter.Interpreter ['boxed_run', 'call_function', 'call_method', 'call_module', 'fetch_args_kwargs_from_env', 'fetch_attr', 'get_attr', 'map_nodes_to_values', 'output', 'placeholder', 'run', 'run_node']
 torch.fx.interpreter.Transformer ['call_function', 'call_module', 'get_attr', 'placeholder', 'transform']
-torch.fx.node.Node ['all_input_nodes', 'append', 'args', 'format_node', 'is_impure', 'kwargs', 'next', 'normalized_arguments', 'prepend', 'prev', 'replace_all_uses_with', 'replace_input_with', 'stack_trace', 'update_arg', 'update_kwarg']
+torch.fx.node.Node ['all_input_nodes', 'append', 'args', 'format_node', 'insert_arg', 'is_impure', 'kwargs', 'next', 'normalized_arguments', 'prepend', 'prev', 'replace_all_uses_with', 'replace_input_with', 'stack_trace', 'update_arg', 'update_kwarg']
 torch.fx.passes.shape_prop.ShapeProp ['propagate', 'run_node']
 torch.fx.passes.shape_prop.TensorMetadata ['dtype', 'is_quantized', 'memory_format', 'qparams', 'requires_grad', 'shape', 'stride']
 torch.fx.passes.split_module.Partition []

test/expect/TestFXAPIBackwardCompatibility.test_function_back_compat-fx_backcompat_function_signatures.expect

@@ -53,6 +53,7 @@ torch.fx.interpreter.Transformer.transform(self) -> torch.fx.graph_module.GraphM
 torch.fx.node.Node.__init__(self, graph: 'Graph', name: str, op: str, target: 'Target', args: Tuple[Argument, ...], kwargs: Dict[str, Argument], return_type: Optional[Any] = None) -> None
 torch.fx.node.Node.append(self, x: 'Node') -> None
 torch.fx.node.Node.format_node(self, placeholder_names: Optional[List[str]] = None, maybe_return_typename: Optional[List[str]] = None) -> Optional[str]
+torch.fx.node.Node.insert_arg(self, idx: int, arg: torch.fx.node.Argument) -> None
 torch.fx.node.Node.prepend(self, x: 'Node') -> None
 torch.fx.node.Node.replace_all_uses_with(self, replace_with: 'Node', delete_user_cb: Callable[[Node], bool] = <function <lambda>>, propagate_meta = False) -> List[Node]
 torch.fx.node.Node.replace_input_with(self, old_input: 'Node', new_input: 'Node')

test/test_fx.py

@@ -3813,6 +3813,25 @@ def foo(x):
         traced = torch.fx.symbolic_trace(foo)
         self.assertEqual(foo([]), traced([]))
 
+    def test_insert_arg(self):
+        m = symbolic_trace(SimpleTest())
+        m.register_buffer("buf", torch.tensor(0))
+        output_node = next(iter(reversed(m.graph.nodes)))
+        with m.graph.inserting_before(output_node):
+            a = m.graph.get_attr("buf")
+        r = len(output_node.args)
+        output_node.insert_arg(0, a)
+        self.assertEqual(len(output_node.args), r + 1)
+        self.assertEqual(len(a.users), 1)
+        self.assertIs(output_node.args[0], a)
+        self.assertIs(list(a.users.keys())[0], output_node)
+        output_node.insert_arg(2, a)
+        self.assertEqual(len(output_node.args), r + 2)
+        self.assertEqual(len(a.users), 1)
+        self.assertIs(output_node.args[2], a)
+        self.assertIs(list(a.users.keys())[0], output_node)
+        m.graph.lint()
+
 
 
 def run_getitem_target():

torch/fx/node.py

@@ -359,6 +359,30 @@ def update_arg(self, idx : int, arg : Argument) -> None:
         args[idx] = arg
         self.args = tuple(args)
 
+    @compatibility(is_backward_compatible=True)
+    def insert_arg(self, idx : int, arg : Argument) -> None:
+        """
+        Insert an positional argument to the argument list with given index.
+
+        Args:
+
+            idx (int): The index of the element in ``self.args`` to be inserted before.
+            arg (Argument): The new argument value to insert into ``args``
+        """
+        assert 0 <= idx <= len(self.args), "insert_args index must be between 0 and len(self.args)"
+        args_left = self.args[:idx]
+        args_right = self.args[idx:]
+
+        self._args = args_left + (arg,) + args_right
+
+        _new_input_nodes = {}
+        map_arg(arg, lambda n: _new_input_nodes.setdefault(n))
+
+        for new_use in _new_input_nodes.keys():
+            if new_use not in self._input_nodes:
+                self._input_nodes.setdefault(new_use)
+                new_use.users.setdefault(self)
+
     @compatibility(is_backward_compatible=True)
     def update_kwarg(self, key : str, arg : Argument) -> None:
         """