luma_tensor/device/cpu/ops/
bool_ops.rs

1//! `impl BoolOps for Cpu`. Bool storage is a plain `Vec<bool>`, so no per-dtype
2//! dispatch is needed; kernels operate on it directly.
3
4use std::borrow::Cow;
5
6use super::kernels::{elementwise as ew, reduce, shape as shape_k};
7use super::{Cpu, CpuBoolStorage, CpuFloatStorage, CpuIntStorage};
8use crate::dtype::{BoolDType, FloatDType, IntDType};
9use crate::{BoolOps, Device, Error, Layout, Result, Shape};
10
11impl BoolOps<Cpu> for Cpu {
12    fn b_falses(shape: &Shape, _device: &Cpu, _dtype: BoolDType) -> Result<<Cpu as Device>::BoolStorage> {
13        Ok(CpuBoolStorage(vec![false; shape.element_count()]))
14    }
15
16    fn b_trues(shape: &Shape, _device: &Cpu, _dtype: BoolDType) -> Result<<Cpu as Device>::BoolStorage> {
17        Ok(CpuBoolStorage(vec![true; shape.element_count()]))
18    }
19
20    fn b_from_bool<'a>(data: impl Into<Cow<'a, [bool]>>, _device: &Cpu) -> Result<<Cpu as Device>::BoolStorage> {
21        let data = data.into();
22        Ok(match data {
23            Cow::Owned(v) => CpuBoolStorage(v),
24            Cow::Borrowed(s) => CpuBoolStorage(s.to_vec()),
25        })
26    }
27
28    fn b_from_bytes<'a>(
29        bytes: impl Into<Cow<'a, [u8]>>,
30        _shape: &Shape,
31        _device: &Cpu,
32        _dtype: BoolDType,
33    ) -> Result<<Cpu as Device>::BoolStorage> {
34        let bytes = bytes.into();
35        Ok(CpuBoolStorage(bytes.iter().map(|&x| x != 0).collect()))
36    }
37
38    fn b_contiguous(x: &<Cpu as Device>::BoolStorage, l: &Layout) -> Result<<Cpu as Device>::BoolStorage> {
39        Ok(CpuBoolStorage(super::kernels::iter::gather(&x.0, l)))
40    }
41
42    fn b_cast_float(x: &CpuBoolStorage, layout: &Layout, to: FloatDType) -> Result<CpuFloatStorage> {
43        let s = match to {
44            FloatDType::F32 => CpuFloatStorage::F32(layout.storage_indices().map(|i| if x.0[i] { 1.0 } else { 0.0 }).collect()),
45            FloatDType::F64 => CpuFloatStorage::F64(layout.storage_indices().map(|i| if x.0[i] { 1.0 } else { 0.0 }).collect()),
46        };
47        Ok(s)
48    }
49
50    fn b_cast_int(x: &CpuBoolStorage, layout: &Layout, to: IntDType) -> Result<CpuIntStorage> {
51        let s = match to {
52            IntDType::I32 => CpuIntStorage::I32(layout.storage_indices().map(|i| x.0[i] as i32).collect()),
53            IntDType::U32 => CpuIntStorage::U32(layout.storage_indices().map(|i| x.0[i] as u32).collect()),
54            IntDType::U8 => CpuIntStorage::U8(layout.storage_indices().map(|i| x.0[i] as u8).collect()),
55        };
56        Ok(s)
57    }
58
59    fn b_cast_bool(x: &CpuBoolStorage, layout: &Layout, _to: BoolDType) -> Result<CpuBoolStorage> {
60        Ok(CpuBoolStorage(layout.storage_indices().map(|i| x.0[i]).collect()))
61    }
62
63    fn b_to_vec(x: &<Cpu as Device>::BoolStorage, layout: &Layout) -> Result<Vec<bool>> {
64        Ok(layout.storage_indices().map(|i| x.0[i]).collect())
65    }
66
67    fn b_to_bytes<'a>(x: &'a <Cpu as Device>::BoolStorage, layout: &Layout) -> Result<Cow<'a, [u8]>> {
68        // bool is not Pod, so bytemuck doesn't work — convert manually.
69        if layout.is_contiguous() {
70            let bytes: Vec<u8> = x.0.iter().map(|&b| b as u8).collect();
71            Ok(Cow::Owned(bytes))
72        } else {
73            let contig = Self::b_contiguous(x, layout)?;
74            let bytes: Vec<u8> = contig.0.iter().map(|&b| b as u8).collect();
75            Ok(Cow::Owned(bytes))
76        }
77    }
78
79    fn b_and(
80        lhs: &<Cpu as Device>::BoolStorage,
81        lhs_l: &Layout,
82        rhs: &<Cpu as Device>::BoolStorage,
83        rhs_l: &Layout,
84    ) -> Result<<Cpu as Device>::BoolStorage> {
85        Ok(CpuBoolStorage(ew::binary(&lhs.0, lhs_l, &rhs.0, rhs_l, |a, b| a & b)))
86    }
87
88    fn b_or(
89        lhs: &<Cpu as Device>::BoolStorage,
90        lhs_l: &Layout,
91        rhs: &<Cpu as Device>::BoolStorage,
92        rhs_l: &Layout,
93    ) -> Result<<Cpu as Device>::BoolStorage> {
94        Ok(CpuBoolStorage(ew::binary(&lhs.0, lhs_l, &rhs.0, rhs_l, |a, b| a | b)))
95    }
96
97    fn b_xor(
98        lhs: &<Cpu as Device>::BoolStorage,
99        lhs_l: &Layout,
100        rhs: &<Cpu as Device>::BoolStorage,
101        rhs_l: &Layout,
102    ) -> Result<<Cpu as Device>::BoolStorage> {
103        Ok(CpuBoolStorage(ew::binary(&lhs.0, lhs_l, &rhs.0, rhs_l, |a, b| a ^ b)))
104    }
105
106    fn b_not(x: &<Cpu as Device>::BoolStorage, l: &Layout) -> Result<<Cpu as Device>::BoolStorage> {
107        Ok(CpuBoolStorage(ew::unary(&x.0, l, |v| !v)))
108    }
109
110    fn b_reduce_all(
111        x: &<Cpu as Device>::BoolStorage,
112        l: &Layout,
113        dims: &[usize],
114        keepdim: bool,
115    ) -> Result<(<Cpu as Device>::BoolStorage, Shape)> {
116        // reduce as u8 with min (all true == min == 1), then back to bool.
117        let as_u8: Vec<u8> = super::kernels::iter::gather(&x.0, l).iter().map(|&b| b as u8).collect();
118        let contig = Layout::contiguous(l.shape().clone());
119        let (v, shape) = reduce::reduce_dims(&as_u8, &contig, dims, keepdim, reduce::Reducer::Min)?;
120        Ok((CpuBoolStorage(v.into_iter().map(|u| u != 0).collect()), shape))
121    }
122
123    fn b_reduce_any(
124        x: &<Cpu as Device>::BoolStorage,
125        l: &Layout,
126        dims: &[usize],
127        keepdim: bool,
128    ) -> Result<(<Cpu as Device>::BoolStorage, Shape)> {
129        let as_u8: Vec<u8> = super::kernels::iter::gather(&x.0, l).iter().map(|&b| b as u8).collect();
130        let contig = Layout::contiguous(l.shape().clone());
131        let (v, shape) = reduce::reduce_dims(&as_u8, &contig, dims, keepdim, reduce::Reducer::Max)?;
132        Ok((CpuBoolStorage(v.into_iter().map(|u| u != 0).collect()), shape))
133    }
134
135    fn b_true_count(x: &<Cpu as Device>::BoolStorage, l: &Layout) -> Result<usize> {
136        Ok(l.storage_indices().filter(|&i| x.0[i]).count())
137    }
138
139    fn b_cat(srcs: &[(&<Cpu as Device>::BoolStorage, &Layout)], dim: usize) -> Result<(<Cpu as Device>::BoolStorage, Shape)> {
140        if srcs.is_empty() {
141            return Err(Error::OpRequiresAtLeastOneTensor { op: "cat" });
142        }
143        let views: Vec<(&[bool], &Layout)> = srcs.iter().map(|(s, l)| (s.0.as_slice(), *l)).collect();
144        let (v, shape) = shape_k::cat(&views, dim)?;
145        Ok((CpuBoolStorage(v), shape))
146    }
147
148    fn b_pick(
149        mask: &<Cpu as Device>::BoolStorage,
150        mask_l: &Layout,
151        on_true: &<Cpu as Device>::BoolStorage,
152        true_l: &Layout,
153        on_false: &<Cpu as Device>::BoolStorage,
154        false_l: &Layout,
155    ) -> Result<<Cpu as Device>::BoolStorage> {
156        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
157        let tv: Vec<bool> = true_l.storage_indices().map(|i| on_true.0[i]).collect();
158        let fv: Vec<bool> = false_l.storage_indices().map(|i| on_false.0[i]).collect();
159        Ok(CpuBoolStorage(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { fv[i] }).collect()))
160    }
161
162    fn b_pick_true(
163        mask: &<Cpu as Device>::BoolStorage,
164        mask_l: &Layout,
165        value: bool,
166        on_false: &<Cpu as Device>::BoolStorage,
167        false_l: &Layout,
168    ) -> Result<<Cpu as Device>::BoolStorage> {
169        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
170        let fv: Vec<bool> = false_l.storage_indices().map(|i| on_false.0[i]).collect();
171        Ok(CpuBoolStorage(m.iter().enumerate().map(|(i, &c)| if c { value } else { fv[i] }).collect()))
172    }
173
174    fn b_pick_false(
175        mask: &<Cpu as Device>::BoolStorage,
176        mask_l: &Layout,
177        on_true: &<Cpu as Device>::BoolStorage,
178        true_l: &Layout,
179        value: bool,
180    ) -> Result<<Cpu as Device>::BoolStorage> {
181        let m: Vec<bool> = mask_l.storage_indices().map(|i| mask.0[i]).collect();
182        let tv: Vec<bool> = true_l.storage_indices().map(|i| on_true.0[i]).collect();
183        Ok(CpuBoolStorage(m.iter().enumerate().map(|(i, &c)| if c { tv[i] } else { value }).collect()))
184    }
185
186    fn b_allclose(a: &CpuBoolStorage, a_l: &Layout, b: &CpuBoolStorage, b_l: &Layout) -> Result<bool> {
187        Ok(a_l.storage_indices().zip(b_l.storage_indices()).all(|(ai, bi)| a.0[ai] == b.0[bi]))
188    }
189}
luma_tensor/device/cpu/ops/bool_ops.rs

luma_tensor/device/cpu/ops/
bool_ops.rs
