Repository: pytorch
Issue: torch.fft.rfft: Unsupported dtype BFloat16
Labels: module: cuda,triaged,module: bfloat16,needs design,module: fft

### 🐛 Describe the bug

```py
import torch
x = torch.rand((16, 64, 256), dtype=torch.bfloat16, device='cuda')
f = torch.fft.rfft(x)  # RuntimeError: Unsupported dtype BFloat16
```

In my use-case this fft operation is on very long sequences, so it would save a lot of memory!

Using FP16 leads to instability issues, so I would prefer using BF16 instead.

Are there any plans to support this? [It looks like cuFFT supports BF16 already...](https://docs.nvidia.com/cuda/cufft/index.html#bfloat16-precision-cufft-transforms)

### Versions

<details>
<summary>Versions</summary>

```
PyTorch version: 2.1.2+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.6
CMake version: version 3.28.1
Libc version: glibc-2.35

Python version: 3.10.13 | packaged by conda-forge | (main, Dec 23 2023, 15:36:39) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-6.5.0-14-generic-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.1.105
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA GeForce GTX 1080 Ti
GPU 1: NVIDIA GeForce GTX 1080 Ti
GPU 2: NVIDIA GeForce RTX 3090

Nvidia driver version: 535.146.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.5
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.5
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                       x86_64
CPU op-mode(s):                     32-bit, 64-bit
Address sizes:                      43 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             24
On-line CPU(s) list:                0-23
Vendor ID:                          AuthenticAMD
Model name:                         AMD Ryzen Threadripper 1920X 12-Core Processor
CPU family:                         23
Model:                              1
Thread(s) per core:                 2
Core(s) per socket:                 12
Socket(s):                          1
Stepping:                           1
Frequency boost:                    enabled
CPU max MHz:                        3500.0000
CPU min MHz:                        2200.0000
BogoMIPS:                           6999.09
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid amd_dcm aperfmperf rapl pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb hw_pstate ssbd ibpb vmmcall fsgsbase bmi1 avx2 smep bmi2 rdseed adx smap clflushopt sha_ni xsaveopt xsavec xgetbv1 clzero irperf xsaveerptr arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif overflow_recov succor smca sev
Virtualization:                     AMD-V
L1d cache:                          384 KiB (12 instances)
L1i cache:                          768 KiB (12 instances)
L2 cache:                           6 MiB (12 instances)
L3 cache:                           32 MiB (4 instances)
NUMA node(s):                       1
NUMA node0 CPU(s):                  0-23
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Vulnerable
Vulnerability Spec rstack overflow: Vulnerable
Vulnerability Spec store bypass:    Vulnerable
Vulnerability Spectre v1:           Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2:           Vulnerable, IBPB: disabled, STIBP: disabled, PBRSB-eIBRS: Not affected
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] ema-pytorch==0.3.2
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.23.5
[pip3] pytorch-lightning==2.1.3
[pip3] torch==2.1.2
[pip3] torch-dct==0.1.6
[pip3] torchaudio==2.1.2
[pip3] torchdata==0.7.1
[pip3] torchextractor==0.3.0
[pip3] torchmetrics==1.2.1
[pip3] torchtext==0.16.2
[pip3] torchtnt==0.2.1
[pip3] torchvision==0.16.2
[pip3] triton==2.2.0
[pip3] vector-quantize-pytorch==1.12.5
[conda] ema-pytorch               0.3.2                    pypi_0    pypi
[conda] numpy                     1.23.5                   pypi_0    pypi
[conda] pytorch-lightning         2.1.3                    pypi_0    pypi
[conda] torch                     2.1.2                    pypi_0    pypi
[conda] torch-dct                 0.1.6                    pypi_0    pypi
[conda] torchaudio                2.1.2                    pypi_0    pypi
[conda] torchdata                 0.7.1                    pypi_0    pypi
[conda] torchextractor            0.3.0                    pypi_0    pypi
[conda] torchmetrics              1.2.1                    pypi_0    pypi
[conda] torchtext                 0.16.2                   pypi_0    pypi
[conda] torchtnt                  0.2.1                    pypi_0    pypi
[conda] torchvision               0.16.2                   pypi_0    pypi
[conda] triton                    2.2.0                    pypi_0    pypi
[conda] vector-quantize-pytorch   1.12.5                   pypi_0    pypi
```

</details>

cc @ptrblck @mruberry @peterbell10