cubek-reduce 0.2.0-pre.2

CubeK: Reduce Kernels
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

use cubecl::prelude::*;

// Using plane operations, return the lowest coordinate for each line element
// for which the item equal the target.
#[cube]
pub(crate) fn lowest_coordinate_matching<E: CubePrimitive>(
    target: Line<E>,
    item: Line<E>,
    coordinate: Line<u32>,
) -> Line<u32> {
    let line_size = item.size();
    let is_candidate = item.equal(target);
    let candidate_coordinate = select_many(
        is_candidate,
        coordinate,
        Line::empty(line_size).fill(u32::MAX),
    );
    plane_min(candidate_coordinate)
}