Add CUDA decompression kernels for OnPair - #8920
Conversation
Merging this PR will not alter performance
|
|
I made it use the same machinery like #8921 to avoid triple kernel dispatch |
0ax1
left a comment
0ax1
left a comment
There was a problem hiding this comment.
lgtm. but worth double checking all index types
|
I am reworking this to make sure all the work happens in one place - there's too many conflation. Will make sure that all non trivial compute is on gpu |
|
@joseph-isaacs have another look |
There was a problem hiding this comment.
can we sort these function by usage a bit
Signed-off-by: Robert Kruszewski <github@robertk.io>
| let Some(decoded) = decode_onpair_bytes(onpair.as_view(), ctx).await? else { | ||
| // Zero decoded bytes: all-zero offsets and an empty values heap. | ||
| let offsets = ctx.copy_to_device(vec![0i32; len + 1])?.await?; | ||
| let allocation = CudaDeviceBuffer::new(ctx.device_alloc::<u8>(1)?); | ||
| let values = BufferHandle::new_device(allocation.slice(0..0)); | ||
| return Ok(DecodedVarBin { | ||
| dtype, | ||
| len, | ||
| offsets, | ||
| values, | ||
| validity, | ||
| }); | ||
| }; | ||
|
|
||
| let OnPairDecoded { | ||
| bytes, | ||
| total_size, | ||
| lengths, | ||
| } = decoded; | ||
|
|
||
| // Build the Arrow i32 offsets from the lengths on device; this also | ||
| // rejects windows beyond Arrow's i32 offset range. | ||
| let I32Offsets { | ||
| buffer: offsets, | ||
| total, | ||
| } = i32_offsets_from_lengths(lengths, ctx).await?; | ||
| ensure_lengths_match(u64::try_from(total)?, total_size)?; |
There was a problem hiding this comment.
can these can not be in concurrent?
There was a problem hiding this comment.
I think you're right here that we could move the device copies and executions to be concurrent but I think we need to let functions call back into the execute method with a different stream?
There was a problem hiding this comment.
I guess per Alex's response we could create new cuda execution context?
09207a9 to
1732f2a
Compare
| let lengths = onpair | ||
| .uncompressed_lengths() | ||
| .clone() | ||
| .execute_cuda(ctx) | ||
| .await? | ||
| .into_primitive(); | ||
|
|
||
| // No codes at all (e.g. every row empty): the child's length is host | ||
| // metadata, so this early-out costs no device read. | ||
| if onpair.codes().is_empty() { | ||
| ensure_zero_lengths(lengths).await?; | ||
| return Ok(None); | ||
| } | ||
|
|
||
| // Decompress the per-row code boundaries on device; the token window is | ||
| // resolved from them by a kernel, never by host scalar reads. | ||
| let codes_offsets = onpair | ||
| .codes_offsets() | ||
| .clone() | ||
| .execute_cuda(ctx) | ||
| .await? | ||
| .into_primitive(); | ||
|
|
||
| // Decompress the codes child on device. The kernels are instantiated for | ||
| // the two widths OnPair stores — u16 natively, u8 when the compressor | ||
| // narrowed the codes — so no widening pass is needed. | ||
| let codes = onpair | ||
| .codes() | ||
| .clone() | ||
| .execute_cuda(ctx) | ||
| .await? | ||
| .into_primitive(); | ||
| match codes.ptype() { |
There was a problem hiding this comment.
do these concurrently?
There was a problem hiding this comment.
As discussed offline, the kernels will be submitted in sequenced and return once the kernel is submitted to the stream (not launched). We do want overlap / concurrency if the kernel does not occupy all SMs. In that case, we would need to submit the work through diff streams. This is not supported directly as of now by the API through execute. The current approach is that we round robin CUDA streams per CudaExecutionCtx.
As the context holds the session as a field
pub struct CudaExecutionCtx {
stream: VortexCudaStream,
ctx: ExecutionCtx,
cuda_session: CudaSession,
strategy: Arc<dyn LaunchStrategy>,
dispatch_mode: CudaDispatchMode,
}
we could do an execute version that hooks into the stream round robin mechanism.
|
I have added a feature to let you spawn work on another stream here #9234. We can then rebase this |
3 participants
They follow the same pattern as FSST where we let users choose varbin/varbinview