From ef0b8ea4a11839c39cd0127641a08fc8394b9443 Mon Sep 17 00:00:00 2001 From: Ben Howe Date: Tue, 2 Dec 2025 21:56:37 -0800 Subject: [PATCH 1/3] Create cudaq.trace_kernel() Signed-off-by: Ben Howe --- python/cudaq/__init__.py | 1 + python/extension/CMakeLists.txt | 1 + python/extension/CUDAQuantumExtension.cpp | 2 + .../cudaq/algorithms/py_trace_kernel.cpp | 99 ++++++++++++ .../cudaq/algorithms/py_trace_kernel.h | 17 ++ python/tests/backends/test_trace_kernel.py | 148 ++++++++++++++++++ .../qis/managers/BasicExecutionManager.h | 5 +- runtime/nvqir/CircuitSimulator.h | 5 +- 8 files changed, 276 insertions(+), 2 deletions(-) create mode 100644 python/runtime/cudaq/algorithms/py_trace_kernel.cpp create mode 100644 python/runtime/cudaq/algorithms/py_trace_kernel.h create mode 100644 python/tests/backends/test_trace_kernel.py diff --git a/python/cudaq/__init__.py b/python/cudaq/__init__.py index 7a652e30a92..1df8103a578 100644 --- a/python/cudaq/__init__.py +++ b/python/cudaq/__init__.py @@ -245,6 +245,7 @@ def _configure_cuda_library_paths() -> None: get_unitary = cudaq_runtime.get_unitary run = cudaq_runtime.run estimate_resources = cudaq_runtime.estimate_resources +trace_kernel = cudaq_runtime.trace_kernel translate = cudaq_runtime.translate displaySVG = display_trace.displaySVG getSVGstring = display_trace.getSVGstring diff --git a/python/extension/CMakeLists.txt b/python/extension/CMakeLists.txt index 9f366ae908e..41ab4a9dd2d 100644 --- a/python/extension/CMakeLists.txt +++ b/python/extension/CMakeLists.txt @@ -60,6 +60,7 @@ declare_mlir_python_extension(CUDAQuantumPythonSources.Extension ../runtime/cudaq/algorithms/py_run.cpp ../runtime/cudaq/algorithms/py_state.cpp ../runtime/cudaq/algorithms/py_evolve.cpp + ../runtime/cudaq/algorithms/py_trace_kernel.cpp ../runtime/cudaq/algorithms/py_translate.cpp ../runtime/cudaq/algorithms/py_unitary.cpp ../runtime/cudaq/algorithms/py_utils.cpp diff --git a/python/extension/CUDAQuantumExtension.cpp b/python/extension/CUDAQuantumExtension.cpp index 2d8a70f9b1d..7a118b9ed58 100644 --- a/python/extension/CUDAQuantumExtension.cpp +++ b/python/extension/CUDAQuantumExtension.cpp @@ -26,6 +26,7 @@ #include "runtime/cudaq/algorithms/py_run.h" #include "runtime/cudaq/algorithms/py_sample_async.h" #include "runtime/cudaq/algorithms/py_state.h" +#include "runtime/cudaq/algorithms/py_trace_kernel.h" #include "runtime/cudaq/algorithms/py_translate.h" #include "runtime/cudaq/algorithms/py_unitary.h" #include "runtime/cudaq/algorithms/py_utils.h" @@ -121,6 +122,7 @@ PYBIND11_MODULE(_quakeDialects, m) { cudaq::bindPyDataClassRegistry(cudaqRuntime); cudaq::bindPyEvolve(cudaqRuntime); cudaq::bindEvolveResult(cudaqRuntime); + cudaq::bindPyTraceKernel(cudaqRuntime); cudaq::bindPyDraw(cudaqRuntime); cudaq::bindPyUnitary(cudaqRuntime); cudaq::bindPyRun(cudaqRuntime); diff --git a/python/runtime/cudaq/algorithms/py_trace_kernel.cpp b/python/runtime/cudaq/algorithms/py_trace_kernel.cpp new file mode 100644 index 00000000000..aac3aec2336 --- /dev/null +++ b/python/runtime/cudaq/algorithms/py_trace_kernel.cpp @@ -0,0 +1,99 @@ +/******************************************************************************* + * Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates. * + * All rights reserved. * + * * + * This source code and the accompanying materials are made available under * + * the terms of the Apache License 2.0 which accompanies this distribution. * + ******************************************************************************/ + +#include "common/Trace.h" +#include "cudaq/qis/execution_manager.h" +#include "runtime/cudaq/platform/py_alt_launch_kernel.h" +#include "utils/LinkedLibraryHolder.h" +#include "utils/OpaqueArguments.h" + +#include +#include +#include +#include +#include +#include +#include + +namespace py = pybind11; + +namespace cudaq { +void bindPyTraceKernel(py::module &mod) { + + py::class_(mod, "QuditInfo") + .def(py::init()) + .def_readonly("levels", &QuditInfo::levels) + .def_readonly("id", &QuditInfo::id) + .def("__repr__", [](const QuditInfo &q) { + return fmt::format("QuditInfo(levels={}, id={})", q.levels, q.id); + }); + + py::class_(mod, "TraceInstruction") + .def(py::init, + std::vector, std::vector>()) + .def_readonly("name", &Trace::Instruction::name) + .def_readonly("params", &Trace::Instruction::params) + .def_readonly("controls", &Trace::Instruction::controls) + .def_readonly("targets", &Trace::Instruction::targets) + .def("__repr__", [](const Trace::Instruction &i) { + std::vector controlIds; + std::transform(i.controls.begin(), i.controls.end(), + std::back_inserter(controlIds), + [](const QuditInfo &q) { return q.id; }); + std::vector targetIds; + std::transform(i.targets.begin(), i.targets.end(), + std::back_inserter(targetIds), + [](const QuditInfo &q) { return q.id; }); + return fmt::format( + "Instruction(name={}, params=[{}], controls=[{}], targets=[{}])", + i.name, fmt::join(i.params, ", "), fmt::join(controlIds, ", "), + fmt::join(targetIds, ", ")); + }); + + py::class_(mod, "Trace") + .def(py::init<>()) + .def("append_instruction", &Trace::appendInstruction) + .def("get_num_qudits", &Trace::getNumQudits) + .def( + "__iter__", + [](const Trace &t) { return py::make_iterator(t.begin(), t.end()); }, + py::keep_alive<0, 1>()); + + mod.def( + "trace_kernel", + [&](py::object kernel, py::args args) { + if (py::hasattr(kernel, "compile")) + kernel.attr("compile")(); + auto &platform = cudaq::get_platform(); + auto kernelName = kernel.attr("name").cast(); + auto kernelMod = kernel.attr("module").cast(); + args = simplifiedValidateInputArguments(args); + std::unique_ptr argData( + toOpaqueArgs(args, kernelMod, kernelName)); + + auto ctx = std::make_unique("tracer", 1); + ctx->kernelName = kernelName; + platform.set_exec_ctx(ctx.get()); + pyAltLaunchKernel(kernelName, kernelMod, *argData, {}); + platform.reset_exec_ctx(); + + return ctx->kernelTrace; + }, + py::arg("kernel"), py::kw_only(), + R"#(Executes the given kernel and returns a :class:`Trace` object +representing the sequence of operations performed. + +Args: + kernel (:class:`Kernel`): The :class:`Kernel` to trace. + *arguments (Optional[Any]): The concrete values to evaluate the kernel + function at. Leave empty if the kernel doesn't accept any arguments. + +Returns: + :class:`Trace`: The trace of the kernel execution.)#"); +} +} // namespace cudaq diff --git a/python/runtime/cudaq/algorithms/py_trace_kernel.h b/python/runtime/cudaq/algorithms/py_trace_kernel.h new file mode 100644 index 00000000000..74d6bbaa0c4 --- /dev/null +++ b/python/runtime/cudaq/algorithms/py_trace_kernel.h @@ -0,0 +1,17 @@ +/****************************************************************-*- C++ -*-**** + * Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates. * + * All rights reserved. * + * * + * This source code and the accompanying materials are made available under * + * the terms of the Apache License 2.0 which accompanies this distribution. * + ******************************************************************************/ + +#pragma once + +#include + +namespace py = pybind11; + +namespace cudaq { +void bindPyTraceKernel(py::module &mod); +} // namespace cudaq diff --git a/python/tests/backends/test_trace_kernel.py b/python/tests/backends/test_trace_kernel.py new file mode 100644 index 00000000000..a4586dfccda --- /dev/null +++ b/python/tests/backends/test_trace_kernel.py @@ -0,0 +1,148 @@ +import cudaq +import pytest + +def test_trace_kernel(): + @cudaq.kernel + def my_kernel(): + q = cudaq.qubit() + x(q) + h(q) + mz(q) + + trace = cudaq.trace_kernel(my_kernel) + + assert trace.get_num_qudits() == 1 + + instructions = list(trace) + assert len(instructions) == 3 + + assert instructions[0].name == "x" + assert len(instructions[0].controls) == 0 + assert len(instructions[0].targets) == 1 + assert instructions[0].targets[0].id == 0 + + assert instructions[1].name == "h" + + assert instructions[2].name == "mz" + +def test_trace_kernel_with_args(): + @cudaq.kernel + def my_kernel(angle: float): + q = cudaq.qubit() + rx(angle, q) + + trace = cudaq.trace_kernel(my_kernel, 0.5) + + instructions = list(trace) + assert len(instructions) == 1 + assert instructions[0].name == "rx" + assert len(instructions[0].params) == 1 + assert abs(instructions[0].params[0] - 0.5) < 1e-6 + +def test_trace_kernel_cnot(): + @cudaq.kernel + def my_kernel(): + q = cudaq.qvector(2) + x(q[0]) + cx(q[0], q[1]) + + trace = cudaq.trace_kernel(my_kernel) + + assert trace.get_num_qudits() == 2 + + instructions = list(trace) + assert len(instructions) == 2 + + assert instructions[0].name == "x" + + assert instructions[1].name == "x" # CX is often represented as X with control + assert len(instructions[1].controls) == 1 + assert instructions[1].controls[0].id == 0 + assert len(instructions[1].targets) == 1 + assert instructions[1].targets[0].id == 1 + +def test_trace_kernel_loops(): + @cudaq.kernel + def my_kernel(): + q = cudaq.qvector(2) + for i in range(2): + x(q[i]) + + trace = cudaq.trace_kernel(my_kernel) + assert trace.get_num_qudits() == 2 + + instructions = list(trace) + assert len(instructions) == 2 + assert instructions[0].name == "x" + assert instructions[0].targets[0].id == 0 + assert instructions[1].name == "x" + assert instructions[1].targets[0].id == 1 + +def test_trace_kernel_cx_loop(): + @cudaq.kernel + def my_kernel(): + q = cudaq.qvector(3) + # 0 -> 1, 1 -> 2 + for i in range(2): + cx(q[i], q[i+1]) + + trace = cudaq.trace_kernel(my_kernel) + assert trace.get_num_qudits() == 3 + + instructions = list(trace) + assert len(instructions) == 2 + + # cx(0, 1) + assert instructions[0].name == "x" + assert len(instructions[0].controls) == 1 + assert instructions[0].controls[0].id == 0 + assert len(instructions[0].targets) == 1 + assert instructions[0].targets[0].id == 1 + + # cx(1, 2) + assert instructions[1].name == "x" + assert len(instructions[1].controls) == 1 + assert instructions[1].controls[0].id == 1 + assert len(instructions[1].targets) == 1 + assert instructions[1].targets[0].id == 2 + +def test_trace_kernel_subkernels(): + @cudaq.kernel + def subkernel(q: cudaq.qubit): + x(q) + + @cudaq.kernel + def my_kernel(): + q = cudaq.qubit() + h(q) + subkernel(q) + + trace = cudaq.trace_kernel(my_kernel) + assert trace.get_num_qudits() == 1 + + instructions = list(trace) + assert len(instructions) == 2 + assert instructions[0].name == "h" + assert instructions[1].name == "x" + +def test_trace_kernel_subkernels_with_controls(): + @cudaq.kernel + def subkernel(q: cudaq.qubit): + x(q) + + @cudaq.kernel + def my_kernel(): + q = cudaq.qvector(2) + # Apply subkernel controlled by q[0] on q[1] + cudaq.control(subkernel, q[0], q[1]) + + trace = cudaq.trace_kernel(my_kernel) + assert trace.get_num_qudits() == 2 + + instructions = list(trace) + assert len(instructions) == 1 + assert instructions[0].name == "x" + assert len(instructions[0].controls) == 1 + assert instructions[0].controls[0].id == 0 + assert len(instructions[0].targets) == 1 + assert instructions[0].targets[0].id == 1 diff --git a/runtime/cudaq/qis/managers/BasicExecutionManager.h b/runtime/cudaq/qis/managers/BasicExecutionManager.h index ca33e7fab24..6644f8c12c9 100644 --- a/runtime/cudaq/qis/managers/BasicExecutionManager.h +++ b/runtime/cudaq/qis/managers/BasicExecutionManager.h @@ -268,8 +268,11 @@ class BasicExecutionManager : public cudaq::ExecutionManager { int measure(const cudaq::QuditInfo &target, const std::string registerName = "") override { - if (isInTracerMode()) + if (isInTracerMode()) { + synchronize(); + executionContext->kernelTrace.appendInstruction("mz", {}, {}, {target}); return 0; + } // We hit a measure, need to exec / clear instruction queue synchronize(); diff --git a/runtime/nvqir/CircuitSimulator.h b/runtime/nvqir/CircuitSimulator.h index 057109336ce..06163e8c999 100644 --- a/runtime/nvqir/CircuitSimulator.h +++ b/runtime/nvqir/CircuitSimulator.h @@ -1428,8 +1428,11 @@ class CircuitSimulatorBase : public CircuitSimulator { if (handleBasicSampling(qubitIdx, registerName)) return true; - if (isInTracerMode()) + if (isInTracerMode()) { + executionContext->kernelTrace.appendInstruction( + "mz", {}, {}, {cudaq::QuditInfo(2, qubitIdx)}); return true; + } // Get the actual measurement from the subtype measureQubit implementation auto measureResult = measureQubit(qubitIdx); From cbab873f9c65baee57554df4cfa77c77fad0aaca Mon Sep 17 00:00:00 2001 From: Ben Howe Date: Tue, 9 Dec 2025 18:13:01 -0800 Subject: [PATCH 2/3] License and formatting Signed-off-by: Ben Howe --- python/tests/backends/test_trace_kernel.py | 51 ++++++++++++++++------ 1 file changed, 37 insertions(+), 14 deletions(-) diff --git a/python/tests/backends/test_trace_kernel.py b/python/tests/backends/test_trace_kernel.py index a4586dfccda..e627e62bd10 100644 --- a/python/tests/backends/test_trace_kernel.py +++ b/python/tests/backends/test_trace_kernel.py @@ -1,7 +1,17 @@ +# ============================================================================ # +# Copyright (c) 2022 - 2025 NVIDIA Corporation & Affiliates. # +# All rights reserved. # +# # +# This source code and the accompanying materials are made available under # +# the terms of the Apache License 2.0 which accompanies this distribution. # +# ============================================================================ # + import cudaq import pytest + def test_trace_kernel(): + @cudaq.kernel def my_kernel(): q = cudaq.qubit() @@ -10,36 +20,40 @@ def my_kernel(): mz(q) trace = cudaq.trace_kernel(my_kernel) - + assert trace.get_num_qudits() == 1 - + instructions = list(trace) assert len(instructions) == 3 - + assert instructions[0].name == "x" assert len(instructions[0].controls) == 0 assert len(instructions[0].targets) == 1 assert instructions[0].targets[0].id == 0 assert instructions[1].name == "h" - + assert instructions[2].name == "mz" + def test_trace_kernel_with_args(): + @cudaq.kernel def my_kernel(angle: float): q = cudaq.qubit() rx(angle, q) trace = cudaq.trace_kernel(my_kernel, 0.5) - + instructions = list(trace) assert len(instructions) == 1 assert instructions[0].name == "rx" assert len(instructions[0].params) == 1 assert abs(instructions[0].params[0] - 0.5) < 1e-6 + def test_trace_kernel_cnot(): + @cudaq.kernel def my_kernel(): q = cudaq.qvector(2) @@ -47,21 +61,24 @@ def my_kernel(): cx(q[0], q[1]) trace = cudaq.trace_kernel(my_kernel) - + assert trace.get_num_qudits() == 2 - + instructions = list(trace) assert len(instructions) == 2 - + assert instructions[0].name == "x" - - assert instructions[1].name == "x" # CX is often represented as X with control + + assert instructions[ + 1].name == "x" # CX is often represented as X with control assert len(instructions[1].controls) == 1 assert instructions[1].controls[0].id == 0 assert len(instructions[1].targets) == 1 assert instructions[1].targets[0].id == 1 + def test_trace_kernel_loops(): + @cudaq.kernel def my_kernel(): q = cudaq.qvector(2) @@ -78,20 +95,22 @@ def my_kernel(): assert instructions[1].name == "x" assert instructions[1].targets[0].id == 1 + def test_trace_kernel_cx_loop(): + @cudaq.kernel def my_kernel(): q = cudaq.qvector(3) # 0 -> 1, 1 -> 2 for i in range(2): - cx(q[i], q[i+1]) + cx(q[i], q[i + 1]) trace = cudaq.trace_kernel(my_kernel) assert trace.get_num_qudits() == 3 instructions = list(trace) assert len(instructions) == 2 - + # cx(0, 1) assert instructions[0].name == "x" assert len(instructions[0].controls) == 1 @@ -106,11 +125,13 @@ def my_kernel(): assert len(instructions[1].targets) == 1 assert instructions[1].targets[0].id == 2 + def test_trace_kernel_subkernels(): + @cudaq.kernel def subkernel(q: cudaq.qubit): x(q) - + @cudaq.kernel def my_kernel(): q = cudaq.qubit() @@ -125,11 +146,13 @@ def my_kernel(): assert instructions[0].name == "h" assert instructions[1].name == "x" + def test_trace_kernel_subkernels_with_controls(): + @cudaq.kernel def subkernel(q: cudaq.qubit): x(q) - + @cudaq.kernel def my_kernel(): q = cudaq.qvector(2) From 3d7a01ead0f146f7439d9815c5f16c2432fc76e5 Mon Sep 17 00:00:00 2001 From: Ben Howe Date: Tue, 9 Dec 2025 20:26:51 -0800 Subject: [PATCH 3/3] Update truth data for draw to include mz Signed-off-by: Ben Howe --- python/tests/kernel/test_kernel_features.py | 10 ++++---- targettests/execution/draw_test.cpp | 28 ++++++++++----------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/python/tests/kernel/test_kernel_features.py b/python/tests/kernel/test_kernel_features.py index 15502764a59..12557dbbfa8 100644 --- a/python/tests/kernel/test_kernel_features.py +++ b/python/tests/kernel/test_kernel_features.py @@ -1412,11 +1412,11 @@ def kernel(angle: float): print(cudaq.draw(kernel, 0.59)) circuit = cudaq.draw(bell_pair) print(circuit) - expected_str = ''' ╭───╮ -q0 : ┤ h ├──●── - ╰───╯╭─┴─╮ -q1 : ─────┤ x ├ - ╰───╯ + expected_str = ''' ╭───╮ ╭────╮ +q0 : ┤ h ├──●──┤ mz ├ + ╰───╯╭─┴─╮├────┤ +q1 : ─────┤ x ├┤ mz ├ + ╰───╯╰────╯ ''' assert circuit == expected_str diff --git a/targettests/execution/draw_test.cpp b/targettests/execution/draw_test.cpp index b609a19133e..7dfda167497 100644 --- a/targettests/execution/draw_test.cpp +++ b/targettests/execution/draw_test.cpp @@ -33,18 +33,18 @@ int main() { return 0; } -// IONQ: ╭───╮ -// IONQ: q0 : ┤ x ├──●────●── -// IONQ: ├───┤╭─┴─╮ │ -// IONQ: q1 : ┤ x ├┤ x ├──┼── -// IONQ: ╰───╯╰───╯╭─┴─╮ -// IONQ: q2 : ──────────┤ x ├ -// IONQ: ╰───╯ +// IONQ: ╭───╮ ╭────╮ +// IONQ: q0 : ┤ x ├──●────●──┤ mz ├ +// IONQ: ├───┤╭─┴─╮ │ ├────┤ +// IONQ: q1 : ┤ x ├┤ x ├──┼──┤ mz ├ +// IONQ: ╰───╯╰───╯╭─┴─╮├────┤ +// IONQ: q2 : ──────────┤ x ├┤ mz ├ +// IONQ: ╰───╯╰────╯ -// OQC: ╭───╮ -// OQC: q0 : ┤ x ├──●───╳────── -// OQC: ├───┤╭─┴─╮ │ -// OQC: q1 : ┤ x ├┤ x ├─╳───●── -// OQC: ╰───╯╰───╯ ╭─┴─╮ -// OQC: q2 : ─────────────┤ x ├ -// OQC: ╰───╯ +// OQC: ╭───╮ ╭────╮ +// OQC: q0 : ┤ x ├──●───╳─┤ mz ├────── +// OQC: ├───┤╭─┴─╮ │ ╰────╯╭────╮ +// OQC: q1 : ┤ x ├┤ x ├─╳───●───┤ mz ├ +// OQC: ╰───╯╰───╯ ╭─┴─╮ ├────┤ +// OQC: q2 : ─────────────┤ x ├─┤ mz ├ +// OQC: ╰───╯ ╰────╯