diff --git a/encodings/fastlanes/Cargo.toml b/encodings/fastlanes/Cargo.toml index 9085390b67b..4a187f8faae 100644 --- a/encodings/fastlanes/Cargo.toml +++ b/encodings/fastlanes/Cargo.toml @@ -58,6 +58,11 @@ name = "compute_between" harness = false required-features = ["_test-harness"] +[[bench]] +name = "bit_transpose" +harness = false +required-features = ["_test-harness"] + [[bench]] name = "bitpack_compare" harness = false diff --git a/encodings/fastlanes/benches/bit_transpose.rs b/encodings/fastlanes/benches/bit_transpose.rs new file mode 100644 index 00000000000..08c3ffb12e5 --- /dev/null +++ b/encodings/fastlanes/benches/bit_transpose.rs @@ -0,0 +1,310 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use divan::Bencher; +use vortex_fastlanes::bit_transpose::scalar::transpose_bits_scalar; +use vortex_fastlanes::bit_transpose::scalar::untranspose_bits_scalar; + +fn main() { + divan::main(); +} + +/// Generate deterministic test data. +#[expect(clippy::cast_possible_truncation)] +fn generate_test_data(seed: usize) -> [u8; 128] { + let mut data = [0u8; 128]; + for (i, byte) in data.iter_mut().enumerate() { + *byte = seed.wrapping_mul(17).wrapping_add(i).wrapping_mul(31) as u8; + } + data +} + +const BATCH_SIZE: usize = 1000; + +// ============================================================================ +// Transpose: single array +// ============================================================================ + +#[divan::bench] +fn transpose_scalar(bencher: Bencher) { + let input = generate_test_data(42); + + bencher + .with_inputs(|| (&input, [0u8; 128])) + .bench_refs(|(input, output)| { + transpose_bits_scalar(input, output); + }); +} + +// ============================================================================ +// Transpose: throughput (1000 arrays) +// ============================================================================ + +#[divan::bench] +fn transpose_scalar_throughput(bencher: Bencher) { + let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); + + bencher + .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) + .bench_refs(|(inputs, outputs)| { + for (input, output) in inputs.iter().zip(outputs.iter_mut()) { + transpose_bits_scalar(input, output); + } + }); +} + +// ============================================================================ +// Untranspose: single array +// ============================================================================ + +#[divan::bench] +fn untranspose_scalar(bencher: Bencher) { + let input = generate_test_data(42); + + bencher + .with_inputs(|| (&input, [0u8; 128])) + .bench_refs(|(input, output)| { + untranspose_bits_scalar(input, output); + }); +} + +// ============================================================================ +// Untranspose: throughput (1000 arrays) +// ============================================================================ + +#[divan::bench] +fn untranspose_scalar_throughput(bencher: Bencher) { + let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); + + bencher + .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) + .bench_refs(|(inputs, outputs)| { + for (input, output) in inputs.iter().zip(outputs.iter_mut()) { + untranspose_bits_scalar(input, output); + } + }); +} + +// ============================================================================ +// x86_64 benchmarks +// ============================================================================ + +#[cfg(target_arch = "x86_64")] +mod x86 { + use divan::Bencher; + use vortex_fastlanes::bit_transpose::x86::has_bmi2; + use vortex_fastlanes::bit_transpose::x86::has_vbmi; + use vortex_fastlanes::bit_transpose::x86::transpose_bits_bmi2; + use vortex_fastlanes::bit_transpose::x86::transpose_bits_vbmi; + use vortex_fastlanes::bit_transpose::x86::untranspose_bits_bmi2; + use vortex_fastlanes::bit_transpose::x86::untranspose_bits_vbmi; + + use super::BATCH_SIZE; + use super::generate_test_data; + + // --- Transpose: single array --- + + #[divan::bench] + fn transpose_bmi2(bencher: Bencher) { + if !has_bmi2() { + return; + } + + let input = generate_test_data(42); + + bencher + .with_inputs(|| (&input, [0u8; 128])) + .bench_refs(|(input, output)| { + unsafe { transpose_bits_bmi2(input, output) }; + }); + } + + #[divan::bench] + fn transpose_vbmi(bencher: Bencher) { + if !has_vbmi() { + return; + } + + let input = generate_test_data(42); + + bencher + .with_inputs(|| (&input, [0u8; 128])) + .bench_refs(|(input, output)| { + unsafe { transpose_bits_vbmi(input, output) }; + }); + } + + // --- Untranspose: single array --- + + #[divan::bench] + fn untranspose_bmi2(bencher: Bencher) { + if !has_bmi2() { + return; + } + + let input = generate_test_data(42); + + bencher + .with_inputs(|| (&input, [0u8; 128])) + .bench_refs(|(input, output)| { + unsafe { untranspose_bits_bmi2(input, output) }; + }); + } + + #[divan::bench] + fn untranspose_vbmi(bencher: Bencher) { + if !has_vbmi() { + return; + } + + let input = generate_test_data(42); + + bencher + .with_inputs(|| (&input, [0u8; 128])) + .bench_refs(|(input, output)| { + unsafe { untranspose_bits_vbmi(input, output) }; + }); + } + + // --- Transpose: throughput (1000 arrays) --- + + #[divan::bench] + fn transpose_bmi2_throughput(bencher: Bencher) { + if !has_bmi2() { + return; + } + + let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); + + bencher + .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) + .bench_refs(|(inputs, outputs)| { + for (input, output) in inputs.iter().zip(outputs.iter_mut()) { + unsafe { transpose_bits_bmi2(input, output) }; + } + }); + } + + #[divan::bench] + fn transpose_vbmi_throughput(bencher: Bencher) { + if !has_vbmi() { + return; + } + + let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); + + bencher + .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) + .bench_refs(|(inputs, outputs)| { + for (input, output) in inputs.iter().zip(outputs.iter_mut()) { + unsafe { transpose_bits_vbmi(input, output) }; + } + }); + } + + // --- Untranspose: throughput (1000 arrays) --- + + #[divan::bench] + fn untranspose_bmi2_throughput(bencher: Bencher) { + if !has_bmi2() { + return; + } + + let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); + + bencher + .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) + .bench_refs(|(inputs, outputs)| { + for (input, output) in inputs.iter().zip(outputs.iter_mut()) { + unsafe { untranspose_bits_bmi2(input, output) }; + } + }); + } + + #[divan::bench] + fn untranspose_vbmi_throughput(bencher: Bencher) { + if !has_vbmi() { + return; + } + + let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); + + bencher + .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) + .bench_refs(|(inputs, outputs)| { + for (input, output) in inputs.iter().zip(outputs.iter_mut()) { + unsafe { untranspose_bits_vbmi(input, output) }; + } + }); + } +} + +// ============================================================================ +// aarch64 benchmarks +// ============================================================================ + +#[cfg(target_arch = "aarch64")] +mod aarch64 { + use vortex_fastlanes::bit_transpose::aarch64::transpose_bits_neon; + use vortex_fastlanes::bit_transpose::aarch64::untranspose_bits_neon; + + use super::BATCH_SIZE; + use super::Bencher; + use super::generate_test_data; + + // --- Transpose: single array --- + + #[divan::bench] + fn transpose_neon(bencher: Bencher) { + let input = generate_test_data(42); + + bencher + .with_inputs(|| (&input, [0u8; 128])) + .bench_refs(|(input, output)| { + unsafe { transpose_bits_neon(input, output) }; + }); + } + + // --- Untranspose: single array --- + + #[divan::bench] + fn untranspose_neon(bencher: Bencher) { + let input = generate_test_data(42); + + bencher + .with_inputs(|| (&input, [0u8; 128])) + .bench_refs(|(input, output)| { + unsafe { untranspose_bits_neon(input, output) }; + }); + } + + // --- Transpose: throughput (1000 arrays) --- + + #[divan::bench] + fn transpose_neon_throughput(bencher: Bencher) { + let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); + + bencher + .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) + .bench_refs(|(inputs, outputs)| { + for (input, output) in inputs.iter().zip(outputs.iter_mut()) { + unsafe { transpose_bits_neon(input, output) }; + } + }); + } + + // --- Untranspose: throughput (1000 arrays) --- + + #[divan::bench] + fn untranspose_neon_throughput(bencher: Bencher) { + let inputs: Vec<[u8; 128]> = (0..BATCH_SIZE).map(generate_test_data).collect(); + + bencher + .with_inputs(|| (&inputs, vec![[0u8; 128]; BATCH_SIZE])) + .bench_refs(|(inputs, outputs)| { + for (input, output) in inputs.iter().zip(outputs.iter_mut()) { + unsafe { untranspose_bits_neon(input, output) }; + } + }); + } +} diff --git a/encodings/fastlanes/src/bit_transpose/aarch64.rs b/encodings/fastlanes/src/bit_transpose/aarch64.rs new file mode 100644 index 00000000000..85dbb1e0690 --- /dev/null +++ b/encodings/fastlanes/src/bit_transpose/aarch64.rs @@ -0,0 +1,299 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +#![cfg(target_arch = "aarch64")] + +use core::arch::aarch64::uint64x2_t; +use core::arch::aarch64::vandq_u64; +use core::arch::aarch64::vdupq_n_u64; +use core::arch::aarch64::veorq_u64; +use core::arch::aarch64::vgetq_lane_u64; +use core::arch::aarch64::vld1q_u8; +use core::arch::aarch64::vld1q_u8_x4; +use core::arch::aarch64::vorrq_u8; +use core::arch::aarch64::vqtbl4q_u8; +use core::arch::aarch64::vreinterpretq_u8_u64; +use core::arch::aarch64::vreinterpretq_u64_u8; +use core::arch::aarch64::vshlq_n_u64; +use core::arch::aarch64::vshrq_n_u64; +use core::arch::aarch64::vst1q_u8; + +use crate::bit_transpose::BASE_PATTERN_FIRST; +use crate::bit_transpose::BASE_PATTERN_SECOND; +use crate::bit_transpose::TRANSPOSE_2X2; +use crate::bit_transpose::TRANSPOSE_4X4; +use crate::bit_transpose::TRANSPOSE_8X8; + +/// Gather indices for the first half from input[0..64]. +/// Each group needs 4 bytes at stride 16 (the low half of the stride pattern). +/// Layout: [`g0_from_lo(4` bytes), pad(4 bytes), `g1_from_lo(4` bytes), pad(4 bytes), ...] +/// Two groups per 16-byte NEON register. +static GATHER_FIRST_LO: [[u8; 16]; 4] = [ + // Groups 0,1 from BASE_PATTERN_FIRST: bases 0, 8 + [ + 0, 16, 32, 48, 0xFF, 0xFF, 0xFF, 0xFF, 8, 24, 40, 56, 0xFF, 0xFF, 0xFF, 0xFF, + ], + // Groups 2,3: bases 4, 12 + [ + 4, 20, 36, 52, 0xFF, 0xFF, 0xFF, 0xFF, 12, 28, 44, 60, 0xFF, 0xFF, 0xFF, 0xFF, + ], + // Groups 4,5: bases 2, 10 + [ + 2, 18, 34, 50, 0xFF, 0xFF, 0xFF, 0xFF, 10, 26, 42, 58, 0xFF, 0xFF, 0xFF, 0xFF, + ], + // Groups 6,7: bases 6, 14 + [ + 6, 22, 38, 54, 0xFF, 0xFF, 0xFF, 0xFF, 14, 30, 46, 62, 0xFF, 0xFF, 0xFF, 0xFF, + ], +]; + +/// Gather indices for the first half from input[64..128]. +/// These fill in bytes 4-7 of each u64 (the high half of the stride pattern). +static GATHER_FIRST_HI: [[u8; 16]; 4] = [ + // Groups 0,1: bases 0, 8 (offset by -64 since table starts at input[64]) + [ + 0xFF, 0xFF, 0xFF, 0xFF, 0, 16, 32, 48, 0xFF, 0xFF, 0xFF, 0xFF, 8, 24, 40, 56, + ], + // Groups 2,3: bases 4, 12 + [ + 0xFF, 0xFF, 0xFF, 0xFF, 4, 20, 36, 52, 0xFF, 0xFF, 0xFF, 0xFF, 12, 28, 44, 60, + ], + // Groups 4,5: bases 2, 10 + [ + 0xFF, 0xFF, 0xFF, 0xFF, 2, 18, 34, 50, 0xFF, 0xFF, 0xFF, 0xFF, 10, 26, 42, 58, + ], + // Groups 6,7: bases 6, 14 + [ + 0xFF, 0xFF, 0xFF, 0xFF, 6, 22, 38, 54, 0xFF, 0xFF, 0xFF, 0xFF, 14, 30, 46, 62, + ], +]; + +/// Gather indices for the second half from input[0..64]. +/// Uses `BASE_PATTERN_SECOND`: bases [1, 9, 5, 13, 3, 11, 7, 15] +static GATHER_SECOND_LO: [[u8; 16]; 4] = [ + [ + 1, 17, 33, 49, 0xFF, 0xFF, 0xFF, 0xFF, 9, 25, 41, 57, 0xFF, 0xFF, 0xFF, 0xFF, + ], + [ + 5, 21, 37, 53, 0xFF, 0xFF, 0xFF, 0xFF, 13, 29, 45, 61, 0xFF, 0xFF, 0xFF, 0xFF, + ], + [ + 3, 19, 35, 51, 0xFF, 0xFF, 0xFF, 0xFF, 11, 27, 43, 59, 0xFF, 0xFF, 0xFF, 0xFF, + ], + [ + 7, 23, 39, 55, 0xFF, 0xFF, 0xFF, 0xFF, 15, 31, 47, 63, 0xFF, 0xFF, 0xFF, 0xFF, + ], +]; + +/// Gather indices for the second half from input[64..128]. +static GATHER_SECOND_HI: [[u8; 16]; 4] = [ + [ + 0xFF, 0xFF, 0xFF, 0xFF, 1, 17, 33, 49, 0xFF, 0xFF, 0xFF, 0xFF, 9, 25, 41, 57, + ], + [ + 0xFF, 0xFF, 0xFF, 0xFF, 5, 21, 37, 53, 0xFF, 0xFF, 0xFF, 0xFF, 13, 29, 45, 61, + ], + [ + 0xFF, 0xFF, 0xFF, 0xFF, 3, 19, 35, 51, 0xFF, 0xFF, 0xFF, 0xFF, 11, 27, 43, 59, + ], + [ + 0xFF, 0xFF, 0xFF, 0xFF, 7, 23, 39, 55, 0xFF, 0xFF, 0xFF, 0xFF, 15, 31, 47, 63, + ], +]; + +/// 8x8 byte transpose (scatter) permutation split into 4 × 16-byte chunks for NEON TBL. +/// Input layout: [g0b0..g0b7, g1b0..g1b7, ..., g7b0..g7b7] (64 bytes, group-major) +/// Output layout: [g0b0,g1b0,..,g7b0, g0b1,g1b1,..,g7b1, ...] (64 bytes, row-major) +/// Same permutation as x86 `SCATTER_8X8`, split for 16-byte NEON registers. +static SCATTER_8X8_NEON: [[u8; 16]; 4] = [ + [0, 8, 16, 24, 32, 40, 48, 56, 1, 9, 17, 25, 33, 41, 49, 57], + [2, 10, 18, 26, 34, 42, 50, 58, 3, 11, 19, 27, 35, 43, 51, 59], + [4, 12, 20, 28, 36, 44, 52, 60, 5, 13, 21, 29, 37, 45, 53, 61], + [6, 14, 22, 30, 38, 46, 54, 62, 7, 15, 23, 31, 39, 47, 55, 63], +]; + +/// Perform 8x8 bit transpose on two u64s packed in a `uint64x2_t`. +#[expect(unsafe_op_in_unsafe_fn)] +unsafe fn bit_transpose_8x8_neon(mut v: uint64x2_t) -> uint64x2_t { + let mask1 = vdupq_n_u64(TRANSPOSE_2X2); + let t = vandq_u64(veorq_u64(v, vshrq_n_u64::<7>(v)), mask1); + v = veorq_u64(veorq_u64(v, t), vshlq_n_u64::<7>(t)); + + let mask2 = vdupq_n_u64(TRANSPOSE_4X4); + let t = vandq_u64(veorq_u64(v, vshrq_n_u64::<14>(v)), mask2); + v = veorq_u64(veorq_u64(v, t), vshlq_n_u64::<14>(t)); + + let mask3 = vdupq_n_u64(TRANSPOSE_8X8); + let t = vandq_u64(veorq_u64(v, vshrq_n_u64::<28>(v)), mask3); + veorq_u64(veorq_u64(v, t), vshlq_n_u64::<28>(t)) +} + +/// Transpose 1024 bits using ARM NEON with TBL-based vectorized gather and scatter. +/// +/// Uses `vqtbl4q_u8` to gather bytes from the 128-byte input in parallel, +/// avoiding scalar byte-by-byte loads. Then uses `vqtbl4q_u8` again to perform +/// the 8x8 byte transpose for scatter. This is the NEON analog of x86 VBMI's +/// `vpermb`/`vpermi2b` byte permutation instructions. +/// +/// # Safety +/// Requires `AArch64` with NEON (always available on `AArch64`). +#[expect(unsafe_op_in_unsafe_fn)] +#[inline(never)] +pub unsafe fn transpose_bits_neon(input: &[u8; 128], output: &mut [u8; 128]) { + // Load all 128 input bytes into two uint8x16x4_t tables (64 bytes each) + let tbl_lo = vld1q_u8_x4(input.as_ptr()); + let tbl_hi = vld1q_u8_x4(input.as_ptr().add(64)); + + // Load scatter permutation indices (4 × 16 bytes) + let scatter0 = vld1q_u8(SCATTER_8X8_NEON[0].as_ptr()); + let scatter1 = vld1q_u8(SCATTER_8X8_NEON[1].as_ptr()); + let scatter2 = vld1q_u8(SCATTER_8X8_NEON[2].as_ptr()); + let scatter3 = vld1q_u8(SCATTER_8X8_NEON[3].as_ptr()); + + // Process first 64 output bytes (8 groups from BASE_PATTERN_FIRST) + // Gather and bit-transpose all 4 pairs, then scatter the full 64 bytes + let mut buf = [0u8; 64]; + for (i, (gather_lo, gather_high)) in [ + (GATHER_FIRST_LO, GATHER_FIRST_HI), + (GATHER_SECOND_LO, GATHER_SECOND_HI), + ] + .iter() + .enumerate() + { + for pair in 0..4 { + let idx_lo = vld1q_u8(gather_lo[pair].as_ptr()); + let idx_hi = vld1q_u8(gather_high[pair].as_ptr()); + + let from_lo = vqtbl4q_u8(tbl_lo, idx_lo); + let from_hi = vqtbl4q_u8(tbl_hi, idx_hi); + let gathered = vorrq_u8(from_lo, from_hi); + + let v = bit_transpose_8x8_neon(vreinterpretq_u64_u8(gathered)); + vst1q_u8(buf.as_mut_ptr().add(pair * 16), vreinterpretq_u8_u64(v)); + } + + // Load the 64-byte result as a TBL table and apply 8x8 byte transpose + let result_tbl = vld1q_u8_x4(buf.as_ptr()); + vst1q_u8( + output.as_mut_ptr().add(i * 64), + vqtbl4q_u8(result_tbl, scatter0), + ); + vst1q_u8( + output.as_mut_ptr().add(i * 64 + 16), + vqtbl4q_u8(result_tbl, scatter1), + ); + vst1q_u8( + output.as_mut_ptr().add(i * 64 + 32), + vqtbl4q_u8(result_tbl, scatter2), + ); + vst1q_u8( + output.as_mut_ptr().add(i * 64 + 48), + vqtbl4q_u8(result_tbl, scatter3), + ); + } +} + +/// Untranspose 1024 bits using ARM NEON with TBL-based vectorized operations. +/// +/// # Safety +/// Requires `AArch64` with NEON (always available on `AArch64`). +#[expect(unsafe_op_in_unsafe_fn)] +#[inline(never)] +pub unsafe fn untranspose_bits_neon(input: &[u8; 128], output: &mut [u8; 128]) { + // Load scatter indices (SCATTER_8X8 is self-inverse, so same table un-scatters) + let scatter0 = vld1q_u8(SCATTER_8X8_NEON[0].as_ptr()); + let scatter1 = vld1q_u8(SCATTER_8X8_NEON[1].as_ptr()); + let scatter2 = vld1q_u8(SCATTER_8X8_NEON[2].as_ptr()); + let scatter3 = vld1q_u8(SCATTER_8X8_NEON[3].as_ptr()); + + // Each iteration un-scatters the 64-byte input block to group-major order + let mut buf = [0u8; 64]; + for (i, base_pattern) in [BASE_PATTERN_FIRST, BASE_PATTERN_SECOND].iter().enumerate() { + let in_tbl = vld1q_u8_x4(input.as_ptr().add(i * 64)); + vst1q_u8(buf.as_mut_ptr(), vqtbl4q_u8(in_tbl, scatter0)); + vst1q_u8(buf.as_mut_ptr().add(16), vqtbl4q_u8(in_tbl, scatter1)); + vst1q_u8(buf.as_mut_ptr().add(32), vqtbl4q_u8(in_tbl, scatter2)); + vst1q_u8(buf.as_mut_ptr().add(48), vqtbl4q_u8(in_tbl, scatter3)); + + // Bit-transpose each pair and scatter to stride-16 output + for pair in 0..4 { + let base_group_0 = pair * 2; + let base_group_1 = pair * 2 + 1; + + let gathered = vld1q_u8(buf.as_ptr().add(pair * 16)); + let v = bit_transpose_8x8_neon(vreinterpretq_u64_u8(gathered)); + + let result_0 = vgetq_lane_u64::<0>(v); + let result_1 = vgetq_lane_u64::<1>(v); + + let out_base_0 = base_pattern[base_group_0]; + let out_base_1 = base_pattern[base_group_1]; + for i in 0..8 { + output[out_base_0 + i * 16] = (result_0 >> (i * 8)) as u8; + output[out_base_1 + i * 16] = (result_1 >> (i * 8)) as u8; + } + } + } +} + +#[cfg(test)] +mod tests { + use crate::bit_transpose::aarch64::transpose_bits_neon; + use crate::bit_transpose::aarch64::untranspose_bits_neon; + use crate::bit_transpose::generate_test_data; + use crate::bit_transpose::transpose_bits_baseline; + use crate::bit_transpose::untranspose_bits_baseline; + + #[test] + fn test_neon_matches_baseline() { + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut baseline_out = [0u8; 128]; + let mut tbl_out = [0u8; 128]; + + transpose_bits_baseline(&input, &mut baseline_out); + unsafe { transpose_bits_neon(&input, &mut tbl_out) }; + + assert_eq!( + baseline_out, tbl_out, + "NEON TBL transpose doesn't match baseline for seed {seed}" + ); + } + } + + #[test] + fn test_neon_roundtrip() { + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut transposed = [0u8; 128]; + let mut roundtrip = [0u8; 128]; + + unsafe { + transpose_bits_neon(&input, &mut transposed); + untranspose_bits_neon(&transposed, &mut roundtrip); + } + + assert_eq!( + input, roundtrip, + "NEON TBL roundtrip failed for seed {seed}" + ); + } + } + + #[test] + fn test_untranspose_neon_matches_baseline() { + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut baseline_out = [0u8; 128]; + let mut tbl_out = [0u8; 128]; + + untranspose_bits_baseline(&input, &mut baseline_out); + unsafe { untranspose_bits_neon(&input, &mut tbl_out) }; + + assert_eq!( + baseline_out, tbl_out, + "NEON TBL untranspose doesn't match baseline for seed {seed}" + ); + } + } +} diff --git a/encodings/fastlanes/src/bit_transpose/mod.rs b/encodings/fastlanes/src/bit_transpose/mod.rs new file mode 100644 index 00000000000..26e8b595195 --- /dev/null +++ b/encodings/fastlanes/src/bit_transpose/mod.rs @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Fast implementations of the `FastLanes` 1024-bit transpose. +//! +//! The `FastLanes` transpose is a fixed permutation of 1024 bits (128 bytes) that +//! enables SIMD parallelism for encodings like delta and RLE. This module provides +//! optimized implementations for different x86 SIMD instruction sets. +//! +//! The key insight is that each output byte is formed by extracting the SAME bit +//! position from 8 different input bytes at stride 16. The input byte groups follow +//! the `FL_ORDER` permutation pattern. + +#[cfg(feature = "_test-harness")] +pub mod aarch64; +#[cfg(feature = "_test-harness")] +pub mod scalar; +#[cfg(feature = "_test-harness")] +pub mod x86; + +#[cfg(not(feature = "_test-harness"))] +mod aarch64; +#[cfg(not(feature = "_test-harness"))] +mod scalar; +#[cfg(not(feature = "_test-harness"))] +mod x86; + +mod validity; + +pub use validity::*; +use vortex_buffer::CpuKernel; + +/// Signature shared by all 1024-bit transpose kernels. +type TransposeKernel = unsafe fn(&[u8; 128], &mut [u8; 128]); + +/// Base indices for the first 64 output bytes (lanes 0-7). +/// Each entry indicates the starting input byte index for that output byte group. +/// Pattern: [0*2, 4*2, 2*2, 6*2, 1*2, 5*2, 3*2, 7*2] = [0, 8, 4, 12, 2, 10, 6, 14] +const BASE_PATTERN_FIRST: [usize; 8] = [0, 8, 4, 12, 2, 10, 6, 14]; + +/// Base indices for the second 64 output bytes (lanes 8-15). +/// Pattern: first pattern + 1 = [1, 9, 5, 13, 3, 11, 7, 15] +const BASE_PATTERN_SECOND: [usize; 8] = [1, 9, 5, 13, 3, 11, 7, 15]; + +/// Masks for transposing 8x8 bit blocks. +const TRANSPOSE_2X2: u64 = 0x00AA_00AA_00AA_00AA; +const TRANSPOSE_4X4: u64 = 0x0000_CCCC_0000_CCCC; +const TRANSPOSE_8X8: u64 = 0x0000_0000_F0F0_F0F0; + +/// Transpose 1024-bits into FastLanes layout. +/// +/// Dispatch to the best available implementation, selected once on first call. +#[inline] +pub fn transpose_bits(input: &[u8; 128], output: &mut [u8; 128]) { + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + // VBMI is fastest + if x86::has_vbmi() { + return x86::transpose_bits_vbmi; + } + if x86::has_bmi2() { + return x86::transpose_bits_bmi2; + } + } + // NEON is architecturally guaranteed on aarch64, so it needs no probe. + #[cfg(target_arch = "aarch64")] + return aarch64::transpose_bits_neon; + // The aarch64 arm above returns unconditionally, making this portable default + // unreachable there. + #[allow(unreachable_code)] + { + scalar::transpose_bits_scalar + } + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(input, output) } +} + +/// Untranspose 1024-bits from FastLanes layout. +/// +/// Dispatch untranspose to the best available implementation, selected once on first call. +#[inline] +pub fn untranspose_bits(input: &[u8; 128], output: &mut [u8; 128]) { + static KERNEL: CpuKernel = CpuKernel::new(|| { + #[cfg(target_arch = "x86_64")] + { + // VBMI is fastest + if x86::has_vbmi() { + return x86::untranspose_bits_vbmi; + } + if x86::has_bmi2() { + return x86::untranspose_bits_bmi2; + } + } + // NEON is architecturally guaranteed on aarch64, so it needs no probe. + #[cfg(target_arch = "aarch64")] + return aarch64::untranspose_bits_neon; + // The aarch64 arm above returns unconditionally, making this portable default + // unreachable there. + #[allow(unreachable_code)] + { + scalar::untranspose_bits_scalar + } + }); + // SAFETY: the selector only returns kernels that are safe or whose required CPU + // features were probed before selection. + unsafe { KERNEL.get()(input, output) } +} + +#[cfg(test)] +#[expect(clippy::cast_possible_truncation)] +fn generate_test_data(seed: u8) -> [u8; 128] { + let mut data = [0u8; 128]; + for (i, byte) in data.iter_mut().enumerate() { + *byte = seed.wrapping_mul(17).wrapping_add(i as u8).wrapping_mul(31); + } + data +} + +#[cfg(test)] +pub fn transpose_bits_baseline(input: &[u8; 128], output: &mut [u8; 128]) { + for in_bit in 0..1024 { + let out_bit = fastlanes::transpose(in_bit); + let in_byte = in_bit / 8; + let in_bit_pos = in_bit % 8; + let out_byte = out_bit / 8; + let out_bit_pos = out_bit % 8; + let bit_val = (input[in_byte] >> in_bit_pos) & 1; + output[out_byte] |= bit_val << out_bit_pos; + } +} + +#[cfg(test)] +pub fn untranspose_bits_baseline(input: &[u8; 128], output: &mut [u8; 128]) { + for out_bit in 0..1024 { + let in_bit = fastlanes::transpose(out_bit); + let in_byte = in_bit / 8; + let in_bit_pos = in_bit % 8; + let out_byte = out_bit / 8; + let out_bit_pos = out_bit % 8; + let bit_val = (input[in_byte] >> in_bit_pos) & 1; + output[out_byte] |= bit_val << out_bit_pos; + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_transpose_baseline_roundtrip() { + let input = generate_test_data(42); + let mut transposed = [0u8; 128]; + let mut roundtrip = [0u8; 128]; + + transpose_bits_baseline(&input, &mut transposed); + untranspose_bits_baseline(&transposed, &mut roundtrip); + + assert_eq!(input, roundtrip); + } + + #[test] + fn test_dispatch_matches_baseline() { + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut baseline_out = [0u8; 128]; + let mut out = [0u8; 128]; + + transpose_bits_baseline(&input, &mut baseline_out); + transpose_bits(&input, &mut out); + + assert_eq!( + baseline_out, out, + "best dispatch doesn't match baseline for seed {seed}" + ); + } + } + + #[test] + fn test_untranspose_dispatch_matches_baseline() { + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut baseline_out = [0u8; 128]; + let mut out = [0u8; 128]; + + untranspose_bits_baseline(&input, &mut baseline_out); + untranspose_bits(&input, &mut out); + + assert_eq!( + baseline_out, out, + "best untranspose dispatch doesn't match baseline for seed {seed}" + ); + } + } +} diff --git a/encodings/fastlanes/src/bit_transpose/scalar.rs b/encodings/fastlanes/src/bit_transpose/scalar.rs new file mode 100644 index 00000000000..425e74087ab --- /dev/null +++ b/encodings/fastlanes/src/bit_transpose/scalar.rs @@ -0,0 +1,219 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +use crate::bit_transpose::BASE_PATTERN_FIRST; +use crate::bit_transpose::BASE_PATTERN_SECOND; +use crate::bit_transpose::TRANSPOSE_2X2; +use crate::bit_transpose::TRANSPOSE_4X4; +use crate::bit_transpose::TRANSPOSE_8X8; + +/// Fast scalar transpose using the 8x8 bit matrix transpose algorithm. +/// +/// This version uses 64-bit gather + parallel bit operations instead of +/// extracting bits one by one. Typically 5-10x faster than the basic scalar version. +#[inline(never)] +#[allow(dead_code)] +pub fn transpose_bits_scalar(input: &[u8; 128], output: &mut [u8; 128]) { + // Helper to perform 8x8 bit transpose on a u64 (each byte becomes a row) + fn transpose_8x8(mut x: u64) -> u64 { + // Step 1: Transpose 2x2 bit blocks + let t = (x ^ (x >> 7)) & TRANSPOSE_2X2; + x = x ^ t ^ (t << 7); + // Step 2: Transpose 4x4 bit blocks + let t = (x ^ (x >> 14)) & TRANSPOSE_4X4; + x = x ^ t ^ (t << 14); + // Step 3: Transpose 8x8 bit blocks + let t = (x ^ (x >> 28)) & TRANSPOSE_8X8; + x ^ t ^ (t << 28) + } + + // Helper to gather 8 bytes at stride 16 into a u64 + fn gather(input: &[u8; 128], base: usize) -> u64 { + u64::from(input[base]) + | (u64::from(input[base + 16]) << 8) + | (u64::from(input[base + 32]) << 16) + | (u64::from(input[base + 48]) << 24) + | (u64::from(input[base + 64]) << 32) + | (u64::from(input[base + 80]) << 40) + | (u64::from(input[base + 96]) << 48) + | (u64::from(input[base + 112]) << 56) + } + + // Process first half (8 base groups, fully unrolled) + let r0 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[0])); + let r1 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[1])); + let r2 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[2])); + let r3 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[3])); + let r4 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[4])); + let r5 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[5])); + let r6 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[6])); + let r7 = transpose_8x8(gather(input, BASE_PATTERN_FIRST[7])); + + // Write first 64 output bytes (unrolled) + for bit_pos in 0..8 { + output[bit_pos * 8] = (r0 >> (bit_pos * 8)) as u8; + output[bit_pos * 8 + 1] = (r1 >> (bit_pos * 8)) as u8; + output[bit_pos * 8 + 2] = (r2 >> (bit_pos * 8)) as u8; + output[bit_pos * 8 + 3] = (r3 >> (bit_pos * 8)) as u8; + output[bit_pos * 8 + 4] = (r4 >> (bit_pos * 8)) as u8; + output[bit_pos * 8 + 5] = (r5 >> (bit_pos * 8)) as u8; + output[bit_pos * 8 + 6] = (r6 >> (bit_pos * 8)) as u8; + output[bit_pos * 8 + 7] = (r7 >> (bit_pos * 8)) as u8; + } + + // Process second half + let r0 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[0])); + let r1 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[1])); + let r2 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[2])); + let r3 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[3])); + let r4 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[4])); + let r5 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[5])); + let r6 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[6])); + let r7 = transpose_8x8(gather(input, BASE_PATTERN_SECOND[7])); + + for bit_pos in 0..8 { + output[64 + bit_pos * 8] = (r0 >> (bit_pos * 8)) as u8; + output[64 + bit_pos * 8 + 1] = (r1 >> (bit_pos * 8)) as u8; + output[64 + bit_pos * 8 + 2] = (r2 >> (bit_pos * 8)) as u8; + output[64 + bit_pos * 8 + 3] = (r3 >> (bit_pos * 8)) as u8; + output[64 + bit_pos * 8 + 4] = (r4 >> (bit_pos * 8)) as u8; + output[64 + bit_pos * 8 + 5] = (r5 >> (bit_pos * 8)) as u8; + output[64 + bit_pos * 8 + 6] = (r6 >> (bit_pos * 8)) as u8; + output[64 + bit_pos * 8 + 7] = (r7 >> (bit_pos * 8)) as u8; + } +} + +/// Fast scalar untranspose using the 8x8 bit matrix transpose algorithm. +#[inline(never)] +#[allow(dead_code)] +pub fn untranspose_bits_scalar(input: &[u8; 128], output: &mut [u8; 128]) { + fn transpose_8x8(mut x: u64) -> u64 { + let t = (x ^ (x >> 7)) & TRANSPOSE_2X2; + x = x ^ t ^ (t << 7); + let t = (x ^ (x >> 14)) & TRANSPOSE_4X4; + x = x ^ t ^ (t << 14); + let t = (x ^ (x >> 28)) & TRANSPOSE_8X8; + x ^ t ^ (t << 28) + } + + fn gather_transposed(input: &[u8; 128], base_group: usize, offset: usize) -> u64 { + let mut result: u64 = 0; + for bit_pos in 0..8 { + result |= u64::from(input[offset + bit_pos * 8 + base_group]) << (bit_pos * 8); + } + result + } + + fn scatter(output: &mut [u8; 128], base: usize, val: u64) { + output[base] = val as u8; + output[base + 16] = (val >> 8) as u8; + output[base + 32] = (val >> 16) as u8; + output[base + 48] = (val >> 24) as u8; + output[base + 64] = (val >> 32) as u8; + output[base + 80] = (val >> 40) as u8; + output[base + 96] = (val >> 48) as u8; + output[base + 112] = (val >> 56) as u8; + } + + // First half (unrolled) + let r0 = transpose_8x8(gather_transposed(input, 0, 0)); + let r1 = transpose_8x8(gather_transposed(input, 1, 0)); + let r2 = transpose_8x8(gather_transposed(input, 2, 0)); + let r3 = transpose_8x8(gather_transposed(input, 3, 0)); + let r4 = transpose_8x8(gather_transposed(input, 4, 0)); + let r5 = transpose_8x8(gather_transposed(input, 5, 0)); + let r6 = transpose_8x8(gather_transposed(input, 6, 0)); + let r7 = transpose_8x8(gather_transposed(input, 7, 0)); + + scatter(output, BASE_PATTERN_FIRST[0], r0); + scatter(output, BASE_PATTERN_FIRST[1], r1); + scatter(output, BASE_PATTERN_FIRST[2], r2); + scatter(output, BASE_PATTERN_FIRST[3], r3); + scatter(output, BASE_PATTERN_FIRST[4], r4); + scatter(output, BASE_PATTERN_FIRST[5], r5); + scatter(output, BASE_PATTERN_FIRST[6], r6); + scatter(output, BASE_PATTERN_FIRST[7], r7); + + // Second half + let r0 = transpose_8x8(gather_transposed(input, 0, 64)); + let r1 = transpose_8x8(gather_transposed(input, 1, 64)); + let r2 = transpose_8x8(gather_transposed(input, 2, 64)); + let r3 = transpose_8x8(gather_transposed(input, 3, 64)); + let r4 = transpose_8x8(gather_transposed(input, 4, 64)); + let r5 = transpose_8x8(gather_transposed(input, 5, 64)); + let r6 = transpose_8x8(gather_transposed(input, 6, 64)); + let r7 = transpose_8x8(gather_transposed(input, 7, 64)); + + scatter(output, BASE_PATTERN_SECOND[0], r0); + scatter(output, BASE_PATTERN_SECOND[1], r1); + scatter(output, BASE_PATTERN_SECOND[2], r2); + scatter(output, BASE_PATTERN_SECOND[3], r3); + scatter(output, BASE_PATTERN_SECOND[4], r4); + scatter(output, BASE_PATTERN_SECOND[5], r5); + scatter(output, BASE_PATTERN_SECOND[6], r6); + scatter(output, BASE_PATTERN_SECOND[7], r7); +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::bit_transpose::generate_test_data; + + #[test] + fn test_scalar_matches_baseline() { + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut baseline_out = [0u8; 128]; + let mut fast_out = [0u8; 128]; + + transpose_bits_scalar(&input, &mut baseline_out); + transpose_bits_scalar(&input, &mut fast_out); + + assert_eq!( + baseline_out, fast_out, + "scalar_fast transpose doesn't match baseline for seed {seed}" + ); + } + } + + #[test] + fn test_scalar_roundtrip() { + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut transposed = [0u8; 128]; + let mut roundtrip = [0u8; 128]; + + transpose_bits_scalar(&input, &mut transposed); + untranspose_bits_scalar(&transposed, &mut roundtrip); + + assert_eq!( + input, roundtrip, + "scalar_fast roundtrip failed for seed {seed}" + ); + } + } + + #[test] + fn test_all_zeros() { + let input = [0u8; 128]; + let mut output = [0xFFu8; 128]; + + transpose_bits_scalar(&input, &mut output); + assert_eq!(output, [0u8; 128]); + + untranspose_bits_scalar(&input, &mut output); + assert_eq!(output, [0u8; 128]); + } + + #[test] + fn test_all_ones() { + let input = [0xFFu8; 128]; + let mut output = [0u8; 128]; + + transpose_bits_scalar(&input, &mut output); + assert_eq!(output, [0xFFu8; 128]); + + untranspose_bits_scalar(&input, &mut output); + assert_eq!(output, [0xFFu8; 128]); + } +} diff --git a/encodings/fastlanes/src/bit_transpose.rs b/encodings/fastlanes/src/bit_transpose/validity.rs similarity index 64% rename from encodings/fastlanes/src/bit_transpose.rs rename to encodings/fastlanes/src/bit_transpose/validity.rs index 068c79f0762..a6b2ec43cbc 100644 --- a/encodings/fastlanes/src/bit_transpose.rs +++ b/encodings/fastlanes/src/bit_transpose/validity.rs @@ -4,11 +4,38 @@ use std::mem; use std::mem::MaybeUninit; -use vortex_buffer::Alignment; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; -use vortex_buffer::BufferMut; use vortex_buffer::ByteBuffer; +use vortex_buffer::ByteBufferMut; use vortex_error::VortexExpect; +use vortex_error::VortexResult; + +use crate::bit_transpose::transpose_bits; +use crate::bit_transpose::untranspose_bits; + +pub fn transpose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> VortexResult { + match validity { + Validity::Array(mask) => { + let bools = mask + .clone() + .execute::(ctx)? + .into_bool() + .into_bit_buffer(); + + Ok(Validity::Array( + BoolArray::new(transpose_bitbuffer(bools), Validity::NonNullable).into_array(), + )) + } + v @ Validity::AllValid | v @ Validity::AllInvalid | v @ Validity::NonNullable => { + Ok(v.clone()) + } + } +} pub fn transpose_bitbuffer(bits: BitBuffer) -> BitBuffer { let (offset, len, bytes) = bits.into_inner(); @@ -18,19 +45,36 @@ pub fn transpose_bitbuffer(bits: BitBuffer) -> BitBuffer { Ok(mut bytes_mut) => { // We can ignore the spare trailer capacity that can be an artifact of allocator as we requested 128 multiple chunks let (chunks, _) = bytes_mut.as_chunks_mut::<128>(); - let mut tmp = [0u64; 16]; + let mut tmp = [0u8; 128]; for chunk in chunks { - let chunk_u64 = - unsafe { mem::transmute::<&mut [u8; 128], &mut [u64; 16]>(chunk) }; - fastlanes::transpose_bits(chunk_u64, &mut tmp); - chunk_u64.copy_from_slice(&tmp); + transpose_bits(chunk, &mut tmp); + chunk.copy_from_slice(&tmp); } BitBuffer::new_with_offset(bytes_mut.freeze().into_byte_buffer(), len, offset) } - Err(bytes) => bits_op_with_copy(bytes, len, offset, fastlanes::transpose_bits), + Err(bytes) => bits_op_with_copy(bytes, len, offset, transpose_bits), } } else { - bits_op_with_copy(bytes, len, offset, fastlanes::transpose_bits) + bits_op_with_copy(bytes, len, offset, transpose_bits) + } +} + +pub fn untranspose_validity(validity: &Validity, ctx: &mut ExecutionCtx) -> VortexResult { + match validity { + Validity::Array(mask) => { + let bools = mask + .clone() + .execute::(ctx)? + .into_bool() + .into_bit_buffer(); + + Ok(Validity::Array( + BoolArray::new(untranspose_bitbuffer(bools), Validity::NonNullable).into_array(), + )) + } + v @ Validity::AllValid | v @ Validity::AllInvalid | v @ Validity::NonNullable => { + Ok(v.clone()) + } } } @@ -39,65 +83,51 @@ pub fn untranspose_bitbuffer(bits: BitBuffer) -> BitBuffer { bits.inner().len().is_multiple_of(128), "Transpose BitBuffer byte length must be a multiple of 128" ); - assert!( - bits.inner().is_aligned(Alignment::of::()), - "Transposed buffer must be 8 byte aligned" - ); let (offset, len, bytes) = bits.into_inner(); match bytes.try_into_mut() { Ok(mut bytes_mut) => { - let (prefix, middle, trailer) = unsafe { bytes_mut.align_to_mut::() }; - assert!( - prefix.is_empty() && trailer.is_empty(), - "Transposed buffer must be 8 byte aligned" - ); - let (chunks, _) = middle.as_chunks_mut::<16>(); - let mut tmp = [0u64; 16]; + let (chunks, _) = bytes_mut.as_chunks_mut::<128>(); + let mut tmp = [0u8; 128]; for chunk in chunks { - fastlanes::untranspose_bits::(chunk, &mut tmp); + untranspose_bits(chunk, &mut tmp); chunk.copy_from_slice(&tmp); } BitBuffer::new_with_offset(bytes_mut.freeze().into_byte_buffer(), len, offset) } - Err(bytes) => bits_op_with_copy(bytes, len, offset, fastlanes::untranspose_bits::), + Err(bytes) => bits_op_with_copy(bytes, len, offset, untranspose_bits), } } -fn bits_op_with_copy( +fn bits_op_with_copy( bytes: ByteBuffer, len: usize, offset: usize, op: F, ) -> BitBuffer { - let output_len = bytes.len().div_ceil(8).next_multiple_of(16); - let mut output = BufferMut::::with_capacity(output_len); + let output_len = bytes.len().next_multiple_of(128); + let mut output = ByteBufferMut::with_capacity(output_len); let (input_chunks, input_trailer) = bytes.as_chunks::<128>(); // Bound to the requested `output_len`: `spare_capacity_mut` may expose extra over-aligned // capacity, which would otherwise split into spurious trailing chunks and make `last_mut` // below target a chunk past the data we actually initialize. - let (output_chunks, _) = unsafe { - mem::transmute::<&mut [MaybeUninit], &mut [u64]>( - &mut output.spare_capacity_mut()[..output_len], - ) - } - .as_chunks_mut::<16>(); + let (output_chunks, _) = output.spare_capacity_mut()[..output_len].as_chunks_mut::<128>(); for (input, output) in input_chunks.iter().zip(output_chunks.iter_mut()) { - op( - unsafe { mem::transmute::<&[u8; 128], &[u64; 16]>(input) }, - output, - ); + op(input, unsafe { + mem::transmute::<&mut [MaybeUninit; 128], &mut [u8; 128]>(output) + }); } if !input_trailer.is_empty() { let mut padded_input = [0u8; 128]; padded_input[0..input_trailer.len()].clone_from_slice(input_trailer); - op( - unsafe { mem::transmute::<&[u8; 128], &[u64; 16]>(&padded_input) }, - output_chunks - .last_mut() - .vortex_expect("Output wasn't a multiple of 128 bytes"), - ); + op(&padded_input, unsafe { + mem::transmute::<&mut [MaybeUninit; 128], &mut [u8; 128]>( + output_chunks + .last_mut() + .vortex_expect("Output wasn't a multiple of 128 bytes"), + ) + }); } unsafe { output.set_len(output_len) }; @@ -153,7 +183,7 @@ mod tests { } #[test] - fn transpose_bitbuffer_roundtrip_non_aligned() { + fn transpose_validity_roundtrip_non_aligned() { let original_len = 1500; let bits = make_validity_bits(original_len); diff --git a/encodings/fastlanes/src/bit_transpose/x86.rs b/encodings/fastlanes/src/bit_transpose/x86.rs new file mode 100644 index 00000000000..9eaef248d28 --- /dev/null +++ b/encodings/fastlanes/src/bit_transpose/x86.rs @@ -0,0 +1,715 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors +#![cfg(target_arch = "x86_64")] + +use core::arch::x86_64::__m512i; +use core::arch::x86_64::_mm512_and_si512; +use core::arch::x86_64::_mm512_loadu_si512; +use core::arch::x86_64::_mm512_permutex2var_epi8; +use core::arch::x86_64::_mm512_permutexvar_epi8; +use core::arch::x86_64::_mm512_set1_epi64; +use core::arch::x86_64::_mm512_slli_epi64; +use core::arch::x86_64::_mm512_srli_epi64; +use core::arch::x86_64::_mm512_storeu_si512; +use core::arch::x86_64::_mm512_xor_si512; +use core::arch::x86_64::_pdep_u64; +use core::arch::x86_64::_pext_u64; +use std::is_x86_feature_detected; + +use crate::bit_transpose::BASE_PATTERN_FIRST; +use crate::bit_transpose::BASE_PATTERN_SECOND; +use crate::bit_transpose::TRANSPOSE_2X2; +use crate::bit_transpose::TRANSPOSE_4X4; +use crate::bit_transpose::TRANSPOSE_8X8; + +/// Check if BMI2 is available. +#[inline] +#[must_use] +pub fn has_bmi2() -> bool { + is_x86_feature_detected!("bmi2") +} + +/// Check if AVX-512 VBMI is available (for byte permutation). +#[inline] +#[must_use] +pub fn has_vbmi() -> bool { + is_x86_feature_detected!("avx512vbmi") +} + +/// Transpose 1024 bits using BMI2 PEXT instruction. +/// +/// PEXT extracts bits at positions specified by a mask into contiguous low bits. +/// Fully unrolled for ~12% better performance vs looped version. +/// +/// # Safety +/// Requires BMI2 support. Check with `has_bmi2()` before calling. +#[target_feature(enable = "bmi2")] +#[inline(never)] +#[expect(clippy::too_many_lines)] +pub unsafe fn transpose_bits_bmi2(input: &[u8; 128], output: &mut [u8; 128]) { + // Helper to gather 8 bytes at stride 16 into a u64 + fn gather(input: &[u8; 128], base: usize) -> u64 { + (input[base] as u64) + | ((input[base + 16] as u64) << 8) + | ((input[base + 32] as u64) << 16) + | ((input[base + 48] as u64) << 24) + | ((input[base + 64] as u64) << 32) + | ((input[base + 80] as u64) << 40) + | ((input[base + 96] as u64) << 48) + | ((input[base + 112] as u64) << 56) + } + + // Gather all 16 groups (fully unrolled) + // First half: BASE_PATTERN_FIRST = [0, 8, 4, 12, 2, 10, 6, 14] + let g0 = gather(input, 0); + let g1 = gather(input, 8); + let g2 = gather(input, 4); + let g3 = gather(input, 12); + let g4 = gather(input, 2); + let g5 = gather(input, 10); + let g6 = gather(input, 6); + let g7 = gather(input, 14); + // Second half: BASE_PATTERN_SECOND = [1, 9, 5, 13, 3, 11, 7, 15] + let g8 = gather(input, 1); + let g9 = gather(input, 9); + let g10 = gather(input, 5); + let g11 = gather(input, 13); + let g12 = gather(input, 3); + let g13 = gather(input, 11); + let g14 = gather(input, 7); + let g15 = gather(input, 15); + + // Masks for each bit position + let m0: u64 = 0x0101_0101_0101_0101; + let m1: u64 = 0x0202_0202_0202_0202; + let m2: u64 = 0x0404_0404_0404_0404; + let m3: u64 = 0x0808_0808_0808_0808; + let m4: u64 = 0x1010_1010_1010_1010; + let m5: u64 = 0x2020_2020_2020_2020; + let m6: u64 = 0x4040_4040_4040_4040; + let m7: u64 = 0x8080_8080_8080_8080; + + // First half - 64 PEXT operations (fully unrolled) + output[0] = _pext_u64(g0, m0) as u8; + output[1] = _pext_u64(g1, m0) as u8; + output[2] = _pext_u64(g2, m0) as u8; + output[3] = _pext_u64(g3, m0) as u8; + output[4] = _pext_u64(g4, m0) as u8; + output[5] = _pext_u64(g5, m0) as u8; + output[6] = _pext_u64(g6, m0) as u8; + output[7] = _pext_u64(g7, m0) as u8; + output[8] = _pext_u64(g0, m1) as u8; + output[9] = _pext_u64(g1, m1) as u8; + output[10] = _pext_u64(g2, m1) as u8; + output[11] = _pext_u64(g3, m1) as u8; + output[12] = _pext_u64(g4, m1) as u8; + output[13] = _pext_u64(g5, m1) as u8; + output[14] = _pext_u64(g6, m1) as u8; + output[15] = _pext_u64(g7, m1) as u8; + output[16] = _pext_u64(g0, m2) as u8; + output[17] = _pext_u64(g1, m2) as u8; + output[18] = _pext_u64(g2, m2) as u8; + output[19] = _pext_u64(g3, m2) as u8; + output[20] = _pext_u64(g4, m2) as u8; + output[21] = _pext_u64(g5, m2) as u8; + output[22] = _pext_u64(g6, m2) as u8; + output[23] = _pext_u64(g7, m2) as u8; + output[24] = _pext_u64(g0, m3) as u8; + output[25] = _pext_u64(g1, m3) as u8; + output[26] = _pext_u64(g2, m3) as u8; + output[27] = _pext_u64(g3, m3) as u8; + output[28] = _pext_u64(g4, m3) as u8; + output[29] = _pext_u64(g5, m3) as u8; + output[30] = _pext_u64(g6, m3) as u8; + output[31] = _pext_u64(g7, m3) as u8; + output[32] = _pext_u64(g0, m4) as u8; + output[33] = _pext_u64(g1, m4) as u8; + output[34] = _pext_u64(g2, m4) as u8; + output[35] = _pext_u64(g3, m4) as u8; + output[36] = _pext_u64(g4, m4) as u8; + output[37] = _pext_u64(g5, m4) as u8; + output[38] = _pext_u64(g6, m4) as u8; + output[39] = _pext_u64(g7, m4) as u8; + output[40] = _pext_u64(g0, m5) as u8; + output[41] = _pext_u64(g1, m5) as u8; + output[42] = _pext_u64(g2, m5) as u8; + output[43] = _pext_u64(g3, m5) as u8; + output[44] = _pext_u64(g4, m5) as u8; + output[45] = _pext_u64(g5, m5) as u8; + output[46] = _pext_u64(g6, m5) as u8; + output[47] = _pext_u64(g7, m5) as u8; + output[48] = _pext_u64(g0, m6) as u8; + output[49] = _pext_u64(g1, m6) as u8; + output[50] = _pext_u64(g2, m6) as u8; + output[51] = _pext_u64(g3, m6) as u8; + output[52] = _pext_u64(g4, m6) as u8; + output[53] = _pext_u64(g5, m6) as u8; + output[54] = _pext_u64(g6, m6) as u8; + output[55] = _pext_u64(g7, m6) as u8; + output[56] = _pext_u64(g0, m7) as u8; + output[57] = _pext_u64(g1, m7) as u8; + output[58] = _pext_u64(g2, m7) as u8; + output[59] = _pext_u64(g3, m7) as u8; + output[60] = _pext_u64(g4, m7) as u8; + output[61] = _pext_u64(g5, m7) as u8; + output[62] = _pext_u64(g6, m7) as u8; + output[63] = _pext_u64(g7, m7) as u8; + + // Second half - 64 PEXT operations (fully unrolled) + output[64] = _pext_u64(g8, m0) as u8; + output[65] = _pext_u64(g9, m0) as u8; + output[66] = _pext_u64(g10, m0) as u8; + output[67] = _pext_u64(g11, m0) as u8; + output[68] = _pext_u64(g12, m0) as u8; + output[69] = _pext_u64(g13, m0) as u8; + output[70] = _pext_u64(g14, m0) as u8; + output[71] = _pext_u64(g15, m0) as u8; + output[72] = _pext_u64(g8, m1) as u8; + output[73] = _pext_u64(g9, m1) as u8; + output[74] = _pext_u64(g10, m1) as u8; + output[75] = _pext_u64(g11, m1) as u8; + output[76] = _pext_u64(g12, m1) as u8; + output[77] = _pext_u64(g13, m1) as u8; + output[78] = _pext_u64(g14, m1) as u8; + output[79] = _pext_u64(g15, m1) as u8; + output[80] = _pext_u64(g8, m2) as u8; + output[81] = _pext_u64(g9, m2) as u8; + output[82] = _pext_u64(g10, m2) as u8; + output[83] = _pext_u64(g11, m2) as u8; + output[84] = _pext_u64(g12, m2) as u8; + output[85] = _pext_u64(g13, m2) as u8; + output[86] = _pext_u64(g14, m2) as u8; + output[87] = _pext_u64(g15, m2) as u8; + output[88] = _pext_u64(g8, m3) as u8; + output[89] = _pext_u64(g9, m3) as u8; + output[90] = _pext_u64(g10, m3) as u8; + output[91] = _pext_u64(g11, m3) as u8; + output[92] = _pext_u64(g12, m3) as u8; + output[93] = _pext_u64(g13, m3) as u8; + output[94] = _pext_u64(g14, m3) as u8; + output[95] = _pext_u64(g15, m3) as u8; + output[96] = _pext_u64(g8, m4) as u8; + output[97] = _pext_u64(g9, m4) as u8; + output[98] = _pext_u64(g10, m4) as u8; + output[99] = _pext_u64(g11, m4) as u8; + output[100] = _pext_u64(g12, m4) as u8; + output[101] = _pext_u64(g13, m4) as u8; + output[102] = _pext_u64(g14, m4) as u8; + output[103] = _pext_u64(g15, m4) as u8; + output[104] = _pext_u64(g8, m5) as u8; + output[105] = _pext_u64(g9, m5) as u8; + output[106] = _pext_u64(g10, m5) as u8; + output[107] = _pext_u64(g11, m5) as u8; + output[108] = _pext_u64(g12, m5) as u8; + output[109] = _pext_u64(g13, m5) as u8; + output[110] = _pext_u64(g14, m5) as u8; + output[111] = _pext_u64(g15, m5) as u8; + output[112] = _pext_u64(g8, m6) as u8; + output[113] = _pext_u64(g9, m6) as u8; + output[114] = _pext_u64(g10, m6) as u8; + output[115] = _pext_u64(g11, m6) as u8; + output[116] = _pext_u64(g12, m6) as u8; + output[117] = _pext_u64(g13, m6) as u8; + output[118] = _pext_u64(g14, m6) as u8; + output[119] = _pext_u64(g15, m6) as u8; + output[120] = _pext_u64(g8, m7) as u8; + output[121] = _pext_u64(g9, m7) as u8; + output[122] = _pext_u64(g10, m7) as u8; + output[123] = _pext_u64(g11, m7) as u8; + output[124] = _pext_u64(g12, m7) as u8; + output[125] = _pext_u64(g13, m7) as u8; + output[126] = _pext_u64(g14, m7) as u8; + output[127] = _pext_u64(g15, m7) as u8; +} + +/// Untranspose 1024 bits using BMI2 PDEP instruction. +/// +/// Structured per-output-group: for each group of 8 output bytes at stride 16, +/// PDEP 8 input bytes into different bit positions, OR in registers, then +/// scatter-store once. Each output byte is written exactly once (no read-modify-write). +/// +/// # Safety +/// Requires BMI2 support. Check with `has_bmi2()` before calling. +#[target_feature(enable = "bmi2")] +#[inline(never)] +#[expect(clippy::too_many_lines)] +pub unsafe fn untranspose_bits_bmi2(input: &[u8; 128], output: &mut [u8; 128]) { + // Helper: scatter a u64 to 8 output bytes at stride 16 + fn scatter(output: &mut [u8; 128], base: usize, val: u64) { + output[base] = val as u8; + output[base + 16] = (val >> 8) as u8; + output[base + 32] = (val >> 16) as u8; + output[base + 48] = (val >> 24) as u8; + output[base + 64] = (val >> 32) as u8; + output[base + 80] = (val >> 40) as u8; + output[base + 96] = (val >> 48) as u8; + output[base + 112] = (val >> 56) as u8; + } + + // Masks for each bit position + let m0: u64 = 0x0101_0101_0101_0101; + let m1: u64 = 0x0202_0202_0202_0202; + let m2: u64 = 0x0404_0404_0404_0404; + let m3: u64 = 0x0808_0808_0808_0808; + let m4: u64 = 0x1010_1010_1010_1010; + let m5: u64 = 0x2020_2020_2020_2020; + let m6: u64 = 0x4040_4040_4040_4040; + let m7: u64 = 0x8080_8080_8080_8080; + + // For each output group, the input bytes that contribute are at + // input[bit_pos * 8 + group_idx] for bit_pos 0..8. + // PDEP deposits the 8 bits of the input byte into the bit_pos position + // of each byte in the u64. + + // First half: 8 groups using BASE_PATTERN_FIRST + // Group 0 (base=0): input bytes [0, 8, 16, 24, 32, 40, 48, 56] + let v = _pdep_u64(input[0] as u64, m0) + | _pdep_u64(input[8] as u64, m1) + | _pdep_u64(input[16] as u64, m2) + | _pdep_u64(input[24] as u64, m3) + | _pdep_u64(input[32] as u64, m4) + | _pdep_u64(input[40] as u64, m5) + | _pdep_u64(input[48] as u64, m6) + | _pdep_u64(input[56] as u64, m7); + scatter(output, 0, v); + + // Group 1 (base=8) + let v = _pdep_u64(input[1] as u64, m0) + | _pdep_u64(input[9] as u64, m1) + | _pdep_u64(input[17] as u64, m2) + | _pdep_u64(input[25] as u64, m3) + | _pdep_u64(input[33] as u64, m4) + | _pdep_u64(input[41] as u64, m5) + | _pdep_u64(input[49] as u64, m6) + | _pdep_u64(input[57] as u64, m7); + scatter(output, 8, v); + + // Group 2 (base=4) + let v = _pdep_u64(input[2] as u64, m0) + | _pdep_u64(input[10] as u64, m1) + | _pdep_u64(input[18] as u64, m2) + | _pdep_u64(input[26] as u64, m3) + | _pdep_u64(input[34] as u64, m4) + | _pdep_u64(input[42] as u64, m5) + | _pdep_u64(input[50] as u64, m6) + | _pdep_u64(input[58] as u64, m7); + scatter(output, 4, v); + + // Group 3 (base=12) + let v = _pdep_u64(input[3] as u64, m0) + | _pdep_u64(input[11] as u64, m1) + | _pdep_u64(input[19] as u64, m2) + | _pdep_u64(input[27] as u64, m3) + | _pdep_u64(input[35] as u64, m4) + | _pdep_u64(input[43] as u64, m5) + | _pdep_u64(input[51] as u64, m6) + | _pdep_u64(input[59] as u64, m7); + scatter(output, 12, v); + + // Group 4 (base=2) + let v = _pdep_u64(input[4] as u64, m0) + | _pdep_u64(input[12] as u64, m1) + | _pdep_u64(input[20] as u64, m2) + | _pdep_u64(input[28] as u64, m3) + | _pdep_u64(input[36] as u64, m4) + | _pdep_u64(input[44] as u64, m5) + | _pdep_u64(input[52] as u64, m6) + | _pdep_u64(input[60] as u64, m7); + scatter(output, 2, v); + + // Group 5 (base=10) + let v = _pdep_u64(input[5] as u64, m0) + | _pdep_u64(input[13] as u64, m1) + | _pdep_u64(input[21] as u64, m2) + | _pdep_u64(input[29] as u64, m3) + | _pdep_u64(input[37] as u64, m4) + | _pdep_u64(input[45] as u64, m5) + | _pdep_u64(input[53] as u64, m6) + | _pdep_u64(input[61] as u64, m7); + scatter(output, 10, v); + + // Group 6 (base=6) + let v = _pdep_u64(input[6] as u64, m0) + | _pdep_u64(input[14] as u64, m1) + | _pdep_u64(input[22] as u64, m2) + | _pdep_u64(input[30] as u64, m3) + | _pdep_u64(input[38] as u64, m4) + | _pdep_u64(input[46] as u64, m5) + | _pdep_u64(input[54] as u64, m6) + | _pdep_u64(input[62] as u64, m7); + scatter(output, 6, v); + + // Group 7 (base=14) + let v = _pdep_u64(input[7] as u64, m0) + | _pdep_u64(input[15] as u64, m1) + | _pdep_u64(input[23] as u64, m2) + | _pdep_u64(input[31] as u64, m3) + | _pdep_u64(input[39] as u64, m4) + | _pdep_u64(input[47] as u64, m5) + | _pdep_u64(input[55] as u64, m6) + | _pdep_u64(input[63] as u64, m7); + scatter(output, 14, v); + + // Second half: 8 groups using BASE_PATTERN_SECOND + // Group 0 (base=1) + let v = _pdep_u64(input[64] as u64, m0) + | _pdep_u64(input[72] as u64, m1) + | _pdep_u64(input[80] as u64, m2) + | _pdep_u64(input[88] as u64, m3) + | _pdep_u64(input[96] as u64, m4) + | _pdep_u64(input[104] as u64, m5) + | _pdep_u64(input[112] as u64, m6) + | _pdep_u64(input[120] as u64, m7); + scatter(output, 1, v); + + // Group 1 (base=9) + let v = _pdep_u64(input[65] as u64, m0) + | _pdep_u64(input[73] as u64, m1) + | _pdep_u64(input[81] as u64, m2) + | _pdep_u64(input[89] as u64, m3) + | _pdep_u64(input[97] as u64, m4) + | _pdep_u64(input[105] as u64, m5) + | _pdep_u64(input[113] as u64, m6) + | _pdep_u64(input[121] as u64, m7); + scatter(output, 9, v); + + // Group 2 (base=5) + let v = _pdep_u64(input[66] as u64, m0) + | _pdep_u64(input[74] as u64, m1) + | _pdep_u64(input[82] as u64, m2) + | _pdep_u64(input[90] as u64, m3) + | _pdep_u64(input[98] as u64, m4) + | _pdep_u64(input[106] as u64, m5) + | _pdep_u64(input[114] as u64, m6) + | _pdep_u64(input[122] as u64, m7); + scatter(output, 5, v); + + // Group 3 (base=13) + let v = _pdep_u64(input[67] as u64, m0) + | _pdep_u64(input[75] as u64, m1) + | _pdep_u64(input[83] as u64, m2) + | _pdep_u64(input[91] as u64, m3) + | _pdep_u64(input[99] as u64, m4) + | _pdep_u64(input[107] as u64, m5) + | _pdep_u64(input[115] as u64, m6) + | _pdep_u64(input[123] as u64, m7); + scatter(output, 13, v); + + // Group 4 (base=3) + let v = _pdep_u64(input[68] as u64, m0) + | _pdep_u64(input[76] as u64, m1) + | _pdep_u64(input[84] as u64, m2) + | _pdep_u64(input[92] as u64, m3) + | _pdep_u64(input[100] as u64, m4) + | _pdep_u64(input[108] as u64, m5) + | _pdep_u64(input[116] as u64, m6) + | _pdep_u64(input[124] as u64, m7); + scatter(output, 3, v); + + // Group 5 (base=11) + let v = _pdep_u64(input[69] as u64, m0) + | _pdep_u64(input[77] as u64, m1) + | _pdep_u64(input[85] as u64, m2) + | _pdep_u64(input[93] as u64, m3) + | _pdep_u64(input[101] as u64, m4) + | _pdep_u64(input[109] as u64, m5) + | _pdep_u64(input[117] as u64, m6) + | _pdep_u64(input[125] as u64, m7); + scatter(output, 11, v); + + // Group 6 (base=7) + let v = _pdep_u64(input[70] as u64, m0) + | _pdep_u64(input[78] as u64, m1) + | _pdep_u64(input[86] as u64, m2) + | _pdep_u64(input[94] as u64, m3) + | _pdep_u64(input[102] as u64, m4) + | _pdep_u64(input[110] as u64, m5) + | _pdep_u64(input[118] as u64, m6) + | _pdep_u64(input[126] as u64, m7); + scatter(output, 7, v); + + // Group 7 (base=15) + let v = _pdep_u64(input[71] as u64, m0) + | _pdep_u64(input[79] as u64, m1) + | _pdep_u64(input[87] as u64, m2) + | _pdep_u64(input[95] as u64, m3) + | _pdep_u64(input[103] as u64, m4) + | _pdep_u64(input[111] as u64, m5) + | _pdep_u64(input[119] as u64, m6) + | _pdep_u64(input[127] as u64, m7); + scatter(output, 15, v); +} + +// Static permutation tables for VBMI gather operations +static GATHER_FIRST: [u8; 64] = [ + // Gather bytes at stride 16 for first 8 groups (bases from BASE_PATTERN_FIRST) + // Group 0: base=0 + 0, 16, 32, 48, 64, 80, 96, 112, // Group 1: base=8 + 8, 24, 40, 56, 72, 88, 104, 120, // Group 2: base=4 + 4, 20, 36, 52, 68, 84, 100, 116, // Group 3: base=12 + 12, 28, 44, 60, 76, 92, 108, 124, // Group 4: base=2 + 2, 18, 34, 50, 66, 82, 98, 114, // Group 5: base=10 + 10, 26, 42, 58, 74, 90, 106, 122, // Group 6: base=6 + 6, 22, 38, 54, 70, 86, 102, 118, // Group 7: base=14 + 14, 30, 46, 62, 78, 94, 110, 126, +]; + +static GATHER_SECOND: [u8; 64] = [ + // Gather bytes at stride 16 for second 8 groups (bases from BASE_PATTERN_SECOND) + // Group 0: base=1 + 1, 17, 33, 49, 65, 81, 97, 113, // Group 1: base=9 + 9, 25, 41, 57, 73, 89, 105, 121, // Group 2: base=5 + 5, 21, 37, 53, 69, 85, 101, 117, // Group 3: base=13 + 13, 29, 45, 61, 77, 93, 109, 125, // Group 4: base=3 + 3, 19, 35, 51, 67, 83, 99, 115, // Group 5: base=11 + 11, 27, 43, 59, 75, 91, 107, 123, // Group 6: base=7 + 7, 23, 39, 55, 71, 87, 103, 119, // Group 7: base=15 + 15, 31, 47, 63, 79, 95, 111, 127, +]; + +// 8x8 byte transpose permutation for scatter phase +// Input: [g0b0..g0b7, g1b0..g1b7, ..., g7b0..g7b7] (8 groups of 8 bytes) +// Output: [g0b0,g1b0,..,g7b0, g0b1,g1b1,..,g7b1, ...] (8 rows of 8 bytes) +static SCATTER_8X8: [u8; 64] = [ + 0, 8, 16, 24, 32, 40, 48, 56, // byte 0 from each group + 1, 9, 17, 25, 33, 41, 49, 57, // byte 1 from each group + 2, 10, 18, 26, 34, 42, 50, 58, // byte 2 from each group + 3, 11, 19, 27, 35, 43, 51, 59, // byte 3 from each group + 4, 12, 20, 28, 36, 44, 52, 60, // byte 4 from each group + 5, 13, 21, 29, 37, 45, 53, 61, // byte 5 from each group + 6, 14, 22, 30, 38, 46, 54, 62, // byte 6 from each group + 7, 15, 23, 31, 39, 47, 55, 63, // byte 7 from each group +]; + +/// Transpose 1024 bits using AVX-512 VBMI for vectorized gather and scatter. +/// +/// Uses vpermi2b to gather bytes from stride-16 positions in parallel, +/// and vpermb for the final 8x8 byte transpose to output format. +/// +/// # Safety +/// Requires AVX-512F, AVX-512BW, and AVX-512VBMI support. +#[target_feature(enable = "avx512f", enable = "avx512bw", enable = "avx512vbmi")] +#[inline(never)] +#[expect(clippy::cast_possible_wrap)] +#[expect(clippy::cast_ptr_alignment)] +#[expect(unsafe_op_in_unsafe_fn)] +pub unsafe fn transpose_bits_vbmi(input: &[u8; 128], output: &mut [u8; 128]) { + // Load all 128 input bytes into two ZMM registers + let in_lo = _mm512_loadu_si512(input.as_ptr().cast::<__m512i>()); + let in_hi = _mm512_loadu_si512(input.as_ptr().add(64).cast::<__m512i>()); + + // Load permutation indices (static tables) + let idx_first = _mm512_loadu_si512(GATHER_FIRST.as_ptr().cast::<__m512i>()); + let idx_second = _mm512_loadu_si512(GATHER_SECOND.as_ptr().cast::<__m512i>()); + let idx_scatter = _mm512_loadu_si512(SCATTER_8X8.as_ptr().cast::<__m512i>()); + + // Masks for 8x8 bit transpose + let mask1 = _mm512_set1_epi64(TRANSPOSE_2X2 as i64); + let mask2 = _mm512_set1_epi64(TRANSPOSE_4X4 as i64); + let mask3 = _mm512_set1_epi64(TRANSPOSE_8X8 as i64); + + // Process first half + let gathered = _mm512_permutex2var_epi8(in_lo, idx_first, in_hi); + + // 8x8 bit transpose on all 8 groups in parallel + let mut v = gathered; + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<7>(v)), mask1); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<7>(t)); + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<14>(v)), mask2); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<14>(t)); + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<28>(v)), mask3); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<28>(t)); + + // 8x8 byte transpose for scatter using vpermb + let scattered = _mm512_permutexvar_epi8(idx_scatter, v); + _mm512_storeu_si512(output.as_mut_ptr().cast::<__m512i>(), scattered); + + // Process second half + let gathered = _mm512_permutex2var_epi8(in_lo, idx_second, in_hi); + + let mut v = gathered; + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<7>(v)), mask1); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<7>(t)); + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<14>(v)), mask2); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<14>(t)); + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<28>(v)), mask3); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<28>(t)); + + let scattered = _mm512_permutexvar_epi8(idx_scatter, v); + _mm512_storeu_si512(output.as_mut_ptr().add(64).cast::<__m512i>(), scattered); +} + +/// Untranspose 1024 bits using AVX-512 VBMI for vectorized scatter. +/// +/// # Safety +/// Requires AVX-512F, AVX-512BW, and AVX-512VBMI support. +#[target_feature(enable = "avx512f", enable = "avx512bw", enable = "avx512vbmi")] +#[inline(never)] +#[expect(clippy::cast_possible_wrap)] +#[expect(clippy::cast_ptr_alignment)] +#[expect(unsafe_op_in_unsafe_fn)] +pub unsafe fn untranspose_bits_vbmi(input: &[u8; 128], output: &mut [u8; 128]) { + // For untranspose, we gather consecutive bytes from transposed layout, + // then scatter back to stride-16 positions + + // Gather indices for first half - collect 8 bytes per group from transposed layout + // In transposed layout, bytes for group 0 are at: [0, 8, 16, 24, 32, 40, 48, 56] + let gather_indices: [u8; 64] = [ + 0, 8, 16, 24, 32, 40, 48, 56, // Group 0 + 1, 9, 17, 25, 33, 41, 49, 57, // Group 1 + 2, 10, 18, 26, 34, 42, 50, 58, // Group 2 + 3, 11, 19, 27, 35, 43, 51, 59, // Group 3 + 4, 12, 20, 28, 36, 44, 52, 60, // Group 4 + 5, 13, 21, 29, 37, 45, 53, 61, // Group 5 + 6, 14, 22, 30, 38, 46, 54, 62, // Group 6 + 7, 15, 23, 31, 39, 47, 55, 63, // Group 7 + ]; + + let in_first = _mm512_loadu_si512(input.as_ptr().cast::<__m512i>()); + let idx = _mm512_loadu_si512(gather_indices.as_ptr().cast::<__m512i>()); + let gathered = _mm512_permutexvar_epi8(idx, in_first); + + // 8x8 bit transpose + let mask1 = _mm512_set1_epi64(TRANSPOSE_2X2 as i64); + let mask2 = _mm512_set1_epi64(TRANSPOSE_4X4 as i64); + let mask3 = _mm512_set1_epi64(TRANSPOSE_8X8 as i64); + + let mut v = gathered; + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<7>(v)), mask1); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<7>(t)); + + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<14>(v)), mask2); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<14>(t)); + + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<28>(v)), mask3); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<28>(t)); + + // Scatter to output at stride 16 - need to use scalar stores for now + // (AVX-512 scatter is available but complex for this pattern) + let mut result = [0u64; 8]; + _mm512_storeu_si512(result.as_mut_ptr().cast::<__m512i>(), v); + + for base_group in 0..8 { + let out_base = BASE_PATTERN_FIRST[base_group]; + for i in 0..8 { + output[out_base + i * 16] = (result[base_group] >> (i * 8)) as u8; + } + } + + // Second half + let in_second = _mm512_loadu_si512(input.as_ptr().add(64).cast::<__m512i>()); + let gathered = _mm512_permutexvar_epi8(idx, in_second); + + let mut v = gathered; + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<7>(v)), mask1); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<7>(t)); + + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<14>(v)), mask2); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<14>(t)); + + let t = _mm512_and_si512(_mm512_xor_si512(v, _mm512_srli_epi64::<28>(v)), mask3); + v = _mm512_xor_si512(_mm512_xor_si512(v, t), _mm512_slli_epi64::<28>(t)); + + _mm512_storeu_si512(result.as_mut_ptr().cast::<__m512i>(), v); + + for base_group in 0..8 { + let out_base = BASE_PATTERN_SECOND[base_group]; + for i in 0..8 { + output[out_base + i * 16] = (result[base_group] >> (i * 8)) as u8; + } + } +} + +#[cfg(test)] +mod tests { + use crate::bit_transpose::generate_test_data; + use crate::bit_transpose::transpose_bits_baseline; + use crate::bit_transpose::x86::has_bmi2; + use crate::bit_transpose::x86::has_vbmi; + use crate::bit_transpose::x86::transpose_bits_bmi2; + use crate::bit_transpose::x86::transpose_bits_vbmi; + use crate::bit_transpose::x86::untranspose_bits_bmi2; + use crate::bit_transpose::x86::untranspose_bits_vbmi; + + #[test] + fn test_bmi2_matches_baseline() { + if !has_bmi2() { + return; + } + + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut baseline_out = [0u8; 128]; + let mut bmi2_out = [0u8; 128]; + + transpose_bits_baseline(&input, &mut baseline_out); + unsafe { transpose_bits_bmi2(&input, &mut bmi2_out) }; + + assert_eq!( + baseline_out, bmi2_out, + "BMI2 transpose doesn't match baseline for seed {seed}" + ); + } + } + + #[test] + fn test_bmi2_roundtrip() { + if !has_bmi2() { + return; + } + + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut transposed = [0u8; 128]; + let mut roundtrip = [0u8; 128]; + + unsafe { + transpose_bits_bmi2(&input, &mut transposed); + untranspose_bits_bmi2(&transposed, &mut roundtrip); + } + + assert_eq!(input, roundtrip, "BMI2 roundtrip failed for seed {seed}"); + } + } + + #[test] + fn test_vbmi_matches_baseline() { + if !has_vbmi() { + return; + } + + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut baseline_out = [0u8; 128]; + let mut vbmi_out = [0u8; 128]; + + transpose_bits_baseline(&input, &mut baseline_out); + unsafe { transpose_bits_vbmi(&input, &mut vbmi_out) }; + + assert_eq!( + baseline_out, vbmi_out, + "VBMI transpose doesn't match baseline for seed {seed}" + ); + } + } + + #[test] + fn test_vbmi_roundtrip() { + if !has_vbmi() { + return; + } + + for seed in [0, 42, 123, 255] { + let input = generate_test_data(seed); + let mut transposed = [0u8; 128]; + let mut roundtrip = [0u8; 128]; + + unsafe { + transpose_bits_vbmi(&input, &mut transposed); + untranspose_bits_vbmi(&transposed, &mut roundtrip); + } + + assert_eq!(input, roundtrip, "VBMI roundtrip failed for seed {seed}"); + } + } +} diff --git a/encodings/fastlanes/src/delta/array/delta_compress.rs b/encodings/fastlanes/src/delta/array/delta_compress.rs index f5d9e10b1ea..7f00eb172dc 100644 --- a/encodings/fastlanes/src/delta/array/delta_compress.rs +++ b/encodings/fastlanes/src/delta/array/delta_compress.rs @@ -8,21 +8,17 @@ use fastlanes::Delta; use fastlanes::FastLanes; use fastlanes::Transpose; use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::BoolArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::primitive::PrimitiveArrayExt; use vortex_array::dtype::NativePType; use vortex_array::match_each_unsigned_integer_ptype; -use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; use crate::FL_CHUNK_SIZE; -use crate::bit_transpose::transpose_bitbuffer; +use crate::bit_transpose::transpose_validity; use crate::fill_forward_nulls; - pub fn delta_compress( array: &PrimitiveArray, ctx: &mut ExecutionCtx, @@ -37,15 +33,8 @@ pub fn delta_compress( // corrupted delta values propagate through the cumulative sum during decompression. let filled = fill_forward_nulls(array.to_buffer::(), &validity, ctx)?; let (bases, deltas) = compress_primitive::(&filled); - let validity = match validity { - Validity::Array(mask) => { - let bits = mask.execute::(ctx)?.into_bit_buffer(); - Validity::Array( - BoolArray::new(transpose_bitbuffer(bits), Validity::NonNullable).into_array(), - ) - } - validity => validity, - }; + // TODO(robert): This can be avoided if we add TransposedBoolArray that performs index translation when necessary. + let validity = transpose_validity(&validity, ctx)?; ( PrimitiveArray::new(bases, array.dtype().nullability().into()), PrimitiveArray::new(deltas, validity), @@ -115,13 +104,10 @@ mod tests { use rstest::rstest; use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; - use vortex_array::arrays::Bool; use vortex_array::arrays::PrimitiveArray; use vortex_array::assert_arrays_eq; - use vortex_array::validity::Validity; use vortex_error::VortexExpect; use vortex_error::VortexResult; - use vortex_error::vortex_bail; use vortex_session::VortexSession; use crate::Delta; @@ -179,11 +165,6 @@ mod tests { (0u8..200).map(|i| (!(50..100).contains(&i)).then_some(i)), ); let (bases, deltas) = delta_compress(&array, &mut ctx)?; - let Validity::Array(storage_validity) = deltas.validity()? else { - vortex_bail!("test input should have array-backed validity") - }; - assert!(storage_validity.is::()); - let bitpacked_deltas = bitpack_encode(&deltas, 1, None, &mut ctx)?; let packed_delta = Delta::try_new( bases.into_array(), diff --git a/encodings/fastlanes/src/delta/array/delta_decompress.rs b/encodings/fastlanes/src/delta/array/delta_decompress.rs index c6da7aa1ea9..7dcbeb44950 100644 --- a/encodings/fastlanes/src/delta/array/delta_decompress.rs +++ b/encodings/fastlanes/src/delta/array/delta_decompress.rs @@ -18,6 +18,7 @@ use vortex_buffer::BufferMut; use vortex_error::VortexResult; use crate::DeltaArray; +use crate::bit_transpose::untranspose_validity; use crate::delta::array::DeltaArrayExt; pub fn delta_decompress( @@ -30,7 +31,8 @@ pub fn delta_decompress( let start = array.offset(); let end = start + array.len(); - let validity = array.validity()?; + let validity = untranspose_validity(&deltas.validity()?, ctx)?; + let validity = validity.slice(start..end)?; let original_ptype = deltas.ptype(); // Signed inputs are processed through their unsigned counterpart; `wrapping_add` on the diff --git a/encodings/fastlanes/src/delta/array/mod.rs b/encodings/fastlanes/src/delta/array/mod.rs index 9a3e6480262..7754bb37a59 100644 --- a/encodings/fastlanes/src/delta/array/mod.rs +++ b/encodings/fastlanes/src/delta/array/mod.rs @@ -13,8 +13,6 @@ use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_ensure; -use crate::Delta; - pub mod delta_compress; pub mod delta_decompress; @@ -90,7 +88,7 @@ impl Display for DeltaData { } } -pub trait DeltaArrayExt: TypedArrayRef { +pub trait DeltaArrayExt: TypedArrayRef { fn bases(&self) -> &ArrayRef { self.as_ref().slots()[BASES_SLOT] .as_ref() @@ -108,7 +106,7 @@ pub trait DeltaArrayExt: TypedArrayRef { } } -impl> DeltaArrayExt for T {} +impl> DeltaArrayExt for T {} impl DeltaData { pub fn try_new(offset: usize) -> VortexResult { diff --git a/encodings/fastlanes/src/delta/vtable/validity.rs b/encodings/fastlanes/src/delta/vtable/validity.rs index c2138750fa3..c8d8b7b4dc0 100644 --- a/encodings/fastlanes/src/delta/vtable/validity.rs +++ b/encodings/fastlanes/src/delta/vtable/validity.rs @@ -2,72 +2,26 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors use vortex_array::ArrayView; -use vortex_array::IntoArray; -use vortex_array::arrays::Bool; -use vortex_array::arrays::bool::BoolArrayExt; +use vortex_array::VortexSessionExecute; +use vortex_array::legacy_session; use vortex_array::validity::Validity; use vortex_array::vtable::ValidityVTable; use vortex_error::VortexResult; -use vortex_error::vortex_bail; use crate::Delta; -use crate::TransposedBool; +use crate::bit_transpose::untranspose_validity; use crate::delta::array::DeltaArrayExt; impl ValidityVTable for Delta { + #[allow(clippy::disallowed_methods)] fn validity(array: ArrayView<'_, Delta>) -> VortexResult { let start = array.offset(); - let stop = start + array.len(); - let validity = match array.deltas().validity()? { - Validity::Array(mask) => { - let Some(mask) = mask.as_opt::() else { - vortex_bail!( - "DeltaArray storage validity must be a BoolArray, got {}", - mask.encoding_id() - ); - }; - Validity::Array(TransposedBool::try_new(mask.to_bit_buffer())?.into_array()) - } - validity => validity, - }; - validity.slice(start..stop) - } -} - -#[cfg(test)] -mod tests { - use vortex_array::VortexSessionExecute; - use vortex_array::array_session; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::PrimitiveArray; - use vortex_array::assert_arrays_eq; - use vortex_array::validity::Validity; - use vortex_error::VortexResult; - use vortex_error::vortex_bail; - - use super::*; - use crate::TransposedBool; - - #[test] - fn validity_is_lazy_for_cross_chunk_slice() -> VortexResult<()> { - let session = array_session(); - crate::initialize(&session); - let mut ctx = session.create_execution_ctx(); - let primitive = PrimitiveArray::from_option_iter( - (0u32..2048).map(|value| (value % 3 != 0).then_some(value)), - ); - let delta = Delta::try_from_primitive_array(&primitive, &mut ctx)?; - let sliced = delta.slice(1000..1050)?; + let end = start + array.len(); - let Validity::Array(validity) = sliced.validity()? else { - vortex_bail!("expected array-backed validity") - }; - assert!(validity.is::()); - assert_arrays_eq!( - validity, - BoolArray::from_iter((1000u32..1050).map(|value| value % 3 != 0)), - &mut ctx - ); - Ok(()) + let validity = untranspose_validity( + &array.deltas().validity()?, + &mut legacy_session().create_execution_ctx(), + )?; + validity.slice(start..end) } } diff --git a/encodings/fastlanes/src/lib.rs b/encodings/fastlanes/src/lib.rs index bcc9859e3eb..e6192a9cdfb 100644 --- a/encodings/fastlanes/src/lib.rs +++ b/encodings/fastlanes/src/lib.rs @@ -30,7 +30,6 @@ pub use bitpacking::*; pub use delta::*; pub use r#for::*; pub use rle::*; -pub use transposed_bool::*; use vortex_array::ExecutionCtx; use vortex_array::arrays::BoolArray; use vortex_array::arrays::bool::BoolArrayExt; @@ -44,7 +43,6 @@ mod bitpacking; mod delta; mod r#for; mod rle; -mod transposed_bool; pub const FL_CHUNK_SIZE: usize = 1024; @@ -72,7 +70,6 @@ pub fn initialize(session: &VortexSession) { session.arrays().register(Delta); session.arrays().register(FoR); session.arrays().register(RLE); - session.arrays().register(TransposedBool); bitpacking::initialize(session); r#for::initialize(session); rle::initialize(session); diff --git a/encodings/fastlanes/src/transposed_bool.rs b/encodings/fastlanes/src/transposed_bool.rs deleted file mode 100644 index 9b987b40936..00000000000 --- a/encodings/fastlanes/src/transposed_bool.rs +++ /dev/null @@ -1,360 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright the Vortex contributors - -use std::fmt::Display; -use std::fmt::Formatter; -use std::hash::Hash; -use std::hash::Hasher; -use std::ops::Range; - -use vortex_array::Array; -use vortex_array::ArrayEq; -use vortex_array::ArrayHash; -use vortex_array::ArrayId; -use vortex_array::ArrayParts; -use vortex_array::ArrayRef; -use vortex_array::ArrayView; -use vortex_array::EqMode; -use vortex_array::ExecutionCtx; -use vortex_array::ExecutionResult; -use vortex_array::IntoArray; -use vortex_array::TypedArrayRef; -use vortex_array::arrays::BoolArray; -use vortex_array::arrays::slice::SliceReduce; -use vortex_array::arrays::slice::SliceReduceAdaptor; -use vortex_array::buffer::BufferHandle; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::optimizer::rules::ParentRuleSet; -use vortex_array::scalar::Scalar; -use vortex_array::serde::ArrayChildren; -use vortex_array::validity::Validity; -use vortex_array::vtable::OperationsVTable; -use vortex_array::vtable::VTable; -use vortex_array::vtable::ValidityVTable; -use vortex_buffer::BitBuffer; -use vortex_buffer::BitBufferView; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; -use vortex_error::vortex_panic; -use vortex_session::VortexSession; -use vortex_session::registry::CachedId; - -use crate::FL_CHUNK_SIZE; -use crate::bit_transpose::untranspose_bitbuffer; - -/// A non-nullable boolean array stored in FastLanes-transposed order. -pub type TransposedBoolArray = Array; - -/// The array encoding for a boolean bitmap stored in FastLanes-transposed order. -#[derive(Clone, Debug)] -pub struct TransposedBool; - -/// Per-array data for a [`TransposedBoolArray`]. -#[derive(Clone, Debug)] -pub struct TransposedBoolData { - bits: BufferHandle, - offset: usize, -} - -impl Display for TransposedBoolData { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "offset: {}", self.offset) - } -} - -impl ArrayHash for TransposedBoolData { - fn array_hash(&self, state: &mut H, accuracy: EqMode) { - self.bits.array_hash(state, accuracy); - self.offset.hash(state); - } -} - -impl ArrayEq for TransposedBoolData { - fn array_eq(&self, other: &Self, accuracy: EqMode) -> bool { - self.bits.array_eq(&other.bits, accuracy) && self.offset == other.offset - } -} - -/// Accessors for a [`TransposedBoolArray`]. -pub trait TransposedBoolArrayExt: TypedArrayRef { - /// Returns the logical offset into the first untransposed chunk. - fn offset(&self) -> usize { - self.deref().offset - } - - /// Borrows the full backing bitmap in FastLanes-transposed order. - /// - /// For a sliced array, this includes the complete chunks surrounding the logical range. - fn bit_buffer_view(&self) -> BitBufferView<'_> { - BitBufferView::new( - self.deref().bits.as_host().as_slice(), - self.deref().bits.len() * 8, - ) - } - - /// Returns the full backing bitmap in FastLanes-transposed order, consuming the array. - fn to_bit_buffer(self) -> BitBuffer - where - Self: Sized, - { - let bits = self.deref().bits.clone(); - let len = bits.len() * 8; - BitBuffer::new(bits.unwrap_host(), len) - } -} - -impl> TransposedBoolArrayExt for T {} - -impl TransposedBool { - /// Creates an array from a bitmap already stored in FastLanes-transposed order. - /// - /// # Errors - /// - /// Returns an error if the bitmap does not contain complete 1,024-bit chunks. - pub fn try_new(bits: BitBuffer) -> VortexResult { - let bits = bits.sliced(); - vortex_ensure!( - bits.len().is_multiple_of(FL_CHUNK_SIZE), - "TransposedBoolArray length {} must be a multiple of {FL_CHUNK_SIZE}", - bits.len() - ); - let (_, len, bits) = bits.into_inner(); - Self::try_new_view(BufferHandle::new_host(bits), 0, len) - } - - fn try_new_view( - bits: BufferHandle, - offset: usize, - len: usize, - ) -> VortexResult { - Array::try_from_parts(ArrayParts::new( - TransposedBool, - DType::Bool(Nullability::NonNullable), - len, - TransposedBoolData { bits, offset }, - )) - } -} - -impl VTable for TransposedBool { - type TypedArrayData = TransposedBoolData; - type OperationsVTable = Self; - type ValidityVTable = Self; - - fn id(&self) -> ArrayId { - static ID: CachedId = CachedId::new("fastlanes.transposed_bool"); - *ID - } - - fn validate( - &self, - data: &Self::TypedArrayData, - dtype: &DType, - len: usize, - slots: &[Option], - ) -> VortexResult<()> { - vortex_ensure!( - dtype == &DType::Bool(Nullability::NonNullable), - "TransposedBoolArray must have non-nullable boolean dtype, got {dtype}" - ); - vortex_ensure!( - slots.is_empty(), - "TransposedBoolArray expects no slots, got {}", - slots.len() - ); - vortex_ensure!( - (data.bits.len() * 8).is_multiple_of(FL_CHUNK_SIZE), - "TransposedBoolArray buffer length {} bits must be a multiple of {FL_CHUNK_SIZE}", - data.bits.len() * 8 - ); - vortex_ensure!( - data.offset < FL_CHUNK_SIZE, - "TransposedBoolArray offset {} must be less than {FL_CHUNK_SIZE}", - data.offset - ); - let end = data - .offset - .checked_add(len) - .ok_or_else(|| vortex_error::vortex_err!("TransposedBoolArray range end overflow"))?; - vortex_ensure!( - end <= data.bits.len() * 8, - "TransposedBoolArray range {}..{} exceeds buffer length {} bits", - data.offset, - end, - data.bits.len() * 8 - ); - Ok(()) - } - - fn nbuffers(_array: ArrayView<'_, Self>) -> usize { - 1 - } - - fn buffer(array: ArrayView<'_, Self>, idx: usize) -> BufferHandle { - match idx { - 0 => array.bits.clone(), - _ => vortex_panic!("TransposedBoolArray buffer index {idx} out of bounds"), - } - } - - fn buffer_name(_array: ArrayView<'_, Self>, idx: usize) -> Option { - match idx { - 0 => Some("bits".to_string()), - _ => None, - } - } - - fn with_buffers( - &self, - array: ArrayView<'_, Self>, - buffers: &[BufferHandle], - ) -> VortexResult> { - vortex_ensure!( - buffers.len() == 1, - "TransposedBoolArray expects one buffer, got {}", - buffers.len() - ); - Ok(ArrayParts::new( - self.clone(), - array.dtype().clone(), - array.len(), - TransposedBoolData { - bits: buffers[0].clone(), - offset: array.offset(), - }, - )) - } - - fn serialize( - _array: ArrayView<'_, Self>, - _session: &VortexSession, - ) -> VortexResult>> { - vortex_bail!("Cannot serialise TransposedBoolArray"); - } - - fn deserialize( - &self, - _dtype: &DType, - _len: usize, - _metadata: &[u8], - _buffers: &[BufferHandle], - _children: &dyn ArrayChildren, - _session: &VortexSession, - ) -> VortexResult> { - vortex_bail!("Cannot deserialise TransposedBoolArray"); - } - - fn slot_name(_array: ArrayView<'_, Self>, idx: usize) -> String { - vortex_panic!("TransposedBoolArray slot index {idx} out of bounds") - } - - fn execute(array: Array, _ctx: &mut ExecutionCtx) -> VortexResult { - let len = array.len(); - let offset = array.offset(); - let untransposed = BoolArray::new( - untranspose_bitbuffer(array.to_bit_buffer()), - Validity::NonNullable, - ); - Ok(ExecutionResult::done( - untransposed.slice(offset..offset + len)?, - )) - } - - fn reduce_parent( - array: ArrayView<'_, Self>, - parent: &ArrayRef, - child_idx: usize, - ) -> VortexResult> { - RULES.evaluate(array, parent, child_idx) - } -} - -impl OperationsVTable for TransposedBool { - fn scalar_at( - array: ArrayView<'_, TransposedBool>, - index: usize, - _ctx: &mut ExecutionCtx, - ) -> VortexResult { - let physical_index = array.offset() + index; - let chunk_start = physical_index / FL_CHUNK_SIZE * FL_CHUNK_SIZE; - let transposed_index = chunk_start + fastlanes::transpose(physical_index % FL_CHUNK_SIZE); - Ok(Scalar::bool( - array.bit_buffer_view().value(transposed_index), - Nullability::NonNullable, - )) - } -} - -impl ValidityVTable for TransposedBool { - fn validity(_array: ArrayView<'_, TransposedBool>) -> VortexResult { - Ok(Validity::NonNullable) - } -} - -impl SliceReduce for TransposedBool { - fn slice(array: ArrayView<'_, Self>, range: Range) -> VortexResult> { - let physical_start = array.offset() + range.start; - let physical_stop = array.offset() + range.end; - let start_chunk = physical_start / FL_CHUNK_SIZE; - let stop_chunk = physical_stop.div_ceil(FL_CHUNK_SIZE); - let bits = array.bits.slice(start_chunk * 128..stop_chunk * 128); - - Ok(Some( - TransposedBool::try_new_view(bits, physical_start % FL_CHUNK_SIZE, range.len())? - .into_array(), - )) - } -} - -static RULES: ParentRuleSet = - ParentRuleSet::new(&[ParentRuleSet::lift(&SliceReduceAdaptor(TransposedBool))]); - -#[cfg(test)] -mod tests { - use vortex_array::VortexSessionExecute; - use vortex_array::array_session; - use vortex_array::assert_arrays_eq; - use vortex_buffer::BitBuffer; - use vortex_error::VortexResult; - - use super::*; - use crate::bit_transpose::transpose_bitbuffer; - - fn test_bits() -> BitBuffer { - BitBuffer::from_iter((0..2 * FL_CHUNK_SIZE).map(|i| i % 3 != 0 && i % 11 != 0)) - } - - #[test] - fn execute_full_array() -> VortexResult<()> { - let expected = test_bits(); - let array = TransposedBool::try_new(transpose_bitbuffer(expected.clone()))?; - let mut ctx = array_session().create_execution_ctx(); - - assert_arrays_eq!(array, BoolArray::from(expected), &mut ctx); - Ok(()) - } - - #[test] - fn slice_stays_lazy_and_translates_scalars() -> VortexResult<()> { - let expected = test_bits(); - let array = TransposedBool::try_new(transpose_bitbuffer(expected.clone()))?; - let sliced = array.slice(1000..1050)?; - assert!(sliced.is::()); - - let mut ctx = array_session().create_execution_ctx(); - for index in [0, 23, 49] { - assert_eq!( - sliced.execute_scalar(index, &mut ctx)?.as_bool().value(), - Some(expected.value(1000 + index)) - ); - } - assert_arrays_eq!( - sliced, - BoolArray::from(expected.slice(1000..1050)), - &mut ctx - ); - Ok(()) - } -}