gpuarray 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

gpuarray/__init__.py

@@ -27,6 +27,10 @@ from .api import (
     concatenate,
     stack,
     random,
+    fft,
+    ifft,
+    fft_freq,
+    fft_magnitude,
 )
 
 __version__ = "0.1.0"
@@ -58,4 +62,8 @@ __all__ = [
     "concatenate",
     "stack",
     "random",
+    "fft",
+    "ifft",
+    "fft_freq",
+    "fft_magnitude",
 ]

gpuarray/api.py

@@ -6,7 +6,7 @@ from typing import Sequence
 
 import numpy as np
 
-from .core import GPUArray, _from_numpy, matmul as _matmul
+from .core import GPUArray, _from_numpy, matmul as _matmul, _complex_magnitude
 from . import shaders
 
 
@@ -180,3 +180,37 @@ class _Random:
 
 
 random = _Random()
+
+
+# ----- FFT functions -----
+
+def fft(a) -> GPUArray:
+    """1D FFT. Input is real or complex (interleaved). Returns complex GPUArray."""
+    if not isinstance(a, GPUArray):
+        a = array(a)
+    return a.fft()
+
+
+def ifft(a) -> GPUArray:
+    """1D inverse FFT. Input is complex interleaved. Returns complex GPUArray."""
+    if not isinstance(a, GPUArray):
+        a = array(a)
+    return a.ifft()
+
+
+def fft_freq(n: int, d: float = 1.0) -> GPUArray:
+    """Return the DFT sample frequencies (like numpy.fft.fftfreq)."""
+    results = np.empty(n, dtype=np.float32)
+    N = (n - 1) // 2 + 1
+    results[:N] = np.arange(0, N, dtype=np.float32)
+    results[N:] = np.arange(-(n // 2), 0, dtype=np.float32)
+    results /= (n * d)
+    return _from_numpy(results)
+
+
+def fft_magnitude(a) -> GPUArray:
+    """Compute magnitude spectrum |FFT(x)|. Returns real GPUArray."""
+    if not isinstance(a, GPUArray):
+        a = array(a)
+    fft_result = a.fft()
+    return _complex_magnitude(fft_result)

gpuarray/core.py

@@ -658,6 +658,224 @@ class GPUArray:
         return self._shape[0]
 
 
+    # ----- FFT methods -----
+
+    def fft(self) -> "GPUArray":
+        """1D FFT. Input is real or complex (interleaved). Returns complex GPUArray (interleaved)."""
+        return _fft_impl(self, inverse=False)
+
+    def ifft(self) -> "GPUArray":
+        """1D inverse FFT. Input is complex interleaved. Returns complex GPUArray."""
+        return _fft_impl(self, inverse=True)
+
+    def fft_magnitude(self) -> "GPUArray":
+        """Compute magnitude spectrum |FFT(x)|. Returns real GPUArray."""
+        fft_result = self.fft()
+        return _complex_magnitude(fft_result)
+
+    def fft_phase(self) -> "GPUArray":
+        """Compute phase spectrum angle(FFT(x)). Returns real GPUArray."""
+        fft_result = self.fft()
+        return _complex_phase(fft_result)
+
+
+def _next_power_of_2(n: int) -> int:
+    """Return the smallest power of 2 >= n."""
+    if n <= 1:
+        return 1
+    p = 1
+    while p < n:
+        p <<= 1
+    return p
+
+
+def _fft_impl(arr: GPUArray, inverse: bool = False) -> GPUArray:
+    """Cooley-Tukey radix-2 FFT/IFFT on GPU."""
+    device = get_device()
+
+    # Determine if input is real or complex (interleaved)
+    # Real: 1D array of N floats -> N-point FFT
+    # Complex interleaved: 1D array of 2*N floats -> N-point FFT
+    # We treat it as real if ndim==1 and size is the number of points
+    # For simplicity: if the array was produced by a previous FFT (even size),
+    # the caller should know. We'll check: if doing IFFT, input must be complex.
+
+    flat = arr.to_numpy().flatten()
+
+    if inverse:
+        # Input must be complex interleaved (even number of floats)
+        assert len(flat) % 2 == 0, "IFFT input must be complex interleaved (even length)"
+        N_complex = len(flat) // 2
+    else:
+        # Assume real input -> convert to complex
+        N_complex = len(flat)
+
+    # Pad to next power of 2
+    N = _next_power_of_2(N_complex)
+    log2N = int(math.log2(N))
+
+    if inverse:
+        # Pad complex data
+        if N > N_complex:
+            padded = np.zeros(2 * N, dtype=np.float32)
+            padded[:len(flat)] = flat
+        else:
+            padded = flat.copy()
+        complex_data = padded
+    else:
+        # Convert real to complex interleaved
+        padded_real = np.zeros(N, dtype=np.float32)
+        padded_real[:len(flat)] = flat
+        complex_data = np.zeros(2 * N, dtype=np.float32)
+        complex_data[0::2] = padded_real  # real parts
+        # imaginary parts are already 0
+
+    # Upload complex data to GPU
+    input_buf = device.create_buffer_with_data(
+        data=complex_data.tobytes(),
+        usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
+    )
+
+    # Step 1: Bit-reversal permutation
+    output_buf = device.create_buffer(
+        size=2 * N * 4,
+        usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
+    )
+
+    _, br_pipeline = _get_pipeline(shaders.FFT_BIT_REVERSE_SHADER)
+    params_data = struct.pack("IIII", N, log2N, 0, 0)
+    params_buf = device.create_buffer_with_data(
+        data=params_data, usage=wgpu.BufferUsage.UNIFORM,
+    )
+    bind_group = device.create_bind_group(
+        layout=br_pipeline.get_bind_group_layout(0),
+        entries=[
+            {"binding": 0, "resource": {"buffer": input_buf, "offset": 0, "size": 2 * N * 4}},
+            {"binding": 1, "resource": {"buffer": output_buf, "offset": 0, "size": 2 * N * 4}},
+            {"binding": 2, "resource": {"buffer": params_buf, "offset": 0, "size": 16}},
+        ],
+    )
+    n_workgroups = math.ceil(N / WORKGROUP_SIZE)
+    encoder = device.create_command_encoder()
+    pass_enc = encoder.begin_compute_pass()
+    pass_enc.set_pipeline(br_pipeline)
+    pass_enc.set_bind_group(0, bind_group)
+    pass_enc.dispatch_workgroups(n_workgroups)
+    pass_enc.end()
+    device.queue.submit([encoder.finish()])
+    input_buf.destroy()
+    params_buf.destroy()
+
+    # Step 2: Butterfly stages
+    butterfly_shader = shaders.IFFT_BUTTERFLY_SHADER if inverse else shaders.FFT_BUTTERFLY_SHADER
+    _, bf_pipeline = _get_pipeline(butterfly_shader)
+
+    for stage in range(log2N):
+        params_data = struct.pack("IIII", N, stage, 0, 0)
+        params_buf = device.create_buffer_with_data(
+            data=params_data, usage=wgpu.BufferUsage.UNIFORM,
+        )
+        bind_group = device.create_bind_group(
+            layout=bf_pipeline.get_bind_group_layout(0),
+            entries=[
+                {"binding": 0, "resource": {"buffer": output_buf, "offset": 0, "size": 2 * N * 4}},
+                {"binding": 1, "resource": {"buffer": params_buf, "offset": 0, "size": 16}},
+            ],
+        )
+        n_workgroups_bf = math.ceil(N // 2 / WORKGROUP_SIZE)
+        if n_workgroups_bf == 0:
+            n_workgroups_bf = 1
+        encoder = device.create_command_encoder()
+        pass_enc = encoder.begin_compute_pass()
+        pass_enc.set_pipeline(bf_pipeline)
+        pass_enc.set_bind_group(0, bind_group)
+        pass_enc.dispatch_workgroups(n_workgroups_bf)
+        pass_enc.end()
+        device.queue.submit([encoder.finish()])
+        params_buf.destroy()
+
+    # Step 3: For IFFT, scale by 1/N
+    if inverse:
+        _, scale_pipeline = _get_pipeline(shaders.FFT_SCALE_SHADER)
+        scale_val = 1.0 / float(N)
+        params_data = struct.pack("f", scale_val) + b"\x00" * 12
+        params_buf = device.create_buffer_with_data(
+            data=params_data, usage=wgpu.BufferUsage.UNIFORM,
+        )
+        bind_group = device.create_bind_group(
+            layout=scale_pipeline.get_bind_group_layout(0),
+            entries=[
+                {"binding": 0, "resource": {"buffer": output_buf, "offset": 0, "size": 2 * N * 4}},
+                {"binding": 1, "resource": {"buffer": params_buf, "offset": 0, "size": 16}},
+            ],
+        )
+        n_workgroups_sc = math.ceil(2 * N / WORKGROUP_SIZE)
+        encoder = device.create_command_encoder()
+        pass_enc = encoder.begin_compute_pass()
+        pass_enc.set_pipeline(scale_pipeline)
+        pass_enc.set_bind_group(0, bind_group)
+        pass_enc.dispatch_workgroups(n_workgroups_sc)
+        pass_enc.end()
+        device.queue.submit([encoder.finish()])
+        params_buf.destroy()
+
+    return GPUArray(output_buf, (2 * N,))
+
+
+def _complex_magnitude(arr: GPUArray) -> GPUArray:
+    """Compute magnitude of complex interleaved array."""
+    device = get_device()
+    N = arr.size // 2
+    _, pipeline = _get_pipeline(shaders.COMPLEX_MAGNITUDE_SHADER)
+    out_buf = device.create_buffer(
+        size=N * 4,
+        usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
+    )
+    bind_group = device.create_bind_group(
+        layout=pipeline.get_bind_group_layout(0),
+        entries=[
+            {"binding": 0, "resource": {"buffer": arr._buffer, "offset": 0, "size": arr.size * 4}},
+            {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": N * 4}},
+        ],
+    )
+    n_workgroups = math.ceil(N / WORKGROUP_SIZE)
+    encoder = device.create_command_encoder()
+    pass_enc = encoder.begin_compute_pass()
+    pass_enc.set_pipeline(pipeline)
+    pass_enc.set_bind_group(0, bind_group)
+    pass_enc.dispatch_workgroups(n_workgroups)
+    pass_enc.end()
+    device.queue.submit([encoder.finish()])
+    return GPUArray(out_buf, (N,))
+
+
+def _complex_phase(arr: GPUArray) -> GPUArray:
+    """Compute phase of complex interleaved array."""
+    device = get_device()
+    N = arr.size // 2
+    _, pipeline = _get_pipeline(shaders.COMPLEX_PHASE_SHADER)
+    out_buf = device.create_buffer(
+        size=N * 4,
+        usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
+    )
+    bind_group = device.create_bind_group(
+        layout=pipeline.get_bind_group_layout(0),
+        entries=[
+            {"binding": 0, "resource": {"buffer": arr._buffer, "offset": 0, "size": arr.size * 4}},
+            {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": N * 4}},
+        ],
+    )
+    n_workgroups = math.ceil(N / WORKGROUP_SIZE)
+    encoder = device.create_command_encoder()
+    pass_enc = encoder.begin_compute_pass()
+    pass_enc.set_pipeline(pipeline)
+    pass_enc.set_bind_group(0, bind_group)
+    pass_enc.dispatch_workgroups(n_workgroups)
+    pass_enc.end()
+    device.queue.submit([encoder.finish()])
+    return GPUArray(out_buf, (N,))
+
+
 def _from_numpy(data: np.ndarray) -> GPUArray:
     """Create a GPUArray from a numpy array."""
     device = get_device()

gpuarray/shaders.py

@@ -879,3 +879,210 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     if (i < arrayLength(&a)) { result[i] = a[i] / b[i / dims.cols]; }
 }
 """
+
+# ---------------------------------------------------------------------------
+# FFT shaders (Cooley-Tukey radix-2)
+# Complex numbers stored as interleaved float32 pairs: [re0, im0, re1, im1, ...]
+# ---------------------------------------------------------------------------
+
+REAL_TO_COMPLEX_SHADER = """
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    let n = arrayLength(&input);
+    if (i < n) {
+        output[i * 2u] = input[i];
+        output[i * 2u + 1u] = 0.0;
+    }
+}
+"""
+
+FFT_BIT_REVERSE_SHADER = """
+struct Params {
+    N: u32,
+    log2N: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+@group(0) @binding(2) var<uniform> params: Params;
+
+fn bit_reverse(x: u32, bits: u32) -> u32 {
+    var v = x;
+    var r: u32 = 0u;
+    for (var i: u32 = 0u; i < bits; i = i + 1u) {
+        r = (r << 1u) | (v & 1u);
+        v = v >> 1u;
+    }
+    return r;
+}
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i >= params.N) {
+        return;
+    }
+    let j = bit_reverse(i, params.log2N);
+    output[j * 2u] = input[i * 2u];
+    output[j * 2u + 1u] = input[i * 2u + 1u];
+}
+"""
+
+FFT_BUTTERFLY_SHADER = """
+struct Params {
+    N: u32,
+    stage: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<storage, read_write> data: array<f32>;
+@group(0) @binding(1) var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    let m = 1u << (params.stage + 1u);
+    let half = m >> 1u;
+
+    // Total number of butterflies is N/2
+    if (i >= params.N / 2u) {
+        return;
+    }
+
+    // Which butterfly group and position within group
+    let group = i / half;
+    let pos = i % half;
+    let j = group * m + pos;
+    let k = j + half;
+
+    let pi = 3.14159265358979323846;
+    let angle = -2.0 * pi * f32(pos) / f32(m);
+    let tw_re = cos(angle);
+    let tw_im = sin(angle);
+
+    // Load X[k]
+    let xk_re = data[k * 2u];
+    let xk_im = data[k * 2u + 1u];
+
+    // Twiddle multiply: W * X[k]
+    let t_re = tw_re * xk_re - tw_im * xk_im;
+    let t_im = tw_re * xk_im + tw_im * xk_re;
+
+    // Load X[j]
+    let xj_re = data[j * 2u];
+    let xj_im = data[j * 2u + 1u];
+
+    // Butterfly
+    data[j * 2u] = xj_re + t_re;
+    data[j * 2u + 1u] = xj_im + t_im;
+    data[k * 2u] = xj_re - t_re;
+    data[k * 2u + 1u] = xj_im - t_im;
+}
+"""
+
+IFFT_BUTTERFLY_SHADER = """
+struct Params {
+    N: u32,
+    stage: u32,
+    _pad1: u32,
+    _pad2: u32,
+}
+
+@group(0) @binding(0) var<storage, read_write> data: array<f32>;
+@group(0) @binding(1) var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    let m = 1u << (params.stage + 1u);
+    let half = m >> 1u;
+
+    if (i >= params.N / 2u) {
+        return;
+    }
+
+    let group = i / half;
+    let pos = i % half;
+    let j = group * m + pos;
+    let k = j + half;
+
+    let pi = 3.14159265358979323846;
+    // Conjugate twiddle for IFFT: positive angle
+    let angle = 2.0 * pi * f32(pos) / f32(m);
+    let tw_re = cos(angle);
+    let tw_im = sin(angle);
+
+    let xk_re = data[k * 2u];
+    let xk_im = data[k * 2u + 1u];
+
+    let t_re = tw_re * xk_re - tw_im * xk_im;
+    let t_im = tw_re * xk_im + tw_im * xk_re;
+
+    let xj_re = data[j * 2u];
+    let xj_im = data[j * 2u + 1u];
+
+    data[j * 2u] = xj_re + t_re;
+    data[j * 2u + 1u] = xj_im + t_im;
+    data[k * 2u] = xj_re - t_re;
+    data[k * 2u + 1u] = xj_im - t_im;
+}
+"""
+
+FFT_SCALE_SHADER = """
+struct Params {
+    scalar: f32,
+    _pad1: u32,
+    _pad2: u32,
+    _pad3: u32,
+}
+
+@group(0) @binding(0) var<storage, read_write> data: array<f32>;
+@group(0) @binding(1) var<uniform> params: Params;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&data)) {
+        data[i] = data[i] * params.scalar;
+    }
+}
+"""
+
+COMPLEX_MAGNITUDE_SHADER = """
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    let n = arrayLength(&output);
+    if (i < n) {
+        let re = input[i * 2u];
+        let im = input[i * 2u + 1u];
+        output[i] = sqrt(re * re + im * im);
+    }
+}
+"""
+
+COMPLEX_PHASE_SHADER = """
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    let n = arrayLength(&output);
+    if (i < n) {
+        let re = input[i * 2u];
+        let im = input[i * 2u + 1u];
+        output[i] = atan2(im, re);
+    }
+}
+"""

{gpuarray-0.3.0.dist-info → gpuarray-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpuarray
-Version: 0.3.0
+Version: 0.4.0
 Summary: NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required.
 Author: Arnav Kewalram
 License: MIT

gpuarray-0.4.0.dist-info/RECORD

@@ -0,0 +1,9 @@
+gpuarray/__init__.py,sha256=YmH2uwL9iez8AUGDE2OCY6sVNIrC49FNsUP5PFxrNmk,864
+gpuarray/api.py,sha256=wiZnz4rs631nGiR_wW01Rr1BNgT0RL8fD2rMRQ7SSpg,5883
+gpuarray/core.py,sha256=D2Q5Lwe8p2PmQjaUqMYEWd2qDSQkI1xt5IfTZ3MyIFY,39048
+gpuarray/device.py,sha256=q80VOrKhy1u5i4AMLUKhWVBwFqQjoJ4tNBm234Wpf-A,452
+gpuarray/shaders.py,sha256=0-71ErmWkBCHqPoROsORiVRgl48c4BsoqAjSLOn2QKg,30566
+gpuarray-0.4.0.dist-info/METADATA,sha256=GRL-qds44XKfiHArZsRivg-WEVWondmQtOEXA0tgofI,3132
+gpuarray-0.4.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
+gpuarray-0.4.0.dist-info/top_level.txt,sha256=iruYli4LmPBqhaZxTVvzT1bT0VeN6dbC0GE29lT1CbI,9
+gpuarray-0.4.0.dist-info/RECORD,,

gpuarray-0.3.0.dist-info/RECORD

@@ -1,9 +0,0 @@
-gpuarray/__init__.py,sha256=MB9k9oXPvA6wRCCZii72HtpRxjXPtkiT2RaJsiVY-wY,752
-gpuarray/api.py,sha256=PlYqtvY7jEV3xIneGdRoqXJjVOP3w2fnHyZgVP6fqis,4874
-gpuarray/core.py,sha256=o7UymuEJoyYQAKao4O4wRtFw911Ut3h7tnQJOT6N0BU,30676
-gpuarray/device.py,sha256=q80VOrKhy1u5i4AMLUKhWVBwFqQjoJ4tNBm234Wpf-A,452
-gpuarray/shaders.py,sha256=oG7vd3iqO1eJ7UMIZx-46UPcb9gNVUWhBa6q8iKOSxo,25467
-gpuarray-0.3.0.dist-info/METADATA,sha256=XrCwPw9lCTIHb5qMSZjHypALq5BLwwHjrV2SVoLDnhE,3132
-gpuarray-0.3.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
-gpuarray-0.3.0.dist-info/top_level.txt,sha256=iruYli4LmPBqhaZxTVvzT1bT0VeN6dbC0GE29lT1CbI,9
-gpuarray-0.3.0.dist-info/RECORD,,