vortex_array/arrays/primitive/compute/
is_constant.rs1use vortex_dtype::half::f16;
2use vortex_dtype::{NativePType, match_each_native_ptype};
3use vortex_error::VortexResult;
4
5use crate::arrays::{PrimitiveArray, PrimitiveVTable};
6use crate::compute::{IsConstantKernel, IsConstantKernelAdapter, IsConstantOpts};
7use crate::register_kernel;
8
9cfg_if::cfg_if! {
10 if #[cfg(target_feature = "avx2")] {
11 pub const IS_CONST_LANE_WIDTH: usize = 32;
12 } else {
13 pub const IS_CONST_LANE_WIDTH: usize = 16;
14 }
15}
16
17impl IsConstantKernel for PrimitiveVTable {
18 fn is_constant(
19 &self,
20 array: &PrimitiveArray,
21 opts: &IsConstantOpts,
22 ) -> VortexResult<Option<bool>> {
23 if opts.is_negligible_cost() {
24 return Ok(None);
25 }
26
27 let is_constant = match_each_native_ptype!(array.ptype(), integral: |P| {
28 compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::<P>()}>(array.as_slice::<P>())
29 }, floating_point: |P| {
30 compute_is_constant::<_, {IS_CONST_LANE_WIDTH / size_of::<P>()}>(unsafe { std::mem::transmute::<&[P], &[<P as EqFloat>::IntType]>(array.as_slice::<P>()) })
31 });
32
33 Ok(Some(is_constant))
34 }
35}
36
37register_kernel!(IsConstantKernelAdapter(PrimitiveVTable).lift());
38
39pub fn compute_is_constant<T: NativePType, const WIDTH: usize>(values: &[T]) -> bool {
42 let first_value = values[0];
43 let first_vec = &[first_value; WIDTH];
44
45 let mut chunks = values[1..].chunks_exact(WIDTH);
46 for chunk in &mut chunks {
47 assert_eq!(chunk.len(), WIDTH); if first_vec != chunk {
49 return false;
50 }
51 }
52
53 for value in chunks.remainder() {
54 if !value.is_eq(first_value) {
55 return false;
56 }
57 }
58
59 true
60}
61
62trait EqFloat {
63 type IntType;
64}
65
66impl EqFloat for f16 {
67 type IntType = u16;
68}
69impl EqFloat for f32 {
70 type IntType = u32;
71}
72impl EqFloat for f64 {
73 type IntType = u64;
74}