AOTAutograd: functionalization should emit foreach_copy_() instead of copy_() when possible

Consider this code:

import torch

@torch.compile
def f(args):
    torch._foreach_mul_(args, 2)


inps = [torch.ones(10) for _ in range(10)]
f(inps)

Which generates this AOT graph:

    def forward(self, arg0_1: "f32[10]", arg1_1: "f32[10]", arg2_1: "f32[10]", arg3_1: "f32[10]", arg4_1: "f32[10]", arg5_1: "f32[10]", arg6_1: "f32[10]", arg7_1: "f32[10]", arg8_1: "f32[10]", arg9_1: "f32[10]"):
        # File: /data/users/hirsheybar/a/pytorch/tmp2.py:5 in f, code: torch._foreach_mul_(args, 2)
        _foreach_mul = torch.ops.aten._foreach_mul.Scalar([arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1, arg8_1, arg9_1], 2)
        getitem: "f32[10]" = _foreach_mul[0]
        getitem_1: "f32[10]" = _foreach_mul[1]
        getitem_2: "f32[10]" = _foreach_mul[2]
        getitem_3: "f32[10]" = _foreach_mul[3]
        getitem_4: "f32[10]" = _foreach_mul[4]
        getitem_5: "f32[10]" = _foreach_mul[5]
        getitem_6: "f32[10]" = _foreach_mul[6]
        getitem_7: "f32[10]" = _foreach_mul[7]
        getitem_8: "f32[10]" = _foreach_mul[8]
        getitem_9: "f32[10]" = _foreach_mul[9];  _foreach_mul = None
        copy_: "f32[10]" = torch.ops.aten.copy_.default(arg0_1, getitem);  arg0_1 = getitem = None
        copy__1: "f32[10]" = torch.ops.aten.copy_.default(arg1_1, getitem_1);  arg1_1 = getitem_1 = None
        copy__2: "f32[10]" = torch.ops.aten.copy_.default(arg2_1, getitem_2);  arg2_1 = getitem_2 = None
        copy__3: "f32[10]" = torch.ops.aten.copy_.default(arg3_1, getitem_3);  arg3_1 = getitem_3 = None
        copy__4: "f32[10]" = torch.ops.aten.copy_.default(arg4_1, getitem_4);  arg4_1 = getitem_4 = None
        copy__5: "f32[10]" = torch.ops.aten.copy_.default(arg5_1, getitem_5);  arg5_1 = getitem_5 = None
        copy__6: "f32[10]" = torch.ops.aten.copy_.default(arg6_1, getitem_6);  arg6_1 = getitem_6 = None
        copy__7: "f32[10]" = torch.ops.aten.copy_.default(arg7_1, getitem_7);  arg7_1 = getitem_7 = None
        copy__8: "f32[10]" = torch.ops.aten.copy_.default(arg8_1, getitem_8);  arg8_1 = getitem_8 = None
        copy__9: "f32[10]" = torch.ops.aten.copy_.default(arg9_1, getitem_9);  arg9_1 = getitem_9 = None
        return ()

Functionalization sees that foreach_mul_() is an input mutation, so it adds copy_() nodes to the graph to apply the input mutations properly.

But if we have 100 tensors in the TensorList input to foreach_mul_, then we will end up with 100 individual copy_() nodes in the graph. A larger graph generally means worse compile times for inductor, unless inductor does extra work to avoid it.

Instead, one thing we can do is recognize when a bunch of inputs were mutated by a single foreach_* op, and emit a single foreach_copy_() node in the graph in their place. We need to take care though that the inputs to the foreach op are preserved in the right order

cc @crcrpar @mcarilli @janeyx99 @ezyang @msaroufim @anijain2305 @zou3519