burn_cubecl/kernel/mask/
mask_fill.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_unchecked, address_type = "dynamic")]
16fn mask_fill_kernel<T: Numeric, B: Int, N: Size>(
17    input: LinearView<'_, Vector<T, N>>,
18    mask: LinearView<'_, Vector<B, N>>,
19    mut output: LinearViewMut<'_, Vector<T, N>>,
20    value: InputScalar,
21    #[define(T, B)] _dtypes: [ElemType; 2],
22) {
23    if !output.is_in_bounds(ABSOLUTE_POS) {
24        terminate!();
25    }
26
27    let mask = Vector::cast_from(mask.read(ABSOLUTE_POS));
28    let input = input.read(ABSOLUTE_POS);
29    let value = Vector::new(value.get::<T>());
30
31    output.write(ABSOLUTE_POS, select_many(mask, value, input));
32}
33
34#[derive(Clone, Copy, Debug)]
35/// Define how to run the mask fill kernel.
36///
37/// # Notes
38///
39/// All assertions should be done before choosing the strategy.
40pub enum MaskFillStrategy {
41    /// Don't mutate any input.
42    Readonly,
43    /// Reuse the input tensor inplace.
44    Inplace,
45}
46
47/// Execute the mask fill kernel with the given strategy.
48pub fn mask_fill(
49    input: CubeTensor,
50    mask: CubeTensor,
51    value: InputScalar,
52    strategy: MaskFillStrategy,
53    dtype_bool: DType,
54) -> CubeTensor {
55    let ndims = input.meta.num_dims();
56    let out_shape = broadcast_shape(&[&input, &mask]);
57    if out_shape.num_elements() == 0 {
58        return empty_device_dtype(
59            input.client.clone(),
60            input.device.clone(),
61            out_shape,
62            input.dtype,
63        );
64    }
65
66    let vector_size = max_vector_size_many(&[&input, &mask], ndims - 1);
67    let working_units = out_shape.num_elements() / vector_size as usize;
68    let cube_dim = CubeDim::new(&input.client, working_units);
69    let cube_count = calculate_cube_count_elemwise(&input.client, working_units, cube_dim);
70
71    let output = match strategy {
72        MaskFillStrategy::Readonly => empty_device_dtype(
73            input.client.clone(),
74            input.device.clone(),
75            out_shape,
76            input.dtype,
77        ),
78        MaskFillStrategy::Inplace => input.clone(),
79    };
80
81    let out_arg = match strategy {
82        MaskFillStrategy::Readonly => output.clone().into_linear_view(),
83        MaskFillStrategy::Inplace => output.as_linear_view_alias(0),
84    };
85
86    let at = address_type!(input, mask, output);
87    let mask = mask.into_linear_view_like(&output);
88
89    unsafe {
90        mask_fill_kernel::launch_unchecked(
91            &output.client,
92            cube_count,
93            cube_dim,
94            at,
95            vector_size,
96            input.into_linear_view_like(&output),
97            mask,
98            out_arg,
99            value,
100            [
101                dtype_to_storage_type(output.dtype),
102                dtype_to_storage_type(dtype_bool),
103            ],
104        );
105    }
106
107    output
108}
burn_cubecl/kernel/mask/mask_fill.rs

burn_cubecl/kernel/mask/
mask_fill.rs
