Skip to content

Reduce monomorphization by dispatching over unsigned integer widths #8110

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
joseph-isaacs wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

Reduce monomorphization by dispatching over unsigned integer widths #8110

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
54 changes: 28 additions & 26 deletions encodings/fastlanes/src/rle/array/rle_decompress.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,35 +19,44 @@ use crate::RLEArray;
use crate::rle::RLEArrayExt;

/// Decompresses an RLE array back into a primitive array.
#[expect(
clippy::cognitive_complexity,
reason = "complexity is from nested match_each_* macros"
)]
pub fn rle_decompress(array: &RLEArray, ctx: &mut ExecutionCtx) -> VortexResult<PrimitiveArray> {
// The per-chunk value-index offsets are tiny (one entry per 1024-element chunk), so cast them
// to `u64` once here instead of monomorphizing the whole decode loop over the offset width.
let values_idx_offsets = array
.values_idx_offsets()
.clone()
.execute::<PrimitiveArray>(ctx)?;
let values_idx_offsets: Vec<u64> =
match_each_unsigned_integer_ptype!(values_idx_offsets.ptype(), |O| {
values_idx_offsets
.as_slice::<O>()
.iter()
.map(|&o| o.as_())
.collect()
});

match_each_native_ptype!(array.values().dtype().as_ptype(), |V| {
match_each_unsigned_integer_ptype!(array.values_idx_offsets().dtype().as_ptype(), |O| {
// RLE indices are always u16 (or u8 if downcasted).
match array.indices().dtype().as_ptype() {
PType::U8 => rle_decode_typed::<V, u8, O>(array, ctx),
PType::U16 => rle_decode_typed::<V, u16, O>(array, ctx),
_ => vortex_panic!(
"Unsupported index type for RLE decoding: {}",
array.indices().dtype().as_ptype()
),
}
})
// RLE indices are always u16 (or u8 if downcasted).
match array.indices().dtype().as_ptype() {
PType::U8 => rle_decode_typed::<V, u8>(array, &values_idx_offsets, ctx),
PType::U16 => rle_decode_typed::<V, u16>(array, &values_idx_offsets, ctx),
_ => vortex_panic!(
"Unsupported index type for RLE decoding: {}",
array.indices().dtype().as_ptype()
),
}
})
}

/// Decompresses an `RLEArray` into to a primitive array of unsigned integers.
fn rle_decode_typed<V, I, O>(
fn rle_decode_typed<V, I>(
array: &RLEArray,
values_idx_offsets: &[u64],
ctx: &mut ExecutionCtx,
) -> VortexResult<PrimitiveArray>
where
V: NativePType + RLE + Clone + Copy,
I: NativePType + Into<usize>,
O: NativePType + AsPrimitive<u64>,
{
let values = array.values().clone().execute::<PrimitiveArray>(ctx)?;
let values = values.as_slice::<V>();
Expand All @@ -64,23 +73,16 @@ where
let mut buffer = BufferMut::<V>::with_capacity(num_chunks * FL_CHUNK_SIZE);
let (out_buf, _) = buffer.spare_capacity_mut().as_chunks_mut::<FL_CHUNK_SIZE>();

let values_idx_offsets = array
.values_idx_offsets()
.clone()
.execute::<PrimitiveArray>(ctx)?;
let values_idx_offsets = values_idx_offsets.as_slice::<O>();

for (chunk_idx, (chunk_indices, chunk_out)) in
indices_sl.iter().zip(out_buf.iter_mut()).enumerate()
{
// Offsets in `values_idx_offsets` are absolute and need to be shifted
// by the offset of the first chunk, respective of the current slice,
// to make them relative.
let value_idx_offset =
(values_idx_offsets[chunk_idx].as_() - values_idx_offsets[0].as_()) as usize;
let value_idx_offset = (values_idx_offsets[chunk_idx] - values_idx_offsets[0]) as usize;

let next_value_idx_offset = if chunk_idx + 1 < num_chunks {
(values_idx_offsets[chunk_idx + 1].as_() - values_idx_offsets[0].as_()) as usize
(values_idx_offsets[chunk_idx + 1] - values_idx_offsets[0]) as usize
} else {
values.len()
};
Expand Down
11 changes: 6 additions & 5 deletions encodings/sparse/src/canonical.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ use vortex_array::arrays::VarBinView;
use vortex_array::arrays::VarBinViewArray;
use vortex_array::arrays::fixed_size_list::FixedSizeListArrayExt;
use vortex_array::arrays::listview::ListViewArrayExt;
use vortex_array::arrays::primitive::PrimitiveArrayExt;
use vortex_array::arrays::struct_::StructArrayExt;
use vortex_array::arrays::varbinview::build_views::BinaryView;
use vortex_array::buffer::BufferHandle;
Expand All @@ -41,6 +42,7 @@ use vortex_array::dtype::StructFields;
use vortex_array::match_each_decimal_value_type;
use vortex_array::match_each_integer_ptype;
use vortex_array::match_each_native_ptype;
use vortex_array::match_each_unsigned_integer_ptype;
use vortex_array::match_smallest_offset_type;
use vortex_array::patches::Patches;
use vortex_array::scalar::DecimalScalar;
Expand Down Expand Up @@ -163,10 +165,6 @@ pub(super) fn execute_sparse(parts: SparseParts, ctx: &mut ExecutionCtx) -> Vort
})
}

