Repository: pytorch
Issue: auto-grad graph replicate split_with_sizes(lengths) X times where X = len(lengths) effecting compile time 
Labels: triaged,module: functionalization,oncall: pt2,module: pt2-dispatcher

auto-grad graph replicate split_with_sizes(lengths) X times where X = len(lengths) effecting compile time.
I think this repitions might be related to the functionalization view replay? unsure yet.

We spend 50% of in split_with_sizes .

strobelight profile:
https://fburl.com/scuba/pyperf_experimental/on_demand/pz6u9po2

for when we have lengths size =1000 we have this replicated 1000 times.

I counted how many times we call split_with_sizes when input is 8 and its like 32 times. 
that means for input 1000 its maybe like 32/8 *1000 times = 16K times?

This is noted while debugging :
https://github.com/pytorch/pytorch/issues/138729

ttile, this pre
```
import time

import torch
import random

if __name__ == "__main__":
    def make_val(n=10):
        values = torch.arange(n , dtype=torch.float)
        return values

    @torch.compile(fullgraph=True, dynamic=True)
    def consolidate(tensors):
        start = 0
        lengths = []
        swaps = []
        origs = []


        def view_and_pad(tensor: torch.Tensor, lengths=lengths) -> torch.Tensor:
            nonlocal start

            origs.append(tensor)
            # (1) View the tensor as a contiguous uint8 buffer.

            swap = tensor.view(-1).view(torch.uint8)

            # (2) Pad the tensor to a multiple of 8 elements.

            # result must always have a multiple of 8 elements
            pad = swap.numel() % 8
            if pad != 0:
                swap = torch.cat([swap, swap.new_zeros(8 - pad)])
            n = swap.numel()

            start = start + n

            # add length of swap to lengths. (not all lengths are always the same)
            lengths.append(n)
            
            # Add the swap to the list of swaps.
            swaps.append(swap)


        for val in list(tensors):
            view_and_pad(val)

        # create empty storage with size filesize.
        filesize = start
        storage = torch.empty(
            filesize,
            dtype=torch.uint8,
        )
        # split storage into tensors of size lengths. lengths are not always the same.
        torch.cat(swaps, out=storage)

        # swaps have the same data but viewed as one tensor. 
        swaps = storage.split(lengths)


        def view_old_as_new(
            v: torch.Tensor, oldv: torch.Tensor
        ) -> torch.Tensor:
            if oldv is None:
                return v
            v = v.view(oldv.dtype)
            if v.numel() > oldv.numel():
                return v[: oldv.numel()].view(oldv.shape)
            return v.view(oldv.shape)

        result = [
            view_old_as_new(v, oldv)
            for (v, oldv) in zip(swaps, origs, strict=True)
        ]
        return result

    for n in [7]:
        torch.compiler.reset()
        tensors = [make_val(10) for i in range(n)]
        t0 = time.time()
        results = consolidate(tensors)
        print(n, time.time()-t0)

    assert len(set(t.untyped_storage().data_ptr() for t in results)) == 1



# [[ X inputs tensor each os size yi]]
# can we write this as tensor of size X*y

```
Dynamo graph:
```




def forward(self, s0 : torch.SymInt, L_tensors_0_ : torch.Tensor, L_tensors_1_ : torch.Tensor, L_tensors_2_ : torch.Tensor, L_tensors_3_ : torch.Tensor, L_tensors_4_ : torch.Tensor, L_tensors_5_ : torch.Tensor, L_tensors_6_ : torch.Tensor):
    l_tensors_0_ = L_tensors_0_
    l_tensors_1_ = L_tensors_1_
    l_tensors_2_ = L_tensors_2_
    l_tensors_3_ = L_tensors_3_
    l_tensors_4_ = L_tensors_4_
    l_tensors_5_ = L_tensors_5_
    l_tensors_6_ = L_tensors_6_
    view = l_tensors_0_.view(-1)
    swap = view.view(torch.uint8);  view = None
    mul = 4 * s0;  s0 = None
    view_2 = l_tensors_1_.view(-1)
    swap_1 = view_2.view(torch.uint8);  view_2 = None
    add_1 = mul + mul
    view_4 = l_tensors_2_.view(-1)
    swap_2 = view_4.view(torch.uint8);  view_4 = None
    add_2 = add_1 + mul;  add_1 = None
    view_6 = l_tensors_3_.view(-1)
    swap_3 = view_6.view(torch.uint8);  view_6 = None
    add_3 = add_2 + mul;  add_2 = None
    view_8 = l_tensors_4_.view(-1)
    swap_4 = view_8.view(torch.uint8);  view_8 = None
    add_4 = add_3 + mul;  add_3 = None
    view_10 = l_tensors_5_.view(-1)
    swap_5 = view_10.view(torch.uint8);  view_10 = None
    add_5 = add_4 + mul;  add_4 = None
    view_12 = l_tensors_6_.view(-1)
    swap_6 = view_12.view(torch.uint8);  view_12 = None
    add_6 = add_5 + mul;  add_5 = None
    storage = torch.empty(add_6, dtype = torch.uint8);  add_6 = None
    cat = torch.cat([swap, swap_1, swap_2, swap_3, swap_4, swap_5, swap_6], out = storage);  swap = swap_1 = swap_2 = swap_3 = swap_4 = swap_5 = swap_6 = cat = None
    split = storage.split([mul, mul, mul, mul, mul, mul, mul]);  storage = mul = None
    v = split[0]
    v_2 = split[1]
    v_4 = split[2]
    v_6 = split[3]
    v_8 = split[4]
    v_10 = split[5]
    v_12 = split[6];  split = None
    v_1 = v.view(torch.float32);  v = None
    size = l_tensors_0_.size();  l_tensors_0_ = None
    view_15 = v_1.view(size);  v_1 = size = None
    v_3 = v_2.view(torch.float32);  v_2 = None
    size_1 = l_tensors_1_.size();  l_tensors_1_ = None
    view_17 = v_3.view(size_1);  v_3 = size_1 = None
    v_5 = v_4.view(torch.float32);  v_4 = None
    size_2 = l_tensors_2_.size();  l_tensors_2_ = None
    view_19 = v_5.view(size_2);  v_5 = size_2 = None
    v_7 = v_6.view(torch.float32);  v_6 = None
    size_3 = l_tensors_3_.size();  l_tensors_3_ = None
    view_21 = v_7.view(size_3);  v_7 = size_3 = None
    v_9 = v_8.view(torch.float32);  v_8 = None
    size_4 = l_tensors_4_.size();  l_tensors_4_ = None
    view_23 = v_9.view(size_4);  v_9 = size_4 = None
    v_11 = v_10.view(torch.float32);  v_10 = None
    size_5 = l_tensors_5_.size();  l_tensors_5_ = None
    view_25 = v_11.view(size_5);  v_11 = size_5 = None
    v_13 = v_12.view(torch.float32);  v_12 = None
    size_6 = l_tensors_6_.size();  l_tensors_6_ = None
    view_27 = v_13.view(size_6);  v_13 = size_6 = None
    return (view_15, view_17, view_19, view_21, view_23, view_25, view_27)
```

