This is an automated email from the ASF dual-hosted git repository.
viiccwen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/mahout.git
The following commit(s) were added to refs/heads/main by this push:
new 678d22d12 [QDP] Fix streaming host-to-device copy size (#1464)
678d22d12 is described below
commit 678d22d12afa82c6b318f386ff3f4f7425d4a59e
Author: Vic Wen <[email protected]>
AuthorDate: Sun Aug 16 14:32:13 2026 +0800
[QDP] Fix streaming host-to-device copy size (#1464)
* [QDP] Fix streaming host-to-device copy size
Signed-off-by: viiccwen <[email protected]>
* [QDP] Guard DLPack host length overflow
Signed-off-by: viiccwen <[email protected]>
---------
Signed-off-by: viiccwen <[email protected]>
---
qdp/qdp-core/src/encoding/mod.rs | 14 +++++++++--
qdp/qdp-core/tests/common/mod.rs | 43 ++++++++++++++++++++++++++++++++
qdp/qdp-core/tests/gpu_angle_encoding.rs | 40 ++++++++++++++++++++++-------
3 files changed, 86 insertions(+), 11 deletions(-)
diff --git a/qdp/qdp-core/src/encoding/mod.rs b/qdp/qdp-core/src/encoding/mod.rs
index 2d09b5146..32cadb059 100644
--- a/qdp/qdp-core/src/encoding/mod.rs
+++ b/qdp/qdp-core/src/encoding/mod.rs
@@ -279,12 +279,22 @@ pub(crate) fn stream_encode<E: ChunkEncoder>(
unsafe {
crate::profile_scope!("GPU::Dispatch");
- // Async copy to device (only if staging buffers are allocated)
+ // Async copy to device (only if staging buffers are
allocated).
+ // `current_len` counts f64 elements; the copy takes bytes.
if dev_staging.is_some() {
+ let copy_bytes = current_len
+ .checked_mul(std::mem::size_of::<f64>())
+ .ok_or_else(|| {
+ MahoutError::MemoryAllocation(format!(
+ "Staging copy size overflow: {} * {}",
+ current_len,
+ std::mem::size_of::<f64>()
+ ))
+ })?;
ctx.async_copy_to_device(
host_buffer.ptr() as *const c_void,
dev_ptr as *mut c_void,
- current_len,
+ copy_bytes,
)?;
ctx.record_copy_done(event_slot)?;
ctx.wait_for_copy(event_slot)?;
diff --git a/qdp/qdp-core/tests/common/mod.rs b/qdp/qdp-core/tests/common/mod.rs
index 241012d0f..85b018312 100644
--- a/qdp/qdp-core/tests/common/mod.rs
+++ b/qdp/qdp-core/tests/common/mod.rs
@@ -231,6 +231,49 @@ pub unsafe fn assert_dlpack_shape_2d_and_delete(
unsafe { take_deleter_and_delete(dlpack_ptr) };
}
+/// Downloads a complex128 DLPack tensor as interleaved real and imaginary
values.
+#[cfg(target_os = "linux")]
+#[allow(dead_code)]
+pub unsafe fn download_dlpack_complex_f64(dlpack_ptr: *mut DLManagedTensor) ->
Vec<f64> {
+ assert!(!dlpack_ptr.is_null(), "DLPack pointer should not be null");
+
+ let tensor = unsafe { &(*dlpack_ptr).dl_tensor };
+ assert!(tensor.ndim > 0, "DLPack tensor should have dimensions");
+ assert!(!tensor.shape.is_null(), "DLPack shape should not be null");
+ assert!(!tensor.data.is_null(), "DLPack data should not be null");
+ assert_eq!(
+ tensor.device.device_type,
+ qdp_core::dlpack::DLDeviceType::kDLCUDA
+ );
+ assert_eq!(tensor.dtype.code, qdp_core::dlpack::DL_COMPLEX);
+ assert_eq!(tensor.dtype.bits, 128);
+ assert_eq!(tensor.dtype.lanes, 1);
+
+ let shape = unsafe { std::slice::from_raw_parts(tensor.shape, tensor.ndim
as usize) };
+ let num_elements = shape.iter().try_fold(1_usize, |count, &dim| {
+ usize::try_from(dim)
+ .ok()
+ .and_then(|dim| count.checked_mul(dim))
+ });
+ let num_elements = num_elements.expect("DLPack shape should fit in usize");
+ let byte_count = num_elements
+ .checked_mul(std::mem::size_of::<qdp_kernels::CuDoubleComplex>())
+ .expect("DLPack byte count should fit in usize");
+ let host_len = num_elements
+ .checked_mul(2)
+ .expect("DLPack complex128 element count should fit in usize");
+ let mut host = vec![0.0_f64; host_len];
+ let device_ptr = (tensor.data as u64)
+ .checked_add(tensor.byte_offset)
+ .expect("DLPack byte offset should fit in a device pointer");
+
+ let ret = unsafe {
+ cudarc::driver::sys::lib().cuMemcpyDtoH_v2(host.as_mut_ptr().cast(),
device_ptr, byte_count)
+ };
+ assert_eq!(ret, cudarc::driver::sys::CUresult::CUDA_SUCCESS);
+ host
+}
+
/// Takes the DLPack deleter from the managed tensor and invokes it exactly
once.
#[cfg(target_os = "linux")]
#[allow(dead_code)]
diff --git a/qdp/qdp-core/tests/gpu_angle_encoding.rs
b/qdp/qdp-core/tests/gpu_angle_encoding.rs
index 874e948c9..937c408b6 100644
--- a/qdp/qdp-core/tests/gpu_angle_encoding.rs
+++ b/qdp/qdp-core/tests/gpu_angle_encoding.rs
@@ -143,11 +143,9 @@ fn test_angle_successful_encoding_from_parquet() {
return;
};
- let num_qubits = 4;
+ let num_qubits = 3;
let num_samples = 3;
- let data: Vec<f64> = (0..num_samples * num_qubits)
- .map(|i| (i as f64) * std::f64::consts::PI / (num_qubits *
num_samples) as f64)
- .collect();
+ let data = vec![0.1_f64, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9];
let path = "/tmp/test_angle_success.parquet";
common::write_fixed_size_list_parquet(path, &data, num_qubits);
@@ -158,11 +156,35 @@ fn test_angle_successful_encoding_from_parquet() {
let _ = std::fs::remove_file(path);
unsafe {
- common::assert_dlpack_shape_2d_and_delete(
- dlpack_ptr,
- num_samples as i64,
- (1 << num_qubits) as i64,
- );
+ let state_len = 1 << num_qubits;
+ common::assert_dlpack_shape_2d(dlpack_ptr, num_samples as i64,
state_len as i64);
+ let actual = common::download_dlpack_complex_f64(dlpack_ptr);
+
+ for sample_index in 0..num_samples {
+ let sample = &data[sample_index * num_qubits..(sample_index + 1) *
num_qubits];
+ for state_index in 0..state_len {
+ let expected = sample
+ .iter()
+ .enumerate()
+ .fold(1.0, |amplitude, (bit, angle)| {
+ amplitude
+ * if state_index & (1 << bit) == 0 {
+ angle.cos()
+ } else {
+ angle.sin()
+ }
+ });
+ let output_index = sample_index * state_len + state_index;
+ assert!(
+ (actual[output_index * 2] - expected).abs() < 1e-12,
+ "state[{sample_index}, {state_index}] real part: expected
{expected}, got {}",
+ actual[output_index * 2]
+ );
+ assert_eq!(actual[output_index * 2 + 1], 0.0);
+ }
+ }
+
+ common::take_deleter_and_delete(dlpack_ptr);
}
}
