Skip to content

Commit

Permalink
Update documentation
Browse files Browse the repository at this point in the history
  • Loading branch information
shwina committed Jan 10, 2025
1 parent 6b0bfa1 commit 210dc51
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 0 deletions.
6 changes: 6 additions & 0 deletions docs/cuda_parallel/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -22,3 +22,9 @@ Iterators
:members:
:undoc-members:
:imported-members:

Utilities
---------

.. automodule:: cuda.parallel.experimental.struct
:members:
32 changes: 32 additions & 0 deletions python/cuda_parallel/tests/test_reduce_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -178,3 +178,35 @@ def square_op(a):
)
assert (d_output == expected_output).all()
# example-end transform-iterator


def test_reduce_struct_type():
from cuda.parallel.experimental.struct import gpu_struct

# example-begin reduce-struct
@gpu_struct
class Pixel:
r: np.int32
g: np.int32
b: np.int32

def max_g_value(x, y):
return x if x.g > y.g else y

d_rgb = cp.random.randint(0, 256, (10, 3), dtype=np.int32).view(Pixel.dtype)
d_out = cp.zeros(1, Pixel.dtype)

h_init = Pixel(0, 0, 0)

reduce_into = algorithms.reduce_into(d_rgb, d_out, max_g_value, h_init)
temp_storage_bytes = reduce_into(None, d_rgb, d_out, len(d_rgb), h_init)

d_temp_storage = cp.zeros(temp_storage_bytes, dtype=np.uint8)
_ = reduce_into(d_temp_storage, d_rgb, d_out, len(d_rgb), h_init)

h_rgb = d_rgb.get()
expected = h_rgb[h_rgb.view("int32")[:, 1].argmax()]

np.testing.assert_equal(expected["g"], d_out.get()["g"])
# example-end reduce-struct

0 comments on commit 210dc51

Please sign in to comment.