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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
gpuarray/__init__.py CHANGED
@@ -20,6 +20,13 @@ from .api import (
20
20
  relu,
21
21
  sigmoid,
22
22
  tanh,
23
+ where,
24
+ clip,
25
+ argmax,
26
+ argmin,
27
+ concatenate,
28
+ stack,
29
+ random,
23
30
  )
24
31
 
25
32
  __version__ = "0.1.0"
@@ -44,4 +51,11 @@ __all__ = [
44
51
  "relu",
45
52
  "sigmoid",
46
53
  "tanh",
54
+ "where",
55
+ "clip",
56
+ "argmax",
57
+ "argmin",
58
+ "concatenate",
59
+ "stack",
60
+ "random",
47
61
  ]
gpuarray/api.py CHANGED
@@ -97,3 +97,86 @@ def sigmoid(a: GPUArray) -> GPUArray:
97
97
 
98
98
  def tanh(a: GPUArray) -> GPUArray:
99
99
  return a._run_unary_op(shaders.TANH_SHADER)
100
+
101
+
102
+ # ----- new operations -----
103
+
104
+ def where(condition: GPUArray, x: GPUArray, y: GPUArray) -> GPUArray:
105
+ """Conditional select: where condition is 1.0, pick x; else pick y."""
106
+ import wgpu
107
+ from .device import get_device
108
+ from .core import _get_pipeline
109
+
110
+ assert condition.size == x.size == y.size
111
+ device = get_device()
112
+ _, pipeline = _get_pipeline(shaders.WHERE_SHADER)
113
+
114
+ out_buf = device.create_buffer(
115
+ size=condition.size * 4,
116
+ usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
117
+ )
118
+ bind_group = device.create_bind_group(
119
+ layout=pipeline.get_bind_group_layout(0),
120
+ entries=[
121
+ {"binding": 0, "resource": {"buffer": condition._buffer, "offset": 0, "size": condition.size * 4}},
122
+ {"binding": 1, "resource": {"buffer": x._buffer, "offset": 0, "size": x.size * 4}},
123
+ {"binding": 2, "resource": {"buffer": y._buffer, "offset": 0, "size": y.size * 4}},
124
+ {"binding": 3, "resource": {"buffer": out_buf, "offset": 0, "size": condition.size * 4}},
125
+ ],
126
+ )
127
+ import math
128
+ n_workgroups = math.ceil(condition.size / 256)
129
+ encoder = device.create_command_encoder()
130
+ pass_enc = encoder.begin_compute_pass()
131
+ pass_enc.set_pipeline(pipeline)
132
+ pass_enc.set_bind_group(0, bind_group)
133
+ pass_enc.dispatch_workgroups(n_workgroups)
134
+ pass_enc.end()
135
+ device.queue.submit([encoder.finish()])
136
+ return GPUArray(out_buf, condition._shape)
137
+
138
+
139
+ def clip(a: GPUArray, min_val: float, max_val: float) -> GPUArray:
140
+ """Clamp values to [min_val, max_val]."""
141
+ return a.clip(min_val, max_val)
142
+
143
+
144
+ def argmax(a: GPUArray) -> int:
145
+ """Return index of maximum value."""
146
+ return a.argmax()
147
+
148
+
149
+ def argmin(a: GPUArray) -> int:
150
+ """Return index of minimum value."""
151
+ return a.argmin()
152
+
153
+
154
+ def concatenate(arrays, axis: int = 0) -> GPUArray:
155
+ """Join arrays along an axis (CPU-side, creates new GPU buffer)."""
156
+ np_arrays = [a.to_numpy() for a in arrays]
157
+ result = np.concatenate(np_arrays, axis=axis)
158
+ return _from_numpy(result)
159
+
160
+
161
+ def stack(arrays, axis: int = 0) -> GPUArray:
162
+ """Stack arrays along a new axis (CPU-side, creates new GPU buffer)."""
163
+ np_arrays = [a.to_numpy() for a in arrays]
164
+ result = np.stack(np_arrays, axis=axis)
165
+ return _from_numpy(result)
166
+
167
+
168
+ # ----- random submodule -----
169
+
170
+ class _Random:
171
+ """Random array generation (CPU-side, uploaded to GPU)."""
172
+
173
+ def rand(self, *shape) -> GPUArray:
174
+ """Uniform random values in [0, 1)."""
175
+ return _from_numpy(np.random.rand(*shape).astype(np.float32))
176
+
177
+ def randn(self, *shape) -> GPUArray:
178
+ """Standard normal random values."""
179
+ return _from_numpy(np.random.randn(*shape).astype(np.float32))
180
+
181
+
182
+ random = _Random()
gpuarray/core.py CHANGED
@@ -13,6 +13,7 @@ from .device import get_device
13
13
  from . import shaders
14
14
 
15
15
  WORKGROUP_SIZE = 256
16
+ REDUCTION_ELEMENTS_PER_WORKGROUP = 1024 # 256 threads * 4 elements each
16
17
 
17
18
  # Global pipeline cache: shader_code -> (shader_module, compute_pipeline)
18
19
  _pipeline_cache: dict[str, tuple] = {}
@@ -174,7 +175,10 @@ class GPUArray:
174
175
  return GPUArray(out_buf, self._shape)
175
176
 
176
177
  def _run_reduction(self, shader_code: str) -> float:
177
- """Multi-pass parallel reduction until 1 element remains."""
178
+ """Multi-pass parallel reduction until 1 element remains.
179
+
180
+ Each workgroup of 256 threads handles 1024 elements (4 per thread).
181
+ """
178
182
  device = get_device()
179
183
  _, pipeline = _get_pipeline(shader_code)
180
184
 
@@ -182,16 +186,16 @@ class GPUArray:
182
186
  current_size = self.size
183
187
 
184
188
  while current_size > 1:
185
- n_workgroups = math.ceil(current_size / WORKGROUP_SIZE)
189
+ n_workgroups = math.ceil(current_size / REDUCTION_ELEMENTS_PER_WORKGROUP)
186
190
  out_buf = device.create_buffer(
187
- size=n_workgroups * 4,
191
+ size=max(n_workgroups * 4, 4),
188
192
  usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
189
193
  )
