burn-cubecl 0.22.0-pre.4

Generic backend that can be compiled just-in-time to any shader language target
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39

use burn_backend::DType;
use cubecl::prelude::InputScalar;

use super::{MaskFillStrategy, mask_where::MaskWhereStrategy};
use crate::tensor::CubeTensor;

/// Execute the mask fill kernel.
pub(crate) fn mask_fill_auto(
    tensor: CubeTensor,
    mask: CubeTensor,
    value: InputScalar,
    dtype_bool: DType,
) -> CubeTensor {
    let strategy = if tensor.can_mut_broadcast(&mask) {
        MaskFillStrategy::Inplace
    } else {
        MaskFillStrategy::Readonly
    };

    super::mask_fill(tensor, mask, value, strategy, dtype_bool)
}

/// Execute the mask where kernel.
pub(crate) fn mask_where_auto(
    tensor: CubeTensor,
    mask: CubeTensor,
    value: CubeTensor,
    dtype_bool: DType,
) -> CubeTensor {
    let strategy = if tensor.can_mut_broadcast(&value) && tensor.can_mut_broadcast(&mask) {
        MaskWhereStrategy::InplaceLhs
    } else if value.can_mut_broadcast(&tensor) && value.can_mut_broadcast(&mask) {
        MaskWhereStrategy::InplaceRhs
    } else {
        MaskWhereStrategy::Readonly
    };

    super::mask_where(tensor, mask, value, strategy, dtype_bool)
}