burn_cubecl/kernel/mask/
mask_where.rs

1use burn_backend::DType;
2use burn_backend::cubecl::dtype_to_storage_type;
3use cubecl::{
4    calculate_cube_count_elemwise,
5    prelude::*,
6    std::tensor::layout::linear::{LinearView, LinearViewMut},
7};
8
9use crate::{
10    kernel::utils::{address_type, broadcast_shape},
11    ops::{max_vector_size_many, numeric::empty_device_dtype},
12    tensor::CubeTensor,
13};
14
15#[cube(launch, address_type = "dynamic")]
16fn mask_where_kernel<T: Numeric, B: Int, N: Size>(
17    input: LinearView<'_, Vector<T, N>>,
18    value: LinearView<'_, Vector<T, N>>,
19    mask: LinearView<'_, Vector<B, N>>,
20    mut output: LinearViewMut<'_, Vector<T, N>>,
21    #[define(T, B)] _dtypes: [ElemType; 2],
22) {
23    let pos = ABSOLUTE_POS;
24    if !output.is_in_bounds(pos) {
25        terminate!();
26    }
27
28    output.write(
29        pos,
30        select_many(
31            Vector::cast_from(mask.read(pos)),
32            value.read(pos),
33            input.read(pos),
34        ),
35    );
36}
37
38#[derive(Clone, Copy, Debug)]
39/// Define how to run the mask where kernel.
40///
41/// # Notes
42///
43/// All assertions should be done before choosing the strategy.
44pub enum MaskWhereStrategy {
45    /// Don't mutate any input.
46    Readonly,
47    /// Reuse the lhs tensor inplace.
48    InplaceLhs,
49    /// Reuse the rhs tensor inplace.
50    InplaceRhs,
51}
52
53/// Execute the mask where kernel with the given strategy.
54pub fn mask_where(
55    input: CubeTensor,
56    mask: CubeTensor,
57    value: CubeTensor,
58    strategy: MaskWhereStrategy,
59    dtype_bool: DType,
60) -> CubeTensor {
61    let out_shape = broadcast_shape(&[&input, &mask, &value]);
62
63    // A zero-sized broadcast output has no elements to compute, and the strategies below assume a
64    // non-empty output. Return the empty output directly.
65    if out_shape.num_elements() == 0 {
66        return empty_device_dtype(
67            input.client.clone(),
68            input.device.clone(),
69            out_shape,
70            input.dtype,
71        );
72    }
73
74    let vector_size = max_vector_size_many(&[&input, &mask, &value], input.meta.num_dims() - 1);
75    let working_units = out_shape.num_elements() / vector_size as usize;
76    let cube_dim = CubeDim::new(&input.client, working_units);
77    let cube_count = calculate_cube_count_elemwise(&input.client, working_units, cube_dim);
78
79    let output = match strategy {
80        MaskWhereStrategy::Readonly => empty_device_dtype(
81            input.client.clone(),
82            input.device.clone(),
83            out_shape,
84            input.dtype,
85        ),
86        MaskWhereStrategy::InplaceLhs => input.clone(),
87        MaskWhereStrategy::InplaceRhs => value.clone(),
88    };
89
90    let out = match strategy {
91        MaskWhereStrategy::Readonly => output.clone().into_linear_view(),
92        MaskWhereStrategy::InplaceLhs => output.as_linear_view_alias(0),
93        MaskWhereStrategy::InplaceRhs => output.as_linear_view_alias(1),
94    };
95
96    mask_where_kernel::launch(
97        &output.client,
98        cube_count,
99        cube_dim,
100        address_type!(input, value, mask, output),
101        vector_size,
102        input.into_linear_view_like(&output),
103        value.into_linear_view_like(&output),
104        mask.into_linear_view_like(&output),
105        out,
106        [
107            dtype_to_storage_type(output.dtype),
108            dtype_to_storage_type(dtype_bool),
109        ],
110    );
111
112    output
113}
burn_cubecl/kernel/mask/mask_where.rs

burn_cubecl/kernel/mask/
mask_where.rs
