This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new 0783cf92f Speed up i256 division and remainder operations (#4303)
0783cf92f is described below
commit 0783cf92f1000b2f0da5e3a1c6c5d18f92af6d66
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Wed May 31 03:57:30 2023 -0700
Speed up i256 division and remainder operations (#4303)
* Get rid of BigInt division
* Add test
* Add benchmark
* Fix merging conflicts
* Add comment
* Fix clippy
* Fix clippy
* Fix MIN case
* For review
* Move tests into test_ops
* Fix doc
---
arrow-buffer/benches/i256.rs | 37 ++++++++++
arrow-buffer/src/bigint.rs | 167 +++++++++++++++++++++++++++++++++++--------
2 files changed, 174 insertions(+), 30 deletions(-)
diff --git a/arrow-buffer/benches/i256.rs b/arrow-buffer/benches/i256.rs
index a04e4cb6c..2c43e0e91 100644
--- a/arrow-buffer/benches/i256.rs
+++ b/arrow-buffer/benches/i256.rs
@@ -17,8 +17,23 @@
use arrow_buffer::i256;
use criterion::*;
+use rand::rngs::StdRng;
+use rand::{Rng, SeedableRng};
use std::str::FromStr;
+/// Returns fixed seedable RNG
+fn seedable_rng() -> StdRng {
+ StdRng::seed_from_u64(42)
+}
+
+fn create_i256_vec(size: usize) -> Vec<i256> {
+ let mut rng = seedable_rng();
+
+ (0..size)
+ .map(|_| i256::from_i128(rng.gen::<i128>()))
+ .collect()
+}
+
fn criterion_benchmark(c: &mut Criterion) {
let numbers = vec![
i256::ZERO,
@@ -38,6 +53,28 @@ fn criterion_benchmark(c: &mut Criterion) {
b.iter(|| i256::from_str(&t).unwrap());
});
}
+
+ c.bench_function("i256_div", |b| {
+ b.iter(|| {
+ for number_a in create_i256_vec(10) {
+ for number_b in create_i256_vec(5) {
+ number_a.checked_div(number_b);
+ number_a.wrapping_div(number_b);
+ }
+ }
+ });
+ });
+
+ c.bench_function("i256_rem", |b| {
+ b.iter(|| {
+ for number_a in create_i256_vec(10) {
+ for number_b in create_i256_vec(5) {
+ number_a.checked_rem(number_b);
+ number_a.wrapping_rem(number_b);
+ }
+ }
+ });
+ });
}
criterion_group!(benches, criterion_benchmark);
diff --git a/arrow-buffer/src/bigint.rs b/arrow-buffer/src/bigint.rs
index fab75b792..3b3994689 100644
--- a/arrow-buffer/src/bigint.rs
+++ b/arrow-buffer/src/bigint.rs
@@ -22,22 +22,33 @@ use std::num::ParseIntError;
use std::ops::{BitAnd, BitOr, BitXor, Neg, Shl, Shr};
use std::str::FromStr;
-/// An opaque error similar to [`std::num::ParseIntError`]
+/// [`i256`] operations return this error type.
#[derive(Debug)]
-pub struct ParseI256Error {}
+pub enum I256Error {
+ /// An opaque error similar to [`std::num::ParseIntError`]
+ ParseError,
+ /// Division by zero
+ DivideByZero,
+ /// Division overflow
+ DivideOverflow,
+}
-impl From<ParseIntError> for ParseI256Error {
+impl From<ParseIntError> for I256Error {
fn from(_: ParseIntError) -> Self {
- Self {}
+ I256Error::ParseError
}
}
-impl std::fmt::Display for ParseI256Error {
+impl std::fmt::Display for I256Error {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
- write!(f, "Failed to parse as i256")
+ match self {
+ I256Error::ParseError => write!(f, "Failed to parse as i256"),
+ I256Error::DivideByZero => write!(f, "Division by zero"),
+ I256Error::DivideOverflow => write!(f, "Division overflow"),
+ }
}
}
-impl std::error::Error for ParseI256Error {}
+impl std::error::Error for I256Error {}
/// A signed 256-bit integer
#[allow(non_camel_case_types)]
@@ -60,7 +71,7 @@ impl std::fmt::Display for i256 {
}
impl FromStr for i256 {
- type Err = ParseI256Error;
+ type Err = I256Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
// i128 can store up to 38 decimal digits
@@ -82,7 +93,7 @@ impl FromStr for i256 {
if !s.as_bytes()[0].is_ascii_digit() {
// Ensures no duplicate sign
- return Err(ParseI256Error {});
+ return Err(I256Error::ParseError);
}
parse_impl(s, negative)
@@ -90,7 +101,7 @@ impl FromStr for i256 {
}
/// Parse `s` with any sign and leading 0s removed
-fn parse_impl(s: &str, negative: bool) -> Result<i256, ParseI256Error> {
+fn parse_impl(s: &str, negative: bool) -> Result<i256, I256Error> {
if s.len() <= 38 {
let low = i128::from_str(s)?;
return Ok(match negative {
@@ -102,7 +113,7 @@ fn parse_impl(s: &str, negative: bool) -> Result<i256,
ParseI256Error> {
let split = s.len() - 38;
if !s.as_bytes()[split].is_ascii_digit() {
// Ensures not splitting codepoint and no sign
- return Err(ParseI256Error {});
+ return Err(I256Error::ParseError);
}
let (hs, ls) = s.split_at(split);
@@ -117,7 +128,7 @@ fn parse_impl(s: &str, negative: bool) -> Result<i256,
ParseI256Error> {
high.checked_mul(i256::from_i128(10_i128.pow(38)))
.and_then(|high| high.checked_add(low))
- .ok_or(ParseI256Error {})
+ .ok_or(I256Error::ParseError)
}
impl PartialOrd for i256 {
@@ -396,42 +407,101 @@ impl i256 {
.then_some(Self { low, high })
}
+ /// Return the least number of bits needed to represent the number
+ #[inline]
+ fn bits_required(&self) -> usize {
+ let le_bytes = self.to_le_bytes();
+ let arr: [u128; 2] = [
+ u128::from_le_bytes(le_bytes[0..16].try_into().unwrap()),
+ u128::from_le_bytes(le_bytes[16..32].try_into().unwrap()),
+ ];
+
+ let iter = arr.iter().rev().take(2 - 1);
+ if self.is_negative() {
+ let ctr = iter.take_while(|&&b| b == ::core::u128::MAX).count();
+ (128 * (2 - ctr)) + 1 - (!arr[2 - ctr - 1]).leading_zeros() as
usize
+ } else {
+ let ctr = iter.take_while(|&&b| b == ::core::u128::MIN).count();
+ (128 * (2 - ctr)) + 1 - arr[2 - ctr - 1].leading_zeros() as usize
+ }
+ }
+
+ /// Division operation, returns (quotient, remainder).
+ /// This basically implements [Long division]:
`<https://en.wikipedia.org/wiki/Division_algorithm>`
+ #[inline]
+ fn div_rem(self, other: Self) -> Result<(Self, Self), I256Error> {
+ if other == Self::ZERO {
+ return Err(I256Error::DivideByZero);
+ }
+ if other == Self::MINUS_ONE && self == Self::MIN {
+ return Err(I256Error::DivideOverflow);
+ }
+
+ if self == Self::MIN || other == Self::MIN {
+ let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
+ let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
+ let d = i256::from_bigint_with_overflow(&l / &r).0;
+ let r = i256::from_bigint_with_overflow(&l % &r).0;
+ return Ok((d, r));
+ }
+
+ let mut me = self.checked_abs().unwrap();
+ let mut you = other.checked_abs().unwrap();
+ let mut ret = [0u128; 2];
+ if me < you {
+ return Ok((Self::from_parts(ret[0], ret[1] as i128), self));
+ }
+
+ let shift = me.bits_required() - you.bits_required();
+ you = you.shl(shift as u8);
+ for i in (0..=shift).rev() {
+ if me >= you {
+ ret[i / 128] |= 1 << (i % 128);
+ me = me.checked_sub(you).unwrap();
+ }
+ you = you.shr(1);
+ }
+
+ Ok((
+ if self.is_negative() == other.is_negative() {
+ Self::from_parts(ret[0], ret[1] as i128)
+ } else {
+ -Self::from_parts(ret[0], ret[1] as i128)
+ },
+ if self.is_negative() { -me } else { me },
+ ))
+ }
+
/// Performs wrapping division
#[inline]
pub fn wrapping_div(self, other: Self) -> Self {
- let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
- let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
- Self::from_bigint_with_overflow(l / r).0
+ match self.div_rem(other) {
+ Ok((v, _)) => v,
+ Err(I256Error::DivideByZero) => panic!("attempt to divide by
zero"),
+ Err(_) => Self::MIN,
+ }
}
/// Performs checked division
#[inline]
pub fn checked_div(self, other: Self) -> Option<Self> {
- let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
- let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
- let (val, overflow) = Self::from_bigint_with_overflow(l / r);
- (!overflow).then_some(val)
+ self.div_rem(other).map(|(v, _)| v).ok()
}
/// Performs wrapping remainder
#[inline]
pub fn wrapping_rem(self, other: Self) -> Self {
- let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
- let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
- Self::from_bigint_with_overflow(l % r).0
+ match self.div_rem(other) {
+ Ok((_, v)) => v,
+ Err(I256Error::DivideByZero) => panic!("attempt to divide by
zero"),
+ Err(_) => Self::ZERO,
+ }
}
/// Performs checked remainder
#[inline]
pub fn checked_rem(self, other: Self) -> Option<Self> {
- if other == Self::ZERO {
- return None;
- }
-
- let l = BigInt::from_signed_bytes_le(&self.to_le_bytes());
- let r = BigInt::from_signed_bytes_le(&other.to_le_bytes());
- let (val, overflow) = Self::from_bigint_with_overflow(l % r);
- (!overflow).then_some(val)
+ self.div_rem(other).map(|(_, v)| v).ok()
}
/// Performs checked exponentiation
@@ -853,6 +923,43 @@ mod tests {
),
}
+ // Division
+ if ir != i256::ZERO {
+ let actual = il.wrapping_div(ir);
+ let expected = bl.clone() / br.clone();
+ let checked = il.checked_div(ir);
+
+ if ir == i256::MINUS_ONE && il == i256::MIN {
+ // BigInt produces an integer over i256::MAX
+ assert_eq!(actual, i256::MIN);
+ assert!(checked.is_none());
+ } else {
+ assert_eq!(actual.to_string(), expected.to_string());
+ assert_eq!(checked.unwrap().to_string(), expected.to_string());
+ }
+ } else {
+ // `wrapping_div` panics on division by zero
+ assert!(il.checked_div(ir).is_none());
+ }
+
+ // Remainder
+ if ir != i256::ZERO {
+ let actual = il.wrapping_rem(ir);
+ let expected = bl.clone() % br.clone();
+ let checked = il.checked_rem(ir);
+
+ assert_eq!(actual.to_string(), expected.to_string());
+
+ if ir == i256::MINUS_ONE && il == i256::MIN {
+ assert!(checked.is_none());
+ } else {
+ assert_eq!(checked.unwrap().to_string(), expected.to_string());
+ }
+ } else {
+ // `wrapping_rem` panics on division by zero
+ assert!(il.checked_rem(ir).is_none());
+ }
+
// Exponentiation
for exp in vec![0, 1, 2, 3, 8, 100].into_iter() {
let actual = il.wrapping_pow(exp);
