[WIP] [Rust] Add explicit SIMD vectorization for arithmetic ops in "array_ops"

rust/arrow/Cargo.toml

@@ -45,6 +45,7 @@ csv = "1.0.0"
 num = "0.2"
 regex = "1.1"
 lazy_static = "1.2"
+packed_simd = "0.3.1"
 
 [dev-dependencies]
 criterion = "0.2"
@@ -56,4 +57,8 @@ harness = false
 
 [[bench]]
 name = "builder"
-harness = false
+harness = false
+
+[[bench]]
+name = "arithmetic_kernels"
+harness = false

rust/arrow/benches/arithmetic_kernels.rs

@@ -0,0 +1,63 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#[macro_use]
+extern crate criterion;
+use criterion::Criterion;
+
+extern crate arrow;
+
+use arrow::array::*;
+use arrow::builder::*;
+use arrow::compute::arithmetic_kernels::*;
+use arrow::compute::array_ops::*;
+
+fn create_array(size: usize) -> Float32Array {
+    let mut builder = Float32Builder::new(size);
+    for _i in 0..size {
+        builder.append_value(1.0).unwrap();
+    }
+    builder.finish()
+}
+
+fn primitive_array_add(size: usize) {
+    let arr_a = create_array(size);
+    let arr_b = create_array(size);
+    criterion::black_box(math_op(&arr_a, &arr_b, |a, b| Ok(a + b)).unwrap());
+}
+
+fn primitive_array_add_simd(size: usize) {
+    let arr_a = create_array(size);
+    let arr_b = create_array(size);
+    criterion::black_box(add(&arr_a, &arr_b).unwrap());
+}
+
+fn add_benchmark(c: &mut Criterion) {
+    c.bench_function("add 128", |b| b.iter(|| primitive_array_add(128)));
+    c.bench_function("add 128 simd", |b| b.iter(|| primitive_array_add_simd(128)));
+    c.bench_function("add 256", |b| b.iter(|| primitive_array_add(256)));
+    c.bench_function("add 256 simd", |b| b.iter(|| primitive_array_add_simd(256)));
+    c.bench_function("add 512", |b| b.iter(|| primitive_array_add(512)));
+    c.bench_function("add 512 simd", |b| b.iter(|| primitive_array_add_simd(512)));
+    c.bench_function("add 1024", |b| b.iter(|| primitive_array_add(1024)));
+    c.bench_function("add 1024 simd", |b| {
+        b.iter(|| primitive_array_add_simd(1024))
+    });
+}
+
+criterion_group!(benches, add_benchmark);
+criterion_main!(benches);

rust/arrow/src/array.rs

@@ -236,8 +236,10 @@ impl<T: ArrowNumericType> PrimitiveArray<T> {
     ///
     /// Note this doesn't do any bound checking, for performance reason.
     pub fn value_slice(&self, offset: usize, len: usize) -> &[T::Native] {
-        let raw = unsafe { std::slice::from_raw_parts(self.raw_values(), self.len()) };
-        &raw[offset..offset + len]
+        let raw = unsafe {
+            std::slice::from_raw_parts(self.raw_values().offset(offset as isize), len)
+        };
+        &raw[..]
     }
 
     // Returns a new primitive array builder

rust/arrow/src/buffer.rs

@@ -19,11 +19,16 @@
 //! fixed size aligned at a 64-byte boundary. `MutableBuffer` is like `Buffer`, but it can
 //! be mutated and grown.
 
+use packed_simd::u8x64;
+
 use std::cmp;
 use std::io::{Error as IoError, ErrorKind, Result as IoResult, Write};
 use std::mem;
+use std::ops::BitAnd;
+use std::slice::{from_raw_parts, from_raw_parts_mut};
 use std::sync::Arc;
 
+use crate::builder::{BufferBuilderTrait, UInt8BufferBuilder};
 use crate::error::Result;
 use crate::memory;
 use crate::util::bit_util;
@@ -140,7 +145,63 @@ impl<T: AsRef<[u8]>> From<T> for Buffer {
     }
 }
 
