mbutrovich commented on code in PR #11580:
URL: https://github.com/apache/arrow-rs/pull/11580#discussion_r4314721331
##########
arrow-buffer/src/bigint/div.rs:
##########
@@ -72,6 +95,104 @@ fn div_rem_small<const N: usize>(numerator: &[u64; N],
divisor: u64) -> ([u64; N
(numerator, rem_padded)
}
+/// Like [`div_rem_small_word`], but divides each digit by multiplying with a
+/// reciprocal of the divisor, computed once
+#[cfg(any(
+ test,
+ all(target_pointer_width = "64", any(not(target_arch = "x86_64"), miri))
+))]
+fn div_rem_small_reciprocal<const N: usize>(
+ numerator: &[u64; N],
+ divisor: u64,
+) -> ([u64; N], [u64; N]) {
+ // Shifting both operands so the divisor's top bit is set leaves the
quotient unchanged
+ let shift = divisor.leading_zeros();
+ let divisor = divisor << shift;
+ let reciprocal = reciprocal_word(divisor);
+ let numerator = full_shl(numerator, shift);
+
+ // The top digit is below 2^shift, and so below the divisor
+ let mut rem = numerator[N];
+ let mut quotient = [0; N];
+ for i in (0..N).rev() {
+ let (q, r) = div_rem_word_reciprocal(rem, numerator[i], divisor,
reciprocal);
+ quotient[i] = q;
+ rem = r;
+ }
+
+ let mut rem_padded = [0; N];
+ rem_padded[0] = rem >> shift;
+ (quotient, rem_padded)
+}
+
+/// Returns `floor((2^128 - 1) / divisor) - 2^64` for a divisor with its top
bit set
+///
+/// This is algorithm 2 of [Möller and Granlund
(2011)](https://gmplib.org/~tege/division-paper.pdf),
+/// which avoids a `u128` division
+#[cfg(any(
+ test,
+ all(target_pointer_width = "64", any(not(target_arch = "x86_64"), miri))
+))]
+fn reciprocal_word(divisor: u64) -> u64 {
+ debug_assert_eq!(divisor >> 63, 1);
+
+ // `floor((2^19 - 3 * 2^8) / d9)` for each value of the top 9 bits `d9`
+ const TABLE: [u16; 256] = {
+ let mut table = [0; 256];
+ let mut i = 0;
+ while i < 256 {
+ table[i] = (((1 << 19) - 3 * (1 << 8)) / (i as u32 + 256)) as u16;
+ i += 1;
+ }
+ table
+ };
+
+ let d0 = divisor & 1;
+ let d9 = divisor >> 55;
+ let d40 = (divisor >> 24) + 1;
+ let d63 = (divisor >> 1) + d0;
+
+ // `d9` is in 256..512, so masking gives `d9 - 256` without a bounds check
+ let v0 = u64::from(TABLE[(d9 & 0xff) as usize]);
+ // Two Newton iterations, each about doubling the bits of precision
+ let v1 = (v0 << 11) - ((v0 * v0 * d40) >> 40) - 1;
+ let v2 = (v1 << 13) + ((v1 * ((1 << 60) - v1 * d40)) >> 47);
+ // `2^96 - v2 * d63 + floor(v2 / 2) * d0`, which the paper shows fits in
64 bits
+ let e = ((v2 >> 1) & d0.wrapping_neg()).wrapping_sub(v2.wrapping_mul(d63));
+ let v3 = (v2 << 31).wrapping_add(((u128::from(v2) * u128::from(e)) >> 65)
as u64);
+ // `floor((v3 + 2^64 + 1) * divisor / 2^64)`, without overflowing a `u128`
+ let p = (((u128::from(v3) + 1) * u128::from(divisor)) >> 64) as u64;
+ v3.wrapping_sub(p.wrapping_add(divisor))
+}
+
+/// Divides `hi:lo` by a divisor with its top bit set, given its
[`reciprocal_word`]
+///
+/// This is algorithm 4 of [Möller and Granlund
(2011)](https://gmplib.org/~tege/division-paper.pdf)
Review Comment:
> Maybe add the paper name in the comment in case the link ever die
Done here too.
##########
arrow-buffer/src/bigint/div.rs:
##########
@@ -300,3 +421,116 @@ impl<T, const N: usize> std::ops::DerefMut for
ArrayPlusOne<T, N> {
unsafe { std::slice::from_raw_parts_mut(x.cast::<T>(), N + 1) }
}
}
+
+#[cfg(test)]
+mod tests {
+ use super::*;
+ use rand::{RngExt, rng};
+
+ fn check_div_rem_small<const N: usize>(numerator: [u64; N], divisor: u64) {
+ let expected = div_rem_small_word(&numerator, divisor);
+ assert_eq!(
+ div_rem_small_reciprocal(&numerator, divisor),
+ expected,
+ "{numerator:?} / {divisor}"
+ );
+ }
+
+ #[test]
+ fn test_reciprocal_word() {
+ let mut divisors = vec![1 << 63, (1 << 63) + 1, u64::MAX - 1,
u64::MAX];
+ // Each boundary between entries of the lookup table
+ for i in 257..512u64 {
+ divisors.extend([(i << 55) - 1, i << 55, (i << 55) + 1]);
+ }
+ let mut rng = rng();
+ divisors.extend((0..100_000).map(|_| rng.random::<u64>() | (1 << 63)));
+ for divisor in divisors {
+ let n = (u128::from(!divisor) << 64) | u128::from(u64::MAX);
+ let expected = (n / u128::from(divisor)) as u64;
+ assert_eq!(reciprocal_word(divisor), expected, "{divisor}");
+ }
+ }
+
+ #[test]
+ fn test_div_rem_word_reciprocal() {
+ // The second correction, which random inputs rarely reach, runs for
this input
+ let (hi, lo, divisor) = (
+ 3787258589690185862,
+ 17800796494353475223,
+ 9357586864501027680,
+ );
+ let n = (u128::from(hi) << 64) | u128::from(lo);
+ let d = u128::from(divisor);
+ assert_eq!(
+ div_rem_word_reciprocal(hi, lo, divisor, reciprocal_word(divisor)),
+ ((n / d) as u64, (n % d) as u64)
+ );
+
+ let mut rng = rng();
Review Comment:
> can you make it seeded rng to be deterministic
All three tests now use `StdRng::seed_from_u64(42)`, like the other seeded
tests in `arrow-buffer`.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
