Skip to content

Commit 0d703d8

Browse files
a10ya10y
committed
save
Signed-off-by: Andrew Duffy <andrew@a10y.dev>
1 parent e325039 commit 0d703d8

File tree

9 files changed

+152
-126
lines changed

9 files changed

+152
-126
lines changed

‎encodings/fastlanes/src/bitpacking/vtable/mod.rs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -297,8 +297,8 @@ impl VTable for BitPacked {
297297
})
298298
}
299299

300-
fn execute(array: &Self::Array, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
301-
Ok(ExecutionResult::done(unpack_array(array)?.into_array()))
300+
fn execute(array: Arc<Self::Array>, _ctx: &mut ExecutionCtx) -> VortexResult<ExecutionResult> {
301+
Ok(ExecutionResult::done(unpack_array(array.as_ref())?.into_array()))
302302
}
303303

304304
fn execute_parent(

‎encodings/fastlanes/src/delta/array/delta_compress.rs‎

Lines changed: 15 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -7,18 +7,18 @@ use std::mem::MaybeUninit;
77
use fastlanes::Delta;
88
use fastlanes::FastLanes;
99
use fastlanes::Transpose;
10-
use vortex_array::ExecutionCtx;
1110
use vortex_array::arrays::PrimitiveArray;
1211
use vortex_array::dtype::NativePType;
1312
use vortex_array::match_each_unsigned_integer_ptype;
1413
use vortex_array::vtable::ValidityHelper;
14+
use vortex_array::ExecutionCtx;
1515
use vortex_buffer::Buffer;
1616
use vortex_buffer::BufferMut;
1717
use vortex_error::VortexResult;
1818

19-
use crate::FL_CHUNK_SIZE;
2019
use crate::bit_transpose::transpose_validity;
2120
use crate::fill_forward_nulls;
21+
use crate::FL_CHUNK_SIZE;
2222

2323
pub fn delta_compress(
2424
array: &PrimitiveArray,
@@ -96,19 +96,19 @@ mod tests {
9696
use std::sync::LazyLock;
9797

9898
use rstest::rstest;
99-
use vortex_array::IntoArray;
100-
use vortex_array::ToCanonical;
101-
use vortex_array::VortexSessionExecute;
10299
use vortex_array::arrays::PrimitiveArray;
103100
use vortex_array::assert_arrays_eq;
104101
use vortex_array::session::ArraySession;
102+
use vortex_array::IntoArray;
103+
use vortex_array::ToCanonical;
104+
use vortex_array::VortexSessionExecute;
105105
use vortex_error::VortexResult;
106106
use vortex_session::VortexSession;
107107

108-
use crate::DeltaArray;
109-
use crate::bitpack_compress::bitpack_encode;
108+
use crate::bitpack_compress::BitPackedEncoder;
110109
use crate::delta::array::delta_decompress::delta_decompress;
111110
use crate::delta_compress;
111+
use crate::DeltaArray;
112112

113113
static SESSION: LazyLock<VortexSession> =
114114
LazyLock::new(|| VortexSession::empty().with::<ArraySession>());
@@ -137,14 +137,14 @@ mod tests {
137137
(0u8..200).map(|i| (!(50..100).contains(&i)).then_some(i)),
138138
);
139139
let (bases, deltas) = delta_compress(&array, &mut SESSION.create_execution_ctx()).unwrap();
140-
let bitpacked_deltas = bitpack_encode(&deltas, 1, None).unwrap();
141-
let packed_delta = DeltaArray::try_new(
142-
bases.into_array(),
143-
bitpacked_deltas.into_array(),
144-
0,
145-
array.len(),
146-
)
147-
.unwrap();
140+
let bitpacked_deltas = BitPackedEncoder::new(&deltas)
141+
.with_bit_width(1)
142+
.pack()
143+
.unwrap()
144+
.into_array()
145+
.unwrap();
146+
let packed_delta =
147+
DeltaArray::try_new(bases.into_array(), bitpacked_deltas, 0, array.len()).unwrap();
148148
assert_arrays_eq!(packed_delta.to_primitive(), array);
149149
}
150150
}

‎vortex-array/src/arrays/patched/array.rs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,19 @@ use crate::patches::Patches;
2626
use crate::stats::ArrayStats;
2727
use crate::validity::Validity;
2828

29+
/// Shredded components of the [`PatchedArray`].
30+
///
31+
/// This is created when you consume the arrary using [`PatchedArray::into_parts`].
32+
pub struct PatchedArrayParts {
33+
pub inner: ArrayRef,
34+
pub n_chunks: usize,
35+
pub n_lanes: usize,
36+
pub offset: usize,
37+
pub lane_offsets: BufferHandle,
38+
pub indices: BufferHandle,
39+
pub values: ArrayRef,
40+
}
41+
2942
/// An array that partially "patches" another array with new values.
3043
#[derive(Debug, Clone)]
3144
pub struct PatchedArray {
@@ -127,6 +140,7 @@ impl PatchedArray {
127140

128141
start..stop
129142
}
143+
}
130144

131145
/// Slice the array to just the patches and inner values that are within the chunk range.
132146
pub(crate) fn slice_chunks(&self, chunks: Range<usize>) -> VortexResult<Self> {

‎vortex-cuda/benches/dynamic_dispatch_cuda.rs‎

Lines changed: 0 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,19 +17,10 @@ use cudarc::driver::LaunchConfig;
1717
use cudarc::driver::PushKernelArg;
1818
use cudarc::driver::sys::CUevent_flags;
1919
use vortex::array::IntoArray;
20-
use vortex::array::ToCanonical;
21-
use vortex::array::arrays::DictArray;
2220
use vortex::array::arrays::PrimitiveArray;
23-
use vortex::array::scalar::Scalar;
2421
use vortex::array::validity::Validity::NonNullable;
2522
use vortex::buffer::Buffer;
2623
use vortex::dtype::PType;
27-
use vortex::encodings::alp::ALPArray;
28-
use vortex::encodings::alp::ALPFloat;
29-
use vortex::encodings::alp::Exponents;
30-
use vortex::encodings::alp::alp_encode;
31-
use vortex::encodings::fastlanes::BitPackedArray;
32-
use vortex::encodings::fastlanes::FoRArray;
3324
use vortex::encodings::runend::RunEndArray;
3425
use vortex::error::VortexExpect;
3526
use vortex::error::VortexResult;

‎vortex-cuda/src/dynamic_dispatch/mod.rs‎

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -277,6 +277,16 @@ impl MaterializedPlan {
277277
mod tests {
278278
use std::sync::Arc;
279279

280+
use super::DynamicDispatchPlan;
281+
use super::SMEM_TILE_SIZE;
282+
use super::ScalarOp;
283+
use super::SourceOp;
284+
use super::Stage;
285+
use super::UnmaterializedPlan;
286+
use crate::CudaBufferExt;
287+
use crate::CudaDeviceBuffer;
288+
use crate::CudaExecutionCtx;
289+
use crate::session::CudaSession;
280290
use cudarc::driver::DevicePtr;
281291
use cudarc::driver::LaunchConfig;
282292
use cudarc::driver::PushKernelArg;
@@ -295,31 +305,26 @@ mod tests {
295305
use vortex::encodings::alp::alp_encode;
296306
use vortex::encodings::fastlanes::BitPackedArray;
297307
use vortex::encodings::fastlanes::FoRArray;
308+
use vortex::encodings::fastlanes::bitpack_compress::BitPackedEncoder;
298309
use vortex::encodings::runend::RunEndArray;
299310
use vortex::encodings::zigzag::ZigZagArray;
300311
use vortex::error::VortexExpect;
301312
use vortex::error::VortexResult;
302313
use vortex::session::VortexSession;
303-
304-
use super::CudaDispatchPlan;
305-
use super::SMEM_TILE_SIZE;
306-
use super::ScalarOp;
307-
use super::SourceOp;
308-
use super::Stage;
309-
use super::UnmaterializedPlan;
310-
use crate::CudaBufferExt;
311-
use crate::CudaDeviceBuffer;
312-
use crate::CudaExecutionCtx;
313-
use crate::session::CudaSession;
314+
use vortex_array::ArrayRef;
314315

315316
fn make_bitpacked_array_u32(bit_width: u8, len: usize) -> BitPackedArray {
316317
let max_val = (1u64 << bit_width).saturating_sub(1);
317318
let values: Vec<u32> = (0..len)
318319
.map(|i| ((i as u64) % (max_val + 1)) as u32)
319320
.collect();
320321
let primitive = PrimitiveArray::new(Buffer::from(values), NonNullable);
321-
BitPackedArray::encode(&primitive.into_array(), bit_width)
322-
.vortex_expect("failed to create BitPacked array")
322+
BitPackedEncoder::new(&primitive)
323+
.with_bit_width(bit_width)
324+
.pack()
325+
.unwrap()
326+
.into_array()
327+
.unwrap()
323328
}
324329

325330
#[crate::test]
@@ -546,7 +551,7 @@ mod tests {
546551
let expected: Vec<u32> = raw.iter().map(|&v| v + reference).collect();
547552

548553
let bp = make_bitpacked_array_u32(bit_width, len);
549-
let for_arr = FoRArray::try_new(bp.into_array(), Scalar::from(reference))?;
554+
let for_arr = FoRArray::try_new(bp, Scalar::from(reference))?;
550555

551556
let cuda_ctx = CudaSession::create_execution_ctx(&VortexSession::empty())?;
552557
let plan = UnmaterializedPlan::new(&for_arr.into_array())?.materialize(&cuda_ctx)?;

‎vortex-cuda/src/dynamic_dispatch/plan_builder.rs‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -105,10 +105,7 @@ fn is_dyn_dispatch_compatible(array: &ArrayRef) -> bool {
105105
return false;
106106
}
107107
if id == BitPacked::ID {
108-
if let Ok(a) = array.clone().try_into::<BitPacked>() {
109-
return a.patches().is_none();
110-
}
111-
return false;
108+
return true;
112109
}
113110
if id == Dict::ID {
114111
if let Ok(a) = array.clone().try_into::<Dict>() {

‎vortex-cuda/src/hybrid_dispatch/mod.rs‎

Lines changed: 78 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -152,6 +152,7 @@ mod tests {
152152
use vortex::buffer::Buffer;
153153
use vortex::encodings::fastlanes::BitPackedArray;
154154
use vortex::encodings::fastlanes::FoRArray;
155+
use vortex::encodings::fastlanes::bitpack_compress::BitPackedEncoder;
155156
use vortex::error::VortexExpect;
156157
use vortex::error::VortexResult;
157158
use vortex::mask::Mask;
@@ -167,12 +168,11 @@ mod tests {
167168
let mut ctx =
168169
CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx");
169170
let values: Vec<u32> = (0..2048).map(|i| (i % 128) as u32).collect();
170-
let bp = BitPackedArray::encode(
171-
&PrimitiveArray::new(Buffer::from(values), NonNullable).into_array(),
172-
7,
173-
)
174-
.vortex_expect("bp");
175-
let arr = FoRArray::try_new(bp.into_array(), 1000u32.into()).vortex_expect("for");
171+
let bp = BitPackedEncoder::new(&PrimitiveArray::from_iter(values))
172+
.with_bit_width(7)
173+
.pack()?
174+
.into_array()?;
175+
let arr = FoRArray::try_new(bp, 1000u32.into()).vortex_expect("for");
176176

177177
let cpu = arr.to_canonical()?.into_array();
178178
let gpu = arr
@@ -196,13 +196,12 @@ mod tests {
196196
let mut ctx =
197197
CudaSession::create_execution_ctx(&VortexSession::empty()).vortex_expect("ctx");
198198
let encoded: Vec<i32> = (0i32..2048).map(|i| i % 500).collect();
199-
let bp = BitPackedArray::encode(
200-
&PrimitiveArray::new(Buffer::from(encoded), NonNullable).into_array(),
201-
9,
202-
)
203-
.vortex_expect("bp");
199+
let bp = BitPackedEncoder::new(&PrimitiveArray::from_iter(encoded))
200+
.with_bit_width(9)
201+
.pack()?
202+
.into_array()?;
204203
let alp = ALPArray::try_new(
205-
FoRArray::try_new(bp.into_array(), 0i32.into())
204+
FoRArray::try_new(bp, 0i32.into())
206205
.vortex_expect("for")
207206
.into_array(),
208207
Exponents { e: 0, f: 2 },
@@ -259,72 +258,73 @@ mod tests {
259258
Ok(())
260259
}
261260

262-
/// Dict(values=ZstdBuffers(FoR(BP)), codes=FoR(BP)) — ZstdBuffers is
263-
/// executed separately, then Dict+FoR+BP fuses with its output as a LOAD.
264-
/// 3 launches: nvcomp + fused FoR+BP + fused LOAD+FoR+BP+DICT.
265-
#[cfg(feature = "unstable_encodings")]
266-
#[crate::test]
267-
async fn test_partial_fusion() -> VortexResult<()> {
268-
use vortex::array::arrays::DictArray;
269-
use vortex::array::session::ArraySessionExt;
270-
use vortex::encodings::fastlanes;
271-
use vortex::encodings::zstd::ZstdBuffers;
272-
use vortex::encodings::zstd::ZstdBuffersArray;
273-
274-
let mut session = VortexSession::empty();
275-
fastlanes::initialize(&mut session);
276-
session.arrays().register(ZstdBuffers);
277-
let mut ctx = CudaSession::create_execution_ctx(&session).vortex_expect("ctx");
278-
279-
let num_values: u32 = 64;
280-
let len: u32 = 2048;
281-
282-
// values = ZstdBuffers(FoR(BitPacked))
283-
let vals = PrimitiveArray::new(
284-
Buffer::from((0..num_values).collect::<Vec<_>>()),
285-
NonNullable,
286-
)
287-
.into_array();
288-
let vals = FoRArray::try_new(
289-
BitPackedArray::encode(&vals, 6)
290-
.vortex_expect("bp")
291-
.into_array(),
292-
0u32.into(),
293-
)
294-
.vortex_expect("for");
295-
let vals = ZstdBuffersArray::compress(&vals.into_array(), 3).vortex_expect("zstd");
296-
297-
// codes = FoR(BitPacked)
298-
let codes = PrimitiveArray::new(
299-
Buffer::from((0..len).map(|i| i % num_values).collect::<Vec<_>>()),
300-
NonNullable,
301-
)
302-
.into_array();
303-
let codes = FoRArray::try_new(
304-
BitPackedArray::encode(&codes, 6)
305-
.vortex_expect("bp")
306-
.into_array(),
307-
0u32.into(),
308-
)
309-
.vortex_expect("for");
310-
311-
let dict = DictArray::try_new(codes.into_array(), vals.into_array()).vortex_expect("dict");
312-
313-
let cpu = PrimitiveArray::new(
314-
Buffer::from((0..len).map(|i| i % num_values).collect::<Vec<_>>()),
315-
NonNullable,
316-
)
317-
.into_array();
318-
let gpu = dict
319-
.into_array()
320-
.execute_cuda(&mut ctx)
321-
.await?
322-
.into_host()
323-
.await?
324-
.into_array();
325-
assert_arrays_eq!(cpu, gpu);
326-
Ok(())
327-
}
261+
// TODO(aduffy): bring this back
262+
// /// Dict(values=ZstdBuffers(FoR(BP)), codes=FoR(BP)) — ZstdBuffers is
263+
// /// executed separately, then Dict+FoR+BP fuses with its output as a LOAD.
264+
// /// 3 launches: nvcomp + fused FoR+BP + fused LOAD+FoR+BP+DICT.
265+
// #[cfg(feature = "unstable_encodings")]
266+
// #[crate::test]
267+
// async fn test_partial_fusion() -> VortexResult<()> {
268+
// use vortex::array::arrays::DictArray;
269+
// use vortex::array::session::ArraySessionExt;
270+
// use vortex::encodings::fastlanes;
271+
// use vortex::encodings::zstd::ZstdBuffers;
272+
// use vortex::encodings::zstd::ZstdBuffersArray;
273+
//
274+
// let mut session = VortexSession::empty();
275+
// fastlanes::initialize(&mut session);
276+
// session.arrays().register(ZstdBuffers);
277+
// let mut ctx = CudaSession::create_execution_ctx(&session).vortex_expect("ctx");
278+
//
279+
// let num_values: u32 = 64;
280+
// let len: u32 = 2048;
281+
//
282+
// // values = ZstdBuffers(FoR(BitPacked))
283+
// let vals = PrimitiveArray::new(
284+
// Buffer::from((0..num_values).collect::<Vec<_>>()),
285+
// NonNullable,
286+
// )
287+
// .into_array();
288+
// let vals = FoRArray::try_new(
289+
// BitPackedArray::encode(&vals, 6)
290+
// .vortex_expect("bp")
291+
// .into_array(),
292+
// 0u32.into(),
293+
// )
294+
// .vortex_expect("for");
295+
// let vals = ZstdBuffersArray::compress(&vals.into_array(), 3).vortex_expect("zstd");
296+
//
297+
// // codes = FoR(BitPacked)
298+
// let codes = PrimitiveArray::new(
299+
// Buffer::from((0..len).map(|i| i % num_values).collect::<Vec<_>>()),
300+
// NonNullable,
301+
// )
302+
// .into_array();
303+
// let codes = FoRArray::try_new(
304+
// BitPackedArray::encode(&codes, 6)
305+
// .vortex_expect("bp")
306+
// .into_array(),
307+
// 0u32.into(),
308+
// )
309+
// .vortex_expect("for");
310+
//
311+
// let dict = DictArray::try_new(codes.into_array(), vals.into_array()).vortex_expect("dict");
312+
//
313+
// let cpu = PrimitiveArray::new(
314+
// Buffer::from((0..len).map(|i| i % num_values).collect::<Vec<_>>()),
315+
// NonNullable,
316+
// )
317+
// .into_array();
318+
// let gpu = dict
319+
// .into_array()
320+
// .execute_cuda(&mut ctx)
321+
// .await?
322+
// .into_host()
323+
// .await?
324+
// .into_array();
325+
// assert_arrays_eq!(cpu, gpu);
326+
// Ok(())
327+
// }
328328

329329
/// Filter(FoR(BP), mask) — FoR+BP fuses via dyn dispatch, then CUB filters the result.
330330
#[crate::test]

0 commit comments

Comments
 (0)