Repository: pytorch
Issue: More economical use of FX nodes when generating symbolic shapes related manipulations
Labels: high priority,triaged,oncall: pt2,module: dynamic shapes

### 🐛 Describe the bug

Internal xref: https://www.internalfb.com/phabricator/paste/view/P1398677554

This is high priority because it potentially affects our E2E performance, which is one of our priorities for H2 2024.

@PaulZhang12 shared with me some export code they're working on and the runtime asserts want to make me gouge my eyes out. Here's some highlights.

**Overaggressive canonicalization leads to unnatural inequalities.**

```
    mul_37 = -1 * item
    le_4 = mul_37 <= 0;  mul_37 = None
    _assert_scalar_default = torch.ops.aten._assert_scalar.default(le_4, "Runtime assertion failed for expression -u71 <= 0 on node 'le_4'");  le_4 = None
```

A much more natural inequality to test here is `u71 >= 0`, which also uses one less FX node. The stilted inequality here arises from `_canonicalize_bool_expr_impl` which eliminates `>=` relations in favor of `<=` relations only.

**tolist desugars into too many nodes.**

```
    unbind = torch.ops.aten.unbind.int(sum_1);  sum_1 = None
    getitem = unbind[0]
    getitem_1 = unbind[1]
    getitem_2 = unbind[2]
    getitem_3 = unbind[3]
    getitem_4 = unbind[4]
    getitem_5 = unbind[5]
    getitem_6 = unbind[6]
    getitem_7 = unbind[7]
    getitem_8 = unbind[8]
    getitem_9 = unbind[9]
    getitem_10 = unbind[10]
    getitem_11 = unbind[11]
    getitem_12 = unbind[12]
    getitem_13 = unbind[13]
    getitem_14 = unbind[14]
    getitem_15 = unbind[15]
    getitem_16 = unbind[16]
    getitem_17 = unbind[17]
    getitem_18 = unbind[18]
    getitem_19 = unbind[19]
    getitem_20 = unbind[20]
    getitem_21 = unbind[21]
    getitem_22 = unbind[22]
    getitem_23 = unbind[23]
    getitem_24 = unbind[24]
    getitem_25 = unbind[25]
    getitem_26 = unbind[26]
    getitem_27 = unbind[27]
    getitem_28 = unbind[28]
    getitem_29 = unbind[29]
    getitem_30 = unbind[30]
    getitem_31 = unbind[31]
    getitem_32 = unbind[32]
    getitem_33 = unbind[33]
    getitem_34 = unbind[34];  unbind = None
    item = torch.ops.aten.item.default(getitem);  getitem = None
    item_1 = torch.ops.aten.item.default(getitem_1);  getitem_1 = None
...
```

Some of the difficulty here is due to FX lack of multi-return, so any N-return op ends up generating for getitem. But we shouldn't have to add another N nodes afterwards for each of the individual item ops; instead, we should just have a consolidated op that does this all in one go. tolist is very, very common in this use case.

**More compact form for check_size tests.** `_assert_scalar` is very general, but it's very common to test if something is non-negative. If we fuse the relation into the `assert_scalar`, this saves us another FX node. Similarly, we often end up doing this assert on every output of a tolist(), and so we can save O(N) FX nodes by fusing it in.

**Useless split/cat.**

Continuing the item calls from above:

