Skip to content

Commit ced8ec9

Browse files
TanishGudiseclaude0ax1
authored
fix(cuda): copy encoded device-resident children to the host in into_host (#10916)
## Summary - Tracking Issue: #10914 `compress-bench --gpu-decompress --gpu-verify` panics with `expected host buffer` on develop for CMSprovider and Euro2016. A run-end encoded validity mask stays encoded on the device after `execute_cuda`. `CanonicalCudaExt::into_host` then executes it on the host while its buffers are still on the device (`validity_into_host` → `child_into_host` → `RunEnd::execute` → `runend_decode_bools` → `BufferHandle::as_host`). A nullable dictionary's validity, a lazy `take` through device-resident codes, fails the same way. ## Changes - `canonical.rs`: `child_into_host` copies every device buffer of a child that isn't canonical to the host before executing it there. - Children that are already canonical take the existing path. - The copies run concurrently, and an array's buffers are replaced only when one of them moves (Sparse rejects buffer replacement). - The decode path is unchanged. - This doesn't add mixed CPU/GPU execution (#7838): the child moves to the host whole, so nothing decodes from device buffers on the host. - Tests (4), each building the child on the device and comparing `into_host` with the expected result: - a primitive array whose validity is a run-end encoded mask; - one whose validity is a sparse mask (Sparse keeps its fill value in a host buffer it doesn't allow to be replaced, so only its children move); - a string dictionary with nullable values and device-resident codes; - a primitive dictionary with nullable values and device-resident codes. Evidence (A40, CUDA 13.0): | Check | develop | this PR | |---|---|---| | The 4 tests | all fail: `expected host buffer` (run-end mask), `as_slice must be called on host buffer` (sparse mask, both dictionaries) | 4 pass | | `--gpu-verify` on CMSprovider and Euro2016, today's preset | both panic | both pass (234 and 22 fields verified) | | `vortex-cuda` lib suite | — | 711 passed, 0 failed | Everything in the table ran on this PR alone. With all three PRs, `--gpu-verify` passes on all 17 GPU-suite datasets. ## AI assistance Developed with Claude Code (Claude Opus 5.5), which wrote most of the code, tests and this description under my direction. I reviewed the diff and the measurements. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Signed-off-by: Tanish Gudise <[email protected]> Co-authored-by: Claude Opus 5.5 <[email protected]> Co-authored-by: Alexander Droste <[email protected]>
1 parent c8ba781 commit ced8ec9

1 file changed

Lines changed: 185 additions & 1 deletion

File tree

‎vortex-cuda/src/canonical.rs‎

Lines changed: 185 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
use std::sync::Arc;
55

66
use async_trait::async_trait;
7+
use futures::future::BoxFuture;
78
use futures::future::try_join_all;
89
use vortex::array::ArrayRef;
910
use vortex::array::Canonical;
@@ -29,18 +30,64 @@ use vortex::array::validity::Validity;
2930
use vortex::buffer::BitBuffer;
3031
use vortex::buffer::Buffer;
3132
use vortex::buffer::ByteBuffer;
33+
use vortex::error::VortexError;
3234
use vortex::error::VortexResult;
3335
use vortex::error::vortex_bail;
3436

35-
/// Copy a canonical child array to the host.
37+
/// Copy a child array to the host, executing it there if CUDA execution left it encoded.
3638
async fn child_into_host(child: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<ArrayRef> {
39+
let child = if child.is_canonical() {
40+
child
41+
} else {
42+
buffers_into_host(child).await?
43+
};
3744
Ok(child
3845
.execute::<Canonical>(ctx)?
3946
.into_host(ctx)
4047
.await?
4148
.into_array())
4249
}
4350

51+
/// Copies every device buffer of `array` and of its children to the host, keeping the encoding.
52+
///
53+
/// A child that CUDA execution left encoded, such as a validity mask that is still a lazy `take`
54+
/// through device-resident dictionary codes or a run-end encoded mask, can only be executed on
55+
/// the host once all of its buffers are there.
56+
fn buffers_into_host(array: ArrayRef) -> BoxFuture<'static, VortexResult<ArrayRef>> {
57+
Box::pin(async move {
58+
if array.is_host() {
59+
return Ok(array);
60+
}
61+
let mut array = array;
62+
if array.children_iter().any(|child| !child.is_host()) {
63+
let mut slots = Vec::with_capacity(array.slots().len());
64+
for slot in array.slots().iter() {
65+
slots.push(match slot {
66+
Some(child) => Some(buffers_into_host(child.clone()).await?),
67+
None => None,
68+
});
69+
}
70+
// SAFETY: the slots are byte-for-byte copies; only their placement changes.
71+
array = unsafe { array.with_slots(slots.into())? };
72+
}
73+
let buffers = array.buffer_handles();
74+
// Some encodings, such as Sparse, reject buffer replacement, so only replace when a
75+
// buffer actually moves.
76+
if buffers.iter().all(BufferHandle::is_on_host) {
77+
return Ok(array);
78+
}
79+
let buffers = try_join_all(buffers.into_iter().map(|buffer| async move {
80+
if buffer.is_on_host() {
81+
return Ok::<_, VortexError>(buffer);
82+
}
83+
Ok(BufferHandle::new_host(buffer.try_into_host()?.await?))
84+
}))
85+
.await?;
86+
// SAFETY: the buffers are byte-for-byte copies; only their placement changes.
87+
unsafe { array.with_buffers(buffers) }
88+
})
89+
}
90+
4491
/// Move all canonical data from to_host from device.
4592
#[async_trait]
4693
pub trait CanonicalCudaExt {
@@ -196,11 +243,19 @@ mod tests {
196243
use std::sync::LazyLock;
197244

198245
use vortex::array::VortexSessionExecute;
246+
use vortex::array::arrays::DictArray;
247+
use vortex::array::assert_arrays_eq;
199248
use vortex::dtype::PType;
249+
use vortex::encodings::runend::RunEnd;
250+
use vortex::encodings::sparse::Sparse;
251+
use vortex::scalar::Scalar;
200252
use vortex::session::VortexSession;
201253

202254
use super::*;
203255
use crate::CudaSession;
256+
use crate::executor::CudaExecute;
257+
use crate::executor::CudaExecutionCtx;
258+
use crate::kernel::DictExecutor;
204259

205260
static SESSION: LazyLock<VortexSession> = LazyLock::new(VortexSession::empty);
206261

@@ -225,4 +280,133 @@ mod tests {
225280
assert!(Canonical::List(list).into_host(&mut ctx).await.is_err());
226281
Ok(())
227282
}
283+
284+
/// Moves every buffer of `array` and of its children to the device without decoding it.
285+
fn encoded_on_device(array: ArrayRef, cuda_ctx: &CudaExecutionCtx) -> VortexResult<ArrayRef> {
286+
let mut slots = Vec::new();
287+
for slot in array.slots().iter() {
288+
slots.push(match slot {
289+
Some(child) => Some(encoded_on_device(child.clone(), cuda_ctx)?),
290+
None => None,
291+
});
292+
}
293+
// SAFETY: the slots are byte-for-byte copies; only their placement changes.
294+
let array = unsafe { array.with_slots(slots.into())? };
295+
// Sparse keeps its fill value in a host buffer that it doesn't allow to be replaced.
296+
if array.as_opt::<Sparse>().is_some() {
297+
return Ok(array);
298+
}
299+
let mut buffers = Vec::new();
300+
for buffer in array.buffer_handles() {
301+
buffers.push(cuda_ctx.ensure_on_device_sync(buffer)?);
302+
}
303+
// SAFETY: the buffers are byte-for-byte copies; only their placement changes.
304+
unsafe { array.with_buffers(buffers) }
305+
}
306+
307+
/// A GPU decode can leave a validity mask encoded on the device, here run-end encoded.
308+
/// Copying the values to the host used to execute that mask on the host while its buffers
309+
/// were still on the device, which panicked with "expected host buffer".
310+
#[crate::test]
311+
async fn test_into_host_copies_encoded_device_validity() -> VortexResult<()> {
312+
let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?;
313+
let mut ctx = SESSION.create_execution_ctx();
314+
let mask = RunEnd::new(
315+
PrimitiveArray::new(Buffer::from(vec![2u32, 5, 6]), Validity::NonNullable).into_array(),
316+
BoolArray::from_iter([true, false, true]).into_array(),
317+
&mut ctx,
318+
)
319+
.into_array();
320+
let mask = encoded_on_device(mask, &cuda_ctx)?;
321+
let values = cuda_ctx
322+
.stream()
323+
.copy_to_device(vec![1i32, 2, 3, 4, 5, 6])?
324+
.await?;
325+
let array = PrimitiveArray::from_buffer_handle(values, PType::I32, Validity::Array(mask));
326+
327+
let host = Canonical::Primitive(array)
328+
.into_host(cuda_ctx.execution_ctx())
329+
.await?
330+
.into_array();
331+
let expected =
332+
PrimitiveArray::from_option_iter([Some(1i32), Some(2), None, None, None, Some(6)]);
333+
assert_arrays_eq!(host, expected, &mut ctx);
334+
Ok(())
335+
}
336+
337+
/// A sparse validity mask keeps its fill value in a host buffer, which Sparse doesn't allow to
338+
/// be replaced, while its patch indices and values are on the device. Only the children move
339+
/// to the host.
340+
#[crate::test]
341+
async fn test_into_host_copies_sparse_device_validity() -> VortexResult<()> {
342+
let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?;
343+
let mut ctx = SESSION.create_execution_ctx();
344+
let mask = Sparse::try_new(
345+
PrimitiveArray::new(Buffer::from(vec![1u32, 4]), Validity::NonNullable).into_array(),
346+
BoolArray::from_iter([false, false]).into_array(),
347+
6,
348+
Scalar::from(true),
349+
&mut ctx,
350+
)?
351+
.into_array();
352+
let mask = encoded_on_device(mask, &cuda_ctx)?;
353+
let values = cuda_ctx
354+
.stream()
355+
.copy_to_device(vec![1i32, 2, 3, 4, 5, 6])?
356+
.await?;
357+
let array = PrimitiveArray::from_buffer_handle(values, PType::I32, Validity::Array(mask));
358+
359+
let host = Canonical::Primitive(array)
360+
.into_host(cuda_ctx.execution_ctx())
361+
.await?
362+
.into_array();
363+
let expected =
364+
PrimitiveArray::from_option_iter([Some(1i32), None, Some(3), Some(4), None, Some(6)]);
365+
assert_arrays_eq!(host, expected, &mut ctx);
366+
Ok(())
367+
}
368+
369+
/// A dictionary with nullable values returns its validity as a lazy gather through the codes.
370+
/// With the codes on the device, as a file scan leaves them, copying the decoded dictionary to
371+
/// the host used to panic on a host-buffer read.
372+
async fn assert_device_dict_into_host(
373+
codes: PrimitiveArray,
374+
values: ArrayRef,
375+
) -> VortexResult<()> {
376+
let mut cuda_ctx = CudaSession::create_execution_ctx(&crate::cuda_session())?;
377+
let mut ctx = SESSION.create_execution_ctx();
378+
let expected = DictArray::try_new(codes.clone().into_array(), values.clone())?;
379+
let device_codes = encoded_on_device(codes.into_array(), &cuda_ctx)?;
380+
let dict = DictArray::try_new(device_codes, values)?;
381+
let host = DictExecutor
382+
.execute(dict.into_array(), &mut cuda_ctx)
383+
.await?
384+
.into_host(cuda_ctx.execution_ctx())
385+
.await?
386+
.into_array();
387+
assert_arrays_eq!(host, expected, &mut ctx);
388+
Ok(())
389+
}
390+
391+
#[crate::test]
392+
async fn test_into_host_string_dict_with_nullable_values() -> VortexResult<()> {
393+
let values = VarBinViewArray::from_iter_nullable_str([
394+
Some("a"),
395+
None,
396+
Some("a string longer than twelve bytes"),
397+
]);
398+
let codes = PrimitiveArray::new(
399+
Buffer::from(vec![0u8, 1, 2, 1, 0, 2]),
400+
Validity::NonNullable,
401+
);
402+
assert_device_dict_into_host(codes, values.into_array()).await
403+
}
404+
405+
#[crate::test]
406+
async fn test_into_host_primitive_dict_with_nullable_values() -> VortexResult<()> {
407+
let values = PrimitiveArray::from_option_iter([Some(7i32), None, Some(-3)]);
408+
let codes =
409+
PrimitiveArray::new(Buffer::from(vec![2u32, 0, 1, 1, 0]), Validity::NonNullable);
410+
assert_device_dict_into_host(codes, values.into_array()).await
411+
}
228412
}

0 commit comments

Comments
 (0)