rich7420 commented on code in PR #932:
URL: https://github.com/apache/mahout/pull/932#discussion_r2730300684
##########
qdp/qdp-core/src/lib.rs:
##########
@@ -469,106 +514,169 @@ impl QdpEngine {
}
let state_len = 1usize << num_qubits;
- if sample_size > state_len {
- return Err(MahoutError::InvalidInput(format!(
- "Sample size {} exceeds state vector size {} (2^{} qubits)",
- sample_size, state_len, num_qubits
- )));
- }
-
- // Allocate output state vector
- let batch_state_vector = {
- crate::profile_scope!("GPU::AllocBatch");
- gpu::GpuStateVector::new_batch(&self.device, num_samples,
num_qubits)?
- };
-
- // Compute inverse norms on GPU using warp-reduced kernel
- let inv_norms_gpu = {
- crate::profile_scope!("GPU::BatchNormKernel");
- use cudarc::driver::DevicePtrMut;
-
- let mut buffer =
self.device.alloc_zeros::<f64>(num_samples).map_err(|e| {
- MahoutError::MemoryAllocation(format!("Failed to allocate norm
buffer: {:?}", e))
- })?;
-
- let ret = unsafe {
- qdp_kernels::launch_l2_norm_batch(
- input_batch_d,
- num_samples,
- sample_size,
- *buffer.device_ptr_mut() as *mut f64,
- std::ptr::null_mut(), // default stream
- )
- };
-
- if ret != 0 {
- return Err(MahoutError::KernelLaunch(format!(
- "Norm reduction kernel failed with CUDA error code: {}
({})",
- ret,
- cuda_error_to_string(ret)
- )));
- }
-
- buffer
- };
-
- // Validate norms on host to catch zero or NaN samples early
- {
- crate::profile_scope!("GPU::NormValidation");
- let host_inv_norms = self
- .device
- .dtoh_sync_copy(&inv_norms_gpu)
- .map_err(|e| MahoutError::Cuda(format!("Failed to copy norms
to host: {:?}", e)))?;
-
- if host_inv_norms.iter().any(|v| !v.is_finite() || *v == 0.0) {
- return Err(MahoutError::InvalidInput(
- "One or more samples have zero or invalid
norm".to_string(),
- ));
+ let method = encoding_method.to_ascii_lowercase();
+
+ match method.as_str() {
+ "amplitude" => {
+ if sample_size > state_len {
+ return Err(MahoutError::InvalidInput(format!(
+ "Sample size {} exceeds state vector size {} (2^{}
qubits)",
+ sample_size, state_len, num_qubits
+ )));
+ }
+
+ let batch_state_vector = {
+ crate::profile_scope!("GPU::AllocBatch");
+ gpu::GpuStateVector::new_batch(&self.device, num_samples,
num_qubits)?
+ };
+
+ let inv_norms_gpu = {
+ crate::profile_scope!("GPU::BatchNormKernel");
+ use cudarc::driver::DevicePtrMut;
+
+ let mut buffer =
self.device.alloc_zeros::<f64>(num_samples).map_err(|e| {
+ MahoutError::MemoryAllocation(format!(
+ "Failed to allocate norm buffer: {:?}",
+ e
+ ))
+ })?;
+
+ let ret = unsafe {
+ qdp_kernels::launch_l2_norm_batch(
+ input_batch_d,
+ num_samples,
+ sample_size,
+ *buffer.device_ptr_mut() as *mut f64,
+ std::ptr::null_mut(), // default stream
+ )
+ };
+
+ if ret != 0 {
+ return Err(MahoutError::KernelLaunch(format!(
+ "Norm reduction kernel failed with CUDA error
code: {} ({})",
+ ret,
+ cuda_error_to_string(ret)
+ )));
+ }
+
+ buffer
+ };
+
+ {
+ crate::profile_scope!("GPU::NormValidation");
+ let host_inv_norms =
+ self.device.dtoh_sync_copy(&inv_norms_gpu).map_err(|e|
{
+ MahoutError::Cuda(format!("Failed to copy norms to
host: {:?}", e))
+ })?;
+
+ if host_inv_norms.iter().any(|v| !v.is_finite() || *v ==
0.0) {
+ return Err(MahoutError::InvalidInput(
+ "One or more samples have zero or invalid
norm".to_string(),
+ ));
+ }
+ }
+
+ {
+ crate::profile_scope!("GPU::BatchKernelLaunch");
+ use cudarc::driver::DevicePtr;
+
+ let state_ptr = batch_state_vector.ptr_f64().ok_or_else(||
{
+ MahoutError::InvalidInput(
+ "Batch state vector precision mismatch (expected
float64 buffer)"
+ .to_string(),
+ )
+ })?;
+
+ let ret = unsafe {
+ qdp_kernels::launch_amplitude_encode_batch(
+ input_batch_d,
+ state_ptr as *mut std::ffi::c_void,
+ *inv_norms_gpu.device_ptr() as *const f64,
+ num_samples,
+ sample_size,
+ state_len,
+ std::ptr::null_mut(), // default stream
+ )
+ };
+
+ if ret != 0 {
+ return Err(MahoutError::KernelLaunch(format!(
+ "Batch kernel launch failed with CUDA error code:
{} ({})",
+ ret,
+ cuda_error_to_string(ret)
+ )));
+ }
+ }
+
+ {
+ crate::profile_scope!("GPU::Synchronize");
+ self.device
+ .synchronize()
+ .map_err(|e| MahoutError::Cuda(format!("Sync failed:
{:?}", e)))?;
+ }
+
+ let batch_state_vector =
+ batch_state_vector.to_precision(&self.device,
self.precision)?;
+ Ok(batch_state_vector.to_dlpack())
}
- }
-
- // Launch batch kernel
- {
- crate::profile_scope!("GPU::BatchKernelLaunch");
- use cudarc::driver::DevicePtr;
-
- let state_ptr = batch_state_vector.ptr_f64().ok_or_else(|| {
- MahoutError::InvalidInput(
- "Batch state vector precision mismatch (expected float64
buffer)".to_string(),
- )
- })?;
-
- let ret = unsafe {
- qdp_kernels::launch_amplitude_encode_batch(
- input_batch_d,
- state_ptr as *mut std::ffi::c_void,
- *inv_norms_gpu.device_ptr() as *const f64,
- num_samples,
- sample_size,
- state_len,
- std::ptr::null_mut(), // default stream
- )
- };
-
- if ret != 0 {
- return Err(MahoutError::KernelLaunch(format!(
- "Batch kernel launch failed with CUDA error code: {} ({})",
- ret,
- cuda_error_to_string(ret)
- )));
+ "angle" => {
+ if sample_size != num_qubits {
+ return Err(MahoutError::InvalidInput(format!(
+ "Angle encoding expects sample_size={} (one angle per
qubit), got {}",
+ num_qubits, sample_size
+ )));
+ }
+
+ let batch_state_vector = {
+ crate::profile_scope!("GPU::AllocBatch");
+ gpu::GpuStateVector::new_batch(&self.device, num_samples,
num_qubits)?
+ };
+
+ let state_ptr = batch_state_vector.ptr_f64().ok_or_else(|| {
+ MahoutError::InvalidInput(
+ "Batch state vector precision mismatch (expected
float64 buffer)"
+ .to_string(),
+ )
+ })?;
+
+ {
+ crate::profile_scope!("GPU::BatchKernelLaunch");
+ let ret = unsafe {
+ qdp_kernels::launch_angle_encode_batch(
+ input_batch_d,
+ state_ptr as *mut std::ffi::c_void,
+ num_samples,
+ state_len,
+ num_qubits as u32,
+ std::ptr::null_mut(), // default stream
+ )
+ };
+
+ if ret != 0 {
+ return Err(MahoutError::KernelLaunch(format!(
+ "Batch angle encoding kernel failed: {} ({})",
+ ret,
+ cuda_error_to_string(ret)
+ )));
+ }
+ }
+
+ {
+ crate::profile_scope!("GPU::Synchronize");
+ self.device
+ .synchronize()
+ .map_err(|e| MahoutError::Cuda(format!("Sync failed:
{:?}", e)))?;
+ }
+
+ let batch_state_vector =
+ batch_state_vector.to_precision(&self.device,
self.precision)?;
+ Ok(batch_state_vector.to_dlpack())
}
Review Comment:
yeah I think it should be validated here.
--
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]