190
194
  bind_group = device.create_bind_group(
191
195
  layout=pipeline.get_bind_group_layout(0),
192
196
  entries=[
193
197
  {"binding": 0, "resource": {"buffer": current_buf, "offset": 0, "size": current_size * 4}},
194
- {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": n_workgroups * 4}},
198
+ {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": max(n_workgroups * 4, 4)}},
195
199
  ],
196
200
  )
197
201
  encoder = device.create_command_encoder()
@@ -214,6 +218,29 @@ class GPUArray:
214
218
  current_buf.destroy()
215
219
  return result
216
220
 
221
+ def _run_binary_op_inplace(self, other: "GPUArray", shader_code: str) -> "GPUArray":
222
+ """Run a binary op that writes the result into self's buffer (no allocation)."""
223
+ assert self.size == other.size, f"Size mismatch: {self.size} vs {other.size}"
224
+ device = get_device()
225
+ _, pipeline = _get_pipeline(shader_code)
226
+
227
+ bind_group = device.create_bind_group(
228
+ layout=pipeline.get_bind_group_layout(0),
229
+ entries=[
230
+ {"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
231
+ {"binding": 1, "resource": {"buffer": other._buffer, "offset": 0, "size": other.size * 4}},
232
+ ],
233
+ )
234
+ n_workgroups = math.ceil(self.size / WORKGROUP_SIZE)
235
+ encoder = device.create_command_encoder()
236
+ pass_enc = encoder.begin_compute_pass()
237
+ pass_enc.set_pipeline(pipeline)
238
+ pass_enc.set_bind_group(0, bind_group)
239
+ pass_enc.dispatch_workgroups(n_workgroups)
240
+ pass_enc.end()
241
+ device.queue.submit([encoder.finish()])
242
+ return self
243
+
217
244
  # ----- dot product (two-input reduction) -----
218
245
 
219
246
  def _run_dot_reduction(self, other: "GPUArray") -> float:
@@ -259,11 +286,71 @@ class GPUArray:
259
286
  current_buf.destroy()
260
287
  return result
261
288
 
289
+ # ----- broadcast helper -----
290
+
291
+ def _run_broadcast_op(self, other: "GPUArray", shader_code: str, out_shape: tuple[int, ...]) -> "GPUArray":
292
+ """Run a broadcast binary shader with a uniform buffer for column count."""
293
+ device = get_device()
294
+ _, pipeline = _get_pipeline(shader_code)
295
+ out_size = 1
296
+ for d in out_shape:
297
+ out_size *= d
298
+ out_buf = device.create_buffer(
299
+ size=out_size * 4,
300
+ usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
301
+ )
302
+ cols = out_shape[-1] if len(out_shape) >= 2 else out_shape[0]
303
+ dims_data = struct.pack("IIII", cols, 0, 0, 0)
304
+ dims_buf = device.create_buffer_with_data(
305
+ data=dims_data, usage=wgpu.BufferUsage.UNIFORM,
306
+ )
307
+ bind_group = device.create_bind_group(
308
+ layout=pipeline.get_bind_group_layout(0),
309
+ entries=[
310
+ {"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
311
+ {"binding": 1, "resource": {"buffer": other._buffer, "offset": 0, "size": other.size * 4}},
312
+ {"binding": 2, "resource": {"buffer": out_buf, "offset": 0, "size": out_size * 4}},
313
+ {"binding": 3, "resource": {"buffer": dims_buf, "offset": 0, "size": 16}},
314
+ ],
315
+ )
316
+ n_workgroups = math.ceil(out_size / WORKGROUP_SIZE)
317
+ encoder = device.create_command_encoder()
318
+ pass_enc = encoder.begin_compute_pass()
319
+ pass_enc.set_pipeline(pipeline)
320
+ pass_enc.set_bind_group(0, bind_group)
321
+ pass_enc.dispatch_workgroups(n_workgroups)
322
+ pass_enc.end()
323
+ device.queue.submit([encoder.finish()])
324
+ dims_buf.destroy()
325
+ return GPUArray(out_buf, out_shape)
326
+
327
+ def _binary_with_broadcast(self, other: "GPUArray",
328
+ same_shader: str, row_shader: str, col_shader: str) -> "GPUArray":
329
+ """Dispatch to the right shader based on shape compatibility."""
330
+ if self.shape == other.shape or self.size == other.size:
331
+ return self._run_binary_op(other, same_shader)
332
+ # (M, N) op (N,) or (M, N) op (1, N) — broadcast row
333
+ if self.ndim == 2:
334
+ M, N = self.shape
335
+ if (other.ndim == 1 and other.shape[0] == N) or \
336
+ (other.ndim == 2 and other.shape == (1, N)):
337
+ b = other.flatten() if other.ndim == 2 else other
338
+ return self._run_broadcast_op(b, row_shader, self.shape)
339
+ # (M, N) op (M, 1) — broadcast col
340
+ if other.ndim == 2 and other.shape == (M, 1):
341
+ b_flat = other.flatten()
342
+ return self._run_broadcast_op(b_flat, col_shader, self.shape)
343
+ raise ValueError(f"Incompatible shapes for broadcast: {self.shape} vs {other.shape}")
344
+
262
345
  # ----- operator overloads -----
263
346
 
264
347
  def __add__(self, other):
265
348
  if isinstance(other, (int, float)):
266
349
  return self._run_scalar_op(float(other), shaders.ADD_SCALAR_SHADER)
350
+ if isinstance(other, GPUArray) and self.shape != other.shape:
351
+ return self._binary_with_broadcast(
352
+ other, shaders.ADD_SHADER,
353
+ shaders.BINARY_BROADCAST_ROW_ADD, shaders.BINARY_BROADCAST_COL_ADD)
267
354
  return self._run_binary_op(other, shaders.ADD_SHADER)
268
355
 
269
356
  def __radd__(self, other):
@@ -274,6 +361,10 @@ class GPUArray:
274
361
  def __sub__(self, other):
275
362
  if isinstance(other, (int, float)):
276
363
  return self._run_scalar_op(float(-other), shaders.ADD_SCALAR_SHADER)
364
+ if isinstance(other, GPUArray) and self.shape != other.shape:
365
+ return self._binary_with_broadcast(
366
+ other, shaders.SUB_SHADER,
367
+ shaders.BINARY_BROADCAST_ROW_SUB, shaders.BINARY_BROADCAST_COL_SUB)
277
368
  return self._run_binary_op(other, shaders.SUB_SHADER)
278
369
 
279
370
  def __rsub__(self, other):
@@ -285,6 +376,10 @@ class GPUArray:
285
376
  def __mul__(self, other):
286
377
  if isinstance(other, (int, float)):
287
378
  return self._run_scalar_op(float(other), shaders.MUL_SCALAR_SHADER)
379
+ if isinstance(other, GPUArray) and self.shape != other.shape:
380
+ return self._binary_with_broadcast(
381
+ other, shaders.MUL_SHADER,
382
+ shaders.BINARY_BROADCAST_ROW_MUL, shaders.BINARY_BROADCAST_COL_MUL)
288
383
  return self._run_binary_op(other, shaders.MUL_SHADER)
289
384
 
290
385
  def __rmul__(self, other):
@@ -295,8 +390,32 @@ class GPUArray:
295
390
  def __truediv__(self, other):
296
391
  if isinstance(other, (int, float)):
297
392
  return self._run_scalar_op(1.0 / float(other), shaders.MUL_SCALAR_SHADER)
393
+ if isinstance(other, GPUArray) and self.shape != other.shape:
394
+ return self._binary_with_broadcast(
395
+ other, shaders.DIV_SHADER,
396
+ shaders.BINARY_BROADCAST_ROW_DIV, shaders.BINARY_BROADCAST_COL_DIV)
298
397
  return self._run_binary_op(other, shaders.DIV_SHADER)
299
398
 
399
+ def __iadd__(self, other):
400
+ if isinstance(other, (int, float)):
401
+ return self._run_scalar_op(float(other), shaders.ADD_SCALAR_SHADER)
402
+ return self._run_binary_op_inplace(other, shaders.ADD_INPLACE_SHADER)
403
+
404
+ def __isub__(self, other):
405
+ if isinstance(other, (int, float)):
406
+ return self._run_scalar_op(float(-other), shaders.ADD_SCALAR_SHADER)
407
+ return self._run_binary_op_inplace(other, shaders.SUB_INPLACE_SHADER)
408
+
409
+ def __imul__(self, other):
410
+ if isinstance(other, (int, float)):
411
+ return self._run_scalar_op(float(other), shaders.MUL_SCALAR_SHADER)
412
+ return self._run_binary_op_inplace(other, shaders.MUL_INPLACE_SHADER)
413
+
414
+ def __itruediv__(self, other):
415
+ if isinstance(other, (int, float)):
416
+ return self._run_scalar_op(1.0 / float(other), shaders.MUL_SCALAR_SHADER)
417
+ return self._run_binary_op_inplace(other, shaders.DIV_INPLACE_SHADER)
418
+
300
419
  def __neg__(self):
301
420
  return self._run_unary_op(shaders.NEG_SHADER)
302
421
 
@@ -305,12 +424,108 @@ class GPUArray:
305
424
  return f"GPUArray({data}, shape={self._shape})"
306
425
 
307
426
  def __getitem__(self, key):
427
+ if isinstance(key, GPUArray):
428
+ return self._getitem_mask(key)
308
429
  data = self.to_numpy()
309
430
  result = data[key]
310
431
  if isinstance(result, np.ndarray):
311
- return _from_numpy(result)
432
+ return _from_numpy(result.copy())
312
433
  return float(result)
313
434
 
435
+ def _getitem_mask(self, mask: "GPUArray") -> "GPUArray":
436
+ """Boolean mask indexing via gather shader."""
437
+ mask_np = mask.to_numpy().flatten()
438
+ indices = np.where(mask_np > 0.5)[0].astype(np.uint32)
439
+ if len(indices) == 0:
440
+ return _from_numpy(np.array([], dtype=np.float32))
441
+ device = get_device()
442
+ _, pipeline = _get_pipeline(shaders.GATHER_SHADER)
443
+ idx_buf = device.create_buffer_with_data(
444
+ data=indices.tobytes(),
445
+ usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC,
446
+ )
447
+ out_buf = device.create_buffer(
448
+ size=len(indices) * 4,
449
+ usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
450
+ )
451
+ bind_group = device.create_bind_group(
452
+ layout=pipeline.get_bind_group_layout(0),
453
+ entries=[
454
+ {"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
455
+ {"binding": 1, "resource": {"buffer": idx_buf, "offset": 0, "size": len(indices) * 4}},
456
+ {"binding": 2, "resource": {"buffer": out_buf, "offset": 0, "size": len(indices) * 4}},
457
+ ],
458
+ )
459
+ n_workgroups = math.ceil(len(indices) / WORKGROUP_SIZE)
460
+ encoder = device.create_command_encoder()
461
+ pass_enc = encoder.begin_compute_pass()
462
+ pass_enc.set_pipeline(pipeline)
463
+ pass_enc.set_bind_group(0, bind_group)
464
+ pass_enc.dispatch_workgroups(n_workgroups)
465
+ pass_enc.end()
466
+ device.queue.submit([encoder.finish()])
467
+ idx_buf.destroy()
468
+ return GPUArray(out_buf, (len(indices),))
469
+
470
+ def __setitem__(self, key, value):
471
+ if isinstance(key, GPUArray):
472
+ self._setitem_mask(key, value)
473
+ return
474
+ device = get_device()
475
+ if isinstance(key, (int, np.integer)):
476
+ idx = int(key)
477
+ if idx < 0:
478
+ idx += self.size
479
+ device.queue.write_buffer(self._buffer, idx * 4, struct.pack("f", float(value)))
480
+ elif isinstance(key, slice):
481
+ if self.ndim == 1:
482
+ indices = range(*key.indices(self._shape[0]))
483
+ if isinstance(value, (int, float)):
484
+ data = struct.pack("f", float(value)) * len(indices)
485
+ elif isinstance(value, GPUArray):
486
+ data = value.to_numpy().astype(np.float32).tobytes()
487
+ else:
488
+ data = np.asarray(value, dtype=np.float32).tobytes()
489
+ start = indices.start if len(indices) > 0 else 0
490
+ device.queue.write_buffer(self._buffer, start * 4, data)
491
+ else:
492
+ indices = range(*key.indices(self._shape[0]))
493
+ cols = self._shape[1]
494
+ if isinstance(value, (int, float)):
495
+ row_data = struct.pack("f", float(value)) * cols
496
+ for r in indices:
497
+ device.queue.write_buffer(self._buffer, r * cols * 4, row_data)
498
+ else:
499
+ val_np = value.to_numpy() if isinstance(value, GPUArray) else np.asarray(value, dtype=np.float32)
500
+ for i, r in enumerate(indices):
501
+ row_bytes = (val_np[i] if val_np.ndim >= 2 else val_np).astype(np.float32).tobytes()
502
+ device.queue.write_buffer(self._buffer, r * cols * 4, row_bytes)
503
+ elif isinstance(key, tuple) and len(key) == 2:
504
+ row, col = key
505
+ if isinstance(row, (int, np.integer)) and isinstance(col, (int, np.integer)):
506
+ cols = self._shape[1]
507
+ offset = (int(row) * cols + int(col)) * 4
508
+ device.queue.write_buffer(self._buffer, offset, struct.pack("f", float(value)))
509
+ else:
510
+ arr = self.to_numpy()
511
+ arr[key] = value if not isinstance(value, GPUArray) else value.to_numpy()
512
+ device.queue.write_buffer(self._buffer, 0, np.ascontiguousarray(arr, dtype=np.float32).tobytes())
513
+ else:
514
+ arr = self.to_numpy()
515
+ arr[key] = value if not isinstance(value, GPUArray) else value.to_numpy()
516
+ device.queue.write_buffer(self._buffer, 0, np.ascontiguousarray(arr, dtype=np.float32).tobytes())
517
+
518
+ def _setitem_mask(self, mask: "GPUArray", value):
519
+ """Masked assignment: a[mask] = value."""
520
+ data = self.to_numpy()
521
+ mask_np = mask.to_numpy() > 0.5
522
+ if isinstance(value, GPUArray):
523
+ data[mask_np] = value.to_numpy().flatten()[:mask_np.sum()]
524
+ else:
525
+ data[mask_np] = float(value)
526
+ device = get_device()
527
+ device.queue.write_buffer(self._buffer, 0, np.ascontiguousarray(data, dtype=np.float32).tobytes())
528
+
314
529
  # ----- reduction methods -----
315
530
 
316
531
  def sum(self) -> float:
@@ -343,6 +558,105 @@ class GPUArray:
343
558
  raise ValueError(f"Cannot reshape array of size {self.size} into shape {new_shape}")
344
559
  return GPUArray(self._buffer, new_shape)
345
560
 
561
+ # ----- clamp / clip -----
562
+
563
+ def clip(self, min_val: float, max_val: float) -> "GPUArray":
564
+ """Clamp values to [min_val, max_val]."""
565
+ device = get_device()
566
+ _, pipeline = _get_pipeline(shaders.CLAMP_SHADER)
567
+
568
+ out_buf = device.create_buffer(
569
+ size=self.size * 4,
570
+ usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
571
+ )
572
+ # Pack two f32 values, padded to 16 bytes
573
+ params_data = struct.pack("ff", min_val, max_val) + b"\x00" * 8
574
+ uniform_buf = device.create_buffer_with_data(
575
+ data=params_data,
576
+ usage=wgpu.BufferUsage.UNIFORM,
577
+ )
578
+ bind_group = device.create_bind_group(
579
+ layout=pipeline.get_bind_group_layout(0),
580
+ entries=[
581
+ {"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
582
+ {"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": self.size * 4}},
583
+ {"binding": 2, "resource": {"buffer": uniform_buf, "offset": 0, "size": 16}},
584
+ ],
585
+ )
586
+ n_workgroups = math.ceil(self.size / WORKGROUP_SIZE)
587
+ encoder = device.create_command_encoder()
588
+ pass_enc = encoder.begin_compute_pass()
589
+ pass_enc.set_pipeline(pipeline)
590
+ pass_enc.set_bind_group(0, bind_group)
591
+ pass_enc.dispatch_workgroups(n_workgroups)
592
+ pass_enc.end()
593
+ device.queue.submit([encoder.finish()])
594
+ uniform_buf.destroy()
595
+ return GPUArray(out_buf, self._shape)
596
+
597
+ # ----- argmax / argmin -----
598
+
599
+ def argmax(self) -> int:
600
+ """Return index of maximum value (as Python int)."""
601
+ data = self.to_numpy().flatten()
602
+ return int(np.argmax(data))
603
+
604
+ def argmin(self) -> int:
605
+ """Return index of minimum value (as Python int)."""
606
+ data = self.to_numpy().flatten()
607
+ return int(np.argmin(data))
608
+
609
+ # ----- copy -----
610
+
611
+ def copy(self) -> "GPUArray":
612
+ """Return a copy of this array."""
613
+ return self._run_unary_op(shaders.COPY_SHADER)
614
+
615
+ # ----- comparison operators -----
616
+
617
+ def __eq__(self, other):
618
+ if isinstance(other, (int, float)):
619
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
620
+ return self._run_binary_op(other, shaders.EQ_SHADER)
621
+
622
+ def __gt__(self, other):
623
+ if isinstance(other, (int, float)):
624
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
625
+ return self._run_binary_op(other, shaders.GT_SHADER)
626
+
627
+ def __lt__(self, other):
628
+ if isinstance(other, (int, float)):
629
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
630
+ return self._run_binary_op(other, shaders.LT_SHADER)
631
+
632
+ def __ge__(self, other):
633
+ if isinstance(other, (int, float)):
634
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
635
+ return self._run_binary_op(other, shaders.GE_SHADER)
636
+
637
+ def __le__(self, other):
638
+ if isinstance(other, (int, float)):
639
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
640
+ return self._run_binary_op(other, shaders.LE_SHADER)
641
+
642
+ # ----- power -----
643
+
644
+ def __pow__(self, other):
645
+ if isinstance(other, (int, float)):
646
+ other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
647
+ return self._run_binary_op(other, shaders.POW_SHADER)
648
+
649
+ # ----- flatten -----
650
+
651
+ def flatten(self) -> "GPUArray":
652
+ """Return a 1D view of the array."""
653
+ return GPUArray(self._buffer, (self.size,))
654
+
655
+ # ----- len -----
656
+
657
+ def __len__(self) -> int:
658
+ return self._shape[0]
659
+
346
660
 
347
661
  def _from_numpy(data: np.ndarray) -> GPUArray:
348
662
  """Create a GPUArray from a numpy array."""
@@ -356,7 +670,10 @@ def _from_numpy(data: np.ndarray) -> GPUArray:
356
670
 
357
671
 
358
672
  def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
359
- """Matrix multiplication C = A @ B."""
673
+ """Matrix multiplication C = A @ B.
674
+
675
+ Uses tiled shared-memory shader for matrices >= 64x64, naive for smaller.
676
+ """
360
677
  if a.ndim != 2 or b.ndim != 2:
361
678
  raise ValueError("matmul requires 2D arrays")
362
679
  M, K = a.shape
@@ -365,7 +682,11 @@ def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
365
682
  raise ValueError(f"Incompatible shapes for matmul: {a.shape} and {b.shape}")
366
683
 
367
684
  device = get_device()
368
- _, pipeline = _get_pipeline(shaders.MATMUL_SHADER)
685
+
686
+ # Choose tiled or naive shader based on matrix size
687
+ use_tiled = M >= 64 and N >= 64 and K >= 64
688
+ shader_code = shaders.MATMUL_TILED if use_tiled else shaders.MATMUL_NAIVE
689
+ _, pipeline = _get_pipeline(shader_code)
369
690
 
370
691
  out_size = M * N
371
692
  out_buf = device.create_buffer(
@@ -390,13 +711,27 @@ def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
390
711
  ],
391
712
  )
392
713
 
393
- n_workgroups = math.ceil(out_size / WORKGROUP_SIZE)
394
- encoder = device.create_command_encoder()
395
- pass_enc = encoder.begin_compute_pass()
396
- pass_enc.set_pipeline(pipeline)
397
- pass_enc.set_bind_group(0, bind_group)
398
- pass_enc.dispatch_workgroups(n_workgroups)
399
- pass_enc.end()
714
+ if use_tiled:
715
+ # Tiled shader uses @workgroup_size(16, 16); dispatch 2D grid
716
+ TILE_SIZE = 16
717
+ wg_x = math.ceil(N / TILE_SIZE)
718
+ wg_y = math.ceil(M / TILE_SIZE)
719
+ encoder = device.create_command_encoder()
720
+ pass_enc = encoder.begin_compute_pass()
721
+ pass_enc.set_pipeline(pipeline)
722
+ pass_enc.set_bind_group(0, bind_group)
723
+ pass_enc.dispatch_workgroups(wg_x, wg_y)
724
+ pass_enc.end()
725
+ else:
726
+ # Naive shader uses @workgroup_size(256); dispatch 1D
727
+ n_workgroups = math.ceil(out_size / WORKGROUP_SIZE)
728
+ encoder = device.create_command_encoder()
729
+ pass_enc = encoder.begin_compute_pass()
730
+ pass_enc.set_pipeline(pipeline)
731
+ pass_enc.set_bind_group(0, bind_group)
732
+ pass_enc.dispatch_workgroups(n_workgroups)
733
+ pass_enc.end()
734
+
400
735
  device.queue.submit([encoder.finish()])
401
736
  dims_buf.destroy()
402
737
  return GPUArray(out_buf, (M, N))
gpuarray/shaders.py CHANGED
@@ -74,6 +74,62 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
74
74
  }
75
75
  """
76
76
 
77
+ # ---------------------------------------------------------------------------
78
+ # In-place elementwise binary operations (result written to a)
79
+ # ---------------------------------------------------------------------------
80
+
81
+ ADD_INPLACE_SHADER = """
82
+ @group(0) @binding(0) var<storage, read_write> a: array<f32>;
83
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
84
+
85
+ @compute @workgroup_size(256)
86
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
87
+ let i = gid.x;
88
+ if (i < arrayLength(&a)) {
89
+ a[i] = a[i] + b[i];
90
+ }
91
+ }
92
+ """
93
+
94
+ SUB_INPLACE_SHADER = """
95
+ @group(0) @binding(0) var<storage, read_write> a: array<f32>;
96
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
97
+
98
+ @compute @workgroup_size(256)
99
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
100
+ let i = gid.x;
101
+ if (i < arrayLength(&a)) {
102
+ a[i] = a[i] - b[i];
103
+ }
104
+ }
105
+ """
106
+
107
+ MUL_INPLACE_SHADER = """
108
+ @group(0) @binding(0) var<storage, read_write> a: array<f32>;
109
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
110
+
111
+ @compute @workgroup_size(256)
112
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
113
+ let i = gid.x;
114
+ if (i < arrayLength(&a)) {
115
+ a[i] = a[i] * b[i];
116
+ }
117
+ }
118
+ """
119
+
120
+ DIV_INPLACE_SHADER = """
121
+ @group(0) @binding(0) var<storage, read_write> a: array<f32>;
122
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
123
+
124
+ @compute @workgroup_size(256)
125
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
126
+ let i = gid.x;
127
+ if (i < arrayLength(&a)) {
128
+ a[i] = a[i] / b[i];
129
+ }
130
+ }
131
+ """
132
+
77
133
  # ---------------------------------------------------------------------------
78
134
  # Elementwise unary operations
79
135
  # ---------------------------------------------------------------------------
@@ -236,17 +292,26 @@ var<workgroup> sdata: array<f32, 256>;
236
292
  fn main(
237
293
  @builtin(local_invocation_id) lid: vec3<u32>,
238
294
  @builtin(workgroup_id) wid: vec3<u32>,
239
- @builtin(global_invocation_id) gid: vec3<u32>,
240
295
  ) {
241
296
  let local_idx = lid.x;
242
- let global_idx = gid.x;
243
297
  let n = arrayLength(&input);
298
+ // Each thread loads and reduces 4 elements (workgroup covers 1024 elements)
299
+ let base = wid.x * 1024u + local_idx;
244
300
 
245
- if (global_idx < n) {
246
- sdata[local_idx] = input[global_idx];
247
- } else {
248
- sdata[local_idx] = 0.0;
301
+ var acc: f32 = 0.0;
302
+ if (base < n) {
303
+ acc = input[base];
304
+ }
305
+ if (base + 256u < n) {
306
+ acc = acc + input[base + 256u];
307
+ }
308
+ if (base + 512u < n) {
309
+ acc = acc + input[base + 512u];
249
310
  }
311
+ if (base + 768u < n) {
312
+ acc = acc + input[base + 768u];
313
+ }
314
+ sdata[local_idx] = acc;
250
315
  workgroupBarrier();
251
316
 
252
317
  // Tree reduction
@@ -276,17 +341,25 @@ var<workgroup> sdata: array<f32, 256>;
276
341
  fn main(
277
342
  @builtin(local_invocation_id) lid: vec3<u32>,
278
343
  @builtin(workgroup_id) wid: vec3<u32>,
279
- @builtin(global_invocation_id) gid: vec3<u32>,
280
344
  ) {
281
345
  let local_idx = lid.x;
282
- let global_idx = gid.x;
283
346
  let n = arrayLength(&input);
347
+ let base = wid.x * 1024u + local_idx;
284
348
 
285
- if (global_idx < n) {
286
- sdata[local_idx] = input[global_idx];
287
- } else {
288
- sdata[local_idx] = -3.402823e+38; // -FLT_MAX
349
+ var acc: f32 = -3.402823e+38;
350
+ if (base < n) {
351
+ acc = input[base];
352
+ }
353
+ if (base + 256u < n) {
354
+ acc = max(acc, input[base + 256u]);
355
+ }
356
+ if (base + 512u < n) {
357
+ acc = max(acc, input[base + 512u]);
289
358
  }
359
+ if (base + 768u < n) {
360
+ acc = max(acc, input[base + 768u]);
361
+ }
362
+ sdata[local_idx] = acc;
290
363
  workgroupBarrier();
291
364
 
292
365
  var stride: u32 = 128u;
@@ -315,17 +388,25 @@ var<workgroup> sdata: array<f32, 256>;
315
388
  fn main(
316
389
  @builtin(local_invocation_id) lid: vec3<u32>,
317
390
  @builtin(workgroup_id) wid: vec3<u32>,
318
- @builtin(global_invocation_id) gid: vec3<u32>,
319
391
  ) {
320
392
  let local_idx = lid.x;
321
- let global_idx = gid.x;
322
393
  let n = arrayLength(&input);
394
+ let base = wid.x * 1024u + local_idx;
323
395
 
324
- if (global_idx < n) {
325
- sdata[local_idx] = input[global_idx];
326
- } else {
327
- sdata[local_idx] = 3.402823e+38; // FLT_MAX
396
+ var acc: f32 = 3.402823e+38;
397
+ if (base < n) {
398
+ acc = input[base];
328
399
  }
400
+ if (base + 256u < n) {
401
+ acc = min(acc, input[base + 256u]);
402
+ }
403
+ if (base + 512u < n) {
404
+ acc = min(acc, input[base + 512u]);
405
+ }
406
+ if (base + 768u < n) {
407
+ acc = min(acc, input[base + 768u]);
408
+ }
409
+ sdata[local_idx] = acc;
329
410
  workgroupBarrier();
330
411
 
331
412
  var stride: u32 = 128u;
@@ -394,7 +475,7 @@ fn main(
394
475
  # Each thread computes one element of C.
395
476
  # ---------------------------------------------------------------------------
396
477
 
397
- MATMUL_SHADER = """
478
+ MATMUL_NAIVE = """
398
479
  struct Dims {
399
480
  M: u32,
400
481
  K: u32,
@@ -424,3 +505,377 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
424
505
  c[row * dims.N + col] = sum;
425
506
  }
426
507
  """
