From e23ca18269fae5889bc05675450f6f9bedd3b8fb Mon Sep 17 00:00:00 2001 From: Anna Gringauze Date: Mon, 30 Sep 2024 13:01:51 -0700 Subject: [PATCH] Fix exp_pauli issues on remote simulators and quantum devices (#2226) * Made exp_pauli work on quantum devices and remote sim * Add more tests * Address CR comments --- lib/Optimizer/CodeGen/QuakeToLLVM.cpp | 15 ++--- .../Transforms/DecompositionPatterns.cpp | 41 ++++++++++-- python/tests/backends/test_IQM.py | 15 +++++ python/tests/backends/test_IonQ.py | 14 +++++ python/tests/backends/test_OQC.py | 14 +++++ .../test_Quantinuum_LocalEmulation_builder.py | 23 +++++++ .../test_Quantinuum_LocalEmulation_kernel.py | 27 ++++++++ .../tests/backends/test_Quantinuum_builder.py | 15 ++++- .../tests/backends/test_Quantinuum_kernel.py | 17 ++++- python/tests/kernel/test_kernel_exp_pauli.py | 37 +++++++++++ targettests/execution/exp_pauli.cpp | 62 +++++++++++++++++++ .../ExpPauliToHRyRzCX.qke | 35 ++++++++++- 12 files changed, 301 insertions(+), 14 deletions(-) create mode 100644 python/tests/kernel/test_kernel_exp_pauli.py create mode 100644 targettests/execution/exp_pauli.cpp diff --git a/lib/Optimizer/CodeGen/QuakeToLLVM.cpp b/lib/Optimizer/CodeGen/QuakeToLLVM.cpp index 7ed9b1b25f..7bd8c38bc4 100644 --- a/lib/Optimizer/CodeGen/QuakeToLLVM.cpp +++ b/lib/Optimizer/CodeGen/QuakeToLLVM.cpp @@ -462,11 +462,12 @@ class ExpPauliRewrite : public ConvertOpToLLVMPattern { cudaq::opt::factory::genLlvmI64Constant(loc, rewriter, numElements); // Set the string literal data - auto strPtr = rewriter.create( - loc, LLVM::LLVMPointerType::get(rewriter.getI8Type()), alloca, - ValueRange{zero, zero}); - auto castedPauli = rewriter.create( - loc, cudaq::opt::factory::getPointerType(context), pauliWord); + auto charPtrTy = cudaq::opt::factory::getPointerType(context); + auto strPtrTy = LLVM::LLVMPointerType::get(charPtrTy); + auto strPtr = rewriter.create(loc, strPtrTy, alloca, + ValueRange{zero, zero}); + auto castedPauli = + rewriter.create(loc, charPtrTy, pauliWord); rewriter.create(loc, castedPauli, strPtr); // Set the integer length @@ -476,8 +477,8 @@ class ExpPauliRewrite : public ConvertOpToLLVMPattern { rewriter.create(loc, size, intPtr); // Cast to raw opaque pointer - auto castedStore = rewriter.create( - loc, cudaq::opt::factory::getPointerType(context), alloca); + auto castedStore = + rewriter.create(loc, charPtrTy, alloca); operands.push_back(castedStore); rewriter.replaceOpWithNewOp(instOp, TypeRange{}, symbolRef, operands); diff --git a/lib/Optimizer/Transforms/DecompositionPatterns.cpp b/lib/Optimizer/Transforms/DecompositionPatterns.cpp index 3849329843..54ee23741f 100644 --- a/lib/Optimizer/Transforms/DecompositionPatterns.cpp +++ b/lib/Optimizer/Transforms/DecompositionPatterns.cpp @@ -335,18 +335,51 @@ struct ExpPauliDecomposition : public OpRewritePattern { LogicalResult matchAndRewrite(quake::ExpPauliOp expPauliOp, PatternRewriter &rewriter) const override { auto loc = expPauliOp.getLoc(); + auto module = expPauliOp->getParentOfType(); auto qubits = expPauliOp.getQubits(); auto theta = expPauliOp.getParameter(); auto pauliWord = expPauliOp.getPauli(); + std::optional optPauliWordStr; + if (auto defOp = + pauliWord.getDefiningOp()) { + optPauliWordStr = defOp.getStringLiteral(); + } else { + // Get the pauli word string from a constant global string generated + // during argument synthesis. + auto stringOp = expPauliOp.getOperand(2); + auto stringTy = stringOp.getType(); + if (auto charSpanTy = dyn_cast(stringTy)) { + if (auto vecInit = stringOp.getDefiningOp()) { + auto addrOp = vecInit.getOperand(0); + if (auto cast = addrOp.getDefiningOp()) + addrOp = cast.getOperand(); + if (auto addr = addrOp.getDefiningOp()) { + auto globalName = addr.getGlobalName(); + auto symbol = module.lookupSymbol(globalName); + if (auto global = dyn_cast(symbol)) { + auto attr = global.getValue(); + auto strAttr = cast(attr.value()); + optPauliWordStr = strAttr.getValue(); + } + } + } + } + } + // Assert that we have a constant known pauli word - auto defOp = pauliWord.getDefiningOp(); - if (!defOp) + if (!optPauliWordStr.has_value()) return failure(); + auto pauliWordStr = optPauliWordStr.value(); + + // Remove optional last zero character + auto size = pauliWordStr.size(); + if (size > 0 && pauliWordStr[size - 1] == '\0') + size--; + SmallVector qubitSupport; - StringRef pauliWordStr = defOp.getStringLiteral(); - for (std::size_t i = 0; i < pauliWordStr.size(); i++) { + for (std::size_t i = 0; i < size; i++) { Value index = rewriter.create(loc, i, 64); Value qubitI = rewriter.create(loc, qubits, index); if (pauliWordStr[i] != 'I') diff --git a/python/tests/backends/test_IQM.py b/python/tests/backends/test_IQM.py index d60f00aa5b..de66693485 100644 --- a/python/tests/backends/test_IQM.py +++ b/python/tests/backends/test_IQM.py @@ -200,6 +200,21 @@ def test_IQM_state_preparation_builder(): assert assert_close(counts["11"], 0., 2) +def test_exp_pauli(): + + @cudaq.kernel + def test(): + q = cudaq.qvector(2) + exp_pauli(1.0, q, "XX") + + shots = 10000 + # gives results like { 11:7074 10:0 01:0 00:2926 } + counts = cudaq.sample(test, shots_count=shots) + counts.dump() + assert assert_close(counts["01"], 0., 2) + assert assert_close(counts["10"], 0., 2) + + def test_arbitrary_unitary_synthesis(): cudaq.register_operation("custom_h", diff --git a/python/tests/backends/test_IonQ.py b/python/tests/backends/test_IonQ.py index 90f2537d7a..23813e5cf4 100644 --- a/python/tests/backends/test_IonQ.py +++ b/python/tests/backends/test_IonQ.py @@ -189,6 +189,20 @@ def test_ionq_state_preparation_builder(): assert not '11' in counts +def test_exp_pauli(): + + @cudaq.kernel + def test(): + q = cudaq.qvector(2) + exp_pauli(1.0, q, "XX") + + counts = cudaq.sample(test) + assert '00' in counts + assert '11' in counts + assert not '01' in counts + assert not '10' in counts + + def test_arbitrary_unitary_synthesis(): cudaq.register_operation("custom_h", diff --git a/python/tests/backends/test_OQC.py b/python/tests/backends/test_OQC.py index 3723916abe..317a59b9f2 100644 --- a/python/tests/backends/test_OQC.py +++ b/python/tests/backends/test_OQC.py @@ -188,6 +188,20 @@ def test_OQC_state_preparation_builder(): assert not '11' in counts +def test_exp_pauli(): + + @cudaq.kernel + def test(): + q = cudaq.qvector(2) + exp_pauli(1.0, q, "XX") + + counts = cudaq.sample(test) + assert '00' in counts + assert '11' in counts + assert not '01' in counts + assert not '10' in counts + + def test_arbitrary_unitary_synthesis(): cudaq.register_operation("custom_h", diff --git a/python/tests/backends/test_Quantinuum_LocalEmulation_builder.py b/python/tests/backends/test_Quantinuum_LocalEmulation_builder.py index e62657ea69..71647a9c70 100644 --- a/python/tests/backends/test_Quantinuum_LocalEmulation_builder.py +++ b/python/tests/backends/test_Quantinuum_LocalEmulation_builder.py @@ -149,6 +149,29 @@ def test_quantinuum_state_synthesis(): assert 'Could not successfully apply quake-synth.' in repr(e) +def test_exp_pauli(): + test = cudaq.make_kernel() + q = test.qalloc(2) + test.exp_pauli(1.0, q, "XX") + + counts = cudaq.sample(test) + assert '00' in counts + assert '11' in counts + assert not '01' in counts + assert not '10' in counts + + +def test_exp_pauli_param(): + test, w = cudaq.make_kernel(cudaq.pauli_word) + q = test.qalloc(2) + test.exp_pauli(1.0, q, w) + + # FIXME: should work after new launchKernel becomes default. + with pytest.raises(RuntimeError) as e: + counts = cudaq.sample(test, cudaq.pauli_word("XX")) + assert 'Remote rest platform Quake lowering failed.' in repr(e) + + # leave for gdb debugging if __name__ == "__main__": loc = os.path.abspath(__file__) diff --git a/python/tests/backends/test_Quantinuum_LocalEmulation_kernel.py b/python/tests/backends/test_Quantinuum_LocalEmulation_kernel.py index fd5c0f536c..27b92dc9e1 100644 --- a/python/tests/backends/test_Quantinuum_LocalEmulation_kernel.py +++ b/python/tests/backends/test_Quantinuum_LocalEmulation_kernel.py @@ -231,6 +231,33 @@ def kernel(): assert counts["1"] == 1000 +def test_exp_pauli(): + + @cudaq.kernel + def test(): + q = cudaq.qvector(2) + exp_pauli(1.0, q, "XX") + + counts = cudaq.sample(test) + assert '00' in counts + assert '11' in counts + assert not '01' in counts + assert not '10' in counts + + +def test_exp_pauli_param(): + + @cudaq.kernel + def test_param(w: cudaq.pauli_word): + q = cudaq.qvector(2) + exp_pauli(1.0, q, w) + + # FIXME: should work after new launchKernel becomes default. + with pytest.raises(RuntimeError) as e: + counts = cudaq.sample(test_param, cudaq.pauli_word("XX")) + assert 'Remote rest platform Quake lowering failed.' in repr(e) + + # leave for gdb debugging if __name__ == "__main__": loc = os.path.abspath(__file__) diff --git a/python/tests/backends/test_Quantinuum_builder.py b/python/tests/backends/test_Quantinuum_builder.py index 952b2cd582..b13c493178 100644 --- a/python/tests/backends/test_Quantinuum_builder.py +++ b/python/tests/backends/test_Quantinuum_builder.py @@ -45,7 +45,8 @@ def startUpMockServer(): if not check_server_connection(port): p.terminate() - pytest.exit("Mock server did not start in time, skipping tests.", returncode=1) + pytest.exit("Mock server did not start in time, skipping tests.", + returncode=1) yield credsName @@ -163,6 +164,18 @@ def test_quantinuum_state_preparation(): assert not '11' in counts +def test_exp_pauli(): + test = cudaq.make_kernel() + q = test.qalloc(2) + test.exp_pauli(1.0, q, "XX") + + counts = cudaq.sample(test) + assert '00' in counts + assert '11' in counts + assert not '01' in counts + assert not '10' in counts + + # leave for gdb debugging if __name__ == "__main__": loc = os.path.abspath(__file__) diff --git a/python/tests/backends/test_Quantinuum_kernel.py b/python/tests/backends/test_Quantinuum_kernel.py index 1472a1f3d4..88989be341 100644 --- a/python/tests/backends/test_Quantinuum_kernel.py +++ b/python/tests/backends/test_Quantinuum_kernel.py @@ -45,7 +45,8 @@ def startUpMockServer(): if not check_server_connection(port): p.terminate() - pytest.exit("Mock server did not start in time, skipping tests.", returncode=1) + pytest.exit("Mock server did not start in time, skipping tests.", + returncode=1) yield credsName @@ -189,6 +190,20 @@ def kernel(vec: List[complex]): assert not '11' in counts +def test_exp_pauli(): + + @cudaq.kernel + def test(): + q = cudaq.qvector(2) + exp_pauli(1.0, q, "XX") + + counts = cudaq.sample(test) + assert '00' in counts + assert '11' in counts + assert not '01' in counts + assert not '10' in counts + + # leave for gdb debugging if __name__ == "__main__": loc = os.path.abspath(__file__) diff --git a/python/tests/kernel/test_kernel_exp_pauli.py b/python/tests/kernel/test_kernel_exp_pauli.py new file mode 100644 index 0000000000..9906ba366a --- /dev/null +++ b/python/tests/kernel/test_kernel_exp_pauli.py @@ -0,0 +1,37 @@ +# ============================================================================ # +# Copyright (c) 2022 - 2024 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 + + +def test_exp_pauli(): + + @cudaq.kernel + def test(): + q = cudaq.qvector(2) + exp_pauli(1.0, q, "XX") + + counts = cudaq.sample(test) + assert '00' in counts + assert '11' in counts + assert not '01' in counts + assert not '10' in counts + + +def test_exp_pauli_param(): + + @cudaq.kernel + def test_param(w: cudaq.pauli_word): + q = cudaq.qvector(2) + exp_pauli(1.0, q, w) + + counts = cudaq.sample(test_param, cudaq.pauli_word("XX")) + assert '00' in counts + assert '11' in counts + assert not '01' in counts + assert not '10' in counts diff --git a/targettests/execution/exp_pauli.cpp b/targettests/execution/exp_pauli.cpp new file mode 100644 index 0000000000..bf7ed5bac1 --- /dev/null +++ b/targettests/execution/exp_pauli.cpp @@ -0,0 +1,62 @@ +/******************************************************************************* + * Copyright (c) 2022 - 2024 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. * + ******************************************************************************/ + +// clang-format off +// Simulators +// RUN: nvq++ %cpp_std --enable-mlir %s -o %t && %t | FileCheck %s +// RUN: nvq++ %cpp_std --enable-mlir --target remote-mqpu -fkernel-exec-kind=2 %s -o %t && %t | FileCheck %s +// +// Quantum emulators +// RUN: nvq++ %cpp_std --target quantinuum --emulate -fkernel-exec-kind=2 %s -o %t && %t | FileCheck %s +// RUN: nvq++ %cpp_std --target ionq --emulate -fkernel-exec-kind=2 %s -o %t && %t | FileCheck %s +// 2 different IQM machines for 2 different topologies +// RUN: nvq++ %cpp_std --target iqm --iqm-machine Adonis --emulate -fkernel-exec-kind=2 %s -o %t && %t | FileCheck %s +// RUN: nvq++ %cpp_std --target iqm --iqm-machine Apollo --emulate -fkernel-exec-kind=2 %s -o %t && %t | FileCheck %s +// RUN: nvq++ %cpp_std --target oqc --emulate -fkernel-exec-kind=2 %s -o %t && %t | FileCheck %s +// RUN: nvq++ %cpp_std --target anyon --emulate -fkernel-exec-kind=2 %s -o %t && %t | FileCheck %s +// clang-format on + +#include +#include + +__qpu__ void test() { + cudaq::qvector q(2); + cudaq::exp_pauli(1.0, q, "XX"); +} + +__qpu__ void test_param(cudaq::pauli_word w) { + cudaq::qvector q(2); + cudaq::exp_pauli(1.0, q, w); +} + +void printCounts(cudaq::sample_result& result) { + std::vector values{}; + for (auto &&[bits, counts] : result) { + values.push_back(bits); + } + + std::sort(values.begin(), values.end()); + for (auto &&bits : values) { + std::cout << bits << '\n'; + } +} + +int main() { + auto counts = cudaq::sample(test); + printCounts(counts); + + counts = cudaq::sample(test_param, cudaq::pauli_word{"XY"}); + printCounts(counts); + return 0; +} + +// CHECK: 00 +// CHECK: 11 + +// CHECK: 00 +// CHECK: 11 diff --git a/test/Transforms/DecompositionPatterns/ExpPauliToHRyRzCX.qke b/test/Transforms/DecompositionPatterns/ExpPauliToHRyRzCX.qke index 3cfe622af9..e7d6454a3b 100644 --- a/test/Transforms/DecompositionPatterns/ExpPauliToHRyRzCX.qke +++ b/test/Transforms/DecompositionPatterns/ExpPauliToHRyRzCX.qke @@ -22,7 +22,6 @@ module attributes {quake.mangled_name_map = {__nvqpp__mlirgen__Z4mainE3$_0 = "_Z quake.exp_pauli %4, %1, %5 : (f64, !quake.veq<4>, !cc.ptr>) -> () return } -} // CHECK-LABEL: func.func @__nvqpp__mlirgen__Z4mainE3$_0( // CHECK-SAME: %[[VAL_0:.*]]: f64) attributes {"cudaq-entrypoint", "cudaq-kernel"} { @@ -66,3 +65,37 @@ module attributes {quake.mangled_name_map = {__nvqpp__mlirgen__Z4mainE3$_0 = "_Z // CHECK: quake.h %[[VAL_20]] : (!quake.ref) -> () // CHECK: return // CHECK: } + + func.func @__nvqpp__mlirgen__function_test_param._Z10test_paramN5cudaq10pauli_wordE() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} { + %cst = arith.constant 1.000000e+00 : f64 + %c2_i64 = arith.constant 2 : i64 + %0 = cc.address_of @cstr.585900 : !cc.ptr> + %1 = cc.cast %0 : (!cc.ptr>) -> !cc.ptr + %2 = cc.stdvec_init %1, %c2_i64 : (!cc.ptr, i64) -> !cc.charspan + %3 = quake.alloca !quake.veq<2> + quake.exp_pauli %cst, %3, %2 : (f64, !quake.veq<2>, !cc.charspan) -> () + return + } + llvm.mlir.global private constant @cstr.585900("XY\00") {addr_space = 0 : i32} + +// CHECK-LABEL: func.func @__nvqpp__mlirgen__function_test_param._Z10test_paramN5cudaq10pauli_wordE() attributes {"cudaq-entrypoint", "cudaq-kernel", no_this} { +// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_2:.*]] = arith.constant 1.5707963267948966 : f64 +// CHECK: %[[VAL_3:.*]] = arith.constant -1.5707963267948966 : f64 +// CHECK: %[[VAL_4:.*]] = arith.constant 1.000000e+00 : f64 +// CHECK: %[[VAL_5:.*]] = quake.alloca !quake.veq<2> +// CHECK: %[[VAL_6:.*]] = quake.extract_ref %[[VAL_5]][%[[VAL_0]]] : (!quake.veq<2>, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_6]] : (!quake.ref) -> () +// CHECK: %[[VAL_7:.*]] = quake.extract_ref %[[VAL_5]][%[[VAL_1]]] : (!quake.veq<2>, i64) -> !quake.ref +// CHECK: quake.rx (%[[VAL_2]]) %[[VAL_7]] : (f64, !quake.ref) -> () +// CHECK: quake.x [%[[VAL_6]]] %[[VAL_7]] : (!quake.ref, !quake.ref) -> () +// CHECK: quake.rz (%[[VAL_4]]) %[[VAL_7]] : (f64, !quake.ref) -> () +// CHECK: quake.x [%[[VAL_6]]] %[[VAL_7]] : (!quake.ref, !quake.ref) -> () +// CHECK: %[[VAL_8:.*]] = quake.extract_ref %[[VAL_5]][%[[VAL_1]]] : (!quake.veq<2>, i64) -> !quake.ref +// CHECK: quake.rx (%[[VAL_3]]) %[[VAL_8]] : (f64, !quake.ref) -> () +// CHECK: %[[VAL_9:.*]] = quake.extract_ref %[[VAL_5]][%[[VAL_0]]] : (!quake.veq<2>, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_9]] : (!quake.ref) -> () +// CHECK: return +// CHECK: } +}