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

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:
@@ -455,7 +670,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 +682,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 +711,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,135 @@ 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]; }
+}
+"""

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpuarray
-Version: 0.2.0
+Version: 0.3.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.dist-info/RECORD

@@ -0,0 +1,9 @@
+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,,

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,,