[Operator] Add radix-2 fft

[Bowen12992]

PR Category

Operator

Type of Change

New Feature

Description

This PR implement a native FFT with Cooley-Tukey method，current it only support radix-2 The performance is far worse than pytorch aten library (which use CuFFT library actually) for some reasons

• triton-lang does not support complex dtype so we need to separate the imag and real part
• triton-lang does not support __syncthreads so we must launch much more kernels
• A fft plan is needed for differnet kinds of cases
• sharedMem should be used for small size FFT

Issue

Progress

• [ ] Change is properly reviewed (1 reviewer required, 2 recommended).
• [x] Change is responded to an issue.
• [x] Change is fully covered by a UT.

Performance

benchmark/test_special_perf.py Operator fft.fft Performance Test (torch.complex64)
Size    Torch Latency (ms)    Gems Latency (ms)    Gems Speedup
---------------------------------------------------------------
1024               0.00512              0.04096           0.125
11264             0.012288             0.082944           0.148
21504              0.01536             0.102144            0.15
31744             0.014336             0.100352           0.143
41984              0.01536              0.11776            0.13
52224             0.017408             0.115712            0.15
62464             0.016384              0.11776           0.139
72704             0.017408             0.131072           0.133
82944             0.014496             0.131072           0.111
93184             0.018432             0.131136           0.141

[Bowen12992]

Since fft is difficult for triton now, i checkout a new branch named fft_dev and close this PR