Create cudaq.trace_kernel()

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

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

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);

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 <algorithm>
+#include <fmt/core.h>
+#include <fmt/ranges.h>
+#include <iterator>
+#include <pybind11/functional.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+
+namespace py = pybind11;
+
+namespace cudaq {
+void bindPyTraceKernel(py::module &mod) {
+
+  py::class_<QuditInfo>(mod, "QuditInfo")
+      .def(py::init<std::size_t, std::size_t>())
+      .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_<Trace::Instruction>(mod, "TraceInstruction")
+      .def(py::init<std::string_view, std::vector<double>,
+                    std::vector<QuditInfo>, std::vector<QuditInfo>>())
+      .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<std::size_t> controlIds;
+        std::transform(i.controls.begin(), i.controls.end(),
+                       std::back_inserter(controlIds),
+                       [](const QuditInfo &q) { return q.id; });
+        std::vector<std::size_t> 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_<Trace>(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<std::string>();
+        auto kernelMod = kernel.attr("module").cast<MlirModule>();
+        args = simplifiedValidateInputArguments(args);
+        std::unique_ptr<OpaqueArguments> argData(
+            toOpaqueArgs(args, kernelMod, kernelName));
+
+        auto ctx = std::make_unique<ExecutionContext>("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

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 <pybind11/pybind11.h>
+
+namespace py = pybind11;
+
+namespace cudaq {
+void bindPyTraceKernel(py::module &mod);
+} // namespace cudaq

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

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
 

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();

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);

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:                   ╰───╯ ╰────╯