auto_grad graph:

```
def forward(self, arg0_1, arg1_1, arg2_1, arg3_1, arg4_1, arg5_1, arg6_1, arg7_1):
    view = torch.ops.aten.view.default(arg1_1, [-1]);  arg1_1 = None
    view_1 = torch.ops.aten.view.dtype(view, torch.uint8);  view = None
    alias = torch.ops.aten.alias.default(view_1);  view_1 = None
    alias_1 = torch.ops.aten.alias.default(alias);  alias = None
    mul_1 = 4 * arg0_1
    view_2 = torch.ops.aten.view.default(arg2_1, [-1]);  arg2_1 = None
    view_3 = torch.ops.aten.view.dtype(view_2, torch.uint8);  view_2 = None
    alias_2 = torch.ops.aten.alias.default(view_3);  view_3 = None
    alias_3 = torch.ops.aten.alias.default(alias_2);  alias_2 = None
    add_12 = mul_1 + mul_1
    view_4 = torch.ops.aten.view.default(arg3_1, [-1]);  arg3_1 = None
    view_5 = torch.ops.aten.view.dtype(view_4, torch.uint8);  view_4 = None
    alias_4 = torch.ops.aten.alias.default(view_5);  view_5 = None
    alias_5 = torch.ops.aten.alias.default(alias_4);  alias_4 = None
    add_19 = add_12 + mul_1;  add_12 = None
    view_6 = torch.ops.aten.view.default(arg4_1, [-1]);  arg4_1 = None
    view_7 = torch.ops.aten.view.dtype(view_6, torch.uint8);  view_6 = None
    alias_6 = torch.ops.aten.alias.default(view_7);  view_7 = None
    alias_7 = torch.ops.aten.alias.default(alias_6);  alias_6 = None
    add_26 = add_19 + mul_1;  add_19 = None
    view_8 = torch.ops.aten.view.default(arg5_1, [-1]);  arg5_1 = None
    view_9 = torch.ops.aten.view.dtype(view_8, torch.uint8);  view_8 = None
    alias_8 = torch.ops.aten.alias.default(view_9);  view_9 = None
    alias_9 = torch.ops.aten.alias.default(alias_8);  alias_8 = None
    add_33 = add_26 + mul_1;  add_26 = None
    view_10 = torch.ops.aten.view.default(arg6_1, [-1]);  arg6_1 = None
    view_11 = torch.ops.aten.view.dtype(view_10, torch.uint8);  view_10 = None
    alias_10 = torch.ops.aten.alias.default(view_11);  view_11 = None
    alias_11 = torch.ops.aten.alias.default(alias_10);  alias_10 = None
    add_40 = add_33 + mul_1;  add_33 = None
    view_12 = torch.ops.aten.view.default(arg7_1, [-1]);  arg7_1 = None
    view_13 = torch.ops.aten.view.dtype(view_12, torch.uint8);  view_12 = None
    alias_12 = torch.ops.aten.alias.default(view_13);  view_13 = None
    alias_13 = torch.ops.aten.alias.default(alias_12);  alias_12 = None
    add_47 = add_40 + mul_1;  add_40 = None
    empty = torch.ops.aten.empty.memory_format([add_47], dtype = torch.uint8, device = device(type='cpu'), pin_memory = False);  add_47 = empty = None
    cat = torch.ops.aten.cat.default([alias_1, alias_3, alias_5, alias_7, alias_9, alias_11, alias_13]);  alias_1 = alias_3 = alias_5 = alias_7 = alias_9 = alias_11 = alias_13 = None
    split_with_sizes = torch.ops.aten.split_with_sizes.default(cat, [mul_1, mul_1, mul_1, mul_1, mul_1, mul_1, mul_1])
    getitem = split_with_sizes[0]
    getitem_1 = split_with_sizes[1]
    getitem_2 = split_with_sizes[2]
    getitem_3 = split_with_sizes[3]
    getitem_4 = split_with_sizes[4]
    getitem_5 = split_with_sizes[5]
    getitem_6 = split_with_sizes[6];  split_with_sizes = None
    view_14 = torch.ops.aten.view.dtype(getitem, torch.float32);  getitem = None
    alias_14 = torch.ops.aten.alias.default(view_14);  view_14 = None
    alias_15 = torch.ops.aten.alias.default(alias_14);  alias_14 = None
    view_15 = torch.ops.aten.view.default(alias_15, [arg0_1]);  alias_15 = view_15 = None
    view_16 = torch.ops.aten.view.dtype(getitem_1, torch.float32);  getitem_1 = None
    alias_16 = torch.ops.aten.alias.default(view_16);  view_16 = None
    alias_17 = torch.ops.aten.alias.default(alias_16);  alias_16 = None
    view_17 = torch.ops.aten.view.default(alias_17, [arg0_1]);  alias_17 = view_17 = None
    view_18 = torch.ops.aten.view.dtype(getitem_2, torch.float32);  getitem_2 = None
    alias_18 = torch.ops.aten.alias.default(view_18);  view_18 = None
    alias_19 = torch.ops.aten.alias.default(alias_18);  alias_18 = None
    view_19 = torch.ops.aten.view.default(alias_19, [arg0_1]);  alias_19 = view_19 = None
    view_20 = torch.ops.aten.view.dtype(getitem_3, torch.float32);  getitem_3 = None
    alias_20 = torch.ops.aten.alias.default(view_20);  view_20 = None
    alias_21 = torch.ops.aten.alias.default(alias_20);  alias_20 = None
    view_21 = torch.ops.aten.view.default(alias_21, [arg0_1]);  alias_21 = view_21 = None
    view_22 = torch.ops.aten.view.dtype(getitem_4, torch.float32);  getitem_4 = None
    alias_22 = torch.ops.aten.alias.default(view_22);  view_22 = None
    alias_23 = torch.ops.aten.alias.default(alias_22);  alias_22 = None
    view_23 = torch.ops.aten.view.default(alias_23, [arg0_1]);  alias_23 = view_23 = None
    view_24 = torch.ops.aten.view.dtype(getitem_5, torch.float32);  getitem_5 = None
    alias_24 = torch.ops.aten.alias.default(view_24);  view_24 = None
    alias_25 = torch.ops.aten.alias.default(alias_24);  alias_24 = None
    view_25 = torch.ops.aten.view.default(alias_25, [arg0_1]);  alias_25 = view_25 = None
    view_26 = torch.ops.aten.view.dtype(getitem_6, torch.float32);  getitem_6 = None
    alias_26 = torch.ops.aten.alias.default(view_26);  view_26 = None
    alias_27 = torch.ops.aten.alias.default(alias_26);  alias_26 = None
    view_27 = torch.ops.aten.view.default(alias_27, [arg0_1]);  alias_27 = view_27 = None
    split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(cat, [mul_1, mul_1, mul_1, mul_1, mul_1, mul_1, mul_1])
    getitem_7 = split_with_sizes_1[0]
    getitem_8 = split_with_sizes_1[1];  getitem_8 = None
    getitem_9 = split_with_sizes_1[2];  getitem_9 = None
    getitem_10 = split_with_sizes_1[3];  getitem_10 = None
    getitem_11 = split_with_sizes_1[4];  getitem_11 = None
    getitem_12 = split_with_sizes_1[5];  getitem_12 = None
    getitem_13 = split_with_sizes_1[6];  split_with_sizes_1 = getitem_13 = None
    view_28 = torch.ops.aten.view.dtype(getitem_7, torch.float32);  getitem_7 = None
    alias_28 = torch.ops.aten.alias.default(view_28);  view_28 = None
    alias_29 = torch.ops.aten.alias.default(alias_28);  alias_28 = None
    alias_30 = torch.ops.aten.alias.default(alias_29);  alias_29 = None
    view_29 = torch.ops.aten.view.default(alias_30, [arg0_1]);  alias_30 = None
    split_with_sizes_2 = torch.ops.aten.split_with_sizes.default(cat, [mul_1, mul_1, mul_1, mul_1, mul_1, mul_1, mul_1])
    getitem_14 = split_with_sizes_2[0];  getitem_14 = None
    getitem_15 = split_with_sizes_2[1]
    getitem_16 = split_with_sizes_2[2];  getitem_16 = None
    getitem_17 = split_with_sizes_2[3];  getitem_17 = None
    getitem_18 = split_with_sizes_2[4];  getitem_18 = None
    getitem_19 = split_with_sizes_2[5];  getitem_19 = None
    getitem_20 = split_with_sizes_2[6];  split_with_sizes_2 = getitem_20 = None
    view_30 = torch.ops.aten.view.dtype(getitem_15, torch.float32);  getitem_15 = None
    alias_31 = torch.ops.aten.alias.default(view_30);  view_30 = None
    alias_32 = torch.ops.aten.alias.default(alias_31);  alias_31 = None
    alias_33 = torch.ops.aten.alias.default(alias_32);  alias_32 = None
    view_31 = torch.ops.aten.view.default(alias_33, [arg0_1]);  alias_33 = None
    split_with_sizes_3 = torch.ops.aten.split_with_sizes.default(cat, [mul_1, mul_1, mul_1, mul_1, mul_1, mul_1, mul_1])
    getitem_21 = split_with_sizes_3[0];  getitem_21 = None
    getitem_22 = split_with_sizes_3[1];  getitem_22 = None
    getitem_23 = split_with_sizes_3[2]
    getitem_24 = split_with_sizes_3[3];  getitem_24 = None
    getitem_25 = split_with_sizes_3[4];  getitem_25 = None
    getitem_26 = split_with_sizes_3[5];  getitem_26 = None
    getitem_27 = split_with_sizes_3[6];  split_with_sizes_3 = getitem_27 = None
    view_32 = torch.ops.aten.view.dtype(getitem_23, torch.float32);  getitem_23 = None
    alias_34 = torch.ops.aten.alias.default(view_32);  view_32 = None
    alias_35 = torch.ops.aten.alias.default(alias_34);  alias_34 = None
    alias_36 = torch.ops.aten.alias.default(alias_35);  alias_35 = None
    view_33 = torch.ops.aten.view.default(alias_36, [arg0_1]);  alias_36 = None
    split_with_sizes_4 = torch.ops.aten.split_with_sizes.default(cat, [mul_1, mul_1, mul_1, mul_1, mul_1, mul_1, mul_1])
    getitem_28 = split_with_sizes_4[0];  getitem_28 = None
    getitem_29 = split_with_sizes_4[1];  getitem_29 = None
    getitem_30 = split_with_sizes_4[2];  getitem_30 = None
    getitem_31 = split_with_sizes_4[3]
    getitem_32 = split_with_sizes_4[4];  getitem_32 = None
    getitem_33 = split_with_sizes_4[5];  getitem_33 = None
    getitem_34 = split_with_sizes_4[6];  split_with_sizes_4 = getitem_34 = None
    view_34 = torch.ops.aten.view.dtype(getitem_31, torch.float32);  getitem_31 = None
    alias_37 = torch.ops.aten.alias.default(view_34);  view_34 = None
    alias_38 = torch.ops.aten.alias.default(alias_37);  alias_37 = None
    alias_39 = torch.ops.aten.alias.default(alias_38);  alias_38 = None
    view_35 = torch.ops.aten.view.default(alias_39, [arg0_1]);  alias_39 = None
    split_with_sizes_5 = torch.ops.aten.split_with_sizes.default(cat, [mul_1, mul_1, mul_1, mul_1, mul_1, mul_1, mul_1])
    getitem_35 = split_with_sizes_5[0];  getitem_35 = None
    getitem_36 = split_with_sizes_5[1];  getitem_36 = None
    getitem_37 = split_with_sizes_5[2];  getitem_37 = None
    getitem_38 = split_with_sizes_5[3];  getitem_38 = None
    getitem_39 = split_with_sizes_5[4]
    getitem_40 = split_with_sizes_5[5];  getitem_40 = None
    getitem_41 = split_with_sizes_5[6];  split_with_sizes_5 = getitem_41 = None
    view_36 = torch.ops.aten.view.dtype(getitem_39, torch.float32);  getitem_39 = None
    alias_40 = torch.ops.aten.alias.default(view_36);  view_36 = None
    alias_41 = torch.ops.aten.alias.default(alias_40);  alias_40 = None
    alias_42 = torch.ops.aten.alias.default(alias_41);  alias_41 = None
    view_37 = torch.ops.aten.view.default(alias_42, [arg0_1]);  alias_42 = None
    split_with_sizes_6 = torch.ops.aten.split_with_sizes.default(cat, [mul_1, mul_1, mul_1, mul_1, mul_1, mul_1, mul_1])
    getitem_42 = split_with_sizes_6[0];  getitem_42 = None
    getitem_43 = split_with_sizes_6[1];  getitem_43 = None
    getitem_44 = split_with_sizes_6[2];  getitem_44 = None
    getitem_45 = split_with_sizes_6[3];  getitem_45 = None
    getitem_46 = split_with_sizes_6[4];  getitem_46 = None
    getitem_47 = split_with_sizes_6[5]
    getitem_48 = split_with_sizes_6[6];  split_with_sizes_6 = getitem_48 = None
    view_38 = torch.ops.aten.view.dtype(getitem_47, torch.float32);  getitem_47 = None
    alias_43 = torch.ops.aten.alias.default(view_38);  view_38 = None
    alias_44 = torch.ops.aten.alias.default(alias_43);  alias_43 = None
    alias_45 = torch.ops.aten.alias.default(alias_44);  alias_44 = None
    view_39 = torch.ops.aten.view.default(alias_45, [arg0_1]);  alias_45 = None
    split_with_sizes_7 = torch.ops.aten.split_with_sizes.default(cat, [mul_1, mul_1, mul_1, mul_1, mul_1, mul_1, mul_1]);  cat = mul_1 = None
    getitem_49 = split_with_sizes_7[0];  getitem_49 = None
    getitem_50 = split_with_sizes_7[1];  getitem_50 = None
    getitem_51 = split_with_sizes_7[2];  getitem_51 = None
    getitem_52 = split_with_sizes_7[3];  getitem_52 = None
    getitem_53 = split_with_sizes_7[4];  getitem_53 = None
    getitem_54 = split_with_sizes_7[5];  getitem_54 = None
    getitem_55 = split_with_sizes_7[6];  split_with_sizes_7 = None
    view_40 = torch.ops.aten.view.dtype(getitem_55, torch.float32);  getitem_55 = None
    alias_46 = torch.ops.aten.alias.default(view_40);  view_40 = None
    alias_47 = torch.ops.aten.alias.default(alias_46);  alias_46 = None
    alias_48 = torch.ops.aten.alias.default(alias_47);  alias_47 = None
    view_41 = torch.ops.aten.view.default(alias_48, [arg0_1]);  alias_48 = arg0_1 = None
    return (view_29, view_31, view_33, view_35, view_37, view_39, view_41)
```

where exactly in the code that happen ? in this function:
```


def _create_graph(f, args, *, aot_config: AOTConfig) -> torch.fx.GraphModule:
    # FunctionalTensorMode must be enabled here.
    # See Note [Accessing .grad_fn on FunctionalTensor]
    with enable_python_dispatcher(), FunctionalTensorMode(
        pre_dispatch=aot_config.pre_dispatch,
        export=aot_config.is_export,
        # Allow token discovery for joint fn tracing as tokens can be used in backward.
        _allow_token_discovery=True,
    ):
        fx_g = make_fx(
            f,
            decomposition_table=aot_config.decompositions,
            record_module_stack=True,
            pre_dispatch=aot_config.pre_dispatch,
        )(*args)

    return fx_g
```
 in /home/lsakka/pytorch/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py


cc @bdhirsh @ezyang @zou3519 @penguinwu @yf225 @albanD @gqchen @pearu @nikitaved @soulitzer @Varal7 @xmfan @chauhang