gpuarray 0.1.0__tar.gz → 0.2.0__tar.gz
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-0.1.0 → gpuarray-0.2.0}/PKG-INFO +1 -1
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray/__init__.py +14 -0
- gpuarray-0.2.0/gpuarray/api.py +182 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray/core.py +99 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray/shaders.py +174 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray.egg-info/PKG-INFO +1 -1
- {gpuarray-0.1.0 → gpuarray-0.2.0}/pyproject.toml +1 -1
- {gpuarray-0.1.0 → gpuarray-0.2.0}/tests/test_core.py +208 -0
- gpuarray-0.1.0/gpuarray/api.py +0 -99
- {gpuarray-0.1.0 → gpuarray-0.2.0}/README.md +0 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray/device.py +0 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray.egg-info/SOURCES.txt +0 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray.egg-info/dependency_links.txt +0 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray.egg-info/requires.txt +0 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/gpuarray.egg-info/top_level.txt +0 -0
- {gpuarray-0.1.0 → gpuarray-0.2.0}/setup.cfg +0 -0
|
@@ -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
|
]
|
|
@@ -0,0 +1,182 @@
|
|
|
1
|
+
"""NumPy-compatible module-level functions."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Sequence
|
|
6
|
+
|
|
7
|
+
import numpy as np
|
|
8
|
+
|
|
9
|
+
from .core import GPUArray, _from_numpy, matmul as _matmul
|
|
10
|
+
from . import shaders
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
# ----- creation functions -----
|
|
14
|
+
|
|
15
|
+
def array(data) -> GPUArray:
|
|
16
|
+
"""Create a GPUArray from a list, nested list, or numpy array."""
|
|
17
|
+
if isinstance(data, GPUArray):
|
|
18
|
+
return data
|
|
19
|
+
arr = np.asarray(data, dtype=np.float32)
|
|
20
|
+
if arr.ndim == 0:
|
|
21
|
+
arr = arr.reshape((1,))
|
|
22
|
+
return _from_numpy(arr)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
def zeros(shape) -> GPUArray:
|
|
26
|
+
if isinstance(shape, int):
|
|
27
|
+
shape = (shape,)
|
|
28
|
+
return _from_numpy(np.zeros(shape, dtype=np.float32))
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def ones(shape) -> GPUArray:
|
|
32
|
+
if isinstance(shape, int):
|
|
33
|
+
shape = (shape,)
|
|
34
|
+
return _from_numpy(np.ones(shape, dtype=np.float32))
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def arange(start, stop=None, step=1) -> GPUArray:
|
|
38
|
+
if stop is None:
|
|
39
|
+
stop = start
|
|
40
|
+
start = 0
|
|
41
|
+
return _from_numpy(np.arange(start, stop, step, dtype=np.float32))
|
|
42
|
+
|
|
43
|
+
|
|
44
|
+
def linspace(start, stop, num=50) -> GPUArray:
|
|
45
|
+
return _from_numpy(np.linspace(start, stop, num, dtype=np.float32))
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
# ----- operations -----
|
|
49
|
+
|
|
50
|
+
def dot(a: GPUArray, b: GPUArray):
|
|
51
|
+
return a.dot(b)
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
|
|
55
|
+
return _matmul(a, b)
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def sum(a: GPUArray) -> float:
|
|
59
|
+
return a.sum()
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def mean(a: GPUArray) -> float:
|
|
63
|
+
return a.mean()
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def exp(a: GPUArray) -> GPUArray:
|
|
67
|
+
return a._run_unary_op(shaders.EXP_SHADER)
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def log(a: GPUArray) -> GPUArray:
|
|
71
|
+
return a._run_unary_op(shaders.LOG_SHADER)
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def sqrt(a: GPUArray) -> GPUArray:
|
|
75
|
+
return a._run_unary_op(shaders.SQRT_SHADER)
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def abs(a: GPUArray) -> GPUArray:
|
|
79
|
+
return a._run_unary_op(shaders.ABS_SHADER)
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
def max(a: GPUArray) -> float:
|
|
83
|
+
return a.max()
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def min(a: GPUArray) -> float:
|
|
87
|
+
return a.min()
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def relu(a: GPUArray) -> GPUArray:
|
|
91
|
+
return a._run_unary_op(shaders.RELU_SHADER)
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def sigmoid(a: GPUArray) -> GPUArray:
|
|
95
|
+
return a._run_unary_op(shaders.SIGMOID_SHADER)
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def tanh(a: GPUArray) -> GPUArray:
|
|
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()
|
|
@@ -343,6 +343,105 @@ class GPUArray:
|
|
|
343
343
|
raise ValueError(f"Cannot reshape array of size {self.size} into shape {new_shape}")
|
|
344
344
|
return GPUArray(self._buffer, new_shape)
|
|
345
345
|
|
|
346
|
+
# ----- clamp / clip -----
|
|
347
|
+
|
|
348
|
+
def clip(self, min_val: float, max_val: float) -> "GPUArray":
|
|
349
|
+
"""Clamp values to [min_val, max_val]."""
|
|
350
|
+
device = get_device()
|
|
351
|
+
_, pipeline = _get_pipeline(shaders.CLAMP_SHADER)
|
|
352
|
+
|
|
353
|
+
out_buf = device.create_buffer(
|
|
354
|
+
size=self.size * 4,
|
|
355
|
+
usage=wgpu.BufferUsage.STORAGE | wgpu.BufferUsage.COPY_SRC | wgpu.BufferUsage.COPY_DST,
|
|
356
|
+
)
|
|
357
|
+
# Pack two f32 values, padded to 16 bytes
|
|
358
|
+
params_data = struct.pack("ff", min_val, max_val) + b"\x00" * 8
|
|
359
|
+
uniform_buf = device.create_buffer_with_data(
|
|
360
|
+
data=params_data,
|
|
361
|
+
usage=wgpu.BufferUsage.UNIFORM,
|
|
362
|
+
)
|
|
363
|
+
bind_group = device.create_bind_group(
|
|
364
|
+
layout=pipeline.get_bind_group_layout(0),
|
|
365
|
+
entries=[
|
|
366
|
+
{"binding": 0, "resource": {"buffer": self._buffer, "offset": 0, "size": self.size * 4}},
|
|
367
|
+
{"binding": 1, "resource": {"buffer": out_buf, "offset": 0, "size": self.size * 4}},
|
|
368
|
+
{"binding": 2, "resource": {"buffer": uniform_buf, "offset": 0, "size": 16}},
|
|
369
|
+
],
|
|
370
|
+
)
|
|
371
|
+
n_workgroups = math.ceil(self.size / WORKGROUP_SIZE)
|
|
372
|
+
encoder = device.create_command_encoder()
|
|
373
|
+
pass_enc = encoder.begin_compute_pass()
|
|
374
|
+
pass_enc.set_pipeline(pipeline)
|
|
375
|
+
pass_enc.set_bind_group(0, bind_group)
|
|
376
|
+
pass_enc.dispatch_workgroups(n_workgroups)
|
|
377
|
+
pass_enc.end()
|
|
378
|
+
device.queue.submit([encoder.finish()])
|
|
379
|
+
uniform_buf.destroy()
|
|
380
|
+
return GPUArray(out_buf, self._shape)
|
|
381
|
+
|
|
382
|
+
# ----- argmax / argmin -----
|
|
383
|
+
|
|
384
|
+
def argmax(self) -> int:
|
|
385
|
+
"""Return index of maximum value (as Python int)."""
|
|
386
|
+
data = self.to_numpy().flatten()
|
|
387
|
+
return int(np.argmax(data))
|
|
388
|
+
|
|
389
|
+
def argmin(self) -> int:
|
|
390
|
+
"""Return index of minimum value (as Python int)."""
|
|
391
|
+
data = self.to_numpy().flatten()
|
|
392
|
+
return int(np.argmin(data))
|
|
393
|
+
|
|
394
|
+
# ----- copy -----
|
|
395
|
+
|
|
396
|
+
def copy(self) -> "GPUArray":
|
|
397
|
+
"""Return a copy of this array."""
|
|
398
|
+
return self._run_unary_op(shaders.COPY_SHADER)
|
|
399
|
+
|
|
400
|
+
# ----- comparison operators -----
|
|
401
|
+
|
|
402
|
+
def __eq__(self, other):
|
|
403
|
+
if isinstance(other, (int, float)):
|
|
404
|
+
other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
|
|
405
|
+
return self._run_binary_op(other, shaders.EQ_SHADER)
|
|
406
|
+
|
|
407
|
+
def __gt__(self, other):
|
|
408
|
+
if isinstance(other, (int, float)):
|
|
409
|
+
other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
|
|
410
|
+
return self._run_binary_op(other, shaders.GT_SHADER)
|
|
411
|
+
|
|
412
|
+
def __lt__(self, other):
|
|
413
|
+
if isinstance(other, (int, float)):
|
|
414
|
+
other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
|
|
415
|
+
return self._run_binary_op(other, shaders.LT_SHADER)
|
|
416
|
+
|
|
417
|
+
def __ge__(self, other):
|
|
418
|
+
if isinstance(other, (int, float)):
|
|
419
|
+
other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
|
|
420
|
+
return self._run_binary_op(other, shaders.GE_SHADER)
|
|
421
|
+
|
|
422
|
+
def __le__(self, other):
|
|
423
|
+
if isinstance(other, (int, float)):
|
|
424
|
+
other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
|
|
425
|
+
return self._run_binary_op(other, shaders.LE_SHADER)
|
|
426
|
+
|
|
427
|
+
# ----- power -----
|
|
428
|
+
|
|
429
|
+
def __pow__(self, other):
|
|
430
|
+
if isinstance(other, (int, float)):
|
|
431
|
+
other = _from_numpy(np.full(self.size, other, dtype=np.float32).reshape(self._shape))
|
|
432
|
+
return self._run_binary_op(other, shaders.POW_SHADER)
|
|
433
|
+
|
|
434
|
+
# ----- flatten -----
|
|
435
|
+
|
|
436
|
+
def flatten(self) -> "GPUArray":
|
|
437
|
+
"""Return a 1D view of the array."""
|
|
438
|
+
return GPUArray(self._buffer, (self.size,))
|
|
439
|
+
|
|
440
|
+
# ----- len -----
|
|
441
|
+
|
|
442
|
+
def __len__(self) -> int:
|
|
443
|
+
return self._shape[0]
|
|
444
|
+
|
|
346
445
|
|
|
347
446
|
def _from_numpy(data: np.ndarray) -> GPUArray:
|
|
348
447
|
"""Create a GPUArray from a numpy array."""
|
|
@@ -424,3 +424,177 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
|
424
424
|
c[row * dims.N + col] = sum;
|
|
425
425
|
}
|
|
426
426
|
"""
|
|
427
|
+
|
|
428
|
+
# ---------------------------------------------------------------------------
|
|
429
|
+
# Clamp (requires uniform buffer with min_val, max_val)
|
|
430
|
+
# ---------------------------------------------------------------------------
|
|
431
|
+
|
|
432
|
+
CLAMP_SHADER = """
|
|
433
|
+
struct Params {
|
|
434
|
+
min_val: f32,
|
|
435
|
+
max_val: f32,
|
|
436
|
+
}
|
|
437
|
+
|
|
438
|
+
@group(0) @binding(0) var<storage, read> a: array<f32>;
|
|
439
|
+
@group(0) @binding(1) var<storage, read_write> result: array<f32>;
|
|
440
|
+
@group(0) @binding(2) var<uniform> params: Params;
|
|
441
|
+
|
|
442
|
+
@compute @workgroup_size(256)
|
|
443
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
444
|
+
let i = gid.x;
|
|
445
|
+
if (i < arrayLength(&a)) {
|
|
446
|
+
result[i] = clamp(a[i], params.min_val, params.max_val);
|
|
447
|
+
}
|
|
448
|
+
}
|
|
449
|
+
"""
|
|
450
|
+
|
|
451
|
+
# ---------------------------------------------------------------------------
|
|
452
|
+
# Where / select: result = cond * a + (1-cond) * b
|
|
453
|
+
# ---------------------------------------------------------------------------
|
|
454
|
+
|
|
455
|
+
WHERE_SHADER = """
|
|
456
|
+
@group(0) @binding(0) var<storage, read> cond: array<f32>;
|
|
457
|
+
@group(0) @binding(1) var<storage, read> a: array<f32>;
|
|
458
|
+
@group(0) @binding(2) var<storage, read> b: array<f32>;
|
|
459
|
+
@group(0) @binding(3) var<storage, read_write> result: array<f32>;
|
|
460
|
+
|
|
461
|
+
@compute @workgroup_size(256)
|
|
462
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
463
|
+
let i = gid.x;
|
|
464
|
+
if (i < arrayLength(&cond)) {
|
|
465
|
+
result[i] = cond[i] * a[i] + (1.0 - cond[i]) * b[i];
|
|
466
|
+
}
|
|
467
|
+
}
|
|
468
|
+
"""
|
|
469
|
+
|
|
470
|
+
# ---------------------------------------------------------------------------
|
|
471
|
+
# Comparison shaders (return 0.0 or 1.0)
|
|
472
|
+
# ---------------------------------------------------------------------------
|
|
473
|
+
|
|
474
|
+
EQ_SHADER = """
|
|
475
|
+
@group(0) @binding(0) var<storage, read> a: array<f32>;
|
|
476
|
+
@group(0) @binding(1) var<storage, read> b: array<f32>;
|
|
477
|
+
@group(0) @binding(2) var<storage, read_write> result: array<f32>;
|
|
478
|
+
|
|
479
|
+
@compute @workgroup_size(256)
|
|
480
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
481
|
+
let i = gid.x;
|
|
482
|
+
if (i < arrayLength(&a)) {
|
|
483
|
+
if (a[i] == b[i]) {
|
|
484
|
+
result[i] = 1.0;
|
|
485
|
+
} else {
|
|
486
|
+
result[i] = 0.0;
|
|
487
|
+
}
|
|
488
|
+
}
|
|
489
|
+
}
|
|
490
|
+
"""
|
|
491
|
+
|
|
492
|
+
GT_SHADER = """
|
|
493
|
+
@group(0) @binding(0) var<storage, read> a: array<f32>;
|
|
494
|
+
@group(0) @binding(1) var<storage, read> b: array<f32>;
|
|
495
|
+
@group(0) @binding(2) var<storage, read_write> result: array<f32>;
|
|
496
|
+
|
|
497
|
+
@compute @workgroup_size(256)
|
|
498
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
499
|
+
let i = gid.x;
|
|
500
|
+
if (i < arrayLength(&a)) {
|
|
501
|
+
if (a[i] > b[i]) {
|
|
502
|
+
result[i] = 1.0;
|
|
503
|
+
} else {
|
|
504
|
+
result[i] = 0.0;
|
|
505
|
+
}
|
|
506
|
+
}
|
|
507
|
+
}
|
|
508
|
+
"""
|
|
509
|
+
|
|
510
|
+
LT_SHADER = """
|
|
511
|
+
@group(0) @binding(0) var<storage, read> a: array<f32>;
|
|
512
|
+
@group(0) @binding(1) var<storage, read> b: array<f32>;
|
|
513
|
+
@group(0) @binding(2) var<storage, read_write> result: array<f32>;
|
|
514
|
+
|
|
515
|
+
@compute @workgroup_size(256)
|
|
516
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
517
|
+
let i = gid.x;
|
|
518
|
+
if (i < arrayLength(&a)) {
|
|
519
|
+
if (a[i] < b[i]) {
|
|
520
|
+
result[i] = 1.0;
|
|
521
|
+
} else {
|
|
522
|
+
result[i] = 0.0;
|
|
523
|
+
}
|
|
524
|
+
}
|
|
525
|
+
}
|
|
526
|
+
"""
|
|
527
|
+
|
|
528
|
+
GE_SHADER = """
|
|
529
|
+
@group(0) @binding(0) var<storage, read> a: array<f32>;
|
|
530
|
+
@group(0) @binding(1) var<storage, read> b: array<f32>;
|
|
531
|
+
@group(0) @binding(2) var<storage, read_write> result: array<f32>;
|
|
532
|
+
|
|
533
|
+
@compute @workgroup_size(256)
|
|
534
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
535
|
+
let i = gid.x;
|
|
536
|
+
if (i < arrayLength(&a)) {
|
|
537
|
+
if (a[i] >= b[i]) {
|
|
538
|
+
result[i] = 1.0;
|
|
539
|
+
} else {
|
|
540
|
+
result[i] = 0.0;
|
|
541
|
+
}
|
|
542
|
+
}
|
|
543
|
+
}
|
|
544
|
+
"""
|
|
545
|
+
|
|
546
|
+
LE_SHADER = """
|
|
547
|
+
@group(0) @binding(0) var<storage, read> a: array<f32>;
|
|
548
|
+
@group(0) @binding(1) var<storage, read> b: array<f32>;
|
|
549
|
+
@group(0) @binding(2) var<storage, read_write> result: array<f32>;
|
|
550
|
+
|
|
551
|
+
@compute @workgroup_size(256)
|
|
552
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
553
|
+
let i = gid.x;
|
|
554
|
+
if (i < arrayLength(&a)) {
|
|
555
|
+
if (a[i] <= b[i]) {
|
|
556
|
+
result[i] = 1.0;
|
|
557
|
+
} else {
|
|
558
|
+
result[i] = 0.0;
|
|
559
|
+
}
|
|
560
|
+
}
|
|
561
|
+
}
|
|
562
|
+
"""
|
|
563
|
+
|
|
564
|
+
# ---------------------------------------------------------------------------
|
|
565
|
+
# Fill (fill array with uniform scalar)
|
|
566
|
+
# ---------------------------------------------------------------------------
|
|
567
|
+
|
|
568
|
+
FILL_SHADER = """
|
|
569
|
+
struct Params {
|
|
570
|
+
scalar: f32,
|
|
571
|
+
}
|
|
572
|
+
|
|
573
|
+
@group(0) @binding(0) var<storage, read_write> result: array<f32>;
|
|
574
|
+
@group(0) @binding(1) var<uniform> params: Params;
|
|
575
|
+
|
|
576
|
+
@compute @workgroup_size(256)
|
|
577
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
578
|
+
let i = gid.x;
|
|
579
|
+
if (i < arrayLength(&result)) {
|
|
580
|
+
result[i] = params.scalar;
|
|
581
|
+
}
|
|
582
|
+
}
|
|
583
|
+
"""
|
|
584
|
+
|
|
585
|
+
# ---------------------------------------------------------------------------
|
|
586
|
+
# Copy
|
|
587
|
+
# ---------------------------------------------------------------------------
|
|
588
|
+
|
|
589
|
+
COPY_SHADER = """
|
|
590
|
+
@group(0) @binding(0) var<storage, read> a: array<f32>;
|
|
591
|
+
@group(0) @binding(1) var<storage, read_write> result: array<f32>;
|
|
592
|
+
|
|
593
|
+
@compute @workgroup_size(256)
|
|
594
|
+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
|
|
595
|
+
let i = gid.x;
|
|
596
|
+
if (i < arrayLength(&a)) {
|
|
597
|
+
result[i] = a[i];
|
|
598
|
+
}
|
|
599
|
+
}
|
|
600
|
+
"""
|
|
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "gpuarray"
|
|
7
|
-
version = "0.
|
|
7
|
+
version = "0.2.0"
|
|
8
8
|
description = "NumPy-like GPU array library that works on every GPU — NVIDIA, AMD, Intel, Apple Silicon. No CUDA required."
|
|
9
9
|
requires-python = ">=3.10,<3.14"
|
|
10
10
|
license = {text = "MIT"}
|
|
@@ -225,3 +225,211 @@ class TestMisc:
|
|
|
225
225
|
def test_getitem(self):
|
|
226
226
|
a = gp.array([10, 20, 30])
|
|
227
227
|
assert a[1] == 20.0
|
|
228
|
+
|
|
229
|
+
|
|
230
|
+
class TestClip:
|
|
231
|
+
def test_clip_basic(self):
|
|
232
|
+
a_np = np.array([-2, -1, 0, 1, 2, 3, 4, 5], dtype=np.float32)
|
|
233
|
+
a = gp.array(a_np)
|
|
234
|
+
result = a.clip(0.0, 3.0)
|
|
235
|
+
expected = np.clip(a_np, 0.0, 3.0)
|
|
236
|
+
npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
|
|
237
|
+
|
|
238
|
+
def test_clip_api(self):
|
|
239
|
+
a_np = np.array([1, 5, 10, -3], dtype=np.float32)
|
|
240
|
+
a = gp.array(a_np)
|
|
241
|
+
result = gp.clip(a, 0.0, 7.0)
|
|
242
|
+
expected = np.clip(a_np, 0.0, 7.0)
|
|
243
|
+
npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
|
|
244
|
+
|
|
245
|
+
|
|
246
|
+
class TestArgmaxArgmin:
|
|
247
|
+
def test_argmax(self):
|
|
248
|
+
a_np = np.array([1, 5, 3, 9, 2], dtype=np.float32)
|
|
249
|
+
a = gp.array(a_np)
|
|
250
|
+
assert a.argmax() == int(np.argmax(a_np))
|
|
251
|
+
|
|
252
|
+
def test_argmin(self):
|
|
253
|
+
a_np = np.array([3, 1, 4, 0, 5], dtype=np.float32)
|
|
254
|
+
a = gp.array(a_np)
|
|
255
|
+
assert a.argmin() == int(np.argmin(a_np))
|
|
256
|
+
|
|
257
|
+
def test_argmax_api(self):
|
|
258
|
+
a_np = np.array([10, 20, 5], dtype=np.float32)
|
|
259
|
+
a = gp.array(a_np)
|
|
260
|
+
assert gp.argmax(a) == 1
|
|
261
|
+
|
|
262
|
+
def test_argmin_api(self):
|
|
263
|
+
a_np = np.array([10, 20, 5], dtype=np.float32)
|
|
264
|
+
a = gp.array(a_np)
|
|
265
|
+
assert gp.argmin(a) == 2
|
|
266
|
+
|
|
267
|
+
|
|
268
|
+
class TestComparisons:
|
|
269
|
+
def test_eq(self):
|
|
270
|
+
a = gp.array([1.0, 2.0, 3.0])
|
|
271
|
+
b = gp.array([1.0, 5.0, 3.0])
|
|
272
|
+
result = (a == b).to_numpy()
|
|
273
|
+
expected = np.array([1.0, 0.0, 1.0], dtype=np.float32)
|
|
274
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
275
|
+
|
|
276
|
+
def test_gt(self):
|
|
277
|
+
a = gp.array([1.0, 5.0, 3.0])
|
|
278
|
+
b = gp.array([2.0, 3.0, 3.0])
|
|
279
|
+
result = (a > b).to_numpy()
|
|
280
|
+
expected = np.array([0.0, 1.0, 0.0], dtype=np.float32)
|
|
281
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
282
|
+
|
|
283
|
+
def test_lt(self):
|
|
284
|
+
a = gp.array([1.0, 5.0, 3.0])
|
|
285
|
+
b = gp.array([2.0, 3.0, 3.0])
|
|
286
|
+
result = (a < b).to_numpy()
|
|
287
|
+
expected = np.array([1.0, 0.0, 0.0], dtype=np.float32)
|
|
288
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
289
|
+
|
|
290
|
+
def test_ge(self):
|
|
291
|
+
a = gp.array([1.0, 5.0, 3.0])
|
|
292
|
+
b = gp.array([2.0, 3.0, 3.0])
|
|
293
|
+
result = (a >= b).to_numpy()
|
|
294
|
+
expected = np.array([0.0, 1.0, 1.0], dtype=np.float32)
|
|
295
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
296
|
+
|
|
297
|
+
def test_le(self):
|
|
298
|
+
a = gp.array([1.0, 5.0, 3.0])
|
|
299
|
+
b = gp.array([2.0, 3.0, 3.0])
|
|
300
|
+
result = (a <= b).to_numpy()
|
|
301
|
+
expected = np.array([1.0, 0.0, 1.0], dtype=np.float32)
|
|
302
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
303
|
+
|
|
304
|
+
def test_eq_scalar(self):
|
|
305
|
+
a = gp.array([1.0, 2.0, 3.0])
|
|
306
|
+
result = (a == 2.0).to_numpy()
|
|
307
|
+
expected = np.array([0.0, 1.0, 0.0], dtype=np.float32)
|
|
308
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
309
|
+
|
|
310
|
+
def test_gt_scalar(self):
|
|
311
|
+
a = gp.array([1.0, 2.0, 3.0])
|
|
312
|
+
result = (a > 1.5).to_numpy()
|
|
313
|
+
expected = np.array([0.0, 1.0, 1.0], dtype=np.float32)
|
|
314
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
315
|
+
|
|
316
|
+
|
|
317
|
+
class TestWhere:
|
|
318
|
+
def test_where_basic(self):
|
|
319
|
+
cond = gp.array([1.0, 0.0, 1.0, 0.0])
|
|
320
|
+
x = gp.array([10.0, 20.0, 30.0, 40.0])
|
|
321
|
+
y = gp.array([50.0, 60.0, 70.0, 80.0])
|
|
322
|
+
result = gp.where(cond, x, y)
|
|
323
|
+
expected = np.array([10.0, 60.0, 30.0, 80.0], dtype=np.float32)
|
|
324
|
+
npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
|
|
325
|
+
|
|
326
|
+
def test_where_from_comparison(self):
|
|
327
|
+
a = gp.array([1.0, 5.0, 3.0, 8.0])
|
|
328
|
+
cond = a > 3.0
|
|
329
|
+
x = gp.array([100.0, 100.0, 100.0, 100.0])
|
|
330
|
+
y = gp.array([0.0, 0.0, 0.0, 0.0])
|
|
331
|
+
result = gp.where(cond, x, y)
|
|
332
|
+
expected = np.array([0.0, 100.0, 0.0, 100.0], dtype=np.float32)
|
|
333
|
+
npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
|
|
334
|
+
|
|
335
|
+
|
|
336
|
+
class TestConcatenate:
|
|
337
|
+
def test_concat_1d(self):
|
|
338
|
+
a = gp.array([1.0, 2.0, 3.0])
|
|
339
|
+
b = gp.array([4.0, 5.0, 6.0])
|
|
340
|
+
result = gp.concatenate([a, b])
|
|
341
|
+
expected = np.array([1, 2, 3, 4, 5, 6], dtype=np.float32)
|
|
342
|
+
npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
|
|
343
|
+
assert result.shape == (6,)
|
|
344
|
+
|
|
345
|
+
def test_concat_2d(self):
|
|
346
|
+
a = gp.array([[1.0, 2.0], [3.0, 4.0]])
|
|
347
|
+
b = gp.array([[5.0, 6.0]])
|
|
348
|
+
result = gp.concatenate([a, b], axis=0)
|
|
349
|
+
expected = np.array([[1, 2], [3, 4], [5, 6]], dtype=np.float32)
|
|
350
|
+
npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
|
|
351
|
+
assert result.shape == (3, 2)
|
|
352
|
+
|
|
353
|
+
def test_stack_1d(self):
|
|
354
|
+
a = gp.array([1.0, 2.0, 3.0])
|
|
355
|
+
b = gp.array([4.0, 5.0, 6.0])
|
|
356
|
+
result = gp.stack([a, b])
|
|
357
|
+
expected = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
|
|
358
|
+
npt.assert_allclose(result.to_numpy(), expected, rtol=RTOL)
|
|
359
|
+
assert result.shape == (2, 3)
|
|
360
|
+
|
|
361
|
+
|
|
362
|
+
class TestPower:
|
|
363
|
+
def test_pow_array(self):
|
|
364
|
+
a_np = np.array([2.0, 3.0, 4.0], dtype=np.float32)
|
|
365
|
+
b_np = np.array([2.0, 2.0, 3.0], dtype=np.float32)
|
|
366
|
+
a = gp.array(a_np)
|
|
367
|
+
b = gp.array(b_np)
|
|
368
|
+
result = (a ** b).to_numpy()
|
|
369
|
+
expected = a_np ** b_np
|
|
370
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
371
|
+
|
|
372
|
+
def test_pow_scalar(self):
|
|
373
|
+
a_np = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32)
|
|
374
|
+
a = gp.array(a_np)
|
|
375
|
+
result = (a ** 2).to_numpy()
|
|
376
|
+
expected = a_np ** 2
|
|
377
|
+
npt.assert_allclose(result, expected, rtol=RTOL)
|
|
378
|
+
|
|
379
|
+
|
|
380
|
+
class TestCopy:
|
|
381
|
+
def test_copy(self):
|
|
382
|
+
a = gp.array([1.0, 2.0, 3.0])
|
|
383
|
+
b = a.copy()
|
|
384
|
+
npt.assert_allclose(b.to_numpy(), a.to_numpy(), rtol=RTOL)
|
|
385
|
+
assert b._buffer is not a._buffer
|
|
386
|
+
|
|
387
|
+
def test_copy_shape(self):
|
|
388
|
+
a = gp.array([[1.0, 2.0], [3.0, 4.0]])
|
|
389
|
+
b = a.copy()
|
|
390
|
+
assert b.shape == a.shape
|
|
391
|
+
npt.assert_allclose(b.to_numpy(), a.to_numpy(), rtol=RTOL)
|
|
392
|
+
|
|
393
|
+
|
|
394
|
+
class TestRandom:
|
|
395
|
+
def test_rand_shape(self):
|
|
396
|
+
a = gp.random.rand(5)
|
|
397
|
+
assert a.shape == (5,)
|
|
398
|
+
data = a.to_numpy()
|
|
399
|
+
assert data.min() >= 0.0
|
|
400
|
+
assert data.max() < 1.0
|
|
401
|
+
|
|
402
|
+
def test_rand_2d(self):
|
|
403
|
+
a = gp.random.rand(3, 4)
|
|
404
|
+
assert a.shape == (3, 4)
|
|
405
|
+
assert a.size == 12
|
|
406
|
+
|
|
407
|
+
def test_randn_shape(self):
|
|
408
|
+
a = gp.random.randn(100)
|
|
409
|
+
assert a.shape == (100,)
|
|
410
|
+
data = a.to_numpy()
|
|
411
|
+
# randn should have mean near 0 and std near 1 for large samples
|
|
412
|
+
assert abs(data.mean()) < 0.5 # loose check
|
|
413
|
+
|
|
414
|
+
def test_randn_2d(self):
|
|
415
|
+
a = gp.random.randn(4, 5)
|
|
416
|
+
assert a.shape == (4, 5)
|
|
417
|
+
assert a.size == 20
|
|
418
|
+
|
|
419
|
+
|
|
420
|
+
class TestFlatten:
|
|
421
|
+
def test_flatten_2d(self):
|
|
422
|
+
a = gp.array([[1.0, 2.0], [3.0, 4.0]])
|
|
423
|
+
b = a.flatten()
|
|
424
|
+
assert b.shape == (4,)
|
|
425
|
+
npt.assert_allclose(b.to_numpy(), np.array([1, 2, 3, 4], dtype=np.float32), rtol=RTOL)
|
|
426
|
+
|
|
427
|
+
|
|
428
|
+
class TestLen:
|
|
429
|
+
def test_len_1d(self):
|
|
430
|
+
a = gp.array([1.0, 2.0, 3.0])
|
|
431
|
+
assert len(a) == 3
|
|
432
|
+
|
|
433
|
+
def test_len_2d(self):
|
|
434
|
+
a = gp.array([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
|
|
435
|
+
assert len(a) == 3
|
gpuarray-0.1.0/gpuarray/api.py
DELETED
|
@@ -1,99 +0,0 @@
|
|
|
1
|
-
"""NumPy-compatible module-level functions."""
|
|
2
|
-
|
|
3
|
-
from __future__ import annotations
|
|
4
|
-
|
|
5
|
-
from typing import Sequence
|
|
6
|
-
|
|
7
|
-
import numpy as np
|
|
8
|
-
|
|
9
|
-
from .core import GPUArray, _from_numpy, matmul as _matmul
|
|
10
|
-
from . import shaders
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
# ----- creation functions -----
|
|
14
|
-
|
|
15
|
-
def array(data) -> GPUArray:
|
|
16
|
-
"""Create a GPUArray from a list, nested list, or numpy array."""
|
|
17
|
-
if isinstance(data, GPUArray):
|
|
18
|
-
return data
|
|
19
|
-
arr = np.asarray(data, dtype=np.float32)
|
|
20
|
-
if arr.ndim == 0:
|
|
21
|
-
arr = arr.reshape((1,))
|
|
22
|
-
return _from_numpy(arr)
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
def zeros(shape) -> GPUArray:
|
|
26
|
-
if isinstance(shape, int):
|
|
27
|
-
shape = (shape,)
|
|
28
|
-
return _from_numpy(np.zeros(shape, dtype=np.float32))
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
def ones(shape) -> GPUArray:
|
|
32
|
-
if isinstance(shape, int):
|
|
33
|
-
shape = (shape,)
|
|
34
|
-
return _from_numpy(np.ones(shape, dtype=np.float32))
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
def arange(start, stop=None, step=1) -> GPUArray:
|
|
38
|
-
if stop is None:
|
|
39
|
-
stop = start
|
|
40
|
-
start = 0
|
|
41
|
-
return _from_numpy(np.arange(start, stop, step, dtype=np.float32))
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
def linspace(start, stop, num=50) -> GPUArray:
|
|
45
|
-
return _from_numpy(np.linspace(start, stop, num, dtype=np.float32))
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
# ----- operations -----
|
|
49
|
-
|
|
50
|
-
def dot(a: GPUArray, b: GPUArray):
|
|
51
|
-
return a.dot(b)
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
def matmul(a: GPUArray, b: GPUArray) -> GPUArray:
|
|
55
|
-
return _matmul(a, b)
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
def sum(a: GPUArray) -> float:
|
|
59
|
-
return a.sum()
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
def mean(a: GPUArray) -> float:
|
|
63
|
-
return a.mean()
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
def exp(a: GPUArray) -> GPUArray:
|
|
67
|
-
return a._run_unary_op(shaders.EXP_SHADER)
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
def log(a: GPUArray) -> GPUArray:
|
|
71
|
-
return a._run_unary_op(shaders.LOG_SHADER)
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
def sqrt(a: GPUArray) -> GPUArray:
|
|
75
|
-
return a._run_unary_op(shaders.SQRT_SHADER)
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
def abs(a: GPUArray) -> GPUArray:
|
|
79
|
-
return a._run_unary_op(shaders.ABS_SHADER)
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
def max(a: GPUArray) -> float:
|
|
83
|
-
return a.max()
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
def min(a: GPUArray) -> float:
|
|
87
|
-
return a.min()
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
def relu(a: GPUArray) -> GPUArray:
|
|
91
|
-
return a._run_unary_op(shaders.RELU_SHADER)
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
def sigmoid(a: GPUArray) -> GPUArray:
|
|
95
|
-
return a._run_unary_op(shaders.SIGMOID_SHADER)
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
def tanh(a: GPUArray) -> GPUArray:
|
|
99
|
-
return a._run_unary_op(shaders.TANH_SHADER)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|