burn_cubecl/ops/
bool_tensor.rs

1use crate::{
2    CubeBackend,
3    element::BoolElement,
4    kernel::{self, AndOp, OrOp},
5    tensor::CubeTensor,
6};
7use burn_backend::cubecl::dtype_to_storage_type;
8use burn_backend::{
9    ExecutionError, Slice,
10    ops::BoolTensorOps,
11    tensor::{BoolTensor, Device, FloatTensor, IntTensor},
12};
13use burn_backend::{Scalar, Shape, TensorData};
14use burn_std::{BoolDType, BoolStore, DType, FloatDType, IntDType};
15use cubecl::prelude::InputScalar;
16use cubek::reduce::components::instructions::ReduceOperationConfig;
17use std::ops::Range;
18
19use super::{expand, numeric, permute, unfold};
20
21/// The boolean storage of a cubecl bool tensor. Cubecl backends never use
22/// native bool (see `CubeBackend::supports_dtype`), so it is always `U8`/`U32`.
23fn bool_store(tensor: &CubeTensor) -> BoolDType {
24    match tensor.dtype {
25        DType::Bool(store) => store,
26        other => unreachable!("cubecl bool tensors are always Bool(_): {other:?}"),
27    }
28}
29
30impl BoolTensorOps<Self> for CubeBackend {
31    fn bool_empty(shape: Shape, device: &Device<Self>, dtype: BoolDType) -> BoolTensor<Self> {
32        super::empty(shape, device, dtype.into())
33    }
34
35    fn bool_zeros(shape: Shape, device: &Device<Self>, dtype: BoolDType) -> BoolTensor<Self> {
36        numeric::zeros(device.clone(), shape, dtype.into())
37    }
38
39    fn bool_ones(shape: Shape, device: &Device<Self>, dtype: BoolDType) -> BoolTensor<Self> {
40        numeric::ones(device.clone(), shape, dtype.into())
41    }
42
43    async fn bool_into_data(tensor: BoolTensor<Self>) -> Result<TensorData, ExecutionError> {
44        super::into_data(tensor).await
45    }
46
47    fn bool_from_data(data: TensorData, device: &Device<Self>) -> BoolTensor<Self> {
48        if !matches!(
49            data.dtype(),
50            DType::Bool(BoolStore::U8) | DType::Bool(BoolStore::U32)
51        ) {
52            unimplemented!("Unsupported dtype for `bool_from_data` {:?}", data.dtype());
53        }
54        super::from_data(data, device)
55    }
56
57    fn bool_into_int(tensor: BoolTensor<Self>, out_dtype: IntDType) -> IntTensor<Self> {
58        kernel::bool_cast(tensor, out_dtype.into())
59    }
60
61    fn bool_to_device(tensor: BoolTensor<Self>, device: &Device<Self>) -> BoolTensor<Self> {
62        super::to_device(tensor, device)
63    }
64
65    fn bool_reshape(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
66        super::reshape(tensor, shape)
67    }
68
69    fn bool_slice(tensor: BoolTensor<Self>, slices: &[Slice]) -> BoolTensor<Self> {
70        // Check if all steps are 1
71        let all_steps_one = slices.iter().all(|info| info.step == 1);
72
73        if all_steps_one {
74            // Use optimized slice for step=1
75            let simple_ranges: Vec<Range<usize>> = slices
76                .iter()
77                .enumerate()
78                .map(|(i, slice)| slice.to_range(tensor.meta.shape()[i]))
79                .collect();
80
81            kernel::slice(tensor, &simple_ranges)
82        } else {
83            // Use slice with steps kernel
84            kernel::slice_with_steps(tensor, slices)
85        }
86    }
87
88    fn bool_slice_assign(
89        tensor: BoolTensor<Self>,
90        ranges: &[Slice],
91        value: BoolTensor<Self>,
92    ) -> BoolTensor<Self> {
93        kernel::slice_assign(tensor, ranges, value)
94    }
95
96    fn bool_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
97        let dtype = lhs.dtype;
98        kernel::equal(lhs, rhs, dtype)
99    }
100
101    fn bool_not_equal(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
102        let dtype = lhs.dtype;
103        kernel::not_equal(lhs, rhs, dtype)
104    }
105
106    fn bool_not(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
107        let dtype = tensor.dtype;
108        let storage = dtype_to_storage_type(dtype);
109        let scalar = match dtype {
110            DType::Bool(BoolStore::U32) => InputScalar::new(u32::false_val(), storage),
111            DType::Bool(BoolStore::U8) => InputScalar::new(u8::false_val(), storage),
112            other => unimplemented!("Unsupported dtype for `bool_from_data` {other:?}"),
113        };
114        kernel::equal_elem(tensor, scalar, dtype)
115    }
116
117    fn bool_and(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
118        kernel::launch_binop::<AndOp>(lhs, rhs)
119    }
120
121    fn bool_or(lhs: BoolTensor<Self>, rhs: BoolTensor<Self>) -> BoolTensor<Self> {
122        kernel::launch_binop::<OrOp>(lhs, rhs)
123    }
124
125    fn bool_any(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
126        let store = bool_store(&tensor);
127        kernel::reduce::reduce_logical(tensor, None, ReduceOperationConfig::Any, store)
128    }
129
130    fn bool_any_dim(tensor: BoolTensor<Self>, dim: usize) -> BoolTensor<Self> {
131        let store = bool_store(&tensor);
132        kernel::reduce::reduce_logical(tensor, Some(dim), ReduceOperationConfig::Any, store)
133    }
134
135    fn bool_all(tensor: BoolTensor<Self>) -> BoolTensor<Self> {
136        let store = bool_store(&tensor);
137        kernel::reduce::reduce_logical(tensor, None, ReduceOperationConfig::All, store)
138    }
139
140    fn bool_all_dim(tensor: BoolTensor<Self>, dim: usize) -> BoolTensor<Self> {
141        let store = bool_store(&tensor);
142        kernel::reduce::reduce_logical(tensor, Some(dim), ReduceOperationConfig::All, store)
143    }
144
145    fn bool_into_float(tensor: BoolTensor<Self>, out_dtype: FloatDType) -> FloatTensor<Self> {
146        kernel::bool_cast(tensor, out_dtype.into())
147    }
148
149    fn bool_swap_dims(mut tensor: BoolTensor<Self>, dim1: usize, dim2: usize) -> BoolTensor<Self> {
150        tensor.meta.swap(dim1, dim2);
151
152        tensor
153    }
154
155    fn bool_repeat_dim(tensor: BoolTensor<Self>, dim: usize, times: usize) -> BoolTensor<Self> {
156        kernel::repeat_dim(tensor, dim, times)
157    }
158
159    fn bool_permute(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
160        permute(tensor, axes)
161    }
162
163    fn bool_expand(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
164        expand(tensor, shape)
165    }
166
167    fn bool_select(
168        tensor: BoolTensor<Self>,
169        dim: usize,
170        indices: IntTensor<Self>,
171    ) -> BoolTensor<Self> {
172        kernel::select(tensor, dim, indices)
173    }
174
175    fn bool_select_or(
176        tensor: BoolTensor<Self>,
177        dim: usize,
178        indices: IntTensor<Self>,
179        value: BoolTensor<Self>,
180    ) -> BoolTensor<Self> {
181        kernel::select_assign(tensor, dim, indices, value, true)
182    }
183
184    fn bool_flip(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
185        let dtype = tensor.dtype;
186        kernel::flip(tensor, axes, dtype)
187    }
188
189    fn bool_unfold(
190        tensor: FloatTensor<Self>,
191        dim: usize,
192        size: usize,
193        step: usize,
194    ) -> FloatTensor<Self> {
195        unfold(tensor, dim, size, step)
196    }
197
198    fn bool_mask_where(
199        tensor: BoolTensor<Self>,
200        mask: BoolTensor<Self>,
201        value: BoolTensor<Self>,
202    ) -> BoolTensor<Self> {
203        let dtype = tensor.dtype;
204        kernel::mask_where_auto(tensor, mask, value, dtype)
205    }
206
207    fn bool_mask_fill(
208        tensor: BoolTensor<Self>,
209        mask: BoolTensor<Self>,
210        value: Scalar,
211    ) -> BoolTensor<Self> {
212        let dtype = tensor.dtype;
213        kernel::mask_fill_auto(
214            tensor,
215            mask,
216            InputScalar::new(value, dtype_to_storage_type(dtype)),
217            dtype,
218        )
219    }
220
221    fn bool_gather(
222        dim: usize,
223        tensor: BoolTensor<Self>,
224        indices: IntTensor<Self>,
225    ) -> BoolTensor<Self> {
226        kernel::gather(dim, tensor, indices)
227    }
228
229    fn bool_scatter_or(
230        dim: usize,
231        tensor: BoolTensor<Self>,
232        indices: IntTensor<Self>,
233        value: BoolTensor<Self>,
234    ) -> BoolTensor<Self> {
235        kernel::scatter(dim, tensor, indices, value, true)
236    }
237
238    fn bool_equal_elem(lhs: BoolTensor<Self>, rhs: Scalar) -> BoolTensor<Self> {
239        let dtype = lhs.dtype;
240        kernel::equal_elem(
241            lhs,
242            InputScalar::new(rhs, dtype_to_storage_type(dtype)),
243            dtype,
244        )
245    }
246}
burn_cubecl/ops/bool_tensor.rs

burn_cubecl/ops/
bool_tensor.rs
