@@ -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