From 7a387d9a0fcb3b02acfb9bb4787730a453bb7737 Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Fri, 29 Jan 2021 12:58:07 +0100 Subject: [PATCH 1/4] make unary arithmetic funciton more generic and implement ssid version of power_scalar --- rust/arrow/src/compute/kernels/arithmetic.rs | 36 +++++++++++++++----- 1 file changed, 28 insertions(+), 8 deletions(-) diff --git a/rust/arrow/src/compute/kernels/arithmetic.rs b/rust/arrow/src/compute/kernels/arithmetic.rs index 06d7c90f6a5..51ad41b8e3e 100644 --- a/rust/arrow/src/compute/kernels/arithmetic.rs +++ b/rust/arrow/src/compute/kernels/arithmetic.rs @@ -35,6 +35,7 @@ use crate::datatypes; use crate::datatypes::ArrowNumericType; use crate::error::{ArrowError, Result}; use crate::{array::*, util::bit_util}; +use num::traits::Pow; #[cfg(simd)] use std::borrow::BorrowMut; #[cfg(simd)] @@ -42,13 +43,9 @@ use std::slice::{ChunksExact, ChunksExactMut}; /// Helper function to perform math lambda function on values from single array of signed numeric /// type. If value is null then the output value is also null, so `-null` is `null`. -pub fn signed_unary_math_op( - array: &PrimitiveArray, - op: F, -) -> Result> +pub fn unary_math_op(array: &PrimitiveArray, op: F) -> Result> where - T: datatypes::ArrowSignedNumericType, - T::Native: Neg, + T: datatypes::ArrowNumericType, F: Fn(T::Native) -> T::Native, { let values = array.values().iter().map(|v| op(*v)); @@ -71,7 +68,7 @@ where Ok(PrimitiveArray::::from(Arc::new(data))) } -/// SIMD vectorized version of `signed_unary_math_op` above. +/// SIMD vectorized version of `unary_math_op` above specialized for signed numerical values. #[cfg(simd)] fn simd_signed_unary_math_op( array: &PrimitiveArray, @@ -536,7 +533,22 @@ where #[cfg(simd)] return simd_signed_unary_math_op(array, |x| -x, |x| -x); #[cfg(not(simd))] - return signed_unary_math_op(array, |x| -x); + return unary_math_op(array, |x| -x); +} + +/// Raise array to the power of a scalar. +pub fn pow_scalar( + array: &PrimitiveArray, + raise: T::Native, +) -> Result> +where + T: datatypes::ArrowNumericType, + T::Native: Pow, +{ + #[cfg(simd)] + todo!(); + #[cfg(not(simd))] + return unary_math_op(array, |x| x.pow(raise)); } /// Perform `left * right` operation on two arrays. If either left or right value is null @@ -808,4 +820,12 @@ mod tests { .collect(); assert_eq!(expected, actual); } + + #[test] + fn test_primitive_array_raise_power_scalar() { + let a = Float64Array::from(vec![1.0, 2.0, 3.0]); + let actual = pow_scalar(&a, 2.0).unwrap(); + let expected = Float64Array::from(vec![1.0, 4.0, 9.0]); + assert_eq!(expected, actual); + } } From 0d1b893f12aee69cd7649aff2339e98aa27837d2 Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Fri, 29 Jan 2021 15:33:26 +0100 Subject: [PATCH 2/4] SIMD implementation of pow_scalar --- rust/arrow/src/compute/kernels/arithmetic.rs | 64 +++++++++++++++++++- rust/arrow/src/datatypes.rs | 36 +++++++++++ 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/rust/arrow/src/compute/kernels/arithmetic.rs b/rust/arrow/src/compute/kernels/arithmetic.rs index 51ad41b8e3e..0ca7c9feaf2 100644 --- a/rust/arrow/src/compute/kernels/arithmetic.rs +++ b/rust/arrow/src/compute/kernels/arithmetic.rs @@ -117,6 +117,55 @@ where Ok(PrimitiveArray::::from(Arc::new(data))) } +#[cfg(simd)] +fn simd_float_unary_math_op( + array: &PrimitiveArray, + simd_op: SIMD_OP, + scalar_op: SCALAR_OP, +) -> Result> +where + T: datatypes::ArrowFloatNumericType, + SIMD_OP: Fn(T::Simd) -> T::Simd, + SCALAR_OP: Fn(T::Native) -> T::Native, +{ + let lanes = T::lanes(); + let buffer_size = array.len() * std::mem::size_of::(); + + let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false); + + let mut result_chunks = result.typed_data_mut().chunks_exact_mut(lanes); + let mut array_chunks = array.values().chunks_exact(lanes); + + result_chunks + .borrow_mut() + .zip(array_chunks.borrow_mut()) + .for_each(|(result_slice, input_slice)| { + let simd_input = T::load(input_slice); + let simd_result = T::unary_op(simd_input, &simd_op); + T::write(simd_result, result_slice); + }); + + let result_remainder = result_chunks.into_remainder(); + let array_remainder = array_chunks.remainder(); + + result_remainder.into_iter().zip(array_remainder).for_each( + |(scalar_result, scalar_input)| { + *scalar_result = scalar_op(*scalar_input); + }, + ); + + let data = ArrayData::new( + T::DATA_TYPE, + array.len(), + None, + array.data_ref().null_buffer().cloned(), + 0, + vec![result.into()], + vec![], + ); + Ok(PrimitiveArray::::from(Arc::new(data))) +} + /// Helper function to perform math lambda function on values from two arrays. If either /// left or right value is null then the output value is also null, so `1 + null` is /// `null`. @@ -542,11 +591,18 @@ pub fn pow_scalar( raise: T::Native, ) -> Result> where - T: datatypes::ArrowNumericType, + T: datatypes::ArrowFloatNumericType, T::Native: Pow, { #[cfg(simd)] - todo!(); + { + let raise_vector = T::init(raise); + return simd_float_unary_math_op( + array, + |x| T::pow(x, raise_vector), + |x| x.pow(raise), + ); + } #[cfg(not(simd))] return unary_math_op(array, |x| x.pow(raise)); } @@ -827,5 +883,9 @@ mod tests { let actual = pow_scalar(&a, 2.0).unwrap(); let expected = Float64Array::from(vec![1.0, 4.0, 9.0]); assert_eq!(expected, actual); + let a = Float64Array::from(vec![Some(1.0), None, Some(3.0)]); + let actual = pow_scalar(&a, 2.0).unwrap(); + let expected = Float64Array::from(vec![Some(1.0), None, Some(9.0)]); + assert_eq!(expected, actual); } } diff --git a/rust/arrow/src/datatypes.rs b/rust/arrow/src/datatypes.rs index 2231da81c37..096a9305891 100644 --- a/rust/arrow/src/datatypes.rs +++ b/rust/arrow/src/datatypes.rs @@ -605,6 +605,8 @@ where /// Writes a SIMD result back to a slice fn write(simd_result: Self::Simd, slice: &mut [Self::Native]); + + fn unary_op Self::Simd>(a: Self::Simd, op: F) -> Self::Simd; } #[cfg(not(simd))] @@ -806,6 +808,14 @@ macro_rules! make_numeric_type { fn write(simd_result: Self::Simd, slice: &mut [Self::Native]) { unsafe { simd_result.write_to_slice_unaligned_unchecked(slice) }; } + + #[inline] + fn unary_op Self::Simd>( + a: Self::Simd, + op: F, + ) -> Self::Simd { + op(a) + } } #[cfg(not(simd))] @@ -909,6 +919,32 @@ make_signed_numeric_type!(Int64Type, i64x8); make_signed_numeric_type!(Float32Type, f32x16); make_signed_numeric_type!(Float64Type, f64x8); +#[cfg(simd)] +pub trait ArrowFloatNumericType: ArrowNumericType { + fn pow(base: Self::Simd, raise: Self::Simd) -> Self::Simd; +} + +#[cfg(not(simd))] +pub trait ArrowFloatNumericType: ArrowNumericType {} + +macro_rules! make_float_numeric_type { + ($impl_ty:ty, $simd_ty:ident) => { + #[cfg(simd)] + impl ArrowFloatNumericType for $impl_ty { + #[inline] + fn pow(base: Self::Simd, raise: Self::Simd) -> Self::Simd { + base.powf(raise) + } + } + + #[cfg(not(simd))] + impl ArrowFloatNumericType for $impl_ty {} + }; +} + +make_float_numeric_type!(Float32Type, f32x16); +make_float_numeric_type!(Float64Type, f64x8); + /// A subtype of primitive type that represents temporal values. pub trait ArrowTemporalType: ArrowPrimitiveType {} From 837d563b91892ba8bd81f4f2be9b5e14542f28ad Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Wed, 3 Feb 2021 09:11:15 +0100 Subject: [PATCH 3/4] remove redundant helper function in favor of arity::unary kernel --- rust/arrow/src/compute/kernels/arithmetic.rs | 33 ++------------------ 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/rust/arrow/src/compute/kernels/arithmetic.rs b/rust/arrow/src/compute/kernels/arithmetic.rs index 0ca7c9feaf2..164497aef0c 100644 --- a/rust/arrow/src/compute/kernels/arithmetic.rs +++ b/rust/arrow/src/compute/kernels/arithmetic.rs @@ -30,7 +30,7 @@ use num::{One, Zero}; use crate::buffer::Buffer; #[cfg(simd)] use crate::buffer::MutableBuffer; -use crate::compute::util::combine_option_bitmap; +use crate::compute::{kernels::arity::unary, util::combine_option_bitmap}; use crate::datatypes; use crate::datatypes::ArrowNumericType; use crate::error::{ArrowError, Result}; @@ -41,33 +41,6 @@ use std::borrow::BorrowMut; #[cfg(simd)] use std::slice::{ChunksExact, ChunksExactMut}; -/// Helper function to perform math lambda function on values from single array of signed numeric -/// type. If value is null then the output value is also null, so `-null` is `null`. -pub fn unary_math_op(array: &PrimitiveArray, op: F) -> Result> -where - T: datatypes::ArrowNumericType, - F: Fn(T::Native) -> T::Native, -{ - let values = array.values().iter().map(|v| op(*v)); - // JUSTIFICATION - // Benefit - // ~60% speedup - // Soundness - // `values` is an iterator with a known size. - let buffer = unsafe { Buffer::from_trusted_len_iter(values) }; - - let data = ArrayData::new( - T::DATA_TYPE, - array.len(), - None, - array.data_ref().null_buffer().cloned(), - 0, - vec![buffer], - vec![], - ); - Ok(PrimitiveArray::::from(Arc::new(data))) -} - /// SIMD vectorized version of `unary_math_op` above specialized for signed numerical values. #[cfg(simd)] fn simd_signed_unary_math_op( @@ -582,7 +555,7 @@ where #[cfg(simd)] return simd_signed_unary_math_op(array, |x| -x, |x| -x); #[cfg(not(simd))] - return unary_math_op(array, |x| -x); + return Ok(unary(array, |x| -x)); } /// Raise array to the power of a scalar. @@ -604,7 +577,7 @@ where ); } #[cfg(not(simd))] - return unary_math_op(array, |x| x.pow(raise)); + return Ok(unary(array, |x| x.pow(raise))); } /// Perform `left * right` operation on two arrays. If either left or right value is null From 511eb65246d235213ce94aad404607c2892191ec Mon Sep 17 00:00:00 2001 From: Ritchie Vink Date: Mon, 8 Feb 2021 08:47:30 +0100 Subject: [PATCH 4/4] rename pow_scalar > powf_scalar --- rust/arrow/src/compute/kernels/arithmetic.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/rust/arrow/src/compute/kernels/arithmetic.rs b/rust/arrow/src/compute/kernels/arithmetic.rs index 164497aef0c..067756662cf 100644 --- a/rust/arrow/src/compute/kernels/arithmetic.rs +++ b/rust/arrow/src/compute/kernels/arithmetic.rs @@ -558,8 +558,8 @@ where return Ok(unary(array, |x| -x)); } -/// Raise array to the power of a scalar. -pub fn pow_scalar( +/// Raise array with floating point values to the power of a scalar. +pub fn powf_scalar( array: &PrimitiveArray, raise: T::Native, ) -> Result> @@ -853,11 +853,11 @@ mod tests { #[test] fn test_primitive_array_raise_power_scalar() { let a = Float64Array::from(vec![1.0, 2.0, 3.0]); - let actual = pow_scalar(&a, 2.0).unwrap(); + let actual = powf_scalar(&a, 2.0).unwrap(); let expected = Float64Array::from(vec![1.0, 4.0, 9.0]); assert_eq!(expected, actual); let a = Float64Array::from(vec![Some(1.0), None, Some(3.0)]); - let actual = pow_scalar(&a, 2.0).unwrap(); + let actual = powf_scalar(&a, 2.0).unwrap(); let expected = Float64Array::from(vec![Some(1.0), None, Some(9.0)]); assert_eq!(expected, actual); }