Copilot commented on code in PR #1464:
URL: https://github.com/apache/mahout/pull/1464#discussion_r3724596167
##########
qdp/qdp-core/tests/common/mod.rs:
##########
@@ -231,6 +231,46 @@ 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 mut host = vec![0.0_f64; num_elements * 2];
Review Comment:
`download_dlpack_complex_f64` allocates `host` with `num_elements * 2`
without overflow checking. If `num_elements` is large enough to overflow
`usize`, the vector can be undersized while `byte_count` is still computed from
`num_elements`, and `cuMemcpyDtoH_v2` would write past the end of `host` (UB).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