#[expect(
clippy::cognitive_complexity,
reason = "complexity is from nested match_smallest_offset_type macro"
)]
fn execute_sparse_lists(
resolved: &Patches,
fill_value: &Scalar,
Expand All @@ -175,14 +173,17 @@ fn execute_sparse_lists(
nullability: Nullability,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
// Patch indices are non-negative; reinterpret to unsigned so this dispatches over 4 widths
// instead of 8. `O` is already unsigned (from `match_smallest_offset_type`).
let indices = resolved.indices().as_::<Primitive>().into_owned();
let indices = indices.reinterpret_cast(indices.ptype().to_unsigned());
let values = resolved.values().as_::<ListView>().into_owned();
let fill_list = fill_value.as_list();

let n_filled = len - resolved.num_patches();
let total_canonical_values = values.elements().len() + fill_list.len() * n_filled;

Ok(match_each_integer_ptype!(indices.ptype(), |I| {
Ok(match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
match_smallest_offset_type!(total_canonical_values, |O| {
execute_sparse_lists_inner::<I, O>(
indices.as_slice(),
Expand Down
11 changes: 8 additions & 3 deletions vortex-array/src/arrays/decimal/array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ use crate::array::validity_to_child;
use crate::arrays::Decimal;
use crate::arrays::DecimalArray;
use crate::arrays::PrimitiveArray;
use crate::arrays::primitive::PrimitiveArrayExt;
use crate::buffer::BufferHandle;
use crate::dtype::BigCast;
use crate::dtype::DType;
Expand All @@ -37,7 +38,7 @@ use crate::dtype::IntegerPType;
use crate::dtype::NativeDecimalType;
use crate::dtype::Nullability;
use crate::match_each_decimal_value_type;
use crate::match_each_integer_ptype;
use crate::match_each_unsigned_integer_ptype;
use crate::patches::Patches;
use crate::validity::Validity;

Expand Down Expand Up @@ -556,8 +557,12 @@ impl Array<Decimal> {
assert_eq!(self.decimal_dtype(), patch_values.decimal_dtype());

let data = self.into_data();
let data = match_each_integer_ptype!(patch_indices.ptype(), |I| {
let patch_indices = patch_indices.as_slice::<I>();
// Patch indices are non-negative; reinterpret to unsigned so this dispatches over 4 widths
// instead of 8 (the decimal value-type dimensions are unaffected).
let patch_indices_unsigned =
patch_indices.reinterpret_cast(patch_indices.ptype().to_unsigned());
let data = match_each_unsigned_integer_ptype!(patch_indices_unsigned.ptype(), |I| {
let patch_indices = patch_indices_unsigned.as_slice::<I>();
match_each_decimal_value_type!(patch_values.values_type(), |PatchDVT| {
let patch_values = patch_values.buffer::<PatchDVT>();
match_each_decimal_value_type!(data.values_type(), |ValuesDVT| {
Expand Down
10 changes: 7 additions & 3 deletions vortex-array/src/arrays/fixed_size_list/compute/take.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,10 @@ use crate::arrays::PrimitiveArray;
use crate::arrays::bool::BoolArrayExt;
use crate::arrays::dict::TakeExecute;
use crate::arrays::fixed_size_list::FixedSizeListArrayExt;
use crate::arrays::primitive::PrimitiveArrayExt;
use crate::dtype::IntegerPType;
use crate::executor::ExecutionCtx;
use crate::match_each_integer_ptype;
use crate::match_each_unsigned_integer_ptype;
use crate::match_smallest_offset_type;
use crate::validity::Validity;

Expand All @@ -31,14 +32,15 @@ use crate::validity::Validity;
/// that elements start at offset 0 and be perfectly packed without gaps. We expand list indices
/// into element indices and push them down to the child elements array.
impl TakeExecute for FixedSizeList {
#[expect(clippy::cognitive_complexity)]
fn take(
array: ArrayView<'_, FixedSizeList>,
indices: &ArrayRef,
ctx: &mut ExecutionCtx,
) -> VortexResult<Option<ArrayRef>> {
let max_element_idx = array.elements().len();
match_each_integer_ptype!(indices.dtype().as_ptype(), |I| {
// Indices are non-negative; dispatch over the 4 unsigned widths (the executed array is
// reinterpreted to unsigned in `take_with_indices`). `E` is already unsigned.
match_each_unsigned_integer_ptype!(indices.dtype().as_ptype().to_unsigned(), |I| {
match_smallest_offset_type!(max_element_idx, |E| {
take_with_indices::<I, E>(array, indices, ctx)
})
Expand All @@ -56,6 +58,8 @@ fn take_with_indices<I: IntegerPType, E: IntegerPType>(
let list_size = array.list_size() as usize;

let indices_array = indices.clone().execute::<PrimitiveArray>(ctx)?;
// Reinterpret to unsigned so `as_slice::<I>` (with unsigned `I`) matches; values are unchanged.
let indices_array = indices_array.reinterpret_cast(indices_array.ptype().to_unsigned());

// Make sure to handle degenerate case where lists have size 0 (these can take fast paths).
if list_size == 0 {
Expand Down
29 changes: 19 additions & 10 deletions vortex-array/src/arrays/list/compute/take.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ use crate::builders::PrimitiveBuilder;
use crate::dtype::IntegerPType;
use crate::dtype::Nullability;
use crate::executor::ExecutionCtx;
use crate::match_each_integer_ptype;
use crate::match_each_unsigned_integer_ptype;
use crate::match_smallest_offset_type;

// TODO(connor)[ListView]: Re-revert to the version where we simply convert to a `ListView` and call
Expand All @@ -38,16 +38,22 @@ impl TakeExecute for List {
ctx: &mut ExecutionCtx,
) -> VortexResult<Option<ArrayRef>> {
let indices = indices.clone().execute::<PrimitiveArray>(ctx)?;
let indices = indices.reinterpret_cast(indices.ptype().to_unsigned());
let offsets = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
let offsets = offsets.reinterpret_cast(offsets.ptype().to_unsigned());
// This is an over-approximation of the total number of elements in the resulting array.
let total_approx = array.elements().len().saturating_mul(indices.len());

match_each_integer_ptype!(array.offsets().dtype().as_ptype(), |O| {
match_each_integer_ptype!(indices.ptype(), |I| {
match_each_unsigned_integer_ptype!(offsets.ptype(), |O| {
match_each_unsigned_integer_ptype!(indices.ptype(), |I| {
match_smallest_offset_type!(total_approx, |OutputOffsetType| {
{
let indices = indices.as_view();
_take::<I, O, OutputOffsetType>(array, indices, ctx).map(Some)
}
_take::<I, O, OutputOffsetType>(
array,
offsets.as_view(),
indices.as_view(),
ctx,
)
.map(Some)
})
})
})
Expand All @@ -56,6 +62,7 @@ impl TakeExecute for List {

fn _take<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
array: ArrayView<'_, List>,
offsets_array: ArrayView<'_, Primitive>,
indices_array: ArrayView<'_, Primitive>,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
Expand All @@ -68,10 +75,9 @@ fn _take<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
.execute_mask(indices_array.as_ref().len(), ctx)?;

if !indices_validity.all_true() || !data_validity.all_true() {
return _take_nullable::<I, O, OutputOffsetType>(array, indices_array, ctx);
return _take_nullable::<I, O, OutputOffsetType>(array, offsets_array, indices_array, ctx);
}

let offsets_array = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
let offsets: &[O] = offsets_array.as_slice();
let indices: &[I] = indices_array.as_slice();

Expand Down Expand Up @@ -121,12 +127,15 @@ fn _take<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
.into_array())
}

// Kept out-of-line: as a single-callsite generic helper it would otherwise be inlined into every
// monomorphization of `_take`, duplicating the entire nullable path across all specializations.
#[inline(never)]
fn _take_nullable<I: IntegerPType, O: IntegerPType, OutputOffsetType: IntegerPType>(
array: ArrayView<'_, List>,
offsets_array: ArrayView<'_, Primitive>,
indices_array: ArrayView<'_, Primitive>,
ctx: &mut ExecutionCtx,
) -> VortexResult<ArrayRef> {
let offsets_array = array.offsets().clone().execute::<PrimitiveArray>(ctx)?;
let offsets: &[O] = offsets_array.as_slice();
let indices: &[I] = indices_array.as_slice();
let data_validity = array
Expand Down
44 changes: 29 additions & 15 deletions vortex-array/src/arrays/listview/array.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,10 +31,12 @@ use crate::arrays::ListView;
use crate::arrays::Primitive;
use crate::arrays::PrimitiveArray;
use crate::arrays::bool;
use crate::arrays::primitive::PrimitiveArrayExt;
use crate::dtype::DType;
use crate::dtype::IntegerPType;
use crate::expr::stats::Stat;
use crate::match_each_integer_ptype;
use crate::match_each_unsigned_integer_ptype;
use crate::validity::Validity;

/// The `elements` data array, where each list scalar is a _slice_ of the `elements` array, and
Expand Down Expand Up @@ -241,10 +243,6 @@ impl ListViewData {
sizes.len()
);

// Check that the size type can fit within the offset type to prevent overflows.
let size_ptype = sizes.dtype().as_ptype();
let offset_ptype = offsets.dtype().as_ptype();

// If a validity array is present, it must be the same length as the `ListViewArray`.
if let Some(validity_len) = validity.maybe_len() {
vortex_ensure!(
Expand All @@ -260,10 +258,16 @@ impl ListViewData {
let offsets_primitive = offsets.to_primitive();
#[expect(deprecated)]
let sizes_primitive = sizes.to_primitive();
// Offsets and sizes are non-negative; reinterpret to unsigned to dispatch over 4 widths
// each (4x4 instead of 8x8). This is a read-only validation, so result types are moot.
let offsets_primitive =
offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
let sizes_primitive =
sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());

// Validate the `offsets` and `sizes` arrays.
match_each_integer_ptype!(offset_ptype, |O| {
match_each_integer_ptype!(size_ptype, |S| {
match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
let offsets_slice = offsets_primitive.as_slice::<O>();
let sizes_slice = sizes_primitive.as_slice::<S>();

Expand Down Expand Up @@ -445,8 +449,13 @@ pub trait ListViewArrayExt: TypedArrayRef<ListView> {

let mut buf = BitBufferMut::new_unset(len);

match_each_integer_ptype!(offsets_primitive.ptype(), |O| {
match_each_integer_ptype!(sizes_primitive.ptype(), |S| {
// Offsets/sizes are non-negative; reinterpret to unsigned (4x4 instead of 8x8).
let offsets_primitive =
offsets_primitive.reinterpret_cast(offsets_primitive.ptype().to_unsigned());
let sizes_primitive =
sizes_primitive.reinterpret_cast(sizes_primitive.ptype().to_unsigned());
match_each_unsigned_integer_ptype!(offsets_primitive.ptype(), |O| {
match_each_unsigned_integer_ptype!(sizes_primitive.ptype(), |S| {
fill_referenced_mask::<O, S>(
&mut buf,
offsets_primitive.as_slice::<O>(),
Expand Down Expand Up @@ -721,11 +730,16 @@ fn validate_zctl(
let sizes_dtype = sizes_primitive.dtype();
let len = offsets_primitive.len();

// Offsets/sizes are non-negative; reinterpret to unsigned (4x4 instead of 8x8).
let offsets_unsigned =
offsets_primitive.reinterpret_cast(offsets_dtype.as_ptype().to_unsigned());
let sizes_unsigned = sizes_primitive.reinterpret_cast(sizes_dtype.as_ptype().to_unsigned());

// Check that offset + size values are monotonic (no overlaps)
match_each_integer_ptype!(offsets_dtype.as_ptype(), |O| {
match_each_integer_ptype!(sizes_dtype.as_ptype(), |S| {
let offsets_slice = offsets_primitive.as_slice::<O>();
let sizes_slice = sizes_primitive.as_slice::<S>();
match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
let offsets_slice = offsets_unsigned.as_slice::<O>();
let sizes_slice = sizes_unsigned.as_slice::<S>();

validate_monotonic_ends(offsets_slice, sizes_slice, len)?;
})
Expand Down Expand Up @@ -756,9 +770,9 @@ fn validate_zctl(
}
}

match_each_integer_ptype!(offsets_primitive.ptype(), |O| {
match_each_integer_ptype!(sizes_primitive.ptype(), |S| {
count_references::<O, S>(&mut element_references, offsets_primitive, sizes_primitive);
match_each_unsigned_integer_ptype!(offsets_unsigned.ptype(), |O| {
match_each_unsigned_integer_ptype!(sizes_unsigned.ptype(), |S| {
count_references::<O, S>(&mut element_references, offsets_unsigned, sizes_unsigned);
})
});

Expand Down
Loading
Loading