Repository: pytorch
Issue: [RFC] Set Float32 Precision for CONV/RNN
Labels: triaged,needs design,topic: performance

# RFC: Extend set fp32 precision API to support Convolution and RNN

## Updated on 2024-03-27:
Incorperate suggestions from @albanD. We have 2 design options to enable low precision acceleration for fp32 `matmul/conv/rnn`.

**Option1**: Layering structure. All these flags should have an `Optional[bool]` type such that whenever a level is not specified, it uses the parent value. Use `allow`+dtype to represent the low precision dtype.
![image](https://github.com/zhuhaozhe/pytorch/assets/54701539/63d41e67-d828-4d5e-8732-98db0127e7d8)
 - **Pros**: Advanced configuration per backend/operator. More descriptive naming. Easier extension with newer data types.
 - **Cons**: Increased debugging complexity. Potential for multiple flags controlling the same thing. We may need to notify users when some low precision computation are enabled by some flags. Relative precision impact not immediately clear. May require additional documentation.

**Option2**: 3 flags `set_fp32_conv/rnn/matmul_precision` under `torch.*` for all backends.The value can be select from `highest, high, medium`. Allow backends to have backend-specific flags and setting backend-specific flags will also change the backend irrelevant flags.
 - **Pros**: Support per-operator configuration. Users can configure without detailed backend knowledge. Clear implication of relative precision.
 - **Cons**: Larger searching scope. Interaction between backend-irrelevant and backend-specific flags. e.g., if user disable tf32 for cuda backend, the `low precision` for `mkldnn matmul` are also disabled. Potential confusion for users familiar with other data types like `tf32` or `bf16`. Lack of direct association between `dtypes` and `highest, high, medium`.


## Previous Proposal on 2024-03-13:

## Overview
This RFC proposes the addition of a user-controlled frontend API to configure the internal precision of `float32` operations in convolutional (`CONV`) and recurrent neural networks (`RNN`) within PyTorch. Currently, PyTorch offers [`torch.set_float32_matmul_precision`](https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch-set-float32-matmul-precision) to configure the internal precision of `float32` matrix multiplication. This RFC suggests extending this functionality to include convolution and recurrent neural network operations, providing `torch.set_float32_conv_precision` and `torch.set_float32_rnn_precision`. The proposed APIs will mimic the behavior of `torch.set_float32_matmul_precision`.

## Frontend Changes

Frontend changes involve introducing new APIs:

- `torch.set_float32_conv_precision`, `torch.get_float32_conv_precision`
- `torch.set_float32_rnn_precision`, `torch.get_float32_rnn_precision`

These APIs will function similarly to `torch.set_float32_matmul_precision` and `torch.get_float32_matmul_precision`. Users can set the precision to `highest`, `high`, or `medium`, each with corresponding backend behavior:
- `highest`: Use the highest available precision, avoiding lower precision.
- `high`: Allow backends to use TensorFloat32 (TF32) or treat each `float32` number as the sum of two `bfloat16` numbers.
- `medium`: Allow backends to use BFloat 16 (BF16).

## Backend Changes

Global flags `float32_conv/rnn_precision` will be introduced at [this location](https://github.com/pytorch/pytorch/blob/9c50ecc84b9a6e699a7f058891b889aafbf976c7/aten/src/ATen/Context.h#L358-L361) in the PyTorch repository. This flag can be accessed and modified by the frontend APIs `torch.get/set_float32_conv/rnn_precision`. Backend-related operators will read this flag to control the internal computation data types. For example:
- For CuDNN backend, we should check `float32_conv_precision` in the [CuDNN Conv kernel](https://github.com/pytorch/pytorch/blob/9c50ecc84b9a6e699a7f058891b889aafbf976c7/aten/src/ATen/native/cudnn/ConvShared.cpp#L728). We should also check `float32_rnn_precision` in the [CuDNN RNN kernel](https://github.com/pytorch/pytorch/blob/9c50ecc84b9a6e699a7f058891b889aafbf976c7/aten/src/ATen/native/cudnn/RNN.cpp#L248). If not set to `highest`, the internal computation data type will be TF32.
- For OneDNN backend, we should check `float32_conv_precision` in [OneDNN Conv kernel](https://github.com/pytorch/pytorch/blob/9c50ecc84b9a6e699a7f058891b889aafbf976c7/aten/src/ATen/native/mkldnn/Conv.cpp#L211) and check `float32_rnn_precision` in [OneDNN RNN kernel](https://github.com/pytorch/pytorch/blob/9c50ecc84b9a6e699a7f058891b889aafbf976c7/aten/src/ATen/native/mkldnn/RNN.cpp#L220C44-L220C60). If set to `medium`, the internal data type will be BF16.

## Flag Overrides

The existing CUDNN backend-specific flag [`torch.backends.cudnn.allow_tf32`](https://pytorch.org/docs/stable/notes/cuda.html#tf32-on-ampere) will interact with the proposed backend-irrelevant flag `torch.set_float32_conv/rnn_precision`. These flags will override each other( we follow similar behavior between `torch.backends.cuda.matmul.allow_tf32` and `float32_matmul_precision`):
- Turning on/off TF32 with `torch.backends.cudnn.allow_tf32` will set `float32_rnn/conv_precision` to `high` (TF32 enabled) and `highest` (TF32 disabled).
```python
torch.backends.cudnn.alow_tf32=True
print("float32_conv_precision", torch.get_float32_conv_precision)
print("float32_rnn_precision", torch.get_float32_rnn_precision)
# output:
# float32_conv_precision, high
# float32_rnn_precision, high
torch.backends.cudnn.alow_tf32=False
print("float32_conv_precision", torch.get_float32_conv_precision)
print("float32_rnn_precision", torch.get_float32_rnn_precision)
# output:
# float32_conv_precision, highest
# float32_rnn_precision, highest
```
- Setting both `float32_rnn/conv_precision` to `high` or `medium` will enable `torch.backends.cudnn.allow_tf32`, while setting **one of it** to `highest` will disable it.
```python
torch.backends.cudnn.alow_tf32=True
torch.set_float32_conv_precision("highest")
print("torch.backends.cudnn.alow_tf32", torch.backends.cudnn.alow_tf32)
# output:
# torch.backends.cudnn.alow_tf32, False
torch.set_float32_rnn_precision("highest")
print("torch.backends.cudnn.alow_tf32", torch.backends.cudnn.alow_tf32)
# output:
# torch.backends.cudnn.alow_tf32, False
torch.set_float32_conv_precision("high")
torch.set_float32_rnn_precision("high")
print("torch.backends.cudnn.alow_tf32", torch.backends.cudnn.alow_tf32)
# output:
# torch.backends.cudnn.alow_tf32, True
```

## Additional CuDNN Flag
We discussed how the existing CuDNN flag, `torch.backends.cudnn.allow_tf32`, interacts with `torch.set_float32_conv/rnn_precision`. However, we believe it is cleaner to use separate flags in CuDNN. We suggest deprecating `torch.backends.cudnn.allow_tf32` in favor of `torch.backends.cudnn.conv.allow_tf32` and `torch.backends.cudnn.rnn.allow_tf32`. Then, the CuDNN backend-specific flags and backend-irrelevant flags can have a one-to-one correspondence, such as `torch.backends.cuda.matmul.allow_tf32` and `torch.float32_matmul_precision`
```
torch.backends.cudnn.conv.allow_tf32 <-> torch.float32_conv_precision
torch.backends.cudnn.rnn.allow_tf32 <-> torch.float32_rnn_precision
# below flags are already existing now
torch.backends.cuda.matmul.allow_tf32 <-> torch.float32_matmul_precision
```

## Motivation

Lower-precision computation from different backends can significantly improve performance for deep learning workloads with minimal impact on precision. For example,  [`TF32`](https://blogs.nvidia.com/blog/tensorfloat-32-precision-format/)  from `CUDA/CUDNN` or  [implicit reduced precision arithmetic feature](https://github.com/mgouicem/oneDNN/tree/mgouicem/rfcs/implicit_downconvert/rfcs/20210301-computation-datatype) from `oneDNN`. By providing a user-controlled frontend API, users can easily configure the internal computation data type of convolutional and recurrent neural networks without knowing the detail of different backends. This allows them to leverage the performance benefits of lower precision while ensuring acceptable precision loss. Compared to [`Autocast`](https://pytorch.org/docs/stable/amp.html#autocasting), the proposed flags offer:
- Higher precision control as they only affect convolutional and recurrent neural network internal data types, unlike Autocast, which impacts more operators.
- Ease of use, as users do not need to modify their model scripts to enable autocasting.

## Pitch

Introduce `float32_conv/rnn_precision` and enable users to control the internal data type for convolutional and recurrent neural networks by configuring the value of `float32_conv/rnn_precision`.