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 +14 -0
- gpuarray/api.py +83 -0
- gpuarray/core.py +349 -14
- gpuarray/shaders.py +474 -19
- {gpuarray-0.1.0.dist-info → gpuarray-0.3.0.dist-info}/METADATA +1 -1
- gpuarray-0.3.0.dist-info/RECORD +9 -0
- gpuarray-0.1.0.dist-info/RECORD +0 -9
- {gpuarray-0.1.0.dist-info → gpuarray-0.3.0.dist-info}/WHEEL +0 -0
- {gpuarray-0.1.0.dist-info → gpuarray-0.3.0.dist-info}/top_level.txt +0 -0
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 /
|
|
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
|
-
|
|
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
|
-
|
|
394
|
-
|
|
395
|
-
|
|
396
|
-
|
|
397
|
-
|
|
398
|
-
|
|
399
|
-
|
|
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
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
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
|
-
|
|
286
|
-
|
|
287
|
-
|
|
288
|
-
|
|
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
|
-
|
|
325
|
-
|
|
326
|
-
|
|
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
|
-
|
|
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
|
+
"""
|
|
@@ -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,,
|
gpuarray-0.1.0.dist-info/RECORD
DELETED
|
@@ -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,,
|
|
File without changes
|
|
File without changes
|