Repository: pytorch
Issue: Foreach Functions Tracking Issue
Labels: module: performance,triaged,tracker,module: mta

Foreach functions would improve optimizers performance as described in https://github.com/pytorch/pytorch/issues/38655, but they need fixes & improvements to replace the current optimizers implementation.


# A: Functionality enhancement
To make it possible to use foreach functions in optimizers, we need to support tensors on different devices, dtypes, and layouts even if those inputs are not going through fastpath (Multi Tensor Apply).

- [x] Push inputs incompatible with MTA to slowpath - #56993
- [x] Refactor unittest
    - [x] unary - ~~#58960~~ https://github.com/pytorch/pytorch/commit/52b2ed65c04b0010522c7fc3f70981ce57b8f1d8
    - [x] binary - #59907
    - [x] pointwise & min/max-imum: #61327
- [x] Support implicit broadcasting via slowpath - #62167 

# B: Performance improvement

## B.1: foreach functions
- [x] Let binary fastpath support inputs of tensorlist and complex scalarlist.
    - binary foreach's with two tensorlists and one tensorlist and one scalar support complex. The limitation derives from that [`TensorListScalarListMetadata`](https://github.com/pytorch/pytorch/blob/50ded095e41d6d3edcde9e5e4a608b89b9a374ce/aten/src/ATen/native/cuda/MultiTensorApply.cuh#L40) violates kernel argument size limit when `scalar_vals_t` is `c10::complex<double>` and `n` is `1`. #59907 is working on this
<!-- 
    - https://github.com/pytorch/pytorch/blob/50ded095e41d6d3edcde9e5e4a608b89b9a374ce/aten/src/ATen/native/cuda/ForeachBinaryOpList.cu#L20
    - https://github.com/pytorch/pytorch/blob/50ded095e41d6d3edcde9e5e4a608b89b9a374ce/aten/src/ATen/native/cuda/ForeachBinaryOpScalar.cu#L19
    - https://github.com/pytorch/pytorch/blob/50ded095e41d6d3edcde9e5e4a608b89b9a374ce/aten/src/ATen/native/cuda/ForeachBinaryOpScalarList.cu#L20 
-->
- [x] Reduction MTA Kernels - ref: [APEX L2 norm impl.](https://github.com/NVIDIA/apex/blob/00c1e56d02f2987e71ee40417ea7e938e1493475/csrc/multi_tensor_l2norm_kernel.cu#L293) - #62646, #68107
- [ ] Type promotion in fastpath.
    - [ ] Unary
    - [ ] Binary
    - [ ] Pointwise 
- [x] Separate input `TensorList`s by devices & dtypes into bunch of vectors of Tensors in C++ to utilize MTA kernel as much as possible<!-- and avoid slow Python list creation (see [GradScaler._unscalegrads\_](https://github.com/pytorch/pytorch/blob/50ded095e41d6d3edcde9e5e4a608b89b9a374ce/torch/cuda/amp/grad_scaler.py#L190)) -->.
    - https://github.com/pytorch/pytorch/pull/94344

For Binary & Pointwise support of type promotion, I think we need a new macro that can handle multiple C++ types or one that can be nested like [APEX](https://github.com/NVIDIA/apex/blob/082f999a6e18a3d02306e27482cc7486dab71a50/csrc/type_shim.h).

## B.2: Optimizers Specific
Rel: https://github.com/pytorch/pytorch/issues/68052

`foreach` option has been added

<details>
<summary> list of foreach optimizer PRs </summary>

- [x] Rewrite `torch.optim._multi_tensor` optimizers in functional form <!-- (my draft branch is https://github.com/crcrpar/pytorch/blob/backup/foreach/optimizers/torch/optim/_multi_tensor/_functional.py)
        - there are some pull requests for this: https://github.com/pytorch/pytorch/pull/59155, https://github.com/pytorch/pytorch/pull/59229 -->
    - [x] SGD ~~#60962~~ #70481
    - [x] Rprop ~~#60963~~ #70483
    - [x] RMSprop ~~#60964~~ #70482
    - [x] Adam ~~#60960~~ #70295
    - [x] AdamW ~~#60961~~ #70484
    - [x] Adagrad #59155 #69981
    - [x] Adamax #60735 #69982
    - [x] Adadelta #60735 #69980
    - [x] asgd #60735 #70231
    - [x] radam #59161 #70230
    - [x] nadam #59165 #70229
    - [x] `foreach_zero_` #69936

</details>

# C: autograd
- https://github.com/pytorch/pytorch/pull/89591
- https://github.com/pytorch/pytorch/pull/93901#issuecomment-1440434125
- #95263
- #96405
    - #100445
- #100256 
- [ ] try to generate `THP<codegen'd Foreach Backward TraceableFunction_properties[]` so that `result_tensors[0].grad_fn` will be explicit
- [x] avoid having each input Tensor's grad_fn have the reference to redundant tensors which seems added by the backward backward function struct <-- partially done in #96405
- https://github.com/pytorch/pytorch/pull/100695
- https://github.com/pytorch/pytorch/pull/105504
- https://github.com/pytorch/pytorch/pull/106043

---

__misc__
- #58209
- #93995

**Math in FP32 for 16 bits I/O**
- #60227
- #60454
- #60715

---
Please kindly let me know if I miss anything and/or I make some mistakes
<!-- cc: @ptrblck @mcarilli @ngimel -->

cc @crcrpar @mcarilli @ngimel @VitalyFedyunin @heitorschueroff