508
+
509
+ # Keep backward compat alias
510
+ MATMUL_SHADER = MATMUL_NAIVE
511
+
512
+ # ---------------------------------------------------------------------------
513
+ # Tiled matrix multiplication (16x16 tiles with shared memory)
514
+ # ---------------------------------------------------------------------------
515
+
516
+ MATMUL_TILED = """
517
+ struct Dims {
518
+ M: u32,
519
+ K: u32,
520
+ N: u32,
521
+ _pad: u32,
522
+ }
523
+
524
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
525
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
526
+ @group(0) @binding(2) var<storage, read_write> c: array<f32>;
527
+ @group(0) @binding(3) var<uniform> dims: Dims;
528
+
529
+ const TILE_SIZE = 16u;
530
+ var<workgroup> tile_a: array<f32, 256>;
531
+ var<workgroup> tile_b: array<f32, 256>;
532
+
533
+ @compute @workgroup_size(16, 16)
534
+ fn main(@builtin(local_invocation_id) lid: vec3<u32>,
535
+ @builtin(workgroup_id) wid: vec3<u32>) {
536
+ let row = wid.y * TILE_SIZE + lid.y;
537
+ let col = wid.x * TILE_SIZE + lid.x;
538
+ let local_idx = lid.y * TILE_SIZE + lid.x;
539
+
540
+ let num_tiles = (dims.K + TILE_SIZE - 1u) / TILE_SIZE;
541
+
542
+ var sum: f32 = 0.0;
543
+
544
+ for (var t: u32 = 0u; t < num_tiles; t = t + 1u) {
545
+ // Load tile of A into shared memory
546
+ let a_col = t * TILE_SIZE + lid.x;
547
+ if (row < dims.M && a_col < dims.K) {
548
+ tile_a[local_idx] = a[row * dims.K + a_col];
549
+ } else {
550
+ tile_a[local_idx] = 0.0;
551
+ }
552
+
553
+ // Load tile of B into shared memory
554
+ let b_row = t * TILE_SIZE + lid.y;
555
+ if (b_row < dims.K && col < dims.N) {
556
+ tile_b[local_idx] = b[b_row * dims.N + col];
557
+ } else {
558
+ tile_b[local_idx] = 0.0;
559
+ }
560
+
561
+ workgroupBarrier();
562
+
563
+ // Compute partial dot product for this tile
564
+ for (var k: u32 = 0u; k < TILE_SIZE; k = k + 1u) {
565
+ sum = sum + tile_a[lid.y * TILE_SIZE + k] * tile_b[k * TILE_SIZE + lid.x];
566
+ }
567
+
568
+ workgroupBarrier();
569
+ }
570
+
571
+ if (row < dims.M && col < dims.N) {
572
+ c[row * dims.N + col] = sum;
573
+ }
574
+ }
575
+ """
576
+
577
+ # ---------------------------------------------------------------------------
578
+ # Clamp (requires uniform buffer with min_val, max_val)
579
+ # ---------------------------------------------------------------------------
580
+
581
+ CLAMP_SHADER = """
582
+ struct Params {
583
+ min_val: f32,
584
+ max_val: f32,
585
+ }
586
+
587
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
588
+ @group(0) @binding(1) var<storage, read_write> result: array<f32>;
589
+ @group(0) @binding(2) var<uniform> params: Params;
590
+
591
+ @compute @workgroup_size(256)
592
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
593
+ let i = gid.x;
594
+ if (i < arrayLength(&a)) {
595
+ result[i] = clamp(a[i], params.min_val, params.max_val);
596
+ }
597
+ }
598
+ """
599
+
600
+ # ---------------------------------------------------------------------------
601
+ # Where / select: result = cond * a + (1-cond) * b
602
+ # ---------------------------------------------------------------------------
603
+
604
+ WHERE_SHADER = """
605
+ @group(0) @binding(0) var<storage, read> cond: array<f32>;
606
+ @group(0) @binding(1) var<storage, read> a: array<f32>;
607
+ @group(0) @binding(2) var<storage, read> b: array<f32>;
608
+ @group(0) @binding(3) var<storage, read_write> result: array<f32>;
609
+
610
+ @compute @workgroup_size(256)
611
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
612
+ let i = gid.x;
613
+ if (i < arrayLength(&cond)) {
614
+ result[i] = cond[i] * a[i] + (1.0 - cond[i]) * b[i];
615
+ }
616
+ }
617
+ """
618
+
619
+ # ---------------------------------------------------------------------------
620
+ # Comparison shaders (return 0.0 or 1.0)
621
+ # ---------------------------------------------------------------------------
622
+
623
+ EQ_SHADER = """
624
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
625
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
626
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
627
+
628
+ @compute @workgroup_size(256)
629
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
630
+ let i = gid.x;
631
+ if (i < arrayLength(&a)) {
632
+ if (a[i] == b[i]) {
633
+ result[i] = 1.0;
634
+ } else {
635
+ result[i] = 0.0;
636
+ }
637
+ }
638
+ }
639
+ """
640
+
641
+ GT_SHADER = """
642
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
643
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
644
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
645
+
646
+ @compute @workgroup_size(256)
647
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
648
+ let i = gid.x;
649
+ if (i < arrayLength(&a)) {
650
+ if (a[i] > b[i]) {
651
+ result[i] = 1.0;
652
+ } else {
653
+ result[i] = 0.0;
654
+ }
655
+ }
656
+ }
657
+ """
658
+
659
+ LT_SHADER = """
660
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
661
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
662
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
663
+
664
+ @compute @workgroup_size(256)
665
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
666
+ let i = gid.x;
667
+ if (i < arrayLength(&a)) {
668
+ if (a[i] < b[i]) {
669
+ result[i] = 1.0;
670
+ } else {
671
+ result[i] = 0.0;
672
+ }
673
+ }
674
+ }
675
+ """
676
+
677
+ GE_SHADER = """
678
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
679
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
680
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
681
+
682
+ @compute @workgroup_size(256)
683
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
684
+ let i = gid.x;
685
+ if (i < arrayLength(&a)) {
686
+ if (a[i] >= b[i]) {
687
+ result[i] = 1.0;
688
+ } else {
689
+ result[i] = 0.0;
690
+ }
691
+ }
692
+ }
693
+ """
694
+
695
+ LE_SHADER = """
696
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
697
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
698
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
699
+
700
+ @compute @workgroup_size(256)
701
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
702
+ let i = gid.x;
703
+ if (i < arrayLength(&a)) {
704
+ if (a[i] <= b[i]) {
705
+ result[i] = 1.0;
706
+ } else {
707
+ result[i] = 0.0;
708
+ }
709
+ }
710
+ }
711
+ """
712
+
713
+ # ---------------------------------------------------------------------------
714
+ # Fill (fill array with uniform scalar)
715
+ # ---------------------------------------------------------------------------
716
+
717
+ FILL_SHADER = """
718
+ struct Params {
719
+ scalar: f32,
720
+ }
721
+
722
+ @group(0) @binding(0) var<storage, read_write> result: array<f32>;
723
+ @group(0) @binding(1) var<uniform> params: Params;
724
+
725
+ @compute @workgroup_size(256)
726
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
727
+ let i = gid.x;
728
+ if (i < arrayLength(&result)) {
729
+ result[i] = params.scalar;
730
+ }
731
+ }
732
+ """
733
+
734
+ # ---------------------------------------------------------------------------
735
+ # Copy
736
+ # ---------------------------------------------------------------------------
737
+
738
+ COPY_SHADER = """
739
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
740
+ @group(0) @binding(1) var<storage, read_write> result: array<f32>;
741
+
742
+ @compute @workgroup_size(256)
743
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
744
+ let i = gid.x;
745
+ if (i < arrayLength(&a)) {
746
+ result[i] = a[i];
747
+ }
748
+ }
749
+ """
750
+
751
+ # ---------------------------------------------------------------------------
752
+ # Gather: result[i] = source[indices[i]]
753
+ # ---------------------------------------------------------------------------
754
+
755
+ GATHER_SHADER = """
756
+ @group(0) @binding(0) var<storage, read> source: array<f32>;
757
+ @group(0) @binding(1) var<storage, read> indices: array<u32>;
758
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
759
+
760
+ @compute @workgroup_size(256)
761
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
762
+ let i = gid.x;
763
+ if (i < arrayLength(&indices)) {
764
+ result[i] = source[indices[i]];
765
+ }
766
+ }
767
+ """
768
+
769
+ # ---------------------------------------------------------------------------
770
+ # Broadcast binary ops: (M,N) op (N,) — broadcast 1D across rows
771
+ # Uniform: dims.x = N (number of columns)
772
+ # ---------------------------------------------------------------------------
773
+
774
+ BINARY_BROADCAST_ROW_ADD = """
775
+ struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
776
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
777
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
778
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
779
+ @group(0) @binding(3) var<uniform> dims: Dims;
780
+ @compute @workgroup_size(256)
781
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
782
+ let i = gid.x;
783
+ if (i < arrayLength(&a)) { result[i] = a[i] + b[i % dims.cols]; }
784
+ }
785
+ """
786
+
787
+ BINARY_BROADCAST_ROW_SUB = """
788
+ struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
789
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
790
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
791
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
792
+ @group(0) @binding(3) var<uniform> dims: Dims;
793
+ @compute @workgroup_size(256)
794
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
795
+ let i = gid.x;
796
+ if (i < arrayLength(&a)) { result[i] = a[i] - b[i % dims.cols]; }
797
+ }
798
+ """
799
+
800
+ BINARY_BROADCAST_ROW_MUL = """
801
+ struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
802
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
803
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
804
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
805
+ @group(0) @binding(3) var<uniform> dims: Dims;
806
+ @compute @workgroup_size(256)
807
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
808
+ let i = gid.x;
809
+ if (i < arrayLength(&a)) { result[i] = a[i] * b[i % dims.cols]; }
810
+ }
811
+ """
812
+
813
+ BINARY_BROADCAST_ROW_DIV = """
814
+ struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
815
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
816
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
817
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
818
+ @group(0) @binding(3) var<uniform> dims: Dims;
819
+ @compute @workgroup_size(256)
820
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
821
+ let i = gid.x;
822
+ if (i < arrayLength(&a)) { result[i] = a[i] / b[i % dims.cols]; }
823
+ }
824
+ """
825
+
826
+ # ---------------------------------------------------------------------------
827
+ # Broadcast binary ops: (M,N) op (M,1) — broadcast column vector
828
+ # Uniform: dims.x = N (number of columns)
829
+ # ---------------------------------------------------------------------------
830
+
831
+ BINARY_BROADCAST_COL_ADD = """
832
+ struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
833
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
834
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
835
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
836
+ @group(0) @binding(3) var<uniform> dims: Dims;
837
+ @compute @workgroup_size(256)
838
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
839
+ let i = gid.x;
840
+ if (i < arrayLength(&a)) { result[i] = a[i] + b[i / dims.cols]; }
841
+ }
842
+ """
843
+
844
+ BINARY_BROADCAST_COL_SUB = """
845
+ struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
846
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
847
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
848
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
849
+ @group(0) @binding(3) var<uniform> dims: Dims;
850
+ @compute @workgroup_size(256)
851
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
852
+ let i = gid.x;
853
+ if (i < arrayLength(&a)) { result[i] = a[i] - b[i / dims.cols]; }
854
+ }
855
+ """
856
+
857
+ BINARY_BROADCAST_COL_MUL = """
858
+ struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
859
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
860
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
861
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
862
+ @group(0) @binding(3) var<uniform> dims: Dims;
863
+ @compute @workgroup_size(256)
864
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
865
+ let i = gid.x;
866
+ if (i < arrayLength(&a)) { result[i] = a[i] * b[i / dims.cols]; }
867
+ }
868
+ """
869
+
870
+ BINARY_BROADCAST_COL_DIV = """
871
+ struct Dims { cols: u32, _p1: u32, _p2: u32, _p3: u32, }
872
+ @group(0) @binding(0) var<storage, read> a: array<f32>;
873
+ @group(0) @binding(1) var<storage, read> b: array<f32>;
874
+ @group(0) @binding(2) var<storage, read_write> result: array<f32>;
875
+ @group(0) @binding(3) var<uniform> dims: Dims;
876
+ @compute @workgroup_size(256)
877
+ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
878
+ let i = gid.x;
879
+ if (i < arrayLength(&a)) { result[i] = a[i] / b[i / dims.cols]; }
880
+ }
881
+ """
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: gpuarray
3
- Version: 0.1.0
3
+ Version: 0.3.0
4
4
  Summary: NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required.
