mbutrovich commented on code in PR #12232:
URL: https://github.com/apache/arrow-rs/pull/12232#discussion_r4564701619
##########
arrow-buffer/src/divisor.rs:
##########
@@ -586,6 +996,18 @@ mod tests {
divisors.retain(|d| *d != 0);
let mut values: Vec<$t> = vec![0, 1, -1, <$t>::MIN, <$t>::MAX,
<$t>::MIN + 1];
+ // Each side of every power of two, which includes the 64-bit
limbs of the
Review Comment:
#12167 asks for dividends on each side of 2^64 and 2^128 and of 10^k.
`test_signed!` adds each side of every power of two and every power of ten,
positive and negative, which covers the 64-bit limbs of the high multiply and
the edges where a kernel would narrow to `i64` or `i128`. The new values also
run for `i32` and `i64`.
An exhaustive check, as in #11467, isn't possible at 128 or 256 bits. The
tests keep the check from #11467 at both ends of the last quotient steps on
each side of zero, where the multiplier's error is largest, for every divisor
in the list.
##########
arrow-buffer/src/divisor.rs:
##########
@@ -470,20 +510,390 @@ unsigned_divisor!(
u64,
u128
);
-signed_divisor!(
+/// Expands to the arguments of `signed_divisor!` that compute with `$wide`, a
type twice as wide
+/// as `$t`
+macro_rules! signed_divisor_wide {
+ ($name:ident, $power_of_two:ident, $add:ident, $t:ty, $wide:ty) => {
+ signed_divisor!(
+ $name,
+ $power_of_two,
+ $add,
+ $t,
+ [const],
+ // m = 1 + floor(2^(N + l - 1) / |divisor|) is between 2^(N - 1)
and 2^N, so m - 2^N
+ // fits in the narrow type
+ |magnitude, l| {
+ let m = 1 + (1 << (<$t>::BITS + l - 1)) / magnitude as $wide;
+ (m - (1 << <$t>::BITS)) as $t
+ },
+ // `n + MULSH(m', n)` in figure 5.1
+ |m, n| n + ((m as $wide * n as $wide) >> <$t>::BITS) as $t,
+ []
+ );
+ };
+}
+
+signed_divisor_wide!(
DivisorI32,
PowerOfTwoDivisorI32,
AddIndicatorDivisorI32,
i32,
i64
);
-signed_divisor!(
+signed_divisor_wide!(
DivisorI64,
PowerOfTwoDivisorI64,
AddIndicatorDivisorI64,
i64,
i128
);
+signed_divisor!(
+ DivisorI128,
+ PowerOfTwoDivisorI128,
+ AddIndicatorDivisorI128,
+ i128,
+ [],
+ // m = 1 + floor(2^(127 + l) / |divisor|), as for the narrower types.
2^(127 + l) is at most
+ // 2^254, so `i256` holds it, and the quotient is below 2^128
+ |magnitude, l| {
+ let pow = i256::ONE << (127 + l) as u8;
+ let (q, _) = pow.wrapping_div(i256::from_parts(magnitude,
0)).to_parts();
+ q.wrapping_add(1) as i128
+ },
+ // Figure 5.1 computes `n + MULSH(m', n)`, and `i128` has no wider type to
compute `MULSH`
+ // in. Section 3 of the paper gives `MULSH(x, y) = MULUH(x, y) - AND(x,
XSIGN(y))
+ // - AND(y, XSIGN(x))`, which is program (1) in section 8-3 of Warren,
"Hacker's Delight", 2nd
+ // edition. `m'` is negative for every divisor but 1 and -1, which are
powers of two, so
+ // `AND(n, XSIGN(m'))` is `n` and cancels the add, leaving `MULUH(m', n) -
AND(m', XSIGN(n))`
+ |m, n| (mulx(m as u128, n as u128).1 as i128).wrapping_sub((n >> 127) & m),
+ [
+ ///
+ /// When every value of a batch fits in `i64`, dividing in `i64` with
[`DivisorI64`] gives
+ /// the same results, except for a divisor of -1, because `i64::MIN /
-1` overflows in
+ /// `i64` but not in `i128`. A divisor that doesn't fit in `i64`
exceeds every such value,
+ /// so the quotient is 0 and the remainder is the value.
+ ///
+ /// ```
+ /// # use arrow_buffer::{DivisorI128, DivisorI64};
+ /// let values: Vec<i128> = vec![-12_345, 678, i64::MIN as i128];
+ /// let divisor: i128 = 100;
+ /// let narrow = values.iter().all(|n| i64::try_from(*n).is_ok());
+ /// let quotients: Vec<i128> = match i64::try_from(divisor) {
+ /// Ok(d) if narrow && d != -1 => {
+ /// let d = DivisorI64::new(d).unwrap();
+ /// values.iter().map(|n| d.wrapping_div_rem(*n as i64).0 as
i128).collect()
+ /// }
+ /// _ => {
+ /// let d = DivisorI128::new(divisor).unwrap();
+ /// values.iter().map(|n| d.wrapping_div_rem(*n).0).collect()
+ /// }
+ /// };
+ /// assert_eq!(quotients, [-123, 6, -92233720368547758]);
+ /// ```
+ ]
+);
+
+// `signed_divisor!` can't build `DivisorI256`, because `i256` has no `BITS`
or `unsigned_abs` and
+// its operators can't be called in a `const fn`. Its variants compute on the
parts of the values
+// with the helpers below `impl DivisorI256` instead
+divisor_enum!(
+ DivisorI256,
+ i256,
+ concat!(
+ "# use arrow_buffer::{i256, DivisorI256};\n",
+ "let values: Vec<i256> = [100, 7,
64].map(i256::from_i128).to_vec();\n",
+ "let quotients: Vec<_> = match
DivisorI256::new(i256::from_i128(7)).unwrap() {\n",
+ " DivisorI256::PowerOfTwo(d) => values.iter().map(move |n|
d.wrapping_div_rem(*n).0).collect(),\n",
+ " DivisorI256::AddIndicator(d) => values.iter().map(move |n|
d.wrapping_div_rem(*n).0).collect(),\n",
+ " d => values.iter().map(move |n|
d.wrapping_div_rem(*n).0).collect(),\n",
+ "};\n",
+ "assert_eq!(quotients, [14, 1, 9].map(i256::from_i128));",
+ ),
+ concat!(
+ "# use arrow_buffer::{i256, DivisorI256};\n",
+ "let divisor = DivisorI256::new(i256::from_i128(7)).unwrap();\n",
+ "let (q, r) = divisor.wrapping_div_rem(i256::from_i128(100));\n",
+ "assert_eq!((q, r), (i256::from_i128(14), i256::from_i128(2)));",
+ ),
+ [
+ ///
+ /// Unlike the other divisor types, the variants' methods branch on
the divisor's shift. The
+ /// shift is the same for every value, so the branch is predicted, and
a form without it
+ /// was slower.
+ ///
+ /// When every value of a batch fits in `i128`, dividing in `i128`
with [`DivisorI128`]
+ /// gives the same results, except for a divisor of -1, because
`i128::MIN / -1` overflows
+ /// in `i128` but not in `i256`. A divisor that doesn't fit in `i128`
exceeds every such
+ /// value, so the quotient is 0 and the remainder is the value.
+ ],
+ /// A divisor whose magnitude is a power of two
+ PowerOfTwo(PowerOfTwoDivisorI256),
+ /// Any other divisor
+ AddIndicator(AddIndicatorDivisorI256),
+);
+
+impl DivisorI256 {
+ /// Returns a divisor for `divisor`, or `None` if `divisor` is zero
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_buffer::{i256, DivisorI256};
+ /// assert!(DivisorI256::new(i256::ZERO).is_none());
+ /// ```
+ pub fn new(divisor: i256) -> Option<Self> {
+ if divisor == i256::ZERO {
+ return None;
+ }
+ Some(match PowerOfTwoDivisorI256::new(divisor) {
+ Some(d) => Self::PowerOfTwo(d),
+ None => Self::AddIndicator(AddIndicatorDivisorI256::new(divisor)),
+ })
+ }
+
+ /// Returns `(n / divisor, n % divisor)`, or `None` for `i256::MIN / -1`
+ ///
+ /// It matches on the variant on each call. Only [`Self::PowerOfTwo`] can
return `None`.
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_buffer::{i256, DivisorI256};
+ /// let divisor = DivisorI256::new(i256::from_i128(-7)).unwrap();
+ /// let qr = divisor.checked_div_rem(i256::from_i128(-100));
+ /// assert_eq!(qr, Some((i256::from_i128(14), i256::from_i128(-2))));
+ ///
+ /// let minus_one = DivisorI256::new(i256::MINUS_ONE).unwrap();
+ /// assert_eq!(minus_one.checked_div_rem(i256::MIN), None);
+ /// ```
+ #[inline]
+ pub const fn checked_div_rem(self, n: i256) -> Option<(i256, i256)> {
+ match self {
+ Self::PowerOfTwo(d) => d.checked_div_rem(n),
+ Self::AddIndicator(d) => Some(d.wrapping_div_rem(n)),
+ }
+ }
+}
+
+/// `i256` operations that a `const fn` can call, on the parts of the values.
`shr` is `SRA` from
+/// table 3.1 of Granlund and Montgomery for `0..=255` bits
+///
+/// It branches on the shift, so the `i256` variants branch on the divisor,
unlike the other
+/// divisor types. The shift is the same for every value, so the branch is
predicted. A
+/// branch-free form that shifts both halves and selects was slower
+const fn shr(x: i256, shift: u32) -> i256 {
+ let (low, high) = x.to_parts();
+ if shift == 0 {
+ x
+ } else if shift < 128 {
+ i256::from_parts(
+ (low >> shift) | ((high as u128) << (128 - shift)),
+ high >> shift,
+ )
+ } else {
+ i256::from_parts((high >> (shift - 128)) as u128, high >> 127)
+ }
+}
+
+const fn and(a: i256, b: i256) -> i256 {
+ let ((al, ah), (bl, bh)) = (a.to_parts(), b.to_parts());
+ i256::from_parts(al & bl, ah & bh)
+}
+
+const fn xor(a: i256, b: i256) -> i256 {
+ let ((al, ah), (bl, bh)) = (a.to_parts(), b.to_parts());
+ i256::from_parts(al ^ bl, ah ^ bh)
+}
+
+const fn not(a: i256) -> i256 {
+ let (low, high) = a.to_parts();
+ i256::from_parts(!low, !high)
+}
+
+const fn eq(a: i256, b: i256) -> bool {
+ let ((al, ah), (bl, bh)) = (a.to_parts(), b.to_parts());
+ al == bl && ah == bh
+}
+
+/// The sign of `x` spread to every bit, -1 if `x` is negative, otherwise 0
+const fn sign(x: i256) -> i256 {
+ let s = x.to_parts().1 >> 127;
+ i256::from_parts(s as u128, s)
+}
+
+/// The high 256 bits of the unsigned product of the bits of `a` and `b`
+#[inline]
+const fn mul_high_u256(a: i256, b: i256) -> i256 {
+ let ((a0, a1), (b0, b1)) = (a.to_parts(), b.to_parts());
+ let (a1, b1) = (a1 as u128, b1 as u128);
+ let (_, p00) = mulx(a0, b0);
+ let (p01_low, p01_high) = mulx(a0, b1);
+ let (p10_low, p10_high) = mulx(a1, b0);
+ let (p11_low, p11_high) = mulx(a1, b1);
+ // Bits 128 to 255 of the product, of which only the carries matter
+ let (s, c1) = p00.overflowing_add(p01_low);
+ let (_, c2) = s.overflowing_add(p10_low);
+ let (low, c3) = p01_high.overflowing_add(p10_high);
+ let (low, c4) = low.overflowing_add(p11_low);
+ let (low, c5) = low.overflowing_add(c1 as u128 + c2 as u128);
+ // The product is below 2^512, so this doesn't overflow
+ let high = p11_high + c3 as u128 + c4 as u128 + c5 as u128;
+ i256::from_parts(low, high as i128)
+}
+
+/// An `i256` divisor whose magnitude is a power of two, see [`DivisorI256`]
+///
+/// The quotient and remainder follow sections 10-1 and 10-2 of Warren,
"Hacker's Delight", 2nd
+/// edition. Those sections compute the addend `2^k - 1` for a negative `n`
with two shifts. This
+/// type computes it with one shift that spreads the sign bit of `n` and an
AND with `2^k - 1`,
+/// which it computes once for the divisor.
+#[derive(Debug, Clone, Copy)]
+pub struct PowerOfTwoDivisorI256 {
+ shift: u32,
+ /// `|divisor| - 1`
+ mask: i256,
+ /// -1 if the divisor is negative, otherwise 0
+ sign: i256,
+}
+
+impl PowerOfTwoDivisorI256 {
+ fn new(divisor: i256) -> Option<Self> {
+ // `wrapping_abs` leaves `MIN` as it is, whose bits are those of its
magnitude, 2^255
+ let magnitude = divisor.wrapping_abs();
+ let mask = magnitude.wrapping_sub(i256::ONE);
+ if !eq(and(magnitude, mask), i256::ZERO) {
+ return None;
+ }
+ Some(Self {
+ shift: magnitude.trailing_zeros(),
+ mask,
+ sign: sign(divisor),
+ })
+ }
+
+ /// Returns `(n.wrapping_div(divisor), n.wrapping_rem(divisor))`, so
`i256::MIN / -1` returns
+ /// `(i256::MIN, 0)`
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_buffer::{i256, DivisorI256};
+ /// let DivisorI256::PowerOfTwo(divisor) =
DivisorI256::new(i256::from_i128(-8)).unwrap() else { unreachable!() };
+ /// let (q, r) = divisor.wrapping_div_rem(i256::from_i128(-100));
+ /// assert_eq!((q, r), (i256::from_i128(12), i256::from_i128(-4)));
+ ///
+ /// let DivisorI256::PowerOfTwo(minus_one) =
DivisorI256::new(i256::MINUS_ONE).unwrap() else { unreachable!() };
+ /// assert_eq!(minus_one.wrapping_div_rem(i256::MIN), (i256::MIN,
i256::ZERO));
+ /// ```
+ #[inline]
+ pub const fn wrapping_div_rem(self, n: i256) -> (i256, i256) {
+ // Adding `|divisor| - 1` to a negative `n` makes the shift round
toward zero
+ let t = n.wrapping_add(and(sign(n), self.mask));
+ let q = shr(t, self.shift);
+ // `(x ^ sign) - sign` negates `x` when `sign` is -1, and wraps only
for `MIN / -1`
+ // `t & !mask` is `n` rounded toward zero to a multiple of `|divisor|`
+ (
+ xor(q, self.sign).wrapping_sub(self.sign),
+ n.wrapping_sub(and(t, not(self.mask))),
+ )
+ }
+
+ /// Returns `(n / divisor, n % divisor)`, or `None` for `i256::MIN / -1`
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_buffer::{i256, DivisorI256};
+ /// let DivisorI256::PowerOfTwo(divisor) =
DivisorI256::new(i256::from_i128(-8)).unwrap() else { unreachable!() };
+ /// let qr = divisor.checked_div_rem(i256::from_i128(-100));
+ /// assert_eq!(qr, Some((i256::from_i128(12), i256::from_i128(-4))));
+ ///
+ /// let DivisorI256::PowerOfTwo(minus_one) =
DivisorI256::new(i256::MINUS_ONE).unwrap() else { unreachable!() };
+ /// assert_eq!(minus_one.checked_div_rem(i256::MIN), None);
+ /// ```
+ #[inline]
+ pub const fn checked_div_rem(self, n: i256) -> Option<(i256, i256)> {
+ // -1 is the only divisor with a negative sign and a zero mask
+ if eq(n, i256::MIN) && eq(self.sign, i256::MINUS_ONE) && eq(self.mask,
i256::ZERO) {
+ None
+ } else {
+ Some(self.wrapping_div_rem(n))
+ }
+ }
+}
+
+/// An `i256` divisor whose magnitude isn't a power of two, see [`DivisorI256`]
+///
+/// It uses the signed multiplier from figure 5.1 of
+/// [Granlund and Montgomery, "Division by Invariant Integers using
Multiplication" (1994)](https://gmplib.org/~tege/divcnst-pldi94.pdf),
+/// as [`AddIndicatorDivisorI64`] does.
+#[derive(Debug, Clone, Copy)]
+pub struct AddIndicatorDivisorI256 {
+ divisor: i256,
+ /// `m'` in figure 5.1
+ multiplier: i256,
+ /// `sh_post` in figure 5.1
+ shift: u32,
+}
+
+impl AddIndicatorDivisorI256 {
+ /// `|divisor|` must be at least 3 and not a power of two
+ fn new(divisor: i256) -> Self {
+ let magnitude = divisor.wrapping_abs();
+ // ceil(log2(|divisor|)), between 2 and 255
+ let l = 256 - magnitude.wrapping_sub(i256::ONE).leading_zeros();
+ // m = 1 + floor(2^(255 + l) / |divisor|) is between 2^255 and 2^256,
so its bits fit in
+ // an `i256`, and they are the bits of m - 2^256
+ let bit = 255 + l as usize;
+ let mut pow = [0u64; 8];
+ pow[bit / 64] = 1 << (bit % 64);
+ let (low, high) = magnitude.to_parts();
+ let high = high as u128;
+ let d = [
+ low as u64,
+ (low >> 64) as u64,
+ high as u64,
+ (high >> 64) as u64,
+ ];
+ let (q, _) = div_rem(&pow, &[d[0], d[1], d[2], d[3], 0, 0, 0, 0]);
+ let q = i256::from_parts(
+ q[0] as u128 | ((q[1] as u128) << 64),
+ (q[2] as u128 | ((q[3] as u128) << 64)) as i128,
+ );
+ Self {
+ divisor,
+ multiplier: q.wrapping_add(i256::ONE),
+ shift: l - 1,
+ }
+ }
+
+ /// Returns `(n / divisor, n % divisor)`, which can't overflow
+ ///
+ /// # Example
+ /// ```
+ /// # use arrow_buffer::{i256, DivisorI256};
+ /// let DivisorI256::AddIndicator(divisor) =
DivisorI256::new(i256::from_i128(-7)).unwrap() else { unreachable!() };
+ /// let (q, r) = divisor.wrapping_div_rem(i256::from_i128(-100));
+ /// assert_eq!((q, r), (i256::from_i128(14), i256::from_i128(-2)));
+ /// ```
+ // At `#[inline]`, LLVM calls this out of line for each value, without LTO
and with thin or
+ // fat LTO, at 16 codegen units and at 1
+ #[inline(always)]
Review Comment:
The other divisor methods are `#[inline]`, as in #11467. This one is about
220 instructions on aarch64 and 310 on x86_64, and at `#[inline]` LLVM called
it out of line for each value. I built the `divisor` benchmark with each
attribute for aarch64, x86_64 and `x86-64-v3`, with the profiles `cargo bench`
and the bot use (no LTO and thin LTO, 16 codegen units) and with those of
DataFusion and Comet releases (fat and thin LTO, 1 codegen unit). At
`#[inline]` the loops called the method in all 12 builds. At
`#[inline(always)]` none did.
On an Apple M5 Max, inlining made `i256` `/` 1.51x to 1.52x and `%` 1.18x to
1.21x faster in the `divisor` benchmark.
--
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]
