Repository: pytorch
Issue: Performance improvement in Autograd Forward AD using ZeroTensors
Labels: module: performance,module: autograd,triaged,actionable,module: forward ad

### 🚀 The feature, motivation and pitch

In https://github.com/pytorch/pytorch/pull/64837 we added a new type of Tensors, called ZeroTensors. They are implemented using a DispatchKey and boxed fallback mechanism similar to Conjugate view.

The goal of this project is to replace the undefined tensors used in the forward and backward mode AD which are currently materialized to a tensor full of zeros with a new type of zero tensor that use O(1) memory (tensors with a storage pointing to nullptr). Having undefined tensors is a common occurrence in the forward mode AD. Currently these undefined tensors are materialized to zero tensors. Not only the materialization is wasteful, but the cases where computationally heavy operations like matrix multiplies are called for these tensors along with some non-zero tensors provide opportunities for perf improvement. With zero tensors, expensive operations like matmul on zero tensors will become O(1).

Note that as of writing this, ZeroTensors do not propagate nans or other extreme values. This is rooted in how the kernels are implemented for ZeroTensors. For example, ZT * (non-ZT) = ZT, i.e., we don't check whether the non-ZT has nans for example.

Design:

1. ZeroTensors are regular tensors with TensorOptions (as usual), and a storage pointing to nullptr (aliasing semantic is correctly maintained), and a ZeroTensor dispatch key set.
2. Zero tensors are immutable. 
    -  As a result, in-place operations on these tensors will be disallowed. This is done to prevent data race in the case of multithreading (when two threads try to read the same zero tensor and materialize it in-place).
    - Zero tensors shouldn’t be passed as an output tensor in out= operations.    
3. Similar to Conjugate and Negative dispatch keys, ZeroTensor keys also has a corresponding boxed fallback for dispatch key resolution. Most functions will use this fallback if the ZeroTensor key is set on one or more inputs. A select few functions that are registered as fallthrough kernels or have a special kernel registered in native_functions.yaml will bypass this fallback (mul, matmul, tensor views, etc.).
    - This fallback materializes each ZeroTensor in the following way: `at::zeros({}, tensor.options()).expand(tensor.sizes());` where tensor is an efficient zero tensor (tensor with ZeroTensor dispatch key set)
    - If there’s a mutable ZeroTensor , then we simply error out to avoid the race condition (as mentioned above).
4. ZeroTensors are handled below autograd. This is necessary because we want the materialization to happen below the autograd to ensure gradients are populated on the correct tensors. 
    - Example: Consider add in the case it was using the fallback: zerotensor_a + b. zerotensor_a would be materialized to c=torch.zeros_like(zerotensor_a) after passing through the fallback. If this happens above the autograd, then the gradients would be populated on c instead of zerotensor_a.
5. The grad field is always populated with an honest to goodness tensor. This materialization of zero tensors will happen in:
    - AcccumulateGrad for Backward Mode AD.
    - will never be required for ForwardMode AD. 
        - This is because if all the tangents were undefined (efficient zero tensors), no computation will be performed (this is ensured via an existing pre-check). 
        - If one or more tangent is non-zero (undefined), then even if the Jacobian is zero, we never get an efficient zero tensor as an output (since we don’t compute efficient zero tensors in Jacobian computation, i.e., Jacobian can never be a ZeroTensor). This will not be true anymore if ZeroTensors were exposed in the Python API.
6. torch.mul(efficient_zero_tensor, other) will return an efficient_zero_tensor. We aim to add similar specialized handling for other relevant functions based on: https://fb.quip.com/Dv5hARK0xyTo
7. All instance of zeros_like should be replaced with a call to _efficientzerotensor. This would ensure that we are taking maximally optimizing the performance wherever we can (combined with not using ZeroTensor fallback for functions where the function computation can be short circuited with the added information that one of the inputs is a zero tensor).

Testing:
The only place where ZeroTensors are currently being used is in Forward Mode AD. To test ZeroTensors, we have expanded torch.autograd.gradcheck to verify correct Forward Mode AD behavior with undefined gradients (which are stored as ZeroTensors). The way this is done is, we loop through the tensor inputs, and set the tangent corresponding to the current input to ZeroTensor in one case and an actual zero tensor (torch.zeros(...)) in the other. The JVP computed for each output is verified to be equal in both cases.

Since we run gradcheck for Forward AD for all OpInfos, this is tested for all operators with OpInfos (which don’t have forward mode AD tests disabled). 

Tasks to be done:

As mentioned in (6), we want to add fastpath for Forward AD formulas for ZeroTensors wherever possible. Here are the list of functions whose JVP formulas need to be fastpathed:

- [x] `add.Tensor`
- [ ] [`addbmm`](https://github.com/pytorch/pytorch/issues/72358)
- [x] [`addcmul`](https://github.com/pytorch/pytorch/issues/72347)
- [ ] [`addmm`](https://github.com/pytorch/pytorch/issues/72345)
- [ ]  [`addmv`](https://github.com/pytorch/pytorch/issues/72359)
- [x]  [`addr`](https://github.com/pytorch/pytorch/issues/72360)
- [ ]  [`baddbmm`](https://github.com/pytorch/pytorch/issues/72361)
- [ ]  [`bmm`](https://github.com/pytorch/pytorch/issues/72362)
- [x] `cat` https://github.com/pytorch/pytorch/pull/75590
- [x] `stack` https://github.com/pytorch/pytorch/pull/75590
- [ ]  `cross`
- [x] `dot`
- [x] `vdot`
- [ ]  `index_add_`
- [ ]  `index_copy_`
- [ ]  `index_fill_.int_Tensor`
- [ ]  `index_put_`
- [ ]  `_index_put_impl_`
- [x]  `max`
- [x]  `min`
- [ ] #75586
- [x]  `mul.Tensor`
- [x]  `sub.Tensor`
- [ ]  `remainder.Tensor` 

To be verified: `pow.Tensor_Tensor`, `linalg_solve`

should special case (even zero tensor can’t help — understandably so):

* clamp_min.Tensor: where(self_p >= min_p, self_t, at::scalar_tensor(0., self_p.options()))
    * return zero if self_t is undefined
* clamp_max.Tensor — same story 

Problematic:

We should add a macro to check for zero tensors so that we can keep the formulas inline in `derivatives.yaml`:

* addcdiv: ... maybe_multiply(tensor2_t * (tensor1_p / tensor2_p) / tensor2_p, value)
* div.Tensor-- result: self_t / other_p - other_t * (self_p / other_p) / other_p
    * but can possibly benefit from zero tensor if written in a different way
        * div.Tensor-- result: self_t / other_p - ((other_t * self_p) / other_p) / other_p
* div.Tensor_mode, hypot, 
* lerp
    * result: at::lerp(self_t, end_t, weight_p) + weight_t * (end_p - self_p)
* logaddexp, logaddexp2
    * result: self_t / (1 + pow(2, other_p - self_p)) + other_t / (1 + pow(2, self_p - other_p))
* xlogy.Tensor
    * self_t * at::xlogy((self_p != 0), other_p) + other_t * self_p / other_p
* special_xlog1py
    * result: self_t * other_p.log1p() + other_t * self_p / (other_p + 1)
* maximum
    * result: other_t + at::where(self_p == other_p, 0.5, (self_p > other_p).to(result.scalar_type())) * (self_t - other_t)
* fmax
    * result: other_t + (self_p > other_p).logical_or_(other_p.isnan()) * (self_t - other_t)
* minimum
* fmin

cc @VitalyFedyunin @ngimel @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @Lezcano @Varal7