```
...
    scalar_tensor = torch.ops.aten.scalar_tensor.default(item, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_1 = torch.ops.aten.scalar_tensor.default(item_1, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_2 = torch.ops.aten.scalar_tensor.default(item_2, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_3 = torch.ops.aten.scalar_tensor.default(item_3, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_4 = torch.ops.aten.scalar_tensor.default(item_4, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_5 = torch.ops.aten.scalar_tensor.default(item_5, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_6 = torch.ops.aten.scalar_tensor.default(item_6, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_7 = torch.ops.aten.scalar_tensor.default(item_7, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_8 = torch.ops.aten.scalar_tensor.default(item_8, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_9 = torch.ops.aten.scalar_tensor.default(item_9, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_10 = torch.ops.aten.scalar_tensor.default(item_10, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_11 = torch.ops.aten.scalar_tensor.default(item_11, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_12 = torch.ops.aten.scalar_tensor.default(item_12, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_13 = torch.ops.aten.scalar_tensor.default(item_13, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_14 = torch.ops.aten.scalar_tensor.default(item_14, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_15 = torch.ops.aten.scalar_tensor.default(item_15, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_16 = torch.ops.aten.scalar_tensor.default(item_16, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_17 = torch.ops.aten.scalar_tensor.default(item_17, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_18 = torch.ops.aten.scalar_tensor.default(item_18, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_19 = torch.ops.aten.scalar_tensor.default(item_19, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_20 = torch.ops.aten.scalar_tensor.default(item_20, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_21 = torch.ops.aten.scalar_tensor.default(item_21, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_22 = torch.ops.aten.scalar_tensor.default(item_22, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_23 = torch.ops.aten.scalar_tensor.default(item_23, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_24 = torch.ops.aten.scalar_tensor.default(item_24, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_25 = torch.ops.aten.scalar_tensor.default(item_25, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_26 = torch.ops.aten.scalar_tensor.default(item_26, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_27 = torch.ops.aten.scalar_tensor.default(item_27, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_28 = torch.ops.aten.scalar_tensor.default(item_28, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_29 = torch.ops.aten.scalar_tensor.default(item_29, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_30 = torch.ops.aten.scalar_tensor.default(item_30, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_31 = torch.ops.aten.scalar_tensor.default(item_31, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_32 = torch.ops.aten.scalar_tensor.default(item_32, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_33 = torch.ops.aten.scalar_tensor.default(item_33, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    scalar_tensor_34 = torch.ops.aten.scalar_tensor.default(item_34, dtype = torch.int32, device = device(type='cpu'), pin_memory = False)
    stack = torch.ops.aten.stack.default([scalar_tensor, scalar_tensor_1, scalar_tensor_2, scalar_tensor_3, scalar_tensor_4, scalar_tensor_5, scalar_tensor_6, scalar_tensor_7, scalar_tensor_8, scalar_tensor_9, scalar_tensor_10, scalar_tensor_11, scalar_tensor_12, scalar_tensor_13, scalar_tensor_14, scalar_tensor_15, scalar_tensor_16, scalar_tensor_17, scalar_tensor_18, scalar_tensor_19, scalar_tensor_20, scalar_tensor_21, scalar_tensor_22, scalar_tensor_23, scalar_tensor_24, scalar_tensor_25, scalar_tensor_26, scalar_tensor_27, scalar_tensor_28, scalar_tensor_29, scalar_tensor_30, scalar_tensor_31, scalar_tensor_32, scalar_tensor_33, scalar_tensor_34]);  scalar_tensor = scalar_tensor_1 = scalar_tensor_2 = scalar_tensor_3 = scalar_tensor_4 = scalar_tensor_5 = scalar_tensor_6 = scalar_tensor_7 = scalar_tensor_8 = scalar_tensor_9 = scalar_tensor_10 = scalar_tensor_11 = scalar_tensor_12 = scalar_tensor_13 = scalar_tensor_14 = scalar_tensor_15 = scalar_tensor_16 = scalar_tensor_17 = scalar_tensor_18 = scalar_tensor_19 = scalar_tensor_20 = scalar_tensor_21 = scalar_tensor_22 = scalar_tensor_23 = scalar_tensor_24 = scalar_tensor_25 = scalar_tensor_26 = scalar_tensor_27 = scalar_tensor_28 = scalar_tensor_29 = scalar_tensor_30 = scalar_tensor_31 = scalar_tensor_32 = scalar_tensor_33 = scalar_tensor_34 = None
```

WRYYYY.

The proximal cause for this is revealed below:

```
    _to_copy = torch.ops.aten._to_copy.default(stack, dtype = torch.int32, device = device(type='cpu'));  stack = None
    detach = torch.ops.aten.detach.default(_to_copy);  _to_copy = None
    asynchronous_complete_cumsum = torch.ops.fbgemm.asynchronous_complete_cumsum.default(detach);  detach = None
    unbind_1 = torch.ops.aten.unbind.int(asynchronous_complete_cumsum);  asynchronous_complete_cumsum = None
```

This is the workaround for performing a cumulative sum, but without this the cumulative sum would generate a lot of nodes and be very slow.

**Peephole eliminate trivial operations.**

```
    add_3 = 0 + item_1;  item_1 = None
```

There's no reason to add zero here, delete this FX node entirely.

**Better way to do split after cumsum.**

Eventually, we do this:

```
 add = 0 + item;  item = None
    add_2 = 0 + add;  add = None
    slice_3 = torch.ops.aten.slice.Tensor(train_input_id_score_list_features__weights, 0, 0, add_2);  add_2 = None
    add_3 = 0 + item_1;  item_1 = None
    add_5 = item_36 + add_3;  add_3 = None
    slice_6 = torch.ops.aten.slice.Tensor(train_input_id_score_list_features__weights, 0, item_36, add_5);  item_36 = add_5 = None
    add_6 = 0 + item_2;  item_2 = None
    add_8 = item_37 + add_6;  add_6 = None
    slice_9 = torch.ops.aten.slice.Tensor(train_input_id_score_list_features__weights, 0, item_37, add_8);  item_37 = add_8 = None
    add_9 = 0 + item_3;  item_3 = None
    add_11 = item_38 + add_9;  add_9 = None
    slice_12 = torch.ops.aten.slice.Tensor(train_input_id_score_list_features__weights, 0, item_38, add_11);  item_38 = add_11 = None
    add_12 = 0 + item_4;  item_4 = None
    add_14 = item_39 + add_12;  add_12 = None
    slice_15 = torch.ops.aten.slice.Tensor(train_input_id_score_list_features__weights, 0, item_39, add_14);  item_39 = add_14 = None
    add_15 = 0 + item_5;  item_5 = None
    add_17 = item_40 + add_15;  add_15 = None
    slice_18 = torch.ops.aten.slice.Tensor(train_input_id_score_list_features__weights, 0, item_40, add_17);  item_40 = add_17 = None
    add_18 = 0 + item_6;  item_6 = None
```

It shouldn't be necessary to materialize the offsets here. Using views is unnecessary except to avoid unnecessary copies. The only unbacked SymInt we actually need is the lengths of the slices.

**Use foreach hop.**

On each of the slices, we repeatedly call the same operations like index_hash/bucketize. Use foreach hop for this case.

cc @gchanan @zou3519 @kadeng @msaroufim @chauhang @penguinwu @bobrenjc93 @bdhirsh @anijain2305 @lezcano @shazqadeer @mlazos 

### Versions

main