+impl BitAnd for Buffer {
+    type Output = Self;
+
+    fn bitand(self, rhs: Self) -> Self {
+        assert_eq!(
+            self.len(),
+            rhs.len(),
+            "Buffers must be the same size to apply Bitwise OR."
+        );
+
+        // SIMD implementation if available
+        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+        {
+            let mut result =
+                MutableBuffer::new(self.len()).with_bitset(self.len(), false);
+            let lanes = u8x64::lanes();
+            for i in (0..self.len()).step_by(lanes) {
+                let left_data =
+                    unsafe { from_raw_parts(self.raw_data().offset(i as isize), lanes) };
+                let left_simd =
+                    unsafe { u8x64::from_slice_unaligned_unchecked(left_data) };
+                let right_data =
+                    unsafe { from_raw_parts(rhs.raw_data().offset(i as isize), lanes) };
+                let right_simd =
+                    unsafe { u8x64::from_slice_unaligned_unchecked(right_data) };
+                let simd_result = left_simd & right_simd;
+                let result_slice: &mut [u8] = unsafe {
+                    from_raw_parts_mut(
+                        (result.data_mut().as_mut_ptr() as *mut u8).offset(i as isize),
+                        lanes,
+                    )
+                };
+                unsafe { simd_result.write_to_slice_unaligned_unchecked(result_slice) };
+            }
+            return result.freeze();
+        }
+
+        // Default implementation
+        #[allow(unreachable_code)]
+        {
+            let mut builder = UInt8BufferBuilder::new(self.len());
+            for i in 0..self.len() {
+                unsafe {
+                    builder
+                        .append(
+                            self.data().get_unchecked(i) & rhs.data().get_unchecked(i),
+                        )
+                        .unwrap();
+                }
+            }
+            builder.finish()
+        }
+    }
+}
+
 unsafe impl Sync for Buffer {}
+
 unsafe impl Send for Buffer {}
 
 /// Similar to `Buffer`, but is growable and can be mutated. A mutable buffer can be
@@ -349,6 +410,22 @@ mod tests {
         assert_ne!(buf1, buf2);
     }
 
