From 2123b47038725da822997ac35d157344130148d4 Mon Sep 17 00:00:00 2001 From: Pradnya Khalate Date: Mon, 9 Sep 2024 16:53:13 -0700 Subject: [PATCH 1/4] * WIP: Python bindings for 'get_state' API with 'photonics' target --- .../examples/python/providers/photonics.py | 3 +++ .../examples/python/providers/photonics_tbi.py | 3 +++ python/runtime/cudaq/algorithms/py_state.cpp | 16 ++++++++++++++++ 3 files changed, 22 insertions(+) diff --git a/docs/sphinx/examples/python/providers/photonics.py b/docs/sphinx/examples/python/providers/photonics.py index 1075c4874a..0bc171cb21 100644 --- a/docs/sphinx/examples/python/providers/photonics.py +++ b/docs/sphinx/examples/python/providers/photonics.py @@ -14,3 +14,6 @@ def photonicsKernel(): counts = cudaq.sample(photonicsKernel) print(counts) + +state = cudaq.get_state(photonicsKernel) +print(state) diff --git a/docs/sphinx/examples/python/providers/photonics_tbi.py b/docs/sphinx/examples/python/providers/photonics_tbi.py index aff5679b5e..0f31de03f5 100644 --- a/docs/sphinx/examples/python/providers/photonics_tbi.py +++ b/docs/sphinx/examples/python/providers/photonics_tbi.py @@ -43,3 +43,6 @@ def TBI( loop_lengths, shots_count=1000000) counts.dump() + +state = cudaq.get_state(TBI, bs_angles, ps_angles, input_state, loop_lengths) +state.dump() diff --git a/python/runtime/cudaq/algorithms/py_state.cpp b/python/runtime/cudaq/algorithms/py_state.cpp index 08665ba78c..9244bb7a59 100644 --- a/python/runtime/cudaq/algorithms/py_state.cpp +++ b/python/runtime/cudaq/algorithms/py_state.cpp @@ -137,6 +137,20 @@ state pyGetStateRemote(py::object kernel, py::args args) { size, returnOffset)); } +state pyGetStateLibraryMode(py::object kernel, py::args args) { + + // cudaq::info("Size of arguments = {}", args.size()); + /// TODO: Unpack arguments + + return details::extractState([&]() mutable { + if (0 == args.size()) + cudaq::invokeKernel(std::forward(kernel)); + else + cudaq::invokeKernel(std::forward(kernel), + std::forward(args)); + }); +} + /// @brief Bind the get_state cudaq function void bindPyState(py::module &mod, LinkedLibraryHolder &holder) { @@ -629,6 +643,8 @@ index pair. if (holder.getTarget().name == "remote-mqpu" || holder.getTarget().name == "nvqc") return pyGetStateRemote(kernel, args); + if (holder.getTarget().name == "photonics") + return pyGetStateLibraryMode(kernel, args); return pyGetState(kernel, args); }, R"#(Return the :class:`State` of the system after execution of the provided `kernel`. From 85ee7b05ad4b9051988bd666df43b0750362dec7 Mon Sep 17 00:00:00 2001 From: Pradnya Khalate Date: Tue, 10 Sep 2024 14:33:09 -0700 Subject: [PATCH 2/4] * Tests for kernel that takes arguments, basic test for state * WIP: Get the state object --- python/runtime/cudaq/algorithms/py_state.cpp | 5 ++--- python/tests/handlers/test_photonics_kernel.py | 17 +++++++++++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/python/runtime/cudaq/algorithms/py_state.cpp b/python/runtime/cudaq/algorithms/py_state.cpp index 9244bb7a59..82ae97e647 100644 --- a/python/runtime/cudaq/algorithms/py_state.cpp +++ b/python/runtime/cudaq/algorithms/py_state.cpp @@ -139,9 +139,8 @@ state pyGetStateRemote(py::object kernel, py::args args) { state pyGetStateLibraryMode(py::object kernel, py::args args) { - // cudaq::info("Size of arguments = {}", args.size()); - /// TODO: Unpack arguments - + cudaq::info("Size of arguments = {}", args.size()); + /// TODO: Pack / unpack arguments return details::extractState([&]() mutable { if (0 == args.size()) cudaq::invokeKernel(std::forward(kernel)); diff --git a/python/tests/handlers/test_photonics_kernel.py b/python/tests/handlers/test_photonics_kernel.py index ddfbaaf37a..99dfd2824e 100644 --- a/python/tests/handlers/test_photonics_kernel.py +++ b/python/tests/handlers/test_photonics_kernel.py @@ -32,6 +32,10 @@ def kernel(): assert len(counts) == 1 assert '3' in counts + state = cudaq.get_state(kernel) + state.dump() + # TODO: Add check for 'state' object + def test_qudit_list(): @@ -78,6 +82,19 @@ def kernel(): counts.dump() +def test_kernel_with_args(): + + @cudaq.kernel + def kernel(theta: float): + q = qudit(4) + plus(q) + phase_shift(q, theta) + mz(q) + + result = cudaq.sample(kernel, 0.5) + result.dump() + + def test_target_change(): @cudaq.kernel From c61a21a5e6477771a4ddb45273d9cdbd8da5da8b Mon Sep 17 00:00:00 2001 From: Pradnya Khalate Date: Tue, 10 Sep 2024 16:14:28 -0700 Subject: [PATCH 3/4] * Override more methods in the `PhotonicsState` class * Add a test for state retrieval with kernel that accepts arguments --- python/runtime/cudaq/algorithms/py_state.cpp | 2 +- .../tests/handlers/test_photonics_kernel.py | 5 ++- .../photonics/PhotonicsExecutionManager.cpp | 35 +++++++++++++------ 3 files changed, 29 insertions(+), 13 deletions(-) diff --git a/python/runtime/cudaq/algorithms/py_state.cpp b/python/runtime/cudaq/algorithms/py_state.cpp index 82ae97e647..62fe055af6 100644 --- a/python/runtime/cudaq/algorithms/py_state.cpp +++ b/python/runtime/cudaq/algorithms/py_state.cpp @@ -138,8 +138,8 @@ state pyGetStateRemote(py::object kernel, py::args args) { } state pyGetStateLibraryMode(py::object kernel, py::args args) { - cudaq::info("Size of arguments = {}", args.size()); + /// TODO: Pack / unpack arguments return details::extractState([&]() mutable { if (0 == args.size()) diff --git a/python/tests/handlers/test_photonics_kernel.py b/python/tests/handlers/test_photonics_kernel.py index 99dfd2824e..8c2d300324 100644 --- a/python/tests/handlers/test_photonics_kernel.py +++ b/python/tests/handlers/test_photonics_kernel.py @@ -34,7 +34,7 @@ def kernel(): state = cudaq.get_state(kernel) state.dump() - # TODO: Add check for 'state' object + assert 4 == state.__len__() def test_qudit_list(): @@ -93,6 +93,9 @@ def kernel(theta: float): result = cudaq.sample(kernel, 0.5) result.dump() + + state = cudaq.get_state(kernel, 0.5) + state.dump() def test_target_change(): diff --git a/runtime/cudaq/qis/managers/photonics/PhotonicsExecutionManager.cpp b/runtime/cudaq/qis/managers/photonics/PhotonicsExecutionManager.cpp index ca0cb35e23..ef45f620c2 100644 --- a/runtime/cudaq/qis/managers/photonics/PhotonicsExecutionManager.cpp +++ b/runtime/cudaq/qis/managers/photonics/PhotonicsExecutionManager.cpp @@ -29,8 +29,9 @@ struct PhotonicsState : public cudaq::SimulationState { PhotonicsState(qpp::ket &&data, std::size_t lvl) : state(std::move(data)), levels(lvl) {} + /// TODO: Rename the API to be generic std::size_t getNumQubits() const override { - throw "not supported for this photonics simulator"; + return (std::log2(state.size()) / std::log2(levels)); } std::complex overlap(const cudaq::SimulationState &other) override { @@ -39,7 +40,11 @@ struct PhotonicsState : public cudaq::SimulationState { std::complex getAmplitude(const std::vector &basisState) override { - /// TODO: Check basisState.size() matches qudit count + if (getNumQubits() != basisState.size()) + throw std::runtime_error(fmt::format( + "[photonics] getAmplitude with an invalid number of bits in the " + "basis state: expected {}, provided {}.", + getNumQubits(), basisState.size())); // Convert the basis state to an index value const std::size_t idx = std::accumulate( @@ -50,27 +55,35 @@ struct PhotonicsState : public cudaq::SimulationState { } Tensor getTensor(std::size_t tensorIdx = 0) const override { - throw "not supported for this photonics simulator"; + if (tensorIdx != 0) + throw std::runtime_error("[photonics] invalid tensor requested."); + return Tensor{ + reinterpret_cast( + const_cast *>(state.data())), + std::vector{static_cast(state.size())}, + getPrecision()}; } - std::vector getTensors() const override { - throw "not supported for this photonics simulator"; - } + // /// @brief Return all tensors that represent this state + std::vector getTensors() const override { return {getTensor()}; } - std::size_t getNumTensors() const override { - throw "not supported for this photonics simulator"; - } + // /// @brief Return the number of tensors that represent this state. + std::size_t getNumTensors() const override { return 1; } std::complex operator()(std::size_t tensorIdx, const std::vector &indices) override { - throw "not supported for this photonics simulator"; + if (tensorIdx != 0) + throw std::runtime_error("[photonics] invalid tensor requested."); + if (indices.size() != 1) + throw std::runtime_error("[photonics] invalid element extraction."); + + return state[indices[0]]; } std::unique_ptr createFromSizeAndPtr(std::size_t size, void *ptr, std::size_t) override { throw "not supported for this photonics simulator"; - ; } void dump(std::ostream &os) const override { os << state << "\n"; } From c13b93419844067fff4406a5b145c9b1e51c95fd Mon Sep 17 00:00:00 2001 From: Pradnya Khalate Date: Mon, 23 Sep 2024 11:39:34 -0700 Subject: [PATCH 4/4] * Unpack arguments based on the caller's execution context * Make tests stable by enforcing garbage collection as part of set up --- python/cudaq/handlers/photonics_kernel.py | 12 ++++- python/cudaq/kernel/kernel_decorator.py | 12 ++++- python/runtime/common/py_ExecutionContext.cpp | 4 ++ python/runtime/cudaq/algorithms/py_state.cpp | 14 +++--- .../tests/handlers/test_photonics_kernel.py | 44 +++++++++++++++++-- .../photonics/PhotonicsExecutionManager.cpp | 5 +-- 6 files changed, 75 insertions(+), 16 deletions(-) diff --git a/python/cudaq/handlers/photonics_kernel.py b/python/cudaq/handlers/photonics_kernel.py index fc676c33a9..5738234314 100644 --- a/python/cudaq/handlers/photonics_kernel.py +++ b/python/cudaq/handlers/photonics_kernel.py @@ -13,6 +13,8 @@ from ..mlir._mlir_libs._quakeDialects import cudaq_runtime +_TARGET_NAME = 'photonics' + # The qudit level must be explicitly defined globalQuditLevel = None @@ -33,7 +35,13 @@ class PyQudit: id: int def __del__(self): - cudaq_runtime.photonics.release_qudit(self.level, self.id) + try: + cudaq_runtime.photonics.release_qudit(self.level, self.id) + except Exception as e: + if _TARGET_NAME == cudaq_runtime.get_target().name: + raise e + else: + pass def _is_qudit_type(q: any) -> bool: @@ -194,7 +202,7 @@ class PhotonicsHandler(object): def __init__(self, function): - if 'photonics' != cudaq_runtime.get_target().name: + if _TARGET_NAME != cudaq_runtime.get_target().name: raise RuntimeError( "A photonics kernel can only be used with 'photonics' target.") diff --git a/python/cudaq/kernel/kernel_decorator.py b/python/cudaq/kernel/kernel_decorator.py index b43c54a6c8..2af09bb891 100644 --- a/python/cudaq/kernel/kernel_decorator.py +++ b/python/cudaq/kernel/kernel_decorator.py @@ -360,7 +360,17 @@ def __call__(self, *args): raise RuntimeError( "The 'photonics' target must be used with a valid function." ) - PhotonicsHandler(self.kernelFunction)(*args) + # NOTE: Since this handler does not support MLIR mode (yet), just + # invoke the kernel. If calling from a bound function, need to + # unpack the arguments, for example, see `pyGetStateLibraryMode` + try: + context_name = cudaq_runtime.getExecutionContextName() + except RuntimeError: + context_name = None + callable_args = args + if "extract-state" == context_name and len(args) == 1: + callable_args = args[0] + PhotonicsHandler(self.kernelFunction)(*callable_args) return # Prepare captured state storage for the run diff --git a/python/runtime/common/py_ExecutionContext.cpp b/python/runtime/common/py_ExecutionContext.cpp index baa1ca518c..b1fc025e1e 100644 --- a/python/runtime/common/py_ExecutionContext.cpp +++ b/python/runtime/common/py_ExecutionContext.cpp @@ -48,5 +48,9 @@ void bindExecutionContext(py::module &mod) { auto &platform = cudaq::get_platform(); return platform.supports_conditional_feedback(); }); + mod.def("getExecutionContextName", []() { + auto &self = cudaq::get_platform(); + return self.get_exec_ctx()->name; + }); } } // namespace cudaq diff --git a/python/runtime/cudaq/algorithms/py_state.cpp b/python/runtime/cudaq/algorithms/py_state.cpp index 62fe055af6..77a8e4a36d 100644 --- a/python/runtime/cudaq/algorithms/py_state.cpp +++ b/python/runtime/cudaq/algorithms/py_state.cpp @@ -138,15 +138,17 @@ state pyGetStateRemote(py::object kernel, py::args args) { } state pyGetStateLibraryMode(py::object kernel, py::args args) { - cudaq::info("Size of arguments = {}", args.size()); - - /// TODO: Pack / unpack arguments return details::extractState([&]() mutable { if (0 == args.size()) cudaq::invokeKernel(std::forward(kernel)); - else - cudaq::invokeKernel(std::forward(kernel), - std::forward(args)); + else { + std::vector argsData; + for (size_t i = 0; i < args.size(); i++) { + py::object arg = args[i]; + argsData.emplace_back(std::forward(arg)); + } + cudaq::invokeKernel(std::forward(kernel), argsData); + } }); } diff --git a/python/tests/handlers/test_photonics_kernel.py b/python/tests/handlers/test_photonics_kernel.py index 8c2d300324..771769b4ad 100644 --- a/python/tests/handlers/test_photonics_kernel.py +++ b/python/tests/handlers/test_photonics_kernel.py @@ -7,6 +7,10 @@ # ============================================================================ # import pytest + +import gc +from typing import List + import cudaq @@ -16,6 +20,8 @@ def do_something(): yield cudaq.reset_target() cudaq.__clearKernelRegistries() + # Make the tests stable by enforcing resource release + gc.collect() def test_qudit(): @@ -83,18 +89,48 @@ def kernel(): def test_kernel_with_args(): + """Test that `PhotonicsHandler` supports basic arguments. + The check here is that all the test kernels run successfully.""" @cudaq.kernel - def kernel(theta: float): + def kernel_1f(theta: float): q = qudit(4) plus(q) phase_shift(q, theta) mz(q) - result = cudaq.sample(kernel, 0.5) + result = cudaq.sample(kernel_1f, 0.5) + result.dump() + + state = cudaq.get_state(kernel_1f, 0.5) + state.dump() + + @cudaq.kernel + def kernel_2f(theta: float, phi: float): + quds = [qudit(3) for _ in range(2)] + plus(quds[0]) + phase_shift(quds[0], theta) + beam_splitter(quds[0], quds[1], phi) + mz(quds) + + result = cudaq.sample(kernel_2f, 0.7854, 0.3927) result.dump() - - state = cudaq.get_state(kernel, 0.5) + + state = cudaq.get_state(kernel_2f, 0.7854, 0.3927) + state.dump() + + @cudaq.kernel + def kernel_list(angles: List[float]): + quds = [qudit(2) for _ in range(3)] + plus(quds[0]) + phase_shift(quds[1], angles[0]) + phase_shift(quds[2], angles[1]) + mz(quds) + + result = cudaq.sample(kernel_list, [0.5236, 1.0472]) + result.dump() + + state = cudaq.get_state(kernel_list, [0.5236, 1.0472]) state.dump() diff --git a/runtime/cudaq/qis/managers/photonics/PhotonicsExecutionManager.cpp b/runtime/cudaq/qis/managers/photonics/PhotonicsExecutionManager.cpp index ef45f620c2..bfdcac1617 100644 --- a/runtime/cudaq/qis/managers/photonics/PhotonicsExecutionManager.cpp +++ b/runtime/cudaq/qis/managers/photonics/PhotonicsExecutionManager.cpp @@ -64,10 +64,10 @@ struct PhotonicsState : public cudaq::SimulationState { getPrecision()}; } - // /// @brief Return all tensors that represent this state + /// @brief Return all tensors that represent this state std::vector getTensors() const override { return {getTensor()}; } - // /// @brief Return the number of tensors that represent this state. + /// @brief Return the number of tensors that represent this state. std::size_t getNumTensors() const override { return 1; } std::complex @@ -77,7 +77,6 @@ struct PhotonicsState : public cudaq::SimulationState { throw std::runtime_error("[photonics] invalid tensor requested."); if (indices.size() != 1) throw std::runtime_error("[photonics] invalid element extraction."); - return state[indices[0]]; }