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..e627e62bd10 --- /dev/null +++ b/python/tests/backends/test_trace_kernel.py @@ -0,0 +1,171 @@ +# ============================================================================ # +# 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() + 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/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/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); 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: ╰───╯ ╰────╯