gpuarray 0.2.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

@@ -13,6 +13,7 @@ from .device import get_device
 from . import shaders
 
 WORKGROUP_SIZE = 256
+REDUCTION_ELEMENTS_PER_WORKGROUP = 1024  # 256 threads * 4 elements each
 
 # Global pipeline cache: shader_code -> (shader_module, compute_pipeline)
 _pipeline_cache: dict[str, tuple] = {}
@@ -174,7 +175,10 @@ class GPUArray:
         return GPUArray(out_buf, self._shape)
 
     def _run_reduction(self, shader_code: str) -> float:
-        """Multi-pass parallel reduction until 1 element remains."""
+        """Multi-pass parallel reduction until 1 element remains.
+
+        Each workgroup of 256 threads handles 1024 elements (4 per thread).
+        """
         device = get_device()
         _, pipeline = _get_pipeline(shader_code)
 
@@ -182,16 +186,16 @@ class GPUArray:
         current_size = self.size
 
         while current_size > 1:
-            n_workgroups = math.ceil(current_size / WORKGROUP_SIZE)
+            n_workgroups = math.ceil(current_size / REDUCTION_ELEMENTS_PER_WORKGROUP)
             out_buf = device.create_buffer(
-                size=n_workgroups * 4,
+                size=max(n_workgroups * 4, 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": current_buf, "offset": 0, "size": current_size * 4}},
-                    {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": n_workgroups * 4}},
+                    {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": max(n_workgroups * 4, 4)}},
                 ],
             )
             encoder = device.create_command_encoder()
@@ -214,6 +218,29 @@ class GPUArray:
             current_buf.destroy()
         return result
 
+    def _run_binary_op_inplace(self, other: "GPUArray", shader_code: str) -> "GPUArray":
+        """Run a binary op that writes the result into self's buffer (no allocation)."""
+        assert self.size == other.size, f"Size mismatch: {self.size} vs {other.size}"
+        device = get_device()
+        _, pipeline = _get_pipeline(shader_code)
+
+        bind_group = device.create_bind_group(
+            layout=pipeline.get_bind_group_layout(0),
+            entries=[
+                {"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
+                {"binding": 1, "resource": {"buffer": other._buffer, "offset": 0, "size": other.size * 4}},
+            ],
+        )
+        n_workgroups = math.ceil(self.size / 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 self
+
     # ----- dot product (two-input reduction) -----
 
     def _run_dot_reduction(self, other: "GPUArray") -> float:
@@ -259,11 +286,71 @@ class GPUArray:
         current_buf.destroy()
         return result
 
+    # ----- broadcast helper -----
+
+    def _run_broadcast_op(self, other: "GPUArray", shader_code: str, out_shape: tuple[int, ...]) -> "GPUArray":
+        """Run a broadcast binary shader with a uniform buffer for column count."""
+        device = get_device()
+        _, pipeline = _get_pipeline(shader_code)
+        out_size = 1
+        for d in out_shape:
+            out_size *= d
+        out_buf = device.create_buffer(
+            size=out_size * 4,
+            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
+        )
+        cols = out_shape[-1] if len(out_shape) >= 2 else out_shape[0]
+        dims_data = struct.pack("IIII", cols, 0, 0, 0)
+        dims_buf = device.create_buffer_with_data(
+            data=dims_data, usage=wgpu.BufferUsage.UNIFORM,
+        )
+        bind_group = device.create_bind_group(
+            layout=pipeline.get_bind_group_layout(0),
+            entries=[
+                {"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
+                {"binding": 1, "resource": {"buffer": other._buffer, "offset": 0, "size": other.size * 4}},
+                {"binding": 2, "resource": {"buffer": out_buf, "offset": 0, "size": out_size * 4}},
+                {"binding": 3, "resource": {"buffer": dims_buf, "offset": 0, "size": 16}},
+            ],
+        )
+        n_workgroups = math.ceil(out_size / 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()])
+        dims_buf.destroy()
+        return GPUArray(out_buf, out_shape)
+
+    def _binary_with_broadcast(self, other: "GPUArray",
+                               same_shader: str, row_shader: str, col_shader: str) -> "GPUArray":
+        """Dispatch to the right shader based on shape compatibility."""
+        if self.shape == other.shape or self.size == other.size:
+            return self._run_binary_op(other, same_shader)
+        # (M, N) op (N,) or (M, N) op (1, N) — broadcast row
+        if self.ndim == 2:
+            M, N = self.shape
+            if (other.ndim == 1 and other.shape[0] == N) or \
+               (other.ndim == 2 and other.shape == (1, N)):
+                b = other.flatten() if other.ndim == 2 else other
+                return self._run_broadcast_op(b, row_shader, self.shape)
+            # (M, N) op (M, 1) — broadcast col
+            if other.ndim == 2 and other.shape == (M, 1):
+                b_flat = other.flatten()
+                return self._run_broadcast_op(b_flat, col_shader, self.shape)
+        raise ValueError(f"Incompatible shapes for broadcast: {self.shape} vs {other.shape}")
+
     # ----- operator overloads -----
 
     def __add__(self, other):
         if isinstance(other, (int, float)):
             return self._run_scalar_op(float(other), shaders.ADD_SCALAR_SHADER)
+        if isinstance(other, GPUArray) and self.shape != other.shape:
+            return self._binary_with_broadcast(
+                other, shaders.ADD_SHADER,
+                shaders.BINARY_BROADCAST_ROW_ADD, shaders.BINARY_BROADCAST_COL_ADD)
         return self._run_binary_op(other, shaders.ADD_SHADER)
 
     def __radd__(self, other):
@@ -274,6 +361,10 @@ class GPUArray:
     def __sub__(self, other):
         if isinstance(other, (int, float)):
             return self._run_scalar_op(float(-other), shaders.ADD_SCALAR_SHADER)
+        if isinstance(other, GPUArray) and self.shape != other.shape:
+            return self._binary_with_broadcast(
+                other, shaders.SUB_SHADER,
+                shaders.BINARY_BROADCAST_ROW_SUB, shaders.BINARY_BROADCAST_COL_SUB)
         return self._run_binary_op(other, shaders.SUB_SHADER)
 
     def __rsub__(self, other):
@@ -285,6 +376,10 @@ class GPUArray:
     def __mul__(self, other):
         if isinstance(other, (int, float)):
             return self._run_scalar_op(float(other), shaders.MUL_SCALAR_SHADER)
+        if isinstance(other, GPUArray) and self.shape != other.shape:
+            return self._binary_with_broadcast(
+                other, shaders.MUL_SHADER,
+                shaders.BINARY_BROADCAST_ROW_MUL, shaders.BINARY_BROADCAST_COL_MUL)
         return self._run_binary_op(other, shaders.MUL_SHADER)
 
     def __rmul__(self, other):
@@ -295,8 +390,32 @@ class GPUArray:
     def __truediv__(self, other):
         if isinstance(other, (int, float)):
             return self._run_scalar_op(1.0 / float(other), shaders.MUL_SCALAR_SHADER)
+        if isinstance(other, GPUArray) and self.shape != other.shape:
+            return self._binary_with_broadcast(
+                other, shaders.DIV_SHADER,
+                shaders.BINARY_BROADCAST_ROW_DIV, shaders.BINARY_BROADCAST_COL_DIV)
         return self._run_binary_op(other, shaders.DIV_SHADER)
 
+    def __iadd__(self, other):
+        if isinstance(other, (int, float)):
+            return self._run_scalar_op(float(other), shaders.ADD_SCALAR_SHADER)
+        return self._run_binary_op_inplace(other, shaders.ADD_INPLACE_SHADER)
+
+    def __isub__(self, other):
+        if isinstance(other, (int, float)):
+            return self._run_scalar_op(float(-other), shaders.ADD_SCALAR_SHADER)
+        return self._run_binary_op_inplace(other, shaders.SUB_INPLACE_SHADER)
+
+    def __imul__(self, other):
+        if isinstance(other, (int, float)):
+            return self._run_scalar_op(float(other), shaders.MUL_SCALAR_SHADER)
+        return self._run_binary_op_inplace(other, shaders.MUL_INPLACE_SHADER)
+
+    def __itruediv__(self, other):
+        if isinstance(other, (int, float)):
+            return self._run_scalar_op(1.0 / float(other), shaders.MUL_SCALAR_SHADER)
+        return self._run_binary_op_inplace(other, shaders.DIV_INPLACE_SHADER)
+
     def __neg__(self):
         return self._run_unary_op(shaders.NEG_SHADER)
 
@@ -305,12 +424,108 @@ class GPUArray:
         return f"GPUArray({data}, shape={self._shape})"
 
     def __getitem__(self, key):
+        if isinstance(key, GPUArray):
+            return self._getitem_mask(key)
         data = self.to_numpy()
         result = data[key]
         if isinstance(result, np.ndarray):
-            return _from_numpy(result)
+            return _from_numpy(result.copy())
         return float(result)
 
+    def _getitem_mask(self, mask: "GPUArray") -> "GPUArray":
+        """Boolean mask indexing via gather shader."""
+        mask_np = mask.to_numpy().flatten()
+        indices = np.where(mask_np > 0.5)[0].astype(np.uint32)
+        if len(indices) == 0:
+            return _from_numpy(np.array([], dtype=np.float32))
+        device = get_device()
+        _, pipeline = _get_pipeline(shaders.GATHER_SHADER)
+        idx_buf = device.create_buffer_with_data(
+            data=indices.tobytes(),
+            usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC,
+        )
+        out_buf = device.create_buffer(
+            size=len(indices) * 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": self._buffer, "offset": 0, "size": self.size * 4}},
+                {"binding": 1, "resource": {"buffer": idx_buf, "offset": 0, "size": len(indices) * 4}},
+                {"binding": 2, "resource": {"buffer": out_buf, "offset": 0, "size": len(indices) * 4}},
+            ],
+        )
+        n_workgroups = math.ceil(len(indices) / 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()])
+        idx_buf.destroy()
+        return GPUArray(out_buf, (len(indices),))
+
+    def __setitem__(self, key, value):
+        if isinstance(key, GPUArray):
+            self._setitem_mask(key, value)
+            return
+        device = get_device()
+        if isinstance(key, (int, np.integer)):
+            idx = int(key)
+            if idx < 0:
+                idx += self.size
+            device.queue.write_buffer(self._buffer, idx * 4, struct.pack("f", float(value)))
+        elif isinstance(key, slice):
+            if self.ndim == 1:
+                indices = range(*key.indices(self._shape[0]))
+                if isinstance(value, (int, float)):
+                    data = struct.pack("f", float(value)) * len(indices)
+                elif isinstance(value, GPUArray):
+                    data = value.to_numpy().astype(np.float32).tobytes()
+                else:
+                    data = np.asarray(value, dtype=np.float32).tobytes()
+                start = indices.start if len(indices) > 0 else 0
+                device.queue.write_buffer(self._buffer, start * 4, data)
+            else:
+                indices = range(*key.indices(self._shape[0]))
+                cols = self._shape[1]
+                if isinstance(value, (int, float)):
+                    row_data = struct.pack("f", float(value)) * cols
+                    for r in indices:
+                        device.queue.write_buffer(self._buffer, r * cols * 4, row_data)
+                else:
+                    val_np = value.to_numpy() if isinstance(value, GPUArray) else np.asarray(value, dtype=np.float32)
+                    for i, r in enumerate(indices):
+                        row_bytes = (val_np[i] if val_np.ndim >= 2 else val_np).astype(np.float32).tobytes()
+                        device.queue.write_buffer(self._buffer, r * cols * 4, row_bytes)
+        elif isinstance(key, tuple) and len(key) == 2:
+            row, col = key
+            if isinstance(row, (int, np.integer)) and isinstance(col, (int, np.integer)):
+                cols = self._shape[1]
+                offset = (int(row) * cols + int(col)) * 4
+                device.queue.write_buffer(self._buffer, offset, struct.pack("f", float(value)))
+            else:
+                arr = self.to_numpy()
+                arr[key] = value if not isinstance(value, GPUArray) else value.to_numpy()
+                device.queue.write_buffer(self._buffer, 0, np.ascontiguousarray(arr, dtype=np.float32).tobytes())
+        else:
+            arr = self.to_numpy()
+            arr[key] = value if not isinstance(value, GPUArray) else value.to_numpy()
+            device.queue.write_buffer(self._buffer, 0, np.ascontiguousarray(arr, dtype=np.float32).tobytes())
+
+    def _setitem_mask(self, mask: "GPUArray", value):
+        """Masked assignment: a[mask] = value."""
+        data = self.to_numpy()
+        mask_np = mask.to_numpy() > 0.5
+        if isinstance(value, GPUArray):
+            data[mask_np] = value.to_numpy().flatten()[:mask_np.sum()]
+        else:
+            data[mask_np] = float(value)
+        device = get_device()
+        device.queue.write_buffer(self._buffer, 0, np.ascontiguousarray(data, dtype=np.float32).tobytes())
+
     # ----- reduction methods -----
 
     def sum(self) -> float:
@@ -443,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()
@@ -455,7 +888,10 @@ def _from_numpy(data: np.ndarray) -> GPUArray:
 
 
 def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
-    """Matrix multiplication C = A @ B."""
+    """Matrix multiplication C = A @ B.
+
+    Uses tiled shared-memory shader for matrices >= 64x64, naive for smaller.
+    """
     if a.ndim != 2 or b.ndim != 2:
         raise ValueError("matmul requires 2D arrays")
     M, K = a.shape
@@ -464,7 +900,11 @@ def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
         raise ValueError(f"Incompatible shapes for matmul: {a.shape} and {b.shape}")
 
     device = get_device()
-    _, pipeline = _get_pipeline(shaders.MATMUL_SHADER)
+
+    # Choose tiled or naive shader based on matrix size
+    use_tiled = M >= 64 and N >= 64 and K >= 64
+    shader_code = shaders.MATMUL_TILED if use_tiled else shaders.MATMUL_NAIVE
+    _, pipeline = _get_pipeline(shader_code)
 
     out_size = M * N
     out_buf = device.create_buffer(
@@ -489,13 +929,27 @@ def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
         ],
     )
 
-    n_workgroups = math.ceil(out_size / 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()
+    if use_tiled:
+        # Tiled shader uses @workgroup_size(16, 16); dispatch 2D grid
+        TILE_SIZE = 16
+        wg_x = math.ceil(N / TILE_SIZE)
+        wg_y = math.ceil(M / TILE_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(wg_x, wg_y)
+        pass_enc.end()
+    else:
+        # Naive shader uses @workgroup_size(256); dispatch 1D
+        n_workgroups = math.ceil(out_size / 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()])
     dims_buf.destroy()
     return GPUArray(out_buf, (M, N))

gpuarray/shaders.py

@@ -74,6 +74,62 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 """
 
+# ---------------------------------------------------------------------------
+# In-place elementwise binary operations (result written to a)
+# ---------------------------------------------------------------------------
+
+ADD_INPLACE_SHADER = """
+@group(0) @binding(0) var<storage, read_write> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        a[i] = a[i] + b[i];
+    }
+}
+"""
+
+SUB_INPLACE_SHADER = """
+@group(0) @binding(0) var<storage, read_write> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        a[i] = a[i] - b[i];
+    }
+}
+"""
+
+MUL_INPLACE_SHADER = """
+@group(0) @binding(0) var<storage, read_write> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        a[i] = a[i] * b[i];
+    }
+}
+"""
+
+DIV_INPLACE_SHADER = """
+@group(0) @binding(0) var<storage, read_write> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) {
+        a[i] = a[i] / b[i];
+    }
+}
+"""
+
 # ---------------------------------------------------------------------------
 # Elementwise unary operations
 # ---------------------------------------------------------------------------
@@ -236,17 +292,26 @@ var<workgroup> sdata: array<f32, 256>;
 fn main(
     @builtin(local_invocation_id) lid: vec3<u32>,
     @builtin(workgroup_id) wid: vec3<u32>,
-    @builtin(global_invocation_id) gid: vec3<u32>,
 ) {
     let local_idx = lid.x;
-    let global_idx = gid.x;
     let n = arrayLength(&input);
+    // Each thread loads and reduces 4 elements (workgroup covers 1024 elements)
+    let base = wid.x * 1024u + local_idx;
 
-    if (global_idx < n) {
-        sdata[local_idx] = input[global_idx];
-    } else {
-        sdata[local_idx] = 0.0;
+    var acc: f32 = 0.0;
+    if (base < n) {
+        acc = input[base];
     }
+    if (base + 256u < n) {
+        acc = acc + input[base + 256u];
+    }
+    if (base + 512u < n) {
+        acc = acc + input[base + 512u];
+    }
+    if (base + 768u < n) {
+        acc = acc + input[base + 768u];
+    }
+    sdata[local_idx] = acc;
     workgroupBarrier();
 
     // Tree reduction
@@ -276,17 +341,25 @@ var<workgroup> sdata: array<f32, 256>;
 fn main(
     @builtin(local_invocation_id) lid: vec3<u32>,
     @builtin(workgroup_id) wid: vec3<u32>,
-    @builtin(global_invocation_id) gid: vec3<u32>,
 ) {
     let local_idx = lid.x;
-    let global_idx = gid.x;
     let n = arrayLength(&input);
+    let base = wid.x * 1024u + local_idx;
 
-    if (global_idx < n) {
-        sdata[local_idx] = input[global_idx];
-    } else {
-        sdata[local_idx] = -3.402823e+38;  // -FLT_MAX
+    var acc: f32 = -3.402823e+38;
+    if (base < n) {
+        acc = input[base];
+    }
+    if (base + 256u < n) {
+        acc = max(acc, input[base + 256u]);
+    }
+    if (base + 512u < n) {
+        acc = max(acc, input[base + 512u]);
     }
+    if (base + 768u < n) {
+        acc = max(acc, input[base + 768u]);
+    }
+    sdata[local_idx] = acc;
     workgroupBarrier();
 
     var stride: u32 = 128u;
@@ -315,17 +388,25 @@ var<workgroup> sdata: array<f32, 256>;
 fn main(
     @builtin(local_invocation_id) lid: vec3<u32>,
     @builtin(workgroup_id) wid: vec3<u32>,
-    @builtin(global_invocation_id) gid: vec3<u32>,
 ) {
     let local_idx = lid.x;
-    let global_idx = gid.x;
     let n = arrayLength(&input);
+    let base = wid.x * 1024u + local_idx;
 
-    if (global_idx < n) {
-        sdata[local_idx] = input[global_idx];
-    } else {
-        sdata[local_idx] = 3.402823e+38;  // FLT_MAX
+    var acc: f32 = 3.402823e+38;
+    if (base < n) {
+        acc = input[base];
     }
+    if (base + 256u < n) {
+        acc = min(acc, input[base + 256u]);
+    }
+    if (base + 512u < n) {
+        acc = min(acc, input[base + 512u]);
+    }
+    if (base + 768u < n) {
+        acc = min(acc, input[base + 768u]);
+    }
+    sdata[local_idx] = acc;
     workgroupBarrier();
 
     var stride: u32 = 128u;
@@ -394,7 +475,7 @@ fn main(
 # Each thread computes one element of C.
 # ---------------------------------------------------------------------------
 
-MATMUL_SHADER = """
+MATMUL_NAIVE = """
 struct Dims {
     M: u32,
     K: u32,
@@ -425,6 +506,74 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 """
 
