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

@@ -462,11 +462,12 @@ class ExpPauliRewrite : public ConvertOpToLLVMPattern<quake::ExpPauliOp> {
           cudaq::opt::factory::genLlvmI64Constant(loc, rewriter, numElements);
 
       // Set the string literal data
-      auto strPtr = rewriter.create<LLVM::GEPOp>(
-          loc, LLVM::LLVMPointerType::get(rewriter.getI8Type()), alloca,
-          ValueRange{zero, zero});
-      auto castedPauli = rewriter.create<LLVM::BitcastOp>(
-          loc, cudaq::opt::factory::getPointerType(context), pauliWord);
+      auto charPtrTy = cudaq::opt::factory::getPointerType(context);
+      auto strPtrTy = LLVM::LLVMPointerType::get(charPtrTy);
+      auto strPtr = rewriter.create<LLVM::GEPOp>(loc, strPtrTy, alloca,
+                                                 ValueRange{zero, zero});
+      auto castedPauli =
+          rewriter.create<LLVM::BitcastOp>(loc, charPtrTy, pauliWord);
       rewriter.create<LLVM::StoreOp>(loc, castedPauli, strPtr);
 
       // Set the integer length
@@ -476,8 +477,8 @@ class ExpPauliRewrite : public ConvertOpToLLVMPattern<quake::ExpPauliOp> {
       rewriter.create<LLVM::StoreOp>(loc, size, intPtr);
 
       // Cast to raw opaque pointer
-      auto castedStore = rewriter.create<LLVM::BitcastOp>(
-          loc, cudaq::opt::factory::getPointerType(context), alloca);
+      auto castedStore =
+          rewriter.create<LLVM::BitcastOp>(loc, charPtrTy, alloca);
       operands.push_back(castedStore);
       rewriter.replaceOpWithNewOp<LLVM::CallOp>(instOp, TypeRange{}, symbolRef,
                                                 operands);

lib/Optimizer/Transforms/DecompositionPatterns.cpp

@@ -335,18 +335,51 @@ struct ExpPauliDecomposition : public OpRewritePattern<quake::ExpPauliOp> {
   LogicalResult matchAndRewrite(quake::ExpPauliOp expPauliOp,
                                 PatternRewriter &rewriter) const override {
     auto loc = expPauliOp.getLoc();
+    auto module = expPauliOp->getParentOfType<ModuleOp>();
     auto qubits = expPauliOp.getQubits();
     auto theta = expPauliOp.getParameter();
     auto pauliWord = expPauliOp.getPauli();
 
+    std::optional<StringRef> optPauliWordStr;
+    if (auto defOp =
+            pauliWord.getDefiningOp<cudaq::cc::CreateStringLiteralOp>()) {
+      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<cudaq::cc::CharspanType>(stringTy)) {
+        if (auto vecInit = stringOp.getDefiningOp<cudaq::cc::StdvecInitOp>()) {
+          auto addrOp = vecInit.getOperand(0);
+          if (auto cast = addrOp.getDefiningOp<cudaq::cc::CastOp>())
+            addrOp = cast.getOperand();
+          if (auto addr = addrOp.getDefiningOp<cudaq::cc::AddressOfOp>()) {
+            auto globalName = addr.getGlobalName();
+            auto symbol = module.lookupSymbol(globalName);
+            if (auto global = dyn_cast<LLVM::GlobalOp>(symbol)) {
+              auto attr = global.getValue();
+              auto strAttr = cast<mlir::StringAttr>(attr.value());
+              optPauliWordStr = strAttr.getValue();
+            }
+          }
+        }
+      }
+    }
+
     // Assert that we have a constant known pauli word
-    auto defOp = pauliWord.getDefiningOp<cudaq::cc::CreateStringLiteralOp>();
-    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<Value> 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<arith::ConstantIntOp>(loc, i, 64);
       Value qubitI = rewriter.create<quake::ExtractRefOp>(loc, qubits, index);
       if (pauliWordStr[i] != 'I')

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",

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",

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",

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

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

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

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

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

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 <cudaq.h>
+#include <iostream>
+
+__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<std::string> 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