You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@mahout.apache.org by ji...@apache.org on 2026/04/22 14:16:42 UTC
(mahout) branch main updated: [Feature][QDP] Add AMD GPU support via Triton backend (#1158)
This is an automated email from the ASF dual-hosted git repository.
400Ping 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 bbce3c392 [Feature][QDP] Add AMD GPU support via Triton backend (#1158)
bbce3c392 is described below
commit bbce3c39266263025b01386d2abd15926b6df5aa
Author: Jie-Kai Chang <ji...@apache.org>
AuthorDate: Wed Apr 22 22:16:35 2026 +0800
[Feature][QDP] Add AMD GPU support via Triton backend (#1158)
---
.github/workflows/wheel-build.yml | 18 +-
.pre-commit-config.yaml | 2 +-
docs/qdp/getting-started.md | 23 ++-
pyproject.toml | 1 +
qdp/qdp-core/src/dlpack.rs | 8 +-
qdp/qdp-core/src/gpu/memory.rs | 13 ++
qdp/qdp-core/src/gpu/pool_metrics.rs | 2 +-
qdp/qdp-core/src/readers/parquet.rs | 9 +-
qdp/qdp-python/README.md | 49 ++++-
qdp/qdp-python/TRITON_AMD_BACKEND.md | 90 +++++++++
qdp/qdp-python/benchmark/benchmark_triton_amd.py | 96 ++++++++++
qdp/qdp-python/pyproject.toml | 12 +-
qdp/qdp-python/qumat_qdp/__init__.py | 18 +-
qdp/qdp-python/qumat_qdp/backend.py | 95 ++++++++++
qdp/qdp-python/qumat_qdp/backends/__init__.py | 22 +++
qdp/qdp-python/qumat_qdp/backends/amd.py | 21 +++
qdp/qdp-python/qumat_qdp/backends/cuda.py | 42 +++++
qdp/qdp-python/qumat_qdp/tensor.py | 56 ++++++
qdp/qdp-python/qumat_qdp/triton_amd.py | 210 +++++++++++++++++++++
qdp/qdp-python/src/dlpack.rs | 78 +++++---
qdp/qdp-python/src/engine.rs | 224 ++++++++++++++++++-----
qdp/qdp-python/src/tensor.rs | 5 +-
qdp/qdp-python/tests/test_amd_engine.py | 52 ++++++
qdp/qdp-python/tests/test_backend_routing.py | 129 +++++++++++++
qdp/qdp-python/tests/test_triton_amd_backend.py | 160 ++++++++++++++++
qdp/qdp-python/uv.lock | 4 +
qumat/qdp.py | 35 ++--
testing/qdp_python/test_torch_ref.py | 3 +-
testing/qumat/test_qdp_module.py | 13 +-
uv.lock | 2 +
30 files changed, 1370 insertions(+), 122 deletions(-)
diff --git a/.github/workflows/wheel-build.yml b/.github/workflows/wheel-build.yml
index 32e0f7f20..bdc864d4b 100644
--- a/.github/workflows/wheel-build.yml
+++ b/.github/workflows/wheel-build.yml
@@ -68,13 +68,19 @@ jobs:
dnf install -y gcc-toolset-13-gcc gcc-toolset-13-gcc-c++
export CC=/opt/rh/gcc-toolset-13/root/usr/bin/gcc
export CXX=/opt/rh/gcc-toolset-13/root/usr/bin/g++
- # Install CUDA 12.5 toolkit (nvcc + headers + cudart) for kernel compilation
+ # Try to install the CUDA 12.5 toolkit for kernel compilation.
+ # If the external NVIDIA repo is temporarily unavailable, fall back
+ # to a no-CUDA smoke build instead of failing the PR wheel job.
dnf install -y 'dnf-command(config-manager)'
- dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo
- dnf install -y cuda-nvcc-12-5 cuda-cudart-devel-12-5
- export CUDA_PATH=/usr/local/cuda-12.5
- export PATH=/usr/local/cuda-12.5/bin:$PATH
- nvcc --version
+ if dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo \
+ && dnf install -y cuda-nvcc-12-5 cuda-cudart-devel-12-5; then
+ export CUDA_PATH=/usr/local/cuda-12.5
+ export PATH=/usr/local/cuda-12.5/bin:$PATH
+ nvcc --version
+ else
+ echo "CUDA repo unavailable; continuing with QDP_NO_CUDA=1 smoke build."
+ export QDP_NO_CUDA=1
+ fi
sccache: true
- uses: actions/upload-artifact@v5
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 77ff509f8..5fafcc55e 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -25,7 +25,7 @@ repos:
hooks:
- id: ty
name: ty check
- entry: uv run ty check .
+ entry: bash -c 'if [ -n "${VIRTUAL_ENV:-}" ]; then python -m ty check .; elif [ -x .venv/bin/python ]; then .venv/bin/python -m ty check .; else uv run ty check .; fi'
language: system
pass_filenames: false
types: [python]
diff --git a/docs/qdp/getting-started.md b/docs/qdp/getting-started.md
index a8747cbb2..129fc5b1b 100644
--- a/docs/qdp/getting-started.md
+++ b/docs/qdp/getting-started.md
@@ -8,9 +8,10 @@ QDP (Quantum Data Plane) is a GPU-accelerated library for encoding classical dat
## Prerequisites
-- Linux with NVIDIA GPU
-- CUDA toolkit installed (`nvcc --version` to verify)
- Python 3.10+
+- One of:
+ - NVIDIA GPU with CUDA toolkit installed (`nvcc --version` to verify)
+ - AMD GPU with ROCm and PyTorch ROCm installed (`python -c "import torch; print(torch.version.hip)"`)
## Installation
@@ -28,13 +29,20 @@ source .venv/bin/activate
uv run --active maturin develop --manifest-path qdp/qdp-python/Cargo.toml
```
+For AMD ROCm with the Triton backend:
+
+```bash
+uv sync --group dev --extra qdp
+pip install triton
+```
+
## Quick Start
```python
import torch
from qumat.qdp import QdpEngine
-engine = QdpEngine(0) # GPU device 0
+engine = QdpEngine(device_id=0, backend="cuda")
data = [0.5, 0.5, 0.5, 0.5]
qtensor = engine.encode(data, num_qubits=2, encoding_method="amplitude")
@@ -42,6 +50,14 @@ qtensor = engine.encode(data, num_qubits=2, encoding_method="amplitude")
tensor = torch.from_dlpack(qtensor) # Note: can only be consumed once
```
+AMD ROCm uses the same API with a different backend selector:
+
+```python
+engine = QdpEngine(device_id=0, backend="amd", precision="float32")
+qtensor = engine.encode(data, num_qubits=2, encoding_method="amplitude")
+tensor = torch.from_dlpack(qtensor)
+```
+
## Encoding Methods
| Method | Constraint | Example |
@@ -79,6 +95,7 @@ Notes:
|---------|----------|
| Import fails | Activate root venv: `source mahout/.venv/bin/activate` (or `cd mahout && source .venv/bin/activate`) |
| CUDA errors | Run `cargo clean` in `qdp/` and rebuild |
+| AMD backend unavailable | Verify ROCm PyTorch reports `torch.version.hip` and install `triton` |
| Out of memory | Reduce `num_qubits` or use `precision="float32"` |
## Next Steps
diff --git a/pyproject.toml b/pyproject.toml
index 53b349d16..0db11d1cb 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -83,6 +83,7 @@ python_functions = "test_*"
addopts = ["-v", "--tb=short", "-rs"]
markers = [
"gpu: marks tests as requiring GPU and _qdp extension (auto-skipped if unavailable)",
+ "rocm: marks tests as requiring ROCm/Triton AMD runtime",
"slow: marks tests as slow running",
]
diff --git a/qdp/qdp-core/src/dlpack.rs b/qdp/qdp-core/src/dlpack.rs
index 0108f4f19..a15344ee8 100644
--- a/qdp/qdp-core/src/dlpack.rs
+++ b/qdp/qdp-core/src/dlpack.rs
@@ -19,7 +19,7 @@
#[cfg(target_os = "linux")]
use crate::error::cuda_error_to_string;
use crate::error::{MahoutError, Result};
-use crate::gpu::memory::{BufferStorage, GpuStateVector, Precision};
+use crate::gpu::memory::{BufferStorage, GpuDeviceType, GpuStateVector, Precision};
use std::os::raw::{c_int, c_void};
use std::sync::Arc;
@@ -113,6 +113,7 @@ pub unsafe fn synchronize_stream(_stream: *mut c_void) -> Result<()> {
pub enum DLDeviceType {
kDLCPU = 1,
kDLCUDA = 2,
+ kDLROCM = 10,
// Other types omitted
}
@@ -291,7 +292,10 @@ impl GpuStateVector {
let tensor = DLTensor {
data: self.ptr_void(),
device: DLDevice {
- device_type: DLDeviceType::kDLCUDA,
+ device_type: match self.device_type {
+ GpuDeviceType::Cuda => DLDeviceType::kDLCUDA,
+ GpuDeviceType::Rocm => DLDeviceType::kDLROCM,
+ },
device_id: self.device_id as c_int,
},
ndim,
diff --git a/qdp/qdp-core/src/gpu/memory.rs b/qdp/qdp-core/src/gpu/memory.rs
index baa9a2988..661060b55 100644
--- a/qdp/qdp-core/src/gpu/memory.rs
+++ b/qdp/qdp-core/src/gpu/memory.rs
@@ -35,6 +35,13 @@ pub enum Precision {
Float64,
}
+/// Backend GPU device type for DLPack metadata.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum GpuDeviceType {
+ Cuda,
+ Rocm,
+}
+
#[cfg(target_os = "linux")]
use crate::gpu::cuda_ffi::{cudaFreeHost, cudaHostAlloc, cudaMemGetInfo};
@@ -212,6 +219,8 @@ pub struct GpuStateVector {
pub(crate) num_samples: Option<usize>,
/// CUDA device ordinal
pub device_id: usize,
+ /// GPU backend type used for DLPack device metadata.
+ pub device_type: GpuDeviceType,
}
// Safety: CudaSlice and Arc are both Send + Sync
@@ -290,6 +299,7 @@ impl GpuStateVector {
size_elements: _size_elements,
num_samples: None,
device_id: _device.ordinal(),
+ device_type: GpuDeviceType::Cuda,
})
}
@@ -401,6 +411,7 @@ impl GpuStateVector {
size_elements: total_elements,
num_samples: Some(num_samples),
device_id: _device.ordinal(),
+ device_type: GpuDeviceType::Cuda,
})
}
@@ -487,6 +498,7 @@ impl GpuStateVector {
size_elements: self.size_elements,
num_samples: self.num_samples,
device_id: device.ordinal(),
+ device_type: self.device_type,
})
}
@@ -562,6 +574,7 @@ impl GpuStateVector {
size_elements: self.size_elements,
num_samples: self.num_samples, // Preserve batch information
device_id: device.ordinal(),
+ device_type: self.device_type,
})
}
diff --git a/qdp/qdp-core/src/gpu/pool_metrics.rs b/qdp/qdp-core/src/gpu/pool_metrics.rs
index 60085aa6a..1cf2779af 100644
--- a/qdp/qdp-core/src/gpu/pool_metrics.rs
+++ b/qdp/qdp-core/src/gpu/pool_metrics.rs
@@ -112,7 +112,7 @@ impl PoolMetrics {
} else {
0.0
},
- avg_wait_time_ns: if waits > 0 { wait_time_ns / waits } else { 0 },
+ avg_wait_time_ns: wait_time_ns.checked_div(waits).unwrap_or(0),
}
}
diff --git a/qdp/qdp-core/src/readers/parquet.rs b/qdp/qdp-core/src/readers/parquet.rs
index ddc0a16ed..0c9a3e162 100644
--- a/qdp/qdp-core/src/readers/parquet.rs
+++ b/qdp/qdp-core/src/readers/parquet.rs
@@ -389,11 +389,10 @@ impl StreamingDataReader for ParquetStreamingReader {
let mut written = 0;
let buf_cap = buffer.len();
let calc_limit = |ss: usize| -> usize {
- if ss == 0 {
- buf_cap
- } else {
- (buf_cap / ss) * ss
- }
+ buf_cap
+ .checked_div(ss)
+ .map(|chunks| chunks * ss)
+ .unwrap_or(buf_cap)
};
let mut limit = self.sample_size.map_or(buf_cap, calc_limit);
diff --git a/qdp/qdp-python/README.md b/qdp/qdp-python/README.md
index ee70834ad..8b3d964ac 100644
--- a/qdp/qdp-python/README.md
+++ b/qdp/qdp-python/README.md
@@ -8,7 +8,25 @@ GPU-accelerated quantum state encoding for [Apache Mahout Qumat](https://github.
pip install qumat[qdp]
```
-Requires CUDA-capable GPU.
+Requires one of:
+- NVIDIA GPU (CUDA path via `QdpEngine`)
+- AMD GPU with ROCm (AMD path via `QdpEngine(backend="amd")`)
+
+Recommended environment setup:
+
+```bash
+python -m venv .venv
+source .venv/bin/activate
+
+# Install the GPU runtime for your platform first:
+# - NVIDIA users: CUDA-compatible torch / triton
+# - AMD users: ROCm-compatible torch / triton
+
+uv sync --active --project qdp/qdp-python --group dev
+```
+
+Use `--active` so `uv` reuses the environment that already has the correct GPU
+runtime stack.
## Usage
@@ -16,8 +34,9 @@ Requires CUDA-capable GPU.
import qumat.qdp as qdp
import torch
-# Initialize engine on GPU 0
-engine = qdp.QdpEngine(device_id=0)
+# Initialize the unified QDP engine on GPU 0.
+# Choose the backend explicitly.
+engine = qdp.QdpEngine(device_id=0, backend="cuda")
# Encode data into quantum state
qtensor = engine.encode([1.0, 2.0, 3.0, 4.0], num_qubits=2, encoding_method="amplitude")
@@ -27,6 +46,26 @@ tensor = torch.from_dlpack(qtensor)
print(tensor) # Complex tensor on CUDA
```
+### AMD ROCm Usage
+
+```python
+import qumat.qdp as qdp
+import torch
+
+# Unified AMD engine route
+engine = qdp.QdpEngine(device_id=0, precision="float32", backend="amd")
+qt = engine.encode(torch.randn(8, 4, device="cuda"), 2, "amplitude")
+state = torch.from_dlpack(qt)
+print(state.device, state.dtype) # cuda:0, complex64
+
+```
+
+The public `QdpEngine` is a unified Python facade with explicit backend selection:
+- `backend="cuda"` routes to the Rust `_qdp.QdpEngine`
+- `backend="amd"` routes to the Triton AMD engine directly
+
+See `qdp/qdp-python/TRITON_AMD_BACKEND.md` for Triton AMD setup and validation details.
+
## Encoding Methods
| Method | Description |
@@ -36,6 +75,10 @@ print(tensor) # Complex tensor on CUDA
| `basis` | Encode integer as computational basis state |
| `iqp` | IQP-style encoding with entanglement |
+Backend support boundary:
+- CUDA (`QdpEngine`): `amplitude`, `angle`, `basis`, `iqp`
+- AMD (`QdpEngine(..., backend="amd")`): `amplitude`, `angle`, `basis` (no `iqp` yet)
+
## Input Sources
```python
diff --git a/qdp/qdp-python/TRITON_AMD_BACKEND.md b/qdp/qdp-python/TRITON_AMD_BACKEND.md
new file mode 100644
index 000000000..93e01d466
--- /dev/null
+++ b/qdp/qdp-python/TRITON_AMD_BACKEND.md
@@ -0,0 +1,90 @@
+# Triton AMD Backend
+
+This document describes the Triton-based implementation used by the QDP AMD backend on ROCm.
+
+## Prerequisites
+
+- AMD GPU supported by ROCm
+- ROCm driver/runtime installed
+- PyTorch ROCm build (`torch.version.hip` is not `None`)
+- Triton installed with HIP support
+
+## Install (project environment)
+
+```bash
+uv sync --project qdp/qdp-python --group benchmark --active
+```
+
+This installs the benchmark group, including `triton`.
+
+## Runtime capability checks
+
+Use:
+
+```python
+from qumat_qdp import is_triton_amd_available
+print(is_triton_amd_available())
+```
+
+The check validates:
+- ROCm runtime is visible through PyTorch
+- Triton is importable
+- Triton active backend is HIP (when query is available)
+
+## Usage
+
+### Unified AMD routing (recommended)
+
+```python
+import torch
+from qumat_qdp import QdpEngine
+
+engine = QdpEngine(device_id=0, precision="float32", backend="amd")
+x = torch.randn(64, 1024, device="cuda", dtype=torch.float32)
+qt = engine.encode(x, num_qubits=10, encoding_method="amplitude")
+state = torch.from_dlpack(qt)
+```
+
+The public `QdpEngine(..., backend="amd")` route goes directly to the Triton
+engine and returns the Python `QuantumTensorWrapper` compatibility wrapper.
+The CUDA route remains Rust-owned and returns the extension `QuantumTensor`.
+
+### Triton implementation details
+
+```python
+import torch
+from qumat_qdp.triton_amd import TritonAmdEngine
+
+engine = TritonAmdEngine(device_id=0, precision="float32")
+x = torch.randn(64, 1024, device="cuda", dtype=torch.float32)
+qt = engine.encode(x, num_qubits=10, encoding_method="amplitude")
+state = torch.from_dlpack(qt)
+```
+
+Supported methods:
+- `amplitude`
+- `angle`
+- `basis`
+
+Not supported in the AMD route yet:
+- `iqp` (currently CUDA backend only)
+
+## Correctness tests
+
+Run Triton backend tests:
+
+```bash
+uv run --project qdp/qdp-python pytest qdp/qdp-python/tests/test_triton_amd_backend.py -q
+uv run --project qdp/qdp-python pytest -m rocm qdp/qdp-python/tests -q
+```
+
+Tests include:
+- parity against Torch reference outputs (amplitude/angle/basis)
+- optional parity against CUDA backend reference (when NVIDIA CUDA path is present)
+
+## Baseline benchmark
+
+```bash
+uv run --project qdp/qdp-python python qdp/qdp-python/benchmark/benchmark_triton_amd.py \
+ --qubits 12 --batch-size 64 --batches 200 --encoding-method amplitude
+```
diff --git a/qdp/qdp-python/benchmark/benchmark_triton_amd.py b/qdp/qdp-python/benchmark/benchmark_triton_amd.py
new file mode 100644
index 000000000..6dd11c442
--- /dev/null
+++ b/qdp/qdp-python/benchmark/benchmark_triton_amd.py
@@ -0,0 +1,96 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Baseline benchmark for Triton AMD backend (ROCm)."""
+
+from __future__ import annotations
+
+import argparse
+import time
+
+import torch
+from qumat_qdp import is_triton_amd_available
+from qumat_qdp.triton_amd import TritonAmdEngine
+
+
+def _build_input(
+ batch_size: int,
+ num_qubits: int,
+ encoding_method: str,
+ *,
+ device: str,
+ dtype: torch.dtype,
+) -> torch.Tensor:
+ if encoding_method == "basis":
+ return torch.randint(0, 1 << num_qubits, (batch_size,), device=device)
+ width = num_qubits if encoding_method == "angle" else 1 << num_qubits
+ return torch.randn(batch_size, width, device=device, dtype=dtype)
+
+
+def main() -> int:
+ parser = argparse.ArgumentParser(description=__doc__)
+ parser.add_argument("--device-id", type=int, default=0)
+ parser.add_argument("--num-qubits", type=int, default=8)
+ parser.add_argument("--batch-size", type=int, default=64)
+ parser.add_argument("--iters", type=int, default=100)
+ parser.add_argument(
+ "--encoding",
+ choices=("amplitude", "angle", "basis"),
+ default="amplitude",
+ )
+ parser.add_argument(
+ "--precision", choices=("float32", "float64"), default="float32"
+ )
+ args = parser.parse_args()
+
+ if not is_triton_amd_available():
+ raise SystemExit("Triton AMD backend unavailable.")
+
+ dtype = torch.float32 if args.precision == "float32" else torch.float64
+ device = f"cuda:{args.device_id}"
+ engine = TritonAmdEngine(device_id=args.device_id, precision=args.precision)
+ data = _build_input(
+ args.batch_size,
+ args.num_qubits,
+ args.encoding,
+ device=device,
+ dtype=dtype,
+ )
+
+ for _ in range(5):
+ engine.encode(data, args.num_qubits, args.encoding)
+ torch.cuda.synchronize()
+
+ start = time.perf_counter()
+ for _ in range(args.iters):
+ engine.encode(data, args.num_qubits, args.encoding)
+ torch.cuda.synchronize()
+ duration = time.perf_counter() - start
+
+ total_vectors = args.batch_size * args.iters
+ print("Triton AMD benchmark")
+ print(f" encoding: {args.encoding}")
+ print(f" precision: {args.precision}")
+ print(f" num_qubits: {args.num_qubits}")
+ print(f" batch_size: {args.batch_size}")
+ print(f" iterations: {args.iters}")
+ print(f" duration_sec: {duration:.6f}")
+ print(f" vectors_per_s: {total_vectors / duration:.2f}")
+ return 0
+
+
+if __name__ == "__main__":
+ raise SystemExit(main())
diff --git a/qdp/qdp-python/pyproject.toml b/qdp/qdp-python/pyproject.toml
index 7b5822893..56c25bb8b 100644
--- a/qdp/qdp-python/pyproject.toml
+++ b/qdp/qdp-python/pyproject.toml
@@ -32,13 +32,18 @@ classifiers = [
dynamic = ["version"]
[dependency-groups]
-dev = ["pytest", "torch>=2.2,<=2.9.0"]
+dev = [
+ "pytest",
+ "torch>=2.2,<=2.9.0",
+ "triton>=3.0",
+]
benchmark = [
"numpy>=1.24,<2.0",
"pandas>=2.0",
"pyarrow>=14.0",
"tensorflow>=2.20",
"torch>=2.2,<=2.9.0",
+ "triton>=3.0",
"qiskit>=1.0",
"qiskit-aer>=0.17.2",
"pennylane>=0.35",
@@ -53,11 +58,6 @@ benchmark = [
[tool.uv.sources]
qumat = { path = "../..", editable = true }
-[[tool.uv.index]]
-name = "pytorch"
-url = "https://download.pytorch.org/whl/cu122"
-explicit = true
-
# Invalidate uv cache when Rust or Cargo changes so extension is rebuilt (run_throughput_pipeline_py etc.).
# Ref: https://docs.astral.sh/uv/concepts/cache/#dynamic-metadata
[tool.uv]
diff --git a/qdp/qdp-python/qumat_qdp/__init__.py b/qdp/qdp-python/qumat_qdp/__init__.py
index f33e8ad5e..2b7566141 100644
--- a/qdp/qdp-python/qumat_qdp/__init__.py
+++ b/qdp/qdp-python/qumat_qdp/__init__.py
@@ -17,7 +17,7 @@
"""
QDP (Quantum Data Processing) Python API.
-Public API: QdpEngine, QuantumTensor (Rust extension _qdp),
+Public API: QdpEngine (unified router), QdpTensor/QuantumTensor (DLPack facade),
QdpBenchmark, ThroughputResult, LatencyResult (benchmark API),
QuantumDataLoader (data loader iterator).
@@ -34,12 +34,12 @@ from qumat_qdp._backend import Backend, force_backend, get_backend, get_qdp
_qdp_mod = get_qdp()
if _qdp_mod is not None:
- QdpEngine = getattr(_qdp_mod, "QdpEngine")
- QuantumTensor = getattr(_qdp_mod, "QuantumTensor")
+ RustQdpEngine = getattr(_qdp_mod, "QdpEngine")
+ NativeQuantumTensor = getattr(_qdp_mod, "QuantumTensor")
run_throughput_pipeline_py = getattr(_qdp_mod, "run_throughput_pipeline_py", None)
else:
- QdpEngine = None
- QuantumTensor = None
+ RustQdpEngine = None
+ NativeQuantumTensor = None
run_throughput_pipeline_py = None
BACKEND = get_backend()
@@ -49,17 +49,25 @@ from qumat_qdp.api import (
QdpBenchmark,
ThroughputResult,
)
+from qumat_qdp.backend import QdpEngine
from qumat_qdp.loader import QuantumDataLoader
+from qumat_qdp.tensor import QdpTensor, QuantumTensor
+from qumat_qdp.triton_amd import TritonAmdEngine, is_triton_amd_available
__all__ = [
"BACKEND",
"Backend",
"LatencyResult",
+ "NativeQuantumTensor",
"QdpBenchmark",
"QdpEngine",
+ "QdpTensor",
"QuantumDataLoader",
"QuantumTensor",
+ "RustQdpEngine",
"ThroughputResult",
+ "TritonAmdEngine",
"force_backend",
+ "is_triton_amd_available",
"run_throughput_pipeline_py",
]
diff --git a/qdp/qdp-python/qumat_qdp/backend.py b/qdp/qdp-python/qumat_qdp/backend.py
new file mode 100644
index 000000000..ea889fed8
--- /dev/null
+++ b/qdp/qdp-python/qumat_qdp/backend.py
@@ -0,0 +1,95 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unified Python facade for explicit QDP backend selection."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from qumat_qdp.backends.cuda import CudaBackendEngine
+
+AmdBackendEngine: Any | None = None
+
+
+def _load_amd_backend_engine_class() -> type[Any]:
+ global AmdBackendEngine
+ if AmdBackendEngine is None:
+ from qumat_qdp.backends.amd import AmdBackendEngine as cls
+
+ AmdBackendEngine = cls
+ return AmdBackendEngine
+
+
+def _load_triton_engine_type() -> tuple[type[Any], Any]:
+ from qumat_qdp.triton_amd import is_triton_amd_available
+
+ return _load_amd_backend_engine_class(), is_triton_amd_available
+
+
+def _load_rust_cuda_engine_class() -> type[Any]:
+ return CudaBackendEngine
+
+
+def _load_rust_cuda_engine(*, device_id: int, precision: str) -> Any:
+ engine_class = _load_rust_cuda_engine_class()
+ return engine_class(device_id=device_id, precision=precision)
+
+
+def _select_engine_adapter(
+ backend: str,
+ device_id: int,
+ precision: str,
+) -> tuple[str, Any]:
+ requested_backend = backend.lower().strip()
+
+ if requested_backend == "cuda":
+ return "cuda", _load_rust_cuda_engine(
+ device_id=device_id,
+ precision=precision,
+ )
+
+ if requested_backend in frozenset({"triton_amd", "amd"}):
+ engine_class, _ = _load_triton_engine_type()
+ return "amd", engine_class(device_id=device_id, precision=precision)
+
+ raise ValueError(f"Unsupported backend '{backend}'. Use one of: amd, cuda.")
+
+
+class QdpEngine:
+ """Unified Python facade over the CUDA and Triton engine routes."""
+
+ def __init__(
+ self,
+ device_id: int = 0,
+ precision: str = "float32",
+ backend: str = "cuda",
+ ) -> None:
+ self.device_id = device_id
+ self.precision = precision
+ self.backend, self._engine_adapter = _select_engine_adapter(
+ backend=backend,
+ device_id=device_id,
+ precision=precision,
+ )
+
+ def encode(
+ self,
+ data: Any,
+ num_qubits: int,
+ encoding_method: str = "amplitude",
+ ) -> Any:
+ return self._engine_adapter.encode(data, num_qubits, encoding_method)
diff --git a/qdp/qdp-python/qumat_qdp/backends/__init__.py b/qdp/qdp-python/qumat_qdp/backends/__init__.py
new file mode 100644
index 000000000..0a61015a4
--- /dev/null
+++ b/qdp/qdp-python/qumat_qdp/backends/__init__.py
@@ -0,0 +1,22 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Backend adapters for the unified Python QDP router."""
+
+from qumat_qdp.backends.amd import AmdBackendEngine
+from qumat_qdp.backends.cuda import CudaBackendEngine
+
+__all__ = ["AmdBackendEngine", "CudaBackendEngine"]
diff --git a/qdp/qdp-python/qumat_qdp/backends/amd.py b/qdp/qdp-python/qumat_qdp/backends/amd.py
new file mode 100644
index 000000000..5697cd08a
--- /dev/null
+++ b/qdp/qdp-python/qumat_qdp/backends/amd.py
@@ -0,0 +1,21 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""AMD backend adapter backed by the Triton ROCm engine."""
+
+from qumat_qdp.triton_amd import TritonAmdEngine as AmdBackendEngine
+
+__all__ = ["AmdBackendEngine"]
diff --git a/qdp/qdp-python/qumat_qdp/backends/cuda.py b/qdp/qdp-python/qumat_qdp/backends/cuda.py
new file mode 100644
index 000000000..a938f19a9
--- /dev/null
+++ b/qdp/qdp-python/qumat_qdp/backends/cuda.py
@@ -0,0 +1,42 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""CUDA backend adapter backed by the native ``_qdp`` extension."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any
+
+
+def _load_qdp_module() -> Any | None:
+ try:
+ return import_module("_qdp")
+ except ImportError:
+ return None
+
+
+_qdp_mod = _load_qdp_module()
+
+
+class _MissingCudaBackendEngine:
+ def __init__(self, device_id: int = 0, precision: str = "float32") -> None:
+ raise RuntimeError(
+ "_qdp.QdpEngine is unavailable. Build the extension with: maturin develop"
+ )
+
+
+CudaBackendEngine = getattr(_qdp_mod, "QdpEngine", _MissingCudaBackendEngine)
diff --git a/qdp/qdp-python/qumat_qdp/tensor.py b/qdp/qdp-python/qumat_qdp/tensor.py
new file mode 100644
index 000000000..4c8b26680
--- /dev/null
+++ b/qdp/qdp-python/qumat_qdp/tensor.py
@@ -0,0 +1,56 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unified tensor facade for backend-native QDP results."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+@dataclass
+class QdpTensor:
+ """Thin DLPack facade over backend-native tensor producers."""
+
+ value: Any
+ backend: str
+
+ def __dlpack__(self, stream: int | None = None) -> Any:
+ if not hasattr(self.value, "__dlpack__"):
+ raise RuntimeError(
+ f"Backend '{self.backend}' returned object without __dlpack__ support: "
+ f"{type(self.value)!r}"
+ )
+ if stream is None:
+ return self.value.__dlpack__()
+ return self.value.__dlpack__(stream=stream)
+
+ def __dlpack_device__(self) -> Any:
+ if not hasattr(self.value, "__dlpack_device__"):
+ raise RuntimeError(
+ f"Backend '{self.backend}' returned object without __dlpack_device__ support: "
+ f"{type(self.value)!r}"
+ )
+ return self.value.__dlpack_device__()
+
+ def to_torch(self) -> Any:
+ import torch
+
+ return torch.from_dlpack(self)
+
+
+QuantumTensor = QdpTensor
diff --git a/qdp/qdp-python/qumat_qdp/triton_amd.py b/qdp/qdp-python/qumat_qdp/triton_amd.py
new file mode 100644
index 000000000..a4453b427
--- /dev/null
+++ b/qdp/qdp-python/qumat_qdp/triton_amd.py
@@ -0,0 +1,210 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Triton AMD backend for QDP encodings on ROCm."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from importlib import import_module
+from typing import Any
+
+
+def _load_optional_module(name: str) -> Any | None:
+ try:
+ return import_module(name)
+ except (
+ ImportError
+ ): # pragma: no cover - import failure is surfaced in check_runtime
+ return None
+
+
+torch_mod = _load_optional_module("torch")
+triton_mod = _load_optional_module("triton")
+
+
+def _is_rocm_runtime() -> bool:
+ if torch_mod is None:
+ return False
+ return (
+ bool(getattr(torch_mod.version, "hip", None)) and torch_mod.cuda.is_available()
+ )
+
+
+def is_triton_amd_available() -> bool:
+ if not _is_rocm_runtime() or triton_mod is None:
+ return False
+ try:
+ target = triton_mod.runtime.driver.active.get_current_target()
+ return str(getattr(target, "backend", "")).lower() == "hip"
+ except Exception:
+ return True
+
+
+@dataclass
+class TritonAmdEngine:
+ """AMD backend implementing amplitude/angle/basis encoders."""
+
+ device_id: int = 0
+ precision: str = "float32"
+
+ def __post_init__(self) -> None:
+ p = self.precision.lower()
+ if p in ("float32", "f32", "float"):
+ self.precision = "float32"
+ return
+ if p in ("float64", "f64", "double"):
+ self.precision = "float64"
+ return
+ raise ValueError(
+ f"Unsupported precision '{self.precision}'. Use float32 or float64."
+ )
+
+ def check_runtime(self) -> None:
+ if not _is_rocm_runtime():
+ raise RuntimeError(
+ "Triton AMD backend unavailable: no PyTorch ROCm device detected."
+ )
+ if triton_mod is None:
+ raise RuntimeError(
+ "Triton AMD backend unavailable: install the Triton Python package."
+ )
+
+ def _device(self) -> str:
+ return f"cuda:{self.device_id}"
+
+ def _require_torch(self) -> Any:
+ if torch_mod is None:
+ raise RuntimeError(
+ "Triton AMD backend unavailable: PyTorch is not installed."
+ )
+ return torch_mod
+
+ def _real_dtype(self) -> Any:
+ torch_mod = self._require_torch()
+ return torch_mod.float32 if self.precision == "float32" else torch_mod.float64
+
+ def _complex_dtype(self) -> Any:
+ torch_mod = self._require_torch()
+ return (
+ torch_mod.complex64 if self.precision == "float32" else torch_mod.complex128
+ )
+
+ def _to_2d(self, data: Any, *, dtype: Any) -> Any:
+ torch_mod = self._require_torch()
+ x = torch_mod.as_tensor(data, device=self._device(), dtype=dtype)
+ if x.ndim == 1:
+ x = x.unsqueeze(0)
+ if x.ndim != 2:
+ raise ValueError(f"Expected 1D or 2D input, got {x.ndim}D.")
+ return x.contiguous()
+
+ def encode_amplitude(self, data: Any, num_qubits: int) -> Any:
+ torch_mod = self._require_torch()
+ x = self._to_2d(data, dtype=self._real_dtype())
+ _, sample_size = x.shape
+ state_len = 1 << num_qubits
+ if sample_size != state_len:
+ raise ValueError(
+ f"Amplitude encoding expects sample size {state_len} (=2^num_qubits), got {sample_size}."
+ )
+
+ norms = torch_mod.linalg.vector_norm(x, dim=1, keepdim=True).clamp_min(1e-12)
+ amp = x / norms
+ return torch_mod.complex(amp, torch_mod.zeros_like(amp)).to(
+ self._complex_dtype()
+ )
+
+ def encode_angle(self, data: Any, num_qubits: int) -> Any:
+ torch_mod = self._require_torch()
+ real_dtype = self._real_dtype()
+ angles = self._to_2d(data, dtype=real_dtype)
+ batch, width = angles.shape
+ if width != num_qubits:
+ raise ValueError(
+ f"Angle encoding expects sample size {num_qubits} (=num_qubits), got {width}."
+ )
+
+ state_len = 1 << num_qubits
+ idx = torch_mod.arange(state_len, device=angles.device).reshape(1, state_len)
+ amp = torch_mod.ones((batch, state_len), device=angles.device, dtype=real_dtype)
+ for bit in range(num_qubits):
+ col = angles[:, bit].unsqueeze(1)
+ factor = torch_mod.where(
+ ((idx >> bit) & 1) == 1,
+ torch_mod.sin(col),
+ torch_mod.cos(col),
+ )
+ amp = amp * factor
+
+ return torch_mod.complex(amp, torch_mod.zeros_like(amp)).to(
+ self._complex_dtype()
+ )
+
+ def encode_basis(self, data: Any, num_qubits: int) -> Any:
+ torch_mod = self._require_torch()
+ idx = torch_mod.as_tensor(data, device=self._device(), dtype=torch_mod.int64)
+ if idx.ndim == 2:
+ if idx.shape[1] != 1:
+ raise ValueError(f"Basis 2D input expects width 1, got {idx.shape[1]}.")
+ idx = idx.squeeze(1)
+ elif idx.ndim != 1:
+ raise ValueError(f"Expected 1D or 2D basis input, got {idx.ndim}D.")
+
+ if idx.numel() == 0:
+ raise ValueError("Basis tensor cannot be empty.")
+
+ state_len = 1 << num_qubits
+ if torch_mod.any(idx < 0) or torch_mod.any(idx >= state_len):
+ raise ValueError(
+ f"Basis index out of range. Valid range is [0, {state_len - 1}]."
+ )
+
+ batch = int(idx.numel())
+ out = torch_mod.zeros(
+ (batch, state_len),
+ device=idx.device,
+ dtype=self._complex_dtype(),
+ )
+ out.scatter_(
+ 1,
+ idx.reshape(batch, 1),
+ torch_mod.ones(
+ (batch, 1),
+ device=idx.device,
+ dtype=self._complex_dtype(),
+ ),
+ )
+ return out
+
+ def encode(
+ self,
+ data: Any,
+ num_qubits: int,
+ encoding_method: str = "amplitude",
+ ) -> Any:
+ self.check_runtime()
+
+ method = encoding_method.lower()
+ if method == "amplitude":
+ return self.encode_amplitude(data, num_qubits)
+ if method == "angle":
+ return self.encode_angle(data, num_qubits)
+ if method == "basis":
+ return self.encode_basis(data, num_qubits)
+ raise ValueError(
+ f"Unsupported encoding '{encoding_method}'. triton_amd supports amplitude, angle, basis."
+ )
diff --git a/qdp/qdp-python/src/dlpack.rs b/qdp/qdp-python/src/dlpack.rs
index 4a1ad08ec..27a4b8154 100644
--- a/qdp/qdp-python/src/dlpack.rs
+++ b/qdp/qdp-python/src/dlpack.rs
@@ -20,6 +20,9 @@ use pyo3::prelude::*;
use qdp_core::dlpack::{DL_FLOAT, DLDeviceType, DLManagedTensor};
use std::ffi::c_void;
+const DLTENSOR_NAME: &[u8] = b"dltensor\0";
+const USED_DLTENSOR_NAME: &[u8] = b"used_dltensor\0";
+
/// DLPack tensor information extracted from a PyCapsule
///
/// This struct owns the DLManagedTensor pointer and ensures proper cleanup
@@ -70,28 +73,7 @@ pub fn extract_dlpack_tensor(
// Call tensor.__dlpack__() to get PyCapsule
// Note: PyTorch's __dlpack__ uses the default stream when called without arguments
let capsule = tensor.call_method0("__dlpack__")?;
-
- const DLTENSOR_NAME: &[u8] = b"dltensor\0";
-
- // SAFETY: capsule is a valid PyCapsule from tensor.__dlpack__(). DLTENSOR_NAME is a
- // null-terminated C string for the lifetime of the call. We only read the capsule
- // and call PyCapsule_IsValid / PyCapsule_GetPointer; we do not invalidate the capsule.
- let managed_ptr = unsafe {
- let capsule_ptr = capsule.as_ptr();
- if ffi::PyCapsule_IsValid(capsule_ptr, DLTENSOR_NAME.as_ptr() as *const i8) == 0 {
- return Err(PyRuntimeError::new_err(
- "Invalid DLPack capsule (expected 'dltensor')",
- ));
- }
- let ptr = ffi::PyCapsule_GetPointer(capsule_ptr, DLTENSOR_NAME.as_ptr() as *const i8)
- as *mut DLManagedTensor;
- if ptr.is_null() {
- return Err(PyRuntimeError::new_err(
- "Failed to extract DLManagedTensor from PyCapsule",
- ));
- }
- ptr
- };
+ let managed_ptr = unsafe { extract_managed_tensor(capsule.as_ptr())? };
// SAFETY: managed_ptr is non-null and was returned by PyCapsule_GetPointer for a valid
// "dltensor" capsule, so it points to a valid DLManagedTensor. The capsule (and thus
@@ -106,9 +88,11 @@ pub fn extract_dlpack_tensor(
));
}
- if dl_tensor.device.device_type != DLDeviceType::kDLCUDA {
+ if dl_tensor.device.device_type != DLDeviceType::kDLCUDA
+ && dl_tensor.device.device_type != DLDeviceType::kDLROCM
+ {
return Err(PyRuntimeError::new_err(
- "DLPack tensor must be on CUDA device",
+ "DLPack tensor must be on CUDA/ROCm device",
));
}
@@ -185,10 +169,7 @@ pub fn extract_dlpack_tensor(
}
let device_id = dl_tensor.device.device_id;
-
- const USED_DLTENSOR_NAME: &[u8] = b"used_dltensor\0";
- // SAFETY: capsule is the same PyCapsule we used above; renaming is allowed and does not free it.
- ffi::PyCapsule_SetName(capsule.as_ptr(), USED_DLTENSOR_NAME.as_ptr() as *const i8);
+ mark_capsule_used(capsule.as_ptr());
Ok(DLPackTensorInfo {
managed_ptr,
@@ -198,3 +179,44 @@ pub fn extract_dlpack_tensor(
})
}
}
+
+unsafe fn extract_managed_tensor(
+ capsule_ptr: *mut ffi::PyObject,
+) -> PyResult<*mut DLManagedTensor> {
+ if unsafe { ffi::PyCapsule_IsValid(capsule_ptr, DLTENSOR_NAME.as_ptr() as *const i8) } == 0 {
+ return Err(PyRuntimeError::new_err(
+ "Invalid DLPack capsule (expected 'dltensor')",
+ ));
+ }
+ let managed_ptr =
+ unsafe { ffi::PyCapsule_GetPointer(capsule_ptr, DLTENSOR_NAME.as_ptr() as *const i8) }
+ as *mut DLManagedTensor;
+ if managed_ptr.is_null() {
+ return Err(PyRuntimeError::new_err(
+ "Failed to extract DLManagedTensor from PyCapsule",
+ ));
+ }
+ unsafe { validate_gpu_managed_tensor(managed_ptr)? };
+ Ok(managed_ptr)
+}
+
+unsafe fn validate_gpu_managed_tensor(managed_ptr: *mut DLManagedTensor) -> PyResult<()> {
+ let dl_tensor = unsafe { &(*managed_ptr).dl_tensor };
+ if dl_tensor.data.is_null() {
+ return Err(PyRuntimeError::new_err(
+ "DLPack tensor has null data pointer",
+ ));
+ }
+ if dl_tensor.device.device_type != DLDeviceType::kDLCUDA
+ && dl_tensor.device.device_type != DLDeviceType::kDLROCM
+ {
+ return Err(PyRuntimeError::new_err(
+ "DLPack tensor must be on CUDA/ROCm device",
+ ));
+ }
+ Ok(())
+}
+
+unsafe fn mark_capsule_used(capsule_ptr: *mut ffi::PyObject) {
+ unsafe { ffi::PyCapsule_SetName(capsule_ptr, USED_DLTENSOR_NAME.as_ptr() as *const i8) };
+}
diff --git a/qdp/qdp-python/src/engine.rs b/qdp/qdp-python/src/engine.rs
index 9d6a2a6ae..f0a1a7b6c 100644
--- a/qdp/qdp-python/src/engine.rs
+++ b/qdp/qdp-python/src/engine.rs
@@ -28,12 +28,36 @@ use qdp_core::{Precision, QdpEngine as CoreEngine};
#[cfg(target_os = "linux")]
use crate::loader::{PyQuantumLoader, config_from_args, parse_null_handling, path_from_py};
-/// PyO3 wrapper for QdpEngine
+struct CudaEngineAdapter {
+ engine: CoreEngine,
+}
+
+impl CudaEngineAdapter {
+ fn new(device_id: usize, precision: Precision) -> PyResult<Self> {
+ let engine = CoreEngine::new_with_precision(device_id, precision).map_err(|e| {
+ PyRuntimeError::new_err(format!("Failed to initialize CUDA backend: {}", e))
+ })?;
+ Ok(Self { engine })
+ }
+
+ fn engine(&self) -> &CoreEngine {
+ &self.engine
+ }
+}
+
+/// PyO3 wrapper for the Rust/CUDA QdpEngine.
///
-/// Provides Python bindings for GPU-accelerated quantum state encoding.
+/// The public Python facade routes AMD/Triton directly in `qumat_qdp.backend`.
+/// `_qdp.QdpEngine` stays focused on the Rust CUDA core and its tensor contract.
+
#[pyclass]
pub struct QdpEngine {
- pub engine: CoreEngine,
+ engine: CudaEngineAdapter,
+ #[allow(dead_code)]
+ device_id: usize,
+ #[allow(dead_code)]
+ precision: Precision,
+ backend: String,
}
#[pymethods]
@@ -50,8 +74,8 @@ impl QdpEngine {
/// Raises:
/// RuntimeError: If CUDA device initialization fails
#[new]
- #[pyo3(signature = (device_id=0, precision="float32"))]
- fn new(device_id: usize, precision: &str) -> PyResult<Self> {
+ #[pyo3(signature = (device_id=0, precision="float32", backend="cuda"))]
+ fn new(device_id: usize, precision: &str, backend: &str) -> PyResult<Self> {
let precision = match precision.to_ascii_lowercase().as_str() {
"float32" | "f32" | "float" => Precision::Float32,
"float64" | "f64" | "double" => Precision::Float64,
@@ -63,9 +87,31 @@ impl QdpEngine {
}
};
- let engine = CoreEngine::new_with_precision(device_id, precision)
- .map_err(|e| PyRuntimeError::new_err(format!("Failed to initialize: {}", e)))?;
- Ok(Self { engine })
+ let backend_name = backend.to_ascii_lowercase();
+ let (engine, resolved_backend) = match backend_name.as_str() {
+ "cuda" => (
+ CudaEngineAdapter::new(device_id, precision)?,
+ "cuda".to_string(),
+ ),
+ "amd" | "triton_amd" => {
+ return Err(PyRuntimeError::new_err(
+ "AMD/Triton routing is provided by the Python facade `qumat_qdp.QdpEngine`; `_qdp.QdpEngine` only supports the Rust CUDA backend.",
+ ));
+ }
+ other => {
+ return Err(PyRuntimeError::new_err(format!(
+ "Unsupported backend '{}'. Use 'cuda'.",
+ other
+ )));
+ }
+ };
+
+ Ok(Self {
+ engine,
+ device_id,
+ precision,
+ backend: resolved_backend,
+ })
}
/// Encode classical data into quantum state (auto-detects input type)
@@ -105,6 +151,94 @@ impl QdpEngine {
data: &Bound<'_, PyAny>,
num_qubits: usize,
encoding_method: &str,
+ ) -> PyResult<QuantumTensor> {
+ self.encode_with_core(data, num_qubits, encoding_method)
+ }
+
+ fn backend(&self) -> &str {
+ &self.backend
+ }
+
+ #[cfg(target_os = "linux")]
+ #[pyo3(signature = (total_batches, batch_size, num_qubits, encoding_method, seed=None, null_handling=None))]
+ fn create_synthetic_loader(
+ &self,
+ total_batches: usize,
+ batch_size: usize,
+ num_qubits: u32,
+ encoding_method: &str,
+ seed: Option<u64>,
+ null_handling: Option<&str>,
+ ) -> PyResult<PyQuantumLoader> {
+ self.create_synthetic_loader_impl(
+ total_batches,
+ batch_size,
+ num_qubits,
+ encoding_method,
+ seed,
+ null_handling,
+ )
+ }
+
+ #[cfg(target_os = "linux")]
+ #[allow(clippy::too_many_arguments)]
+ #[pyo3(signature = (path, batch_size, num_qubits, encoding_method, batch_limit=None, null_handling=None))]
+ fn create_file_loader(
+ &self,
+ py: Python<'_>,
+ path: &Bound<'_, PyAny>,
+ batch_size: usize,
+ num_qubits: u32,
+ encoding_method: &str,
+ batch_limit: Option<usize>,
+ null_handling: Option<&str>,
+ ) -> PyResult<PyQuantumLoader> {
+ self.create_file_loader_impl(
+ py,
+ path,
+ batch_size,
+ num_qubits,
+ encoding_method,
+ batch_limit,
+ null_handling,
+ )
+ }
+
+ #[cfg(target_os = "linux")]
+ #[allow(clippy::too_many_arguments)]
+ #[pyo3(signature = (path, batch_size, num_qubits, encoding_method, batch_limit=None, null_handling=None))]
+ fn create_streaming_file_loader(
+ &self,
+ py: Python<'_>,
+ path: &Bound<'_, PyAny>,
+ batch_size: usize,
+ num_qubits: u32,
+ encoding_method: &str,
+ batch_limit: Option<usize>,
+ null_handling: Option<&str>,
+ ) -> PyResult<PyQuantumLoader> {
+ self.create_streaming_file_loader_impl(
+ py,
+ path,
+ batch_size,
+ num_qubits,
+ encoding_method,
+ batch_limit,
+ null_handling,
+ )
+ }
+}
+
+impl QdpEngine {
+ fn core_engine(&self) -> PyResult<&CoreEngine> {
+ Ok(self.engine.engine())
+ }
+
+ fn encode_with_core(
+ &self,
+ data: &Bound<'_, PyAny>,
+ num_qubits: usize,
+ encoding_method: &str,
) -> PyResult<QuantumTensor> {
// Check if it's a string path
if let Ok(path) = data.extract::<String>() {
@@ -158,7 +292,7 @@ impl QdpEngine {
PyRuntimeError::new_err("NumPy array must be contiguous (C-order)")
})?;
let ptr = self
- .engine
+ .core_engine()?
.encode(data_slice, num_qubits, encoding_method)
.map_err(|e| PyRuntimeError::new_err(format!("Encoding failed: {}", e)))?;
Ok(QuantumTensor {
@@ -180,7 +314,7 @@ impl QdpEngine {
PyRuntimeError::new_err("NumPy array must be contiguous (C-order)")
})?;
let ptr = self
- .engine
+ .core_engine()?
.encode_batch(
data_slice,
num_samples,
@@ -210,7 +344,7 @@ impl QdpEngine {
// Validate CUDA tensor for direct GPU encoding
validate_cuda_tensor_for_encoding(
data,
- self.engine.device().ordinal(),
+ self.core_engine()?.device().ordinal(),
encoding_method,
)?;
@@ -239,7 +373,7 @@ impl QdpEngine {
// (held by Python's GIL), and we validated dtype/contiguity/device above.
// The DLPackTensorInfo RAII wrapper will call deleter when dropped.
let ptr = unsafe {
- self.engine
+ self.core_engine()?
.encode_from_gpu_ptr(
dlpack_info.data_ptr,
input_len,
@@ -261,7 +395,7 @@ impl QdpEngine {
let sample_size = dlpack_info.shape[1] as usize;
// SAFETY: Same as above - pointer from validated DLPack tensor
let ptr = unsafe {
- self.engine
+ self.core_engine()?
.encode_batch_from_gpu_ptr(
dlpack_info.data_ptr,
num_samples,
@@ -317,7 +451,7 @@ impl QdpEngine {
)
})?;
let ptr = self
- .engine
+ .core_engine()?
.encode(data_slice, num_qubits, encoding_method)
.map_err(|e| PyRuntimeError::new_err(format!("Encoding failed: {}", e)))?;
Ok(QuantumTensor {
@@ -343,7 +477,7 @@ impl QdpEngine {
)
})?;
let ptr = self
- .engine
+ .core_engine()?
.encode_batch(
data_slice,
num_samples,
@@ -374,7 +508,7 @@ impl QdpEngine {
)
})?;
let ptr = self
- .engine
+ .core_engine()?
.encode(&vec_data, num_qubits, encoding_method)
.map_err(|e| PyRuntimeError::new_err(format!("Encoding failed: {}", e)))?;
Ok(QuantumTensor {
@@ -405,31 +539,31 @@ impl QdpEngine {
};
let ptr = if path.ends_with(".parquet") {
- self.engine
+ self.core_engine()?
.encode_from_parquet(path, num_qubits, encoding_method)
.map_err(|e| {
PyRuntimeError::new_err(format!("Encoding from parquet failed: {}", e))
})?
} else if path.ends_with(".arrow") || path.ends_with(".feather") {
- self.engine
+ self.core_engine()?
.encode_from_arrow_ipc(path, num_qubits, encoding_method)
.map_err(|e| {
PyRuntimeError::new_err(format!("Encoding from Arrow IPC failed: {}", e))
})?
} else if path.ends_with(".npy") {
- self.engine
+ self.core_engine()?
.encode_from_numpy(path, num_qubits, encoding_method)
.map_err(|e| {
PyRuntimeError::new_err(format!("Encoding from NumPy failed: {}", e))
})?
} else if path.ends_with(".pt") || path.ends_with(".pth") {
- self.engine
+ self.core_engine()?
.encode_from_torch(path, num_qubits, encoding_method)
.map_err(|e| {
PyRuntimeError::new_err(format!("Encoding from PyTorch failed: {}", e))
})?
} else if path.ends_with(".pb") {
- self.engine
+ self.core_engine()?
.encode_from_tensorflow(path, num_qubits, encoding_method)
.map_err(|e| {
PyRuntimeError::new_err(format!("Encoding from TensorFlow failed: {}", e))
@@ -461,6 +595,7 @@ impl QdpEngine {
/// >>> engine = QdpEngine(device_id=0)
/// >>> batched = engine.encode_from_tensorflow("data.pb", 16, "amplitude")
/// >>> torch_tensor = torch.from_dlpack(batched) # Shape: [200, 65536]
+ #[allow(dead_code)]
fn encode_from_tensorflow(
&self,
path: &str,
@@ -468,7 +603,7 @@ impl QdpEngine {
encoding_method: &str,
) -> PyResult<QuantumTensor> {
let ptr = self
- .engine
+ .core_engine()?
.encode_from_tensorflow(path, num_qubits, encoding_method)
.map_err(|e| {
PyRuntimeError::new_err(format!("Encoding from TensorFlow failed: {}", e))
@@ -489,7 +624,11 @@ impl QdpEngine {
num_qubits: usize,
encoding_method: &str,
) -> PyResult<QuantumTensor> {
- validate_cuda_tensor_for_encoding(data, self.engine.device().ordinal(), encoding_method)?;
+ validate_cuda_tensor_for_encoding(
+ data,
+ self.core_engine()?.device().ordinal(),
+ encoding_method,
+ )?;
let dtype = data.getattr("dtype")?;
let dtype_str: String = dtype.str()?.extract()?;
@@ -510,7 +649,7 @@ impl QdpEngine {
let ptr = unsafe {
match method.as_str() {
"amplitude" => self
- .engine
+ .core_engine()?
.encode_from_gpu_ptr_f32_with_stream(
data_ptr, input_len, num_qubits, stream_ptr,
)
@@ -521,7 +660,7 @@ impl QdpEngine {
))
})?,
"angle" => self
- .engine
+ .core_engine()?
.encode_angle_from_gpu_ptr_f32_with_stream(
data_ptr, input_len, num_qubits, stream_ptr,
)
@@ -548,7 +687,7 @@ impl QdpEngine {
let data_ptr = data_ptr_u64 as *const f32;
let ptr = unsafe {
- self.engine
+ self.core_engine()?
.encode_batch_from_gpu_ptr_f32_with_stream(
data_ptr,
num_samples,
@@ -582,7 +721,7 @@ impl QdpEngine {
1 => {
let input_len = tensor_info.shape[0] as usize;
let ptr = unsafe {
- self.engine
+ self.core_engine()?
.encode_from_gpu_ptr_with_stream(
tensor_info.data_ptr as *const std::ffi::c_void,
input_len,
@@ -603,7 +742,7 @@ impl QdpEngine {
let num_samples = tensor_info.shape[0] as usize;
let sample_size = tensor_info.shape[1] as usize;
let ptr = unsafe {
- self.engine
+ self.core_engine()?
.encode_batch_from_gpu_ptr_with_stream(
tensor_info.data_ptr as *const std::ffi::c_void,
num_samples,
@@ -632,9 +771,7 @@ impl QdpEngine {
// --- Loader factory methods (Linux only) ---
#[cfg(target_os = "linux")]
- /// Create a synthetic-data pipeline iterator (for QuantumDataLoader.source_synthetic()).
- #[pyo3(signature = (total_batches, batch_size, num_qubits, encoding_method, seed=None, null_handling=None))]
- fn create_synthetic_loader(
+ fn create_synthetic_loader_impl(
&self,
total_batches: usize,
batch_size: usize,
@@ -643,9 +780,10 @@ impl QdpEngine {
seed: Option<u64>,
null_handling: Option<&str>,
) -> PyResult<PyQuantumLoader> {
+ let engine = self.core_engine()?.clone();
let nh = parse_null_handling(null_handling)?;
let config = config_from_args(
- &self.engine,
+ &engine,
batch_size,
num_qubits,
encoding_method,
@@ -654,17 +792,15 @@ impl QdpEngine {
nh,
true,
);
- let iter = qdp_core::PipelineIterator::new_synthetic(self.engine.clone(), config).map_err(
- |e| PyRuntimeError::new_err(format!("create_synthetic_loader failed: {}", e)),
- )?;
+ let iter = qdp_core::PipelineIterator::new_synthetic(engine, config).map_err(|e| {
+ PyRuntimeError::new_err(format!("create_synthetic_loader failed: {}", e))
+ })?;
Ok(PyQuantumLoader::new(Some(iter)))
}
#[cfg(target_os = "linux")]
- /// Create a file-backed pipeline iterator (full read then batch; for QuantumDataLoader.source_file(path)).
#[allow(clippy::too_many_arguments)]
- #[pyo3(signature = (path, batch_size, num_qubits, encoding_method, batch_limit=None, null_handling=None))]
- fn create_file_loader(
+ fn create_file_loader_impl(
&self,
py: Python<'_>,
path: &Bound<'_, PyAny>,
@@ -676,9 +812,10 @@ impl QdpEngine {
) -> PyResult<PyQuantumLoader> {
let path_str = path_from_py(path)?;
let batch_limit = batch_limit.unwrap_or(usize::MAX);
+ let engine = self.core_engine()?.clone();
let nh = parse_null_handling(null_handling)?;
let config = config_from_args(
- &self.engine,
+ &engine,
batch_size,
num_qubits,
encoding_method,
@@ -687,7 +824,6 @@ impl QdpEngine {
nh,
true, // float32_pipeline
);
- let engine = self.engine.clone();
// Resolve remote URLs before detaching from GIL. The _resolved guard keeps the
// temp file alive until after the file is fully read inside py.detach.
#[cfg(feature = "remote-io")]
@@ -710,10 +846,8 @@ impl QdpEngine {
}
#[cfg(target_os = "linux")]
- /// Create a streaming Parquet pipeline iterator (for QuantumDataLoader.source_file(path, streaming=True)).
#[allow(clippy::too_many_arguments)]
- #[pyo3(signature = (path, batch_size, num_qubits, encoding_method, batch_limit=None, null_handling=None))]
- fn create_streaming_file_loader(
+ fn create_streaming_file_loader_impl(
&self,
py: Python<'_>,
path: &Bound<'_, PyAny>,
@@ -725,9 +859,10 @@ impl QdpEngine {
) -> PyResult<PyQuantumLoader> {
let path_str = path_from_py(path)?;
let batch_limit = batch_limit.unwrap_or(usize::MAX);
+ let engine = self.core_engine()?.clone();
let nh = parse_null_handling(null_handling)?;
let config = config_from_args(
- &self.engine,
+ &engine,
batch_size,
num_qubits,
encoding_method,
@@ -736,7 +871,6 @@ impl QdpEngine {
nh,
true, // float32_pipeline
);
- let engine = self.engine.clone();
// Resolve remote URLs before detaching from GIL. The _resolved guard keeps the
// temp file alive; the streaming reader's open fd preserves data after drop.
#[cfg(feature = "remote-io")]
diff --git a/qdp/qdp-python/src/tensor.rs b/qdp/qdp-python/src/tensor.rs
index ab341d1ac..01fa0c713 100644
--- a/qdp/qdp-python/src/tensor.rs
+++ b/qdp/qdp-python/src/tensor.rs
@@ -101,7 +101,7 @@ impl QuantumTensor {
/// Returns DLPack device information
///
/// Returns:
- /// Tuple of (device_type, device_id) where device_type=2 for CUDA
+ /// Tuple of (device_type, device_id) where device_type follows DLPack constants
fn __dlpack_device__(&self) -> PyResult<(i32, i32)> {
if self.ptr.is_null() {
return Err(PyRuntimeError::new_err("Invalid DLPack tensor pointer"));
@@ -109,9 +109,10 @@ impl QuantumTensor {
unsafe {
let tensor = &(*self.ptr).dl_tensor;
- // DLPack device_type: kDLCUDA = 2, kDLCPU = 1
+ // DLPack device_type: kDLCUDA = 2, kDLROCM = 10, kDLCPU = 1
let device_type = match tensor.device.device_type {
qdp_core::dlpack::DLDeviceType::kDLCUDA => 2,
+ qdp_core::dlpack::DLDeviceType::kDLROCM => 10,
qdp_core::dlpack::DLDeviceType::kDLCPU => 1,
};
// Read device_id from DLPack tensor metadata
diff --git a/qdp/qdp-python/tests/test_amd_engine.py b/qdp/qdp-python/tests/test_amd_engine.py
new file mode 100644
index 000000000..c93b3d401
--- /dev/null
+++ b/qdp/qdp-python/tests/test_amd_engine.py
@@ -0,0 +1,52 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+torch = pytest.importorskip("torch")
+_qdp = pytest.importorskip("_qdp")
+
+
+def _rocm_available() -> bool:
+ return bool(getattr(torch.version, "hip", None)) and torch.cuda.is_available()
+
+
+@pytest.mark.skipif(
+ not hasattr(_qdp, "AmdQdpEngine"), reason="AmdQdpEngine not available"
+)
+@pytest.mark.skipif(not _rocm_available(), reason="ROCm runtime not available")
+def test_amd_engine_amplitude_shape_and_dtype() -> None:
+ engine = _qdp.AmdQdpEngine(0, precision="float32")
+ x = torch.randn(2, 4, device="cuda", dtype=torch.float32)
+ qt = engine.encode(x, 2, "amplitude")
+ out = torch.from_dlpack(qt)
+ assert out.shape == (2, 4)
+ assert out.dtype == torch.complex64
+
+
+@pytest.mark.skipif(
+ not hasattr(_qdp, "AmdQdpEngine"), reason="AmdQdpEngine not available"
+)
+@pytest.mark.skipif(not _rocm_available(), reason="ROCm runtime not available")
+def test_amd_engine_basis_shape_and_values() -> None:
+ engine = _qdp.AmdQdpEngine(0, precision="float64")
+ idx = torch.tensor([0, 3], device="cuda", dtype=torch.int64)
+ qt = engine.encode(idx, 2, "basis")
+ out = torch.from_dlpack(qt)
+ assert out.shape == (2, 4)
+ assert out.dtype == torch.complex128
+ assert out[0, 0].real.item() == pytest.approx(1.0)
+ assert out[1, 3].real.item() == pytest.approx(1.0)
diff --git a/qdp/qdp-python/tests/test_backend_routing.py b/qdp/qdp-python/tests/test_backend_routing.py
new file mode 100644
index 000000000..4ae32899a
--- /dev/null
+++ b/qdp/qdp-python/tests/test_backend_routing.py
@@ -0,0 +1,129 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import pytest
+import qumat_qdp
+from qumat_qdp import QdpEngine
+from qumat_qdp import backend as backend_mod
+
+
+def test_backend_routing_rejects_unknown_backend() -> None:
+ with pytest.raises(ValueError):
+ QdpEngine(backend="unknown-backend")
+
+
+def test_create_encoder_engine_instantiates_python_facade(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ class FakeAdapter:
+ def encode(self) -> str:
+ return "ok"
+
+ monkeypatch.setattr(
+ backend_mod,
+ "_select_engine_adapter",
+ lambda *args, **kwargs: ("fake", FakeAdapter()),
+ )
+
+ engine = QdpEngine(backend="cuda")
+ assert isinstance(engine, QdpEngine)
+ assert engine.backend == "fake"
+
+
+def test_cuda_route_uses_rust_engine(monkeypatch: pytest.MonkeyPatch) -> None:
+ seen: dict[str, Any] = {}
+
+ class FakeRustEngine:
+ def __init__(self, device_id: int = 0, precision: str = "float32") -> None:
+ seen["init"] = {"device_id": device_id, "precision": precision}
+
+ def encode(self, data: Any, num_qubits: int, encoding_method: str) -> str:
+ seen["encode"] = {
+ "data": data,
+ "num_qubits": num_qubits,
+ "encoding_method": encoding_method,
+ }
+ return "rust-quantum-tensor"
+
+ monkeypatch.setattr(
+ backend_mod, "_load_rust_cuda_engine_class", lambda: FakeRustEngine
+ )
+
+ engine = QdpEngine(device_id=3, precision="float64", backend="cuda")
+ value = engine.encode([1.0, 0.0], 1, "amplitude")
+
+ assert engine.backend == "cuda"
+ assert value == "rust-quantum-tensor"
+ assert seen["init"] == {"device_id": 3, "precision": "float64"}
+ assert seen["encode"]["encoding_method"] == "amplitude"
+
+
+def test_amd_route_returns_triton_output_directly(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ class FakeValue:
+ def __dlpack__(self, stream: Any | None = None) -> Any:
+ return ("capsule", stream)
+
+ def __dlpack_device__(self) -> Any:
+ return (2, 0)
+
+ class FakeTritonEngine:
+ def __init__(self, device_id: int = 0, precision: str = "float32") -> None:
+ self.device_id = device_id
+ self.precision = precision
+
+ def check_runtime(self) -> None:
+ return None
+
+ def encode(self, data: Any, num_qubits: int, encoding_method: str) -> Any:
+ return FakeValue()
+
+ monkeypatch.setattr(
+ backend_mod,
+ "_load_triton_engine_type",
+ lambda: (FakeTritonEngine, lambda: True),
+ )
+
+ engine = QdpEngine(device_id=1, precision="float32", backend="amd")
+ qt = engine.encode([1.0, 0.0], 1, "amplitude")
+
+ assert engine.backend == "amd"
+ assert isinstance(qt, FakeValue)
+ assert qt.__dlpack_device__() == (2, 0)
+
+
+def test_missing_cuda_backend_fails_cleanly(monkeypatch: pytest.MonkeyPatch) -> None:
+ class MissingCudaBackendEngine:
+ def __init__(self, device_id: int, precision: str) -> None:
+ raise RuntimeError("_qdp.QdpEngine is unavailable.")
+
+ monkeypatch.setattr(backend_mod, "CudaBackendEngine", MissingCudaBackendEngine)
+
+ with pytest.raises(RuntimeError, match="_qdp.QdpEngine is unavailable"):
+ QdpEngine(backend="cuda")
+
+
+def test_quantum_tensor_wrapper_requires_dlpack() -> None:
+ with pytest.raises(RuntimeError):
+ qumat_qdp.QdpTensor(value=object(), backend="fake").__dlpack__()
+
+
+def test_package_exports_quantum_tensor_contract() -> None:
+ assert hasattr(qumat_qdp, "QuantumTensor")
+ assert hasattr(qumat_qdp, "QdpTensor")
diff --git a/qdp/qdp-python/tests/test_triton_amd_backend.py b/qdp/qdp-python/tests/test_triton_amd_backend.py
new file mode 100644
index 000000000..74730999b
--- /dev/null
+++ b/qdp/qdp-python/tests/test_triton_amd_backend.py
@@ -0,0 +1,160 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from qumat_qdp import QdpEngine, is_triton_amd_available
+from qumat_qdp.triton_amd import TritonAmdEngine
+
+
+def _as_torch(value):
+ if isinstance(value, torch.Tensor):
+ return value
+ return torch.from_dlpack(value)
+
+
+def _torch_amplitude_ref(x: torch.Tensor) -> torch.Tensor:
+ x = x.to(torch.float32)
+ norms = torch.linalg.vector_norm(x, dim=1, keepdim=True).clamp_min(1e-12)
+ y = x / norms
+ return torch.complex(y, torch.zeros_like(y))
+
+
+def _torch_angle_ref(angles: torch.Tensor, num_qubits: int) -> torch.Tensor:
+ real_dtype = angles.dtype
+ batch = angles.shape[0]
+ state_len = 1 << num_qubits
+ idx = torch.arange(state_len, device=angles.device).reshape(1, state_len)
+ amp = torch.ones((batch, state_len), device=angles.device, dtype=real_dtype)
+ for bit in range(num_qubits):
+ col = angles[:, bit].unsqueeze(1)
+ factor = torch.where(
+ ((idx >> bit) & 1) == 1,
+ torch.sin(col),
+ torch.cos(col),
+ )
+ amp = amp * factor
+ return torch.complex(amp, torch.zeros_like(amp))
+
+
+def _torch_basis_ref(idx: torch.Tensor, num_qubits: int) -> torch.Tensor:
+ idx = idx.to(torch.int64)
+ batch = idx.numel()
+ state_len = 1 << num_qubits
+ out = torch.zeros((batch, state_len), device=idx.device, dtype=torch.complex64)
+ out.scatter_(
+ 1,
+ idx.reshape(batch, 1),
+ torch.ones((batch, 1), device=idx.device, dtype=torch.complex64),
+ )
+ return out
+
+
+@pytest.mark.skipif(
+ not is_triton_amd_available(), reason="Triton AMD backend unavailable"
+)
+@pytest.mark.rocm
+def test_triton_amd_amplitude_parity() -> None:
+ engine = TritonAmdEngine(device_id=0, precision="float32")
+ x = torch.randn(4, 8, device="cuda", dtype=torch.float32)
+ got = _as_torch(engine.encode(x, 3, "amplitude"))
+ ref = _torch_amplitude_ref(x)
+ assert got.shape == ref.shape
+ assert got.dtype == torch.complex64
+ assert torch.allclose(got, ref, atol=1e-5, rtol=1e-5)
+
+
+@pytest.mark.skipif(
+ not is_triton_amd_available(), reason="Triton AMD backend unavailable"
+)
+@pytest.mark.rocm
+def test_triton_amd_angle_parity() -> None:
+ engine = TritonAmdEngine(device_id=0, precision="float32")
+ angles = torch.randn(3, 5, device="cuda", dtype=torch.float32)
+ got = _as_torch(engine.encode(angles, 5, "angle"))
+ ref = _torch_angle_ref(angles, 5)
+ assert got.shape == ref.shape
+ assert torch.allclose(got, ref, atol=2e-4, rtol=2e-4)
+
+
+@pytest.mark.skipif(
+ not is_triton_amd_available(), reason="Triton AMD backend unavailable"
+)
+@pytest.mark.rocm
+def test_triton_amd_basis_parity() -> None:
+ engine = TritonAmdEngine(device_id=0, precision="float32")
+ idx = torch.tensor([0, 3, 7], device="cuda", dtype=torch.int64)
+ got = _as_torch(engine.encode(idx, 3, "basis"))
+ ref = _torch_basis_ref(idx, 3)
+ assert got.shape == ref.shape
+ assert torch.equal(got, ref)
+
+
+@pytest.mark.skipif(
+ not is_triton_amd_available(), reason="Triton AMD backend unavailable"
+)
+@pytest.mark.rocm
+def test_triton_amd_angle_float64_precision_contract() -> None:
+ engine = TritonAmdEngine(device_id=0, precision="float64")
+ angles = torch.randn(2, 4, device="cuda", dtype=torch.float64)
+ got = _as_torch(engine.encode(angles, 4, "angle"))
+ ref = _torch_angle_ref(angles.to(torch.float64), 4).to(torch.complex128)
+ assert got.dtype == torch.complex128
+ assert torch.allclose(got, ref, atol=1e-10, rtol=1e-10)
+
+
+@pytest.mark.skipif(
+ not torch.cuda.is_available() or getattr(torch.version, "cuda", None) is None,
+ reason="NVIDIA CUDA reference not available",
+)
+@pytest.mark.rocm
+def test_triton_amd_cuda_reference_optional() -> None:
+ _qdp = pytest.importorskip("_qdp")
+ if not is_triton_amd_available():
+ pytest.skip("Triton AMD backend unavailable")
+
+ engine_triton = TritonAmdEngine(device_id=0, precision="float32")
+ engine_cuda = _qdp.QdpEngine(0, precision="float32")
+ x = torch.randn(2, 4, device="cuda", dtype=torch.float32)
+ got = _as_torch(engine_triton.encode(x, 2, "amplitude"))
+ ref = torch.from_dlpack(engine_cuda.encode(x, 2, "amplitude"))
+ assert torch.allclose(got, ref, atol=1e-4, rtol=1e-4)
+
+
+@pytest.mark.skipif(
+ not is_triton_amd_available(), reason="Triton AMD backend unavailable"
+)
+@pytest.mark.rocm
+def test_triton_amd_direct_engine_returns_dlpack_contract() -> None:
+ engine = TritonAmdEngine(device_id=0, precision="float32")
+ x = torch.randn(2, 4, device="cuda", dtype=torch.float32)
+ qt = engine.encode(x, 2, "amplitude")
+ assert hasattr(qt, "__dlpack__")
+ out = torch.from_dlpack(qt)
+ assert out.shape == (2, 4)
+
+
+@pytest.mark.skipif(
+ not is_triton_amd_available(), reason="Triton AMD backend unavailable"
+)
+@pytest.mark.rocm
+def test_unified_router_contract_returns_torch_tensor() -> None:
+ router = QdpEngine(backend="amd", device_id=0, precision="float32")
+ x = torch.randn(2, 4, device="cuda", dtype=torch.float32)
+ qt = router.encode(x, 2, "amplitude")
+ assert isinstance(qt, torch.Tensor)
+ assert qt.shape == (2, 4)
+ assert qt.dtype == torch.complex64
diff --git a/qdp/qdp-python/uv.lock b/qdp/qdp-python/uv.lock
index 50110588d..cb5ed4cf5 100644
--- a/qdp/qdp-python/uv.lock
+++ b/qdp/qdp-python/uv.lock
@@ -1467,10 +1467,12 @@ benchmark = [
{ name = "tensorflow" },
{ name = "torch" },
{ name = "tqdm" },
+ { name = "triton" },
]
dev = [
{ name = "pytest" },
{ name = "torch" },
+ { name = "triton" },
]
[package.metadata]
@@ -1491,10 +1493,12 @@ benchmark = [
{ name = "tensorflow", specifier = ">=2.20" },
{ name = "torch", specifier = ">=2.2,<=2.9.0" },
{ name = "tqdm" },
+ { name = "triton", specifier = ">=3.0" },
]
dev = [
{ name = "pytest" },
{ name = "torch", specifier = ">=2.2,<=2.9.0" },
+ { name = "triton", specifier = ">=3.0" },
]
[[package]]
diff --git a/qumat/qdp.py b/qumat/qdp.py
index f8d82f77a..c81b0aa52 100644
--- a/qumat/qdp.py
+++ b/qumat/qdp.py
@@ -46,20 +46,33 @@ def _make_stub(name: str) -> type[Any]:
return type(name, (), {"__init__": __init__, "__doc__": f"Stub class - {name}"})
+QdpEngine: type[Any]
+QuantumTensor: type[Any]
+
try:
- from _qdp import QdpEngine as QdpEngine
- from _qdp import QuantumTensor as QuantumTensor
+ from qumat_qdp import QdpEngine as _QdpEngine
+ from qumat_qdp import QuantumTensor as _QuantumTensor
+
+ QdpEngine = _QdpEngine
+ QuantumTensor = _QuantumTensor
except ImportError as e:
- import warnings
+ try:
+ from _qdp import QdpEngine as _QdpEngine
+ from _qdp import QuantumTensor as _QuantumTensor
+
+ QdpEngine = _QdpEngine
+ QuantumTensor = _QuantumTensor
+ except ImportError:
+ import warnings
- warnings.warn(
- f"QDP module not available: {e}. "
- "QDP requires the qumat-qdp native extension which needs to be built with maturin. "
- "See qdp/qdp-python/README.md for installation instructions.",
- ImportWarning,
- )
+ warnings.warn(
+ f"QDP module not available: {e}. "
+ "QDP requires the qumat-qdp package or native extension to be built. "
+ "See qdp/qdp-python/README.md for installation instructions.",
+ ImportWarning,
+ )
- QdpEngine: type[Any] = _make_stub("QdpEngine")
- QuantumTensor: type[Any] = _make_stub("QuantumTensor")
+ QdpEngine = _make_stub("QdpEngine")
+ QuantumTensor = _make_stub("QuantumTensor")
__all__: list[str] = ["QdpEngine", "QuantumTensor"]
diff --git a/testing/qdp_python/test_torch_ref.py b/testing/qdp_python/test_torch_ref.py
index a7ccd25e6..1f28d9cfc 100644
--- a/testing/qdp_python/test_torch_ref.py
+++ b/testing/qdp_python/test_torch_ref.py
@@ -371,7 +371,8 @@ class TestCrossValidation:
import _qdp
import numpy as np
- engine = _qdp.QdpEngine(0) # type: ignore[unresolved-attribute]
+ qdp_engine = getattr(_qdp, "QdpEngine")
+ engine = qdp_engine(0)
num_qubits = 3
state_dim = 1 << num_qubits
diff --git a/testing/qumat/test_qdp_module.py b/testing/qumat/test_qdp_module.py
index e249c12cc..e243d1fbf 100644
--- a/testing/qumat/test_qdp_module.py
+++ b/testing/qumat/test_qdp_module.py
@@ -23,19 +23,21 @@ import pytest
def _reload_qdp_without_extension():
"""
- Safely reload qumat.qdp while simulating missing _qdp extension.
+ Safely reload qumat.qdp while simulating missing qumat_qdp and _qdp.
Restores sys.modules after execution to avoid test pollution.
"""
original_qdp = sys.modules.get("qumat.qdp")
+ original_pkg = sys.modules.get("qumat_qdp")
original_ext = sys.modules.get("_qdp")
try:
# Remove cached modules
sys.modules.pop("qumat.qdp", None)
+ sys.modules.pop("qumat_qdp", None)
sys.modules.pop("_qdp", None)
- # Simulate missing compiled extension
- with patch.dict(sys.modules, {"_qdp": None}):
+ # Simulate the pure-Python package and native extension both missing.
+ with patch.dict(sys.modules, {"qumat_qdp": None, "_qdp": None}):
module = importlib.import_module("qumat.qdp")
return importlib.reload(module)
@@ -46,6 +48,11 @@ def _reload_qdp_without_extension():
else:
sys.modules.pop("qumat.qdp", None)
+ if original_pkg is not None:
+ sys.modules["qumat_qdp"] = original_pkg
+ else:
+ sys.modules.pop("qumat_qdp", None)
+
if original_ext is not None:
sys.modules["_qdp"] = original_ext
else:
diff --git a/uv.lock b/uv.lock
index fb968090e..e626a1ebd 100644
--- a/uv.lock
+++ b/uv.lock
@@ -2095,10 +2095,12 @@ benchmark = [
{ name = "tensorflow", specifier = ">=2.20" },
{ name = "torch", specifier = ">=2.2,<=2.9.0" },
{ name = "tqdm" },
+ { name = "triton", specifier = ">=3.0" },
]
dev = [
{ name = "pytest" },
{ name = "torch", specifier = ">=2.2,<=2.9.0" },
+ { name = "triton", specifier = ">=3.0" },
]
[[package]]