+# Keep backward compat alias
+MATMUL_SHADER = MATMUL_NAIVE
+
+# ---------------------------------------------------------------------------
+# Tiled matrix multiplication (16x16 tiles with shared memory)
+# ---------------------------------------------------------------------------
+
+MATMUL_TILED = """
+struct Dims {
+    M: u32,
+    K: u32,
+    N: u32,
+    _pad: u32,
+}
+
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> c: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+
+const TILE_SIZE = 16u;
+var<workgroup> tile_a: array<f32, 256>;
+var<workgroup> tile_b: array<f32, 256>;
+
+@compute @workgroup_size(16, 16)
+fn main(@builtin(local_invocation_id) lid: vec3<u32>,
+        @builtin(workgroup_id) wid: vec3<u32>) {
+    let row = wid.y * TILE_SIZE + lid.y;
+    let col = wid.x * TILE_SIZE + lid.x;
+    let local_idx = lid.y * TILE_SIZE + lid.x;
+
+    let num_tiles = (dims.K + TILE_SIZE - 1u) / TILE_SIZE;
+
+    var sum: f32 = 0.0;
+
+    for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
+        // Load tile of A into shared memory
+        let a_col = t * TILE_SIZE + lid.x;
+        if (row < dims.M && a_col < dims.K) {
+            tile_a[local_idx] = a[row * dims.K + a_col];
+        } else {
+            tile_a[local_idx] = 0.0;
+        }
+
+        // Load tile of B into shared memory
+        let b_row = t * TILE_SIZE + lid.y;
+        if (b_row < dims.K && col < dims.N) {
+            tile_b[local_idx] = b[b_row * dims.N + col];
+        } else {
+            tile_b[local_idx] = 0.0;
+        }
+
+        workgroupBarrier();
+
+        // Compute partial dot product for this tile
+        for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
+            sum = sum + tile_a[lid.y * TILE_SIZE + k] * tile_b[k * TILE_SIZE + lid.x];
+        }
+
+        workgroupBarrier();
+    }
+
+    if (row < dims.M && col < dims.N) {
+        c[row * dims.N + col] = sum;
+    }
+}
+"""
+
 # ---------------------------------------------------------------------------
 # Clamp (requires uniform buffer with min_val, max_val)
 # ---------------------------------------------------------------------------