5
5
  Author: Arnav Kewalram
6
6
  License: MIT
@@ -0,0 +1,9 @@
1
+ gpuarray/__init__.py,sha256=MB9k9oXPvA6wRCCZii72HtpRxjXPtkiT2RaJsiVY-wY,752
2
+ gpuarray/api.py,sha256=PlYqtvY7jEV3xIneGdRoqXJjVOP3w2fnHyZgVP6fqis,4874
3
+ gpuarray/core.py,sha256=o7UymuEJoyYQAKao4O4wRtFw911Ut3h7tnQJOT6N0BU,30676
4
+ gpuarray/device.py,sha256=q80VOrKhy1u5i4AMLUKhWVBwFqQjoJ4tNBm234Wpf-A,452
5
+ gpuarray/shaders.py,sha256=oG7vd3iqO1eJ7UMIZx-46UPcb9gNVUWhBa6q8iKOSxo,25467
6
+ gpuarray-0.3.0.dist-info/METADATA,sha256=XrCwPw9lCTIHb5qMSZjHypALq5BLwwHjrV2SVoLDnhE,3132
7
+ gpuarray-0.3.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
8
+ gpuarray-0.3.0.dist-info/top_level.txt,sha256=iruYli4LmPBqhaZxTVvzT1bT0VeN6dbC0GE29lT1CbI,9
9
+ gpuarray-0.3.0.dist-info/RECORD,,
@@ -1,9 +0,0 @@
1
- gpuarray/__init__.py,sha256=yDS9YEnJUsiPgAnftBGRQJqG2FwL6oRwvOI98p7WZu8,568
2
- gpuarray/api.py,sha256=tlIjtSCmCs-CYbZSAmW5XPxAxZmmSX3-wuGXI-6u1Zs,2090
3
- gpuarray/core.py,sha256=GIs-PE6Y7m0htoBv1B1hbGxImkmLFkP6z_Zo9Gm5Joc,15122
4
- gpuarray/device.py,sha256=q80VOrKhy1u5i4AMLUKhWVBwFqQjoJ4tNBm234Wpf-A,452
5
- gpuarray/shaders.py,sha256=7fuWEtPVVHXnJxFBeT5SoQ467FDCqM6fP14Wqd8r0HA,11411
6
- gpuarray-0.1.0.dist-info/METADATA,sha256=WZHK9RoABh502yZZ7s8U4h4vy4UGknpfOQ1duSeb-3I,3132
7
- gpuarray-0.1.0.dist-info/WHEEL,sha256=K260EYznzXsJYBQGqmI8VTxEdiZYNvDZwW9cBh9-_MA,91
8
- gpuarray-0.1.0.dist-info/top_level.txt,sha256=iruYli4LmPBqhaZxTVvzT1bT0VeN6dbC0GE29lT1CbI,9
9
- gpuarray-0.1.0.dist-info/RECORD,,