diff --git a/qdp/qdp-python/src/engine.rs b/qdp/qdp-python/src/engine.rs index b2b006ff78..9d6a2a6ae6 100644 --- a/qdp/qdp-python/src/engine.rs +++ b/qdp/qdp-python/src/engine.rs @@ -75,10 +75,14 @@ impl QdpEngine { /// - Python list: [1.0, 2.0, 3.0, 4.0] /// - NumPy array: 1D (single sample) or 2D (batch) array /// - PyTorch tensor: CPU tensor (float64 recommended; will be copied to GPU) + /// or CUDA tensor for zero-copy encoding /// - String path: .parquet, .arrow, .feather, .npy, .pt, .pth, .pb file /// - pathlib.Path: Path object (converted via os.fspath()) /// num_qubits: Number of qubits for encoding /// encoding_method: Encoding strategy ("amplitude" default, "angle", or "basis") + /// CUDA tensor note: + /// - amplitude accepts float64 and float32 + /// - angle accepts float64 generally, plus float32 for 1D single-sample tensors /// /// Returns: /// QuantumTensor: DLPack-compatible tensor for zero-copy PyTorch integration @@ -477,7 +481,7 @@ impl QdpEngine { /// Encode directly from a PyTorch CUDA tensor. Internal helper. /// - /// Dispatches to the core f32 GPU pointer API for float32 amplitude encoding, + /// Dispatches to the core f32 GPU pointer APIs for supported float32 CUDA paths, /// or to the float64/basis GPU pointer APIs for other dtypes and methods. fn _encode_from_cuda_tensor( &self, @@ -495,7 +499,7 @@ impl QdpEngine { let ndim: usize = data.call_method0("dim")?.extract()?; let tensor_info = extract_cuda_tensor_info(data)?; - if method.as_str() == "amplitude" && is_f32 { + if is_f32 && matches!(method.as_str(), "amplitude" | "angle") { match ndim { 1 => { let input_len: usize = data.call_method0("numel")?.extract()?; @@ -504,16 +508,31 @@ impl QdpEngine { let data_ptr = data_ptr_u64 as *const f32; let ptr = unsafe { - self.engine - .encode_from_gpu_ptr_f32_with_stream( - data_ptr, input_len, num_qubits, stream_ptr, - ) - .map_err(|e| { - PyRuntimeError::new_err(format!( - "Encoding failed (float32 amplitude): {}", - e - )) - })? + match method.as_str() { + "amplitude" => self + .engine + .encode_from_gpu_ptr_f32_with_stream( + data_ptr, input_len, num_qubits, stream_ptr, + ) + .map_err(|e| { + PyRuntimeError::new_err(format!( + "Encoding failed (float32 amplitude): {}", + e + )) + })?, + "angle" => self + .engine + .encode_angle_from_gpu_ptr_f32_with_stream( + data_ptr, input_len, num_qubits, stream_ptr, + ) + .map_err(|e| { + PyRuntimeError::new_err(format!( + "Encoding failed (float32 angle): {}", + e + )) + })?, + _ => unreachable!("unreachable: unhandled f32 encoding method"), + } }; Ok(QuantumTensor { diff --git a/qdp/qdp-python/src/pytorch.rs b/qdp/qdp-python/src/pytorch.rs index 17185fd548..9b6645105c 100644 --- a/qdp/qdp-python/src/pytorch.rs +++ b/qdp/qdp-python/src/pytorch.rs @@ -151,6 +151,7 @@ pub fn validate_cuda_tensor_for_encoding( encoding_method: &str, ) -> PyResult<()> { let method = encoding_method.to_ascii_lowercase(); + let ndim: usize = tensor.call_method0("dim")?.extract()?; if !CUDA_ENCODING_METHODS.contains(&method.as_str()) { return Err(PyRuntimeError::new_err(format!( @@ -176,7 +177,14 @@ pub fn validate_cuda_tensor_for_encoding( } } "angle" | "iqp" | "iqp-z" => { - if !dtype_str_lower.contains("float64") { + if method == "angle" && dtype_str_lower.contains("float32") { + if ndim != 1 { + return Err(PyRuntimeError::new_err( + "CUDA tensor float32 angle encoding currently supports only 1D single-sample tensors. \ + Use tensor.to(torch.float64) for batch angle encoding.", + )); + } + } else if !dtype_str_lower.contains("float64") { return Err(PyRuntimeError::new_err(format!( "CUDA tensor must have dtype float64 for {} encoding, got {}. \ Use tensor.to(torch.float64)", diff --git a/testing/qdp/test_bindings.py b/testing/qdp/test_bindings.py index c44ee05010..7664cb61e0 100644 --- a/testing/qdp/test_bindings.py +++ b/testing/qdp/test_bindings.py @@ -385,6 +385,25 @@ def test_encode_cuda_tensor_wrong_dtype(): engine.encode(data, 2, "amplitude") +@requires_qdp +@pytest.mark.gpu +def test_encode_cuda_tensor_angle_float16_rejected(): + """Test error when CUDA tensor has wrong dtype for angle float32 fast path.""" + pytest.importorskip("torch") + from _qdp import QdpEngine + + if not torch.cuda.is_available(): + pytest.skip("GPU required for QdpEngine") + + engine = QdpEngine(0) + data = torch.tensor([0.0, torch.pi / 2], dtype=torch.float16, device="cuda:0") + + with pytest.raises( + RuntimeError, match="float64 for angle encoding|supports only 1D" + ): + engine.encode(data, 2, "angle") + + @requires_qdp @pytest.mark.gpu def test_encode_cuda_tensor_non_contiguous(): @@ -643,6 +662,57 @@ def test_encode_cuda_tensor_float32_input_output_dtype(precision, expected_dtype ) +@requires_qdp +@pytest.mark.gpu +@pytest.mark.parametrize( + ("precision", "expected_dtype"), + [ + ("float32", torch.complex64), + ("float64", torch.complex128), + ], +) +def test_angle_encode_cuda_tensor_float32_input_output_dtype(precision, expected_dtype): + """Test that 1D float32 CUDA angle encoding respects engine precision (f32 path).""" + pytest.importorskip("torch") + from _qdp import QdpEngine + + if not torch.cuda.is_available(): + pytest.skip("GPU required for QdpEngine") + + engine = QdpEngine(0, precision=precision) + data = torch.tensor([torch.pi / 2, 0.0], dtype=torch.float32, device="cuda:0") + result = torch.from_dlpack(engine.encode(data, 2, "angle")) + + assert result.dtype == expected_dtype, ( + f"Expected {expected_dtype}, got {result.dtype}" + ) + assert result.shape == (1, 4) + + expected = torch.tensor([[0.0 + 0j, 1.0 + 0j, 0.0 + 0j, 0.0 + 0j]], device="cuda:0") + assert torch.allclose(result, expected.to(result.dtype), atol=1e-6, rtol=1e-6) + + +@requires_qdp +@pytest.mark.gpu +def test_angle_encode_cuda_tensor_float32_batch_rejected(): + """Test that float32 CUDA angle encoding stays limited to 1D single-sample tensors.""" + pytest.importorskip("torch") + from _qdp import QdpEngine + + if not torch.cuda.is_available(): + pytest.skip("GPU required for QdpEngine") + + engine = QdpEngine(0) + data = torch.tensor( + [[0.0, 0.0], [torch.pi / 2, 0.0]], + dtype=torch.float32, + device="cuda:0", + ) + + with pytest.raises(RuntimeError, match="supports only 1D single-sample tensors"): + engine.encode(data, 2, "angle") + + @requires_qdp @pytest.mark.gpu def test_basis_encode_basic(): @@ -1365,7 +1435,7 @@ def test_iqp_fwt_matches_naive_reference(): if not torch.cuda.is_available(): pytest.skip("GPU required for QdpEngine") - engine = QdpEngine(0) + engine = QdpEngine(0, precision="float64") for encoding_method, enable_zz in [("iqp-z", False), ("iqp", True)]: for num_qubits in [4, 5]: diff --git a/testing/qdp_python/test_dlpack_validation.py b/testing/qdp_python/test_dlpack_validation.py index 414e44f7bc..8b90c41911 100644 --- a/testing/qdp_python/test_dlpack_validation.py +++ b/testing/qdp_python/test_dlpack_validation.py @@ -115,6 +115,65 @@ def test_cuda_float32_amplitude_2d_respects_engine_precision() -> None: assert torch.allclose(qt, expected) +@pytest.mark.skipif(not _cuda_available(), reason="CUDA not available") +def test_cuda_float32_angle_supported_single_sample() -> None: + """1D float32 CUDA tensor should be supported for single-sample angle encoding.""" + engine = _engine() + t = torch.tensor([torch.pi / 2, 0.0], dtype=torch.float32, device="cuda") + + result = engine.encode(t, num_qubits=2, encoding_method="angle") + assert result is not None + + qt = torch.from_dlpack(result) + assert qt.is_cuda + assert qt.shape == (1, 4) + assert qt.dtype == torch.complex64 + + expected = torch.tensor( + [[0.0 + 0j, 1.0 + 0j, 0.0 + 0j, 0.0 + 0j]], + dtype=torch.complex64, + device="cuda", + ) + assert torch.allclose(qt, expected, atol=1e-6, rtol=1e-6) + + +@pytest.mark.skipif(not _cuda_available(), reason="CUDA not available") +def test_cuda_float32_angle_2d_rejected() -> None: + """Float32 CUDA angle encoding should remain single-sample only.""" + engine = _engine() + t = torch.tensor( + [[0.0, 0.0], [torch.pi / 2, 0.0]], + dtype=torch.float32, + device="cuda", + ) + + with pytest.raises(RuntimeError, match="1D single-sample"): + engine.encode(t, num_qubits=2, encoding_method="angle") + + +@pytest.mark.skipif(not _cuda_available(), reason="CUDA not available") +def test_cuda_float32_angle_non_contiguous_rejected() -> None: + """1D float32 CUDA angle tensor must still be contiguous.""" + engine = _engine() + t = torch.randn(4, dtype=torch.float32, device="cuda")[::2] + assert t.stride(0) != 1 + + with pytest.raises(RuntimeError, match="contiguous"): + engine.encode(t, num_qubits=2, encoding_method="angle") + + +@pytest.mark.skipif(not _cuda_available(), reason="CUDA not available") +def test_cuda_float16_angle_rejected() -> None: + """Angle encoding should not accept float16 CUDA tensors.""" + engine = _engine() + t = torch.tensor([0.0, torch.pi / 2], dtype=torch.float16, device="cuda") + + with pytest.raises( + RuntimeError, match="float64 for angle encoding|supports only 1D" + ): + engine.encode(t, num_qubits=2, encoding_method="angle") + + @pytest.mark.skipif(not _cuda_available(), reason="CUDA not available") def test_stride_1d_non_contiguous_rejected() -> None: """Non-contiguous 1D CUDA tensor (stride != 1) should fail with contiguous requirement."""