+    #[test]
+    fn test_buffer_bitand() {
+        let buf1 = Buffer::from([0b01101010]);
+        let buf2 = Buffer::from([0b01001110]);
+        let buf3 = buf1 & buf2;
+        assert_eq!(Buffer::from([0b01001010]), buf3);
+    }
+
+    #[test]
+    #[should_panic(expected = "Buffers must be the same size to apply Bitwise OR.")]
+    fn test_buffer_bitand_different_sizes() {
+        let buf1 = Buffer::from([1_u8, 1_u8]);
+        let buf2 = Buffer::from([0b01001110]);
+        let _buf3 = buf1 & buf2;
+    }
+
     #[test]
     fn test_from_raw_parts() {
         let buf = Buffer::from_raw_parts(null_mut(), 0);

rust/arrow/src/compute/arithmetic_kernels.rs

@@ -0,0 +1,154 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+//! Defines basic arithmetic kernels for `PrimitiveArrays`.
+//!
+//! These kernels can leverage SIMD if available on your system.  Currently no runtime
+//! detection is provided, you should enable the specific SIMD intrinsics using
+//! `RUSTFLAGS="-C target-feature=+avx2"` for example.  See the
+//! [here] (https://doc.rust-lang.org/stable/std/arch/) for more information.
+
+use std::mem;
+use std::ops::{Add, Div, Mul, Sub};
+use std::slice::from_raw_parts_mut;
+
+use num::Zero;
+
+use crate::array::*;
+use crate::buffer::MutableBuffer;
+use crate::compute::array_ops::math_op;
+use crate::datatypes;
+use crate::error::{ArrowError, Result};
+
+/// Vectorized version of add operation
+#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+fn simd_bin_op<T, F>(
+    left: &PrimitiveArray<T>,
+    right: &PrimitiveArray<T>,
+    op: F,
+) -> Result<PrimitiveArray<T>>
+where
+    T: datatypes::ArrowNumericType + datatypes::ArrowSIMDType,
+    T::Simd: Add<Output = T::Simd>
+        + Sub<Output = T::Simd>
+        + Mul<Output = T::Simd>
+        + Div<Output = T::Simd>,
+    F: Fn(T::Simd, T::Simd) -> T::Simd,
+{
+    if left.len() != right.len() {
+        return Err(ArrowError::ComputeError(
+            "Cannot perform math operation on arrays of different length".to_string(),
+        ));
+    }
+
+    let lanes = T::lanes();
+    let buffer_size = left.len() * mem::size_of::<T::Native>();
+    let mut result = MutableBuffer::new(buffer_size).with_bitset(buffer_size, false);
+
+    for i in (0..left.len()).step_by(lanes) {
+        let simd_left = T::load(left.value_slice(i, lanes));
+        let simd_right = T::load(right.value_slice(i, lanes));
+        let simd_result = T::bin_op(simd_left, simd_right, &op);
+
+        let result_slice: &mut [T::Native] = unsafe {
+            from_raw_parts_mut(
+                (result.data_mut().as_mut_ptr() as *mut T::Native).offset(i as isize),
+                lanes,
+            )
+        };
+        T::write(simd_result, result_slice);
+    }
+
+    Ok(PrimitiveArray::<T>::new(left.len(), result.freeze(), 0, 0))
+}
+
+/// Perform `left + right` operation on two arrays. If either left or right value is null
+/// then the result is also null.
+pub fn add<T>(
+    left: &PrimitiveArray<T>,
+    right: &PrimitiveArray<T>,
+) -> Result<PrimitiveArray<T>>
+where
+    T: datatypes::ArrowNumericType + datatypes::ArrowSIMDType,
+    T::Native: Add<Output = T::Native>
+        + Sub<Output = T::Native>
+        + Mul<Output = T::Native>
+        + Div<Output = T::Native>
+        + Zero,
+{
+    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+    return simd_bin_op(&left, &right, |a, b| a + b);
+
+    #[allow(unreachable_code)]
+    math_op(left, right, |a, b| Ok(a + b))
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::array::Int32Array;
+
+    #[test]
+    fn test_primitive_array_add() {
+        let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
+        let b = Int32Array::from(vec![6, 7, 8, 9, 8]);
+
+        #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
+        {
+            let c = simd_bin_op(&a, &b, |x, y| x + y).unwrap();
+
+            assert_eq!(11, c.value(0));
+            assert_eq!(13, c.value(1));
+            assert_eq!(15, c.value(2));
+            assert_eq!(17, c.value(3));
+            assert_eq!(17, c.value(4));
+        }
+
+        let d = add(&a, &b).unwrap();
+        assert_eq!(11, d.value(0));
+        assert_eq!(13, d.value(1));
+        assert_eq!(15, d.value(2));
+        assert_eq!(17, d.value(3));
+        assert_eq!(17, d.value(4));
+    }
+
+    #[test]
+    fn test_primitive_array_add_mismatched_length() {
+        let a = Int32Array::from(vec![5, 6, 7, 8, 9]);
+        let b = Int32Array::from(vec![6, 7, 8]);
+        let e = add(&a, &b)
+            .err()
+            .expect("should have failed due to different lengths");
+        assert_eq!(
+            "ComputeError(\"Cannot perform math operation on arrays of different length\")",
+            format!("{:?}", e)
+        );
+    }
+
+    #[ignore]
+    #[test]
+    fn test_primitive_array_add_with_nulls() {
+        let a = Int32Array::from(vec![Some(5), None, Some(7), None]);
+        let b = Int32Array::from(vec![None, None, Some(6), Some(7)]);
+        let c = add(&a, &b).unwrap();
+        assert_eq!(true, c.is_null(0));
+        assert_eq!(true, c.is_null(1));
+        assert_eq!(false, c.is_null(2));
+        assert_eq!(true, c.is_null(3));
+        assert_eq!(13, c.value(2));
+    }
+}