ggml: init fft operations

Signed-off-by: Ilham Syahid S <ilhamsyahids@gmail.com>

include/ggml.h

@@ -521,6 +521,9 @@ extern "C" {
         GGML_OP_OPT_STEP_ADAMW,
 
         GGML_OP_COUNT,
+
+        GGML_OP_FFT,       // Fast Fourier Transform
+        GGML_OP_IFFT,      // Inverse Fast Fourier Transform
     };
 
     enum ggml_unary_op {
@@ -1775,6 +1778,15 @@ extern "C" {
             struct ggml_tensor  * a,
             int                   k);
 
+    // Fast Fourier Transform operations
+    GGML_API struct ggml_tensor * ggml_fft(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
+    GGML_API struct ggml_tensor * ggml_ifft(
+            struct ggml_context * ctx,
+            struct ggml_tensor  * a);
+
 #define GGML_KQ_MASK_PAD 32
 
     // q:    [n_embd, n_batch,     n_head,    1]

src/ggml-cpu/ggml-cpu.c

@@ -23,6 +23,7 @@
 #include <errno.h>
 #include <time.h>
 #include <math.h>
+#include <complex.h>
 #include <stdlib.h>
 #include <string.h>
 #include <stdint.h>
@@ -12716,6 +12717,106 @@ static void ggml_compute_forward_opt_step_adamw(
             }
     }
 }
+
+// FFT
+
+static bool is_power_of_2(int64_t n) {
+    return (n & (n - 1)) == 0;
+}
+
+// Recursive implementation of FFT
+static void fft_recursive(float complex *x, int64_t n, int64_t stride, bool inverse) {
+    if (n <= 1) return;
+
+    // Split into even and odd
+    int64_t half = n / 2;
+    float complex *even = (float complex *)malloc(half * sizeof(float complex));
+    float complex *odd = (float complex *)malloc(half * sizeof(float complex));
+
+    for (int64_t i = 0; i < half; i++) {
+        even[i] = x[2*i*stride];
+        odd[i] = x[(2*i+1)*stride];
+    }
+
+    // Recursive FFT on even and odd parts
+    fft_recursive(even, half, 1, inverse);
+    fft_recursive(odd, half, 1, inverse);
+
+    // Combine results
+    float angle_factor = inverse ? 2.0 * M_PI / n : -2.0 * M_PI / n;
+    for (int64_t k = 0; k < half; k++) {
+        float complex t = cexp(angle_factor * k * I) * odd[k];
+        x[k*stride] = even[k] + t;
+        x[(k+half)*stride] = even[k] - t;
+    }
+
+    // Scale if inverse
+    if (inverse) {
+        for (int64_t k = 0; k < n; k++) {
+            x[k*stride] /= 2.0;
+        }
+    }
+
+    free(even);
+    free(odd);
+}
+
+static void ggml_compute_forward_fft(struct ggml_tensor * dst, const struct ggml_tensor * src) {
+    GGML_ASSERT(src->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_vector(src));  // Only support 1D FFT for now
+    GGML_ASSERT(is_power_of_2(src->ne[0]));  // Length must be power of 2
+
+    // Allocate temporary complex array
+    int64_t n = src->ne[0];
+    float complex *x = (float complex *)malloc(n * sizeof(float complex));
+
+    // Copy input data to complex array
+    float *src_data = (float *)src->data;
+    for (int64_t i = 0; i < n; i++) {
+        x[i] = src_data[i] + 0.0*I;
+    }
+
+    // Perform FFT
+    fft_recursive(x, n, 1, false);
+
+    // Copy result back
+    float *dst_data = (float *)dst->data;
+    for (int64_t i = 0; i < n; i++) {
+        // Store real and imaginary parts in interleaved format
+        dst_data[2*i] = crealf(x[i]);
+        dst_data[2*i+1] = cimagf(x[i]);
+    }
+
+    free(x);
+}
+
+static void ggml_compute_forward_ifft(struct ggml_tensor * dst, const struct ggml_tensor * src) {
+    GGML_ASSERT(src->type == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_vector(src));  // Only support 1D IFFT for now
+    GGML_ASSERT(is_power_of_2(src->ne[0]/2));  // Complex length must be power of 2
+
+    // Allocate temporary complex array
+    int64_t n = src->ne[0]/2;  // Complex length is half of the real array length
+    float complex *x = (float complex *)malloc(n * sizeof(float complex));
+
+    // Copy input data to complex array
+    float *src_data = (float *)src->data;
+    for (int64_t i = 0; i < n; i++) {
+        x[i] = src_data[2*i] + src_data[2*i+1]*I;
+    }
+
+    // Perform IFFT
+    fft_recursive(x, n, 1, true);
+
+    // Copy result back (real part only for IFFT)
+    float *dst_data = (float *)dst->data;
+    for (int64_t i = 0; i < n; i++) {
+        dst_data[i] = crealf(x[i]);
+    }
+
+    free(x);
+}
+
 /////////////////////////////////
 
 static void ggml_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * tensor) {
@@ -13081,6 +13182,16 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
                 ggml_compute_forward_opt_step_adamw(params, tensor);
             }
             break;
+        case GGML_OP_FFT:
+            {
+                ggml_compute_forward_fft(tensor, tensor->src[0]);
+            }
+            break;
+        case GGML_OP_IFFT:
+            {
+                ggml_compute_forward_ifft(tensor, tensor->src[0]);
+            }
+            break;
         case GGML_OP_NONE:
             {
                 // nop

src/ggml.c

@@ -5130,6 +5130,65 @@ struct ggml_tensor * ggml_opt_step_adamw(
     return result;
 }
 
+// ggml_fft
+
+static struct ggml_tensor * ggml_fft_impl(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a,
+        bool                  inplace) {
+    GGML_ASSERT(ggml_is_vector(a));
+    bool is_node = false;
+
+    // Instead of checking a->grad directly, we should check if the tensor
+    // is part of a computation graph and has gradients enabled
+    if (!inplace && (a->flags & GGML_TENSOR_FLAG_PARAM)) {
+        is_node = true;
+    }
+
+    struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a);
+
+    // Set the operation type directly
+    result->op = GGML_OP_FFT;
+    // Initialize op params if needed
+    ggml_set_op_params(result, NULL, 0);
+
+    // Instead of setting grad directly, we should mark the tensor as requiring gradients
+    if (is_node) {
+        result->flags |= GGML_TENSOR_FLAG_PARAM;
+    }
+    result->src[0] = a;
+
+    return result;
+}
+
+struct ggml_tensor * ggml_fft(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a) {
+    return ggml_fft_impl(ctx, a, false);
+}
+
+// ggml_ifft
+
+struct ggml_tensor * ggml_ifft(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * a) {
+    GGML_ASSERT(ggml_is_vector(a));
+    struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, a->ne[0]/2);
+
+    // Set the operation type directly
+    result->op = GGML_OP_IFFT;
+    // Initialize op params if needed
+    ggml_set_op_params(result, NULL, 0);
+
+    // Instead of checking a->grad directly, check if input tensor requires gradients
+    if (a->flags & GGML_TENSOR_FLAG_PARAM) {
+        // Mark the result tensor as requiring gradients
+        result->flags |= GGML_TENSOR_FLAG_PARAM;
+    }
+    result->src[0] = a;
+    return result;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 struct ggml_hash_set ggml_hash_set_new(size_t size) {