@@ -598,3 +747,342 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
     }
 }
 """
+
+# ---------------------------------------------------------------------------
+# Gather: result[i] = source[indices[i]]
+# ---------------------------------------------------------------------------
+
+GATHER_SHADER = """
+@group(0) @binding(0) var<storage, read> source: array<f32>;
+@group(0) @binding(1) var<storage, read> indices: array<u32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&indices)) {
+        result[i] = source[indices[i]];
+    }
+}
+"""
+
+# ---------------------------------------------------------------------------
+# Broadcast binary ops: (M,N) op (N,) — broadcast 1D across rows
+# Uniform: dims.x = N (number of columns)
+# ---------------------------------------------------------------------------
+
+BINARY_BROADCAST_ROW_ADD = """
+struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) { result[i] = a[i] + b[i % dims.cols]; }
+}
+"""
+
+BINARY_BROADCAST_ROW_SUB = """
+struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) { result[i] = a[i] - b[i % dims.cols]; }
+}
+"""
+
+BINARY_BROADCAST_ROW_MUL = """
+struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) { result[i] = a[i] * b[i % dims.cols]; }
+}
+"""
+
+BINARY_BROADCAST_ROW_DIV = """
+struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) { result[i] = a[i] / b[i % dims.cols]; }
+}
+"""
+
+# ---------------------------------------------------------------------------
+# Broadcast binary ops: (M,N) op (M,1) — broadcast column vector
+# Uniform: dims.x = N (number of columns)
+# ---------------------------------------------------------------------------
+
+BINARY_BROADCAST_COL_ADD = """
+struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) { result[i] = a[i] + b[i / dims.cols]; }
+}
+"""
+
+BINARY_BROADCAST_COL_SUB = """
+struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) { result[i] = a[i] - b[i / dims.cols]; }
+}
+"""
+
+BINARY_BROADCAST_COL_MUL = """
+struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    if (i < arrayLength(&a)) { result[i] = a[i] * b[i / dims.cols]; }
+}
+"""
+
+BINARY_BROADCAST_COL_DIV = """
+struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
+@group(0) @binding(0) var<storage, read> a: array<f32>;
+@group(0) @binding(1) var<storage, read> b: array<f32>;
+@group(0) @binding(2) var<storage, read_write> result: array<f32>;
+@group(0) @binding(3) var<uniform> dims: Dims;
+@compute @workgroup_size(256)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let i = gid.x;
+    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.2.0.dist-info → gpuarray-0.4.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpuarray
-Version: 0.2.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.2.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=bvBlQ9YSTRUcYawpnHpWnD-zJFdeMV16mv_rNFxYBGQ,18934
-gpuarray/device.py,sha256=q80VOrKhy1u5i4AMLUKhWVBwFqQjoJ4tNBm234Wpf-A,452
-gpuarray/shaders.py,sha256=LhS3vEZc08Mx_OPnsv06bT9gRPThuHIxBuWO0V2C_SI,16301
-gpuarray-0.2.0.dist-info/METADATA,sha256=8ZL33ZQ7eVBnm4GuvjLxNVCiVl39u2maLCpGXbuI24w,3132
-gpuarray-0.2.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
-gpuarray-0.2.0.dist-info/top_level.txt,sha256=iruYli4LmPBqhaZxTVvzT1bT0VeN6dbC0GE29lT1CbI,9
-gpuarray-0.2.0.dist-info/RECORD,,