Skip to content

Commit 06c924d

Browse files
committed
Update documentation
1 parent 0a6b9dd commit 06c924d

File tree

2 files changed

+37
-0
lines changed

2 files changed

+37
-0
lines changed

docs/cuda_parallel/index.rst

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,3 +22,9 @@ Iterators
2222
:members:
2323
:undoc-members:
2424
:imported-members:
25+
26+
Utilities
27+
---------
28+
29+
.. automodule:: cuda.parallel.experimental.struct
30+
:members:

python/cuda_parallel/tests/test_reduce_api.py

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,3 +37,34 @@ def min_op(a, b):
3737
expected_output = 0
3838
assert (d_output == expected_output).all()
3939
# example-end reduce-min
40+
41+
42+
def test_reduce_struct_type():
43+
from cuda.parallel.experimental.struct import gpu_struct
44+
45+
# example-begin reduce-struct
46+
@gpu_struct
47+
class Pixel:
48+
r: np.int32
49+
g: np.int32
50+
b: np.int32
51+
52+
def max_g_value(x, y):
53+
return x if x.g > y.g else y
54+
55+
d_rgb = cp.random.randint(0, 256, (10, 3), dtype=np.int32).view(Pixel.dtype)
56+
d_out = cp.zeros(1, Pixel.dtype)
57+
58+
h_init = Pixel(0, 0, 0)
59+
60+
reduce_into = algorithms.reduce_into(d_rgb, d_out, max_g_value, h_init)
61+
temp_storage_bytes = reduce_into(None, d_rgb, d_out, len(d_rgb), h_init)
62+
63+
d_temp_storage = cp.zeros(temp_storage_bytes, dtype=np.uint8)
64+
_ = reduce_into(d_temp_storage, d_rgb, d_out, len(d_rgb), h_init)
65+
66+
h_rgb = d_rgb.get()
67+
expected = h_rgb[h_rgb.view("int32")[:, 1].argmax()]
68+
69+
np.testing.assert_equal(expected["g"], d_out.get()["g"])
70+
# example-end reduce-struct

0 commit comments

Comments
 (0)