Add fft 1d OP

benchmark/test_special_perf.py

@@ -630,3 +630,31 @@ def set_more_shapes(self):
 
     bench.set_gems(gems_op)
     bench.run()
+
+
+@pytest.mark.fft_1d
+def test_perf_fft_1d():
+    def fft_1d_input_fn(shape, dtype, device):
+        N = 1 << (shape[0] - 1).bit_length()
+
+        input_tensor = (
+            torch.randn((N,), device="cuda") + torch.randn((N,), device="cuda") * 1j
+        )
+        output_tensor = (
+            torch.empty((N,), device="cuda") + torch.empty((N,), device="cuda") * 1j
+        )
+        yield input_tensor, output_tensor
+
+    def torch_op(input_tensor, output_tensor):
+        output_tensor.copy_(torch.fft.fft(input_tensor))
+
+    gems_op = flag_gems.fft_1d
+
+    bench = GenericBenchmarkExcluse1D(
+        input_fn=fft_1d_input_fn,
+        op_name="fft_1d",
+        torch_op=torch_op,
+        # dtypes=FLOAT_DTYPES,
+    )
+    bench.set_gems(gems_op)
+    bench.run()

src/flag_gems/ops/__init__.py

@@ -88,6 +88,7 @@
 from flag_gems.ops.exponential_ import exponential_
 from flag_gems.ops.eye import eye
 from flag_gems.ops.eye_m import eye_m
+from flag_gems.ops.fft_1d import fft_1d
 from flag_gems.ops.fill import fill_scalar, fill_scalar_, fill_tensor, fill_tensor_
 from flag_gems.ops.flip import flip
 from flag_gems.ops.full import full
@@ -329,6 +330,7 @@
     "exponential_",
     "eye",
     "eye_m",
+    "fft_1d",
     "fill_scalar",
     "fill_scalar_",
     "fill_tensor",

src/flag_gems/ops/fft_1d.py

@@ -0,0 +1,103 @@
+import math
+
+import torch
+import torch.profiler
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def bit_reverse_kernel(real_in, imag_in, real_out, imag_out, n):
+    """do reverse first: input[i] -> output[bit_reverse(i)]"""
+    tid = tl.program_id(0)
+    if tid >= n:
+        return
+
+    # compute bits & reverse
+    temp_n = n
+    idx = tid
+    rev_idx = 0
+    temp_idx = idx
+    while temp_n > 1:
+        temp_n //= 2
+        rev_idx = (rev_idx << 1) | (temp_idx & 1)
+        temp_idx = temp_idx >> 1
+
+    val_real = tl.load(real_in + idx)
+    val_imag = tl.load(imag_in + idx)
+    tl.store(real_out + rev_idx, val_real)
+    tl.store(imag_out + rev_idx, val_imag)
+
+
+@triton.jit
+def fft_stage_kernel(real_ptr, imag_ptr, n, stage):
+    """iterate the FFT stage"""
+    PI = math.pi
+    tid = tl.program_id(0)
+
+    if tid >= n // 2:
+        return
+
+    # compute current parameter
+    half_block = 1 << (stage - 1)  # 2^stage
+
+    # Each thread processes one butterfly pair
+    butterfly_group = tid // half_block
+    pos_in_group = tid % half_block
+
+    # compute the index of two elements in butterfly pair
+    first_idx = butterfly_group * half_block * 2 + pos_in_group
+    second_idx = first_idx + half_block
+
+    if second_idx >= n:
+        return
+
+    # load
+    a_real = tl.load(real_ptr + first_idx)
+    a_imag = tl.load(imag_ptr + first_idx)
+    b_real = tl.load(real_ptr + second_idx)
+    b_imag = tl.load(imag_ptr + second_idx)
+
+    # calculate complex amplitude
+    angle = PI * pos_in_group / half_block
+    w_real = tl.cos(-angle)
+    w_imag = tl.sin(-angle)
+
+    tw_real = b_real * w_real - b_imag * w_imag
+    tw_imag = b_real * w_imag + b_imag * w_real
+
+    # butterfly
+    result_a_real = a_real + tw_real
+    result_a_imag = a_imag + tw_imag
+    result_b_real = a_real - tw_real
+    result_b_imag = a_imag - tw_imag
+
+    # store
+    tl.store(real_ptr + first_idx, result_a_real)
+    tl.store(imag_ptr + first_idx, result_a_imag)
+    tl.store(real_ptr + second_idx, result_b_real)
+    tl.store(imag_ptr + second_idx, result_b_imag)
+
+
+def fft_1d(x: torch.Tensor, output: torch.Tensor) -> torch.Tensor:
+    N = x.shape[0]
+    # make sure N is an integer power of 2
+    assert N > 0 and (N & (N - 1)) == 0
+
+    x_real = x.real.clone()
+    x_imag = x.imag.clone()
+
+    temp_real = torch.zeros_like(x_real)
+    temp_imag = torch.zeros_like(x_imag)
+    bit_reverse_kernel[(N,)](x_real, x_imag, temp_real, temp_imag, N)
+
+    x_real.copy_(temp_real)
+    x_imag.copy_(temp_imag)
+
+    log2n = N.bit_length() - 1
+    for stage in range(1, log2n + 1):
+        fft_stage_kernel[(N // 2,)](x_real, x_imag, N, stage)
+
+    output.real.copy_(x_real)
+    output.imag.copy_(x_imag)
+    return output

tests/test_special_ops.py

@@ -1456,3 +1456,18 @@ def test_accuracy_moe_align_block_size(
     gems_assert_close(sorted_ids, sorted_ids_vllm, dtype=dtype)
     gems_assert_close(expert_ids, expert_ids_vllm, dtype=dtype)
     gems_assert_close(num_tokens_post_pad, num_tokens_post_pad_vllm, dtype=dtype)
+
+
+@pytest.mark.fft_1d
+@pytest.mark.parametrize("N", [4, 8, 16, 32, 128, 256, 512])
+def test_fft_1d(N):
+    # FlagGems
+    input = torch.randn((N,), device="cuda") + torch.randn((N,), device="cuda") * 1j
+    output = torch.empty((N,), device="cuda") + torch.empty((N,), device="cuda") * 1j
+    with flag_gems.use_gems():
+        flag_gems.fft_1d(input, output)
+    # ref: torch
+    ref_x = to_reference(input)
+    ref_out = torch.fft.fft(ref_x)
+    dtype = torch.complex64
+    gems_assert_close(output, ref_out, dtype)