gpuarray 0.3.0__tar.gz → 0.4.0__tar.gz

{gpuarray-0.3.0 → gpuarray-0.4.0}/PKG-INFO

@@ -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.3.0 → gpuarray-0.4.0}/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-0.3.0 → gpuarray-0.4.0}/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-0.3.0 → gpuarray-0.4.0}/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-0.3.0 → gpuarray-0.4.0}/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 → gpuarray-0.4.0}/gpuarray.egg-info/PKG-INFO

@@ -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.3.0 → gpuarray-0.4.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gpuarray"
-version = "0.3.0"
+version = "0.4.0"
 description = "NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required."
 requires-python = ">=3.10,<3.14"
 license = {text = "MIT"}

{gpuarray-0.3.0 → gpuarray-0.4.0}/tests/test_core.py

@@ -557,3 +557,64 @@ class TestBroadcast:
         result = A + b
         expected = A_np + b_np
         npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
+
+
+class TestFFT:
+    def test_fft_simple(self):
+        """FFT of [1, 0, 0, 0] should be [1, 1, 1, 1] (all real)."""
+        signal = np.array([1.0, 0.0, 0.0, 0.0], dtype=np.float32)
+        gpu_result = gp.fft(gp.array(signal)).to_numpy()
+        np_result = np.fft.fft(signal)
+        # gpu_result is interleaved [re0, im0, re1, im1, ...]
+        gpu_re = gpu_result[0::2]
+        gpu_im = gpu_result[1::2]
+        npt.assert_allclose(gpu_re, np_result.real, rtol=1e-3, atol=1e-5)
+        npt.assert_allclose(gpu_im, np_result.imag, atol=1e-5)
+
+    def test_fft_sine(self):
+        """FFT of a sine wave should have peak at the frequency."""
+        N = 256
+        freq = 10
+        t = np.linspace(0, 1, N, endpoint=False, dtype=np.float32)
+        signal = np.sin(2 * np.pi * freq * t).astype(np.float32)
+
+        gpu_mag = gp.fft_magnitude(gp.array(signal)).to_numpy()
+        np_mag = np.abs(np.fft.fft(signal)).astype(np.float32)
+        npt.assert_allclose(gpu_mag, np_mag, rtol=1e-3, atol=1e-2)
+
+    def test_ifft_roundtrip(self):
+        """ifft(fft(x)) should equal x."""
+        signal = np.array([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], dtype=np.float32)
+        gpu_arr = gp.array(signal)
+        fft_result = gp.fft(gpu_arr)
+        ifft_result = gp.ifft(fft_result)
+        # Extract real parts from interleaved complex result
+        recovered = ifft_result.to_numpy()
+        recovered_re = recovered[0::2]
+        # Should match original (padded to power-of-2 if needed, but 8 is already power of 2)
+        npt.assert_allclose(recovered_re[:len(signal)], signal, rtol=1e-3, atol=1e-5)
+
+    def test_fft_power_of_2(self):
+        """Test with exact power of 2 input."""
+        N = 16
+        signal = np.random.randn(N).astype(np.float32)
+        gpu_result = gp.fft(gp.array(signal)).to_numpy()
+        np_result = np.fft.fft(signal)
+        gpu_re = gpu_result[0::2]
+        gpu_im = gpu_result[1::2]
+        npt.assert_allclose(gpu_re, np_result.real.astype(np.float32), rtol=1e-3, atol=1e-4)
+        npt.assert_allclose(gpu_im, np_result.imag.astype(np.float32), rtol=1e-3, atol=1e-4)
+
+    def test_fft_freq(self):
+        """Test frequency bin generation matches numpy."""
+        for n in [8, 16, 32]:
+            gpu_freqs = gp.fft_freq(n).to_numpy()
+            np_freqs = np.fft.fftfreq(n).astype(np.float32)
+            npt.assert_allclose(gpu_freqs, np_freqs, rtol=1e-5)
+
+    def test_fft_magnitude_api(self):
+        """Test fft_magnitude module-level function."""
+        signal = np.array([1.0, 0.0, 1.0, 0.0], dtype=np.float32)
+        gpu_mag = gp.fft_magnitude(gp.array(signal)).to_numpy()
+        np_mag = np.abs(np.fft.fft(signal)).astype(np.float32)
+        npt.assert_allclose(gpu_mag, np_mag, rtol=1e-3, atol=1e-5)