Allow controlled Z gates for profile-specific QIR (#2245)

NVIDIA · Oct 8, 2024 · f3907a3 · f3907a3
1 parent 77a4d82
commit f3907a3
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 9 deletions.
diff --git a/include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h b/include/cudaq/Optimizer/CodeGen/QIRFunctionNames.h
@@ -24,6 +24,7 @@ static constexpr const char QIRMeasureToRegister[] =
 
 static constexpr const char QIRCnot[] = "__quantum__qis__cnot";
 static constexpr const char QIRCphase[] = "__quantum__qis__cphase";
+static constexpr const char QIRCZ[] = "__quantum__qis__cz";
 static constexpr const char QIRReadResultBody[] =
     "__quantum__qis__read_result__body";
 

diff --git a/lib/Optimizer/CodeGen/ConvertToQIRProfile.cpp b/lib/Optimizer/CodeGen/ConvertToQIRProfile.cpp
@@ -434,6 +434,32 @@ struct CallAlloc : public OpRewritePattern<LLVM::CallOp> {
   }
 };
 
+// %1 = address_of @__quantum__qis__z__ctl
+// %2 = call @invokewithControlBits %1, %ctrl, %targ
+// ─────────────────────────────────────────────────
+// %2 = call __quantum__qis__cz %ctrl, %targ
+struct ZCtrlOneTargetToCZ : public OpRewritePattern<LLVM::CallOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(LLVM::CallOp call,
+                                PatternRewriter &rewriter) const override {
+    ValueRange args(call.getArgOperands());
+    if (args.size() == 4 && call.getCallee() &&
+        call.getCallee()->equals(cudaq::opt::NVQIRInvokeWithControlBits)) {
+      if (auto addrOf = dyn_cast_or_null<mlir::LLVM::AddressOfOp>(
+              args[1].getDefiningOp())) {
+        if (addrOf.getGlobalName().startswith(
+                std::string(cudaq::opt::QIRQISPrefix) + "z__ctl")) {
+          rewriter.replaceOpWithNewOp<LLVM::CallOp>(
+              call, TypeRange{}, cudaq::opt::QIRCZ, args.drop_front(2));
+          return success();
+        }
+      }
+    }
+    return failure();
+  }
+};
+
 /// QIR to the Specific QIR Profile
 ///
 /// This pass converts patterns in LLVM-IR dialect using QIR calls, etc. into a
@@ -457,10 +483,11 @@ struct QIRToQIRProfileQIRPass
     RewritePatternSet patterns(context);
     // Note: LoadMeasureResult is not compliant with the Base Profile, so don't
     // add it here unless we're specifically doing the Adaptive Profile.
-    patterns.insert<AddrOfCisToBase, ArrayGetElementPtrConv, CallAlloc,
-                    CalleeConv, EraseArrayAlloc, EraseArrayRelease,
-                    EraseDeadArrayGEP, MeasureCallConv,
-                    MeasureToRegisterCallConv, XCtrlOneTargetToCNot>(context);
+    patterns
+        .insert<AddrOfCisToBase, ArrayGetElementPtrConv, CallAlloc, CalleeConv,
+                EraseArrayAlloc, EraseArrayRelease, EraseDeadArrayGEP,
+                MeasureCallConv, MeasureToRegisterCallConv,
+                XCtrlOneTargetToCNot, ZCtrlOneTargetToCZ>(context);
     if (convertTo.getValue() == "qir-adaptive")
       patterns.insert<LoadMeasureResult>(context);
     if (failed(applyPatternsAndFoldGreedily(op, std::move(patterns))))
@@ -505,6 +532,11 @@ struct QIRProfilePreparationPass
         cudaq::opt::QIRCnot, LLVM::LLVMVoidType::get(ctx),
         {cudaq::opt::getQubitType(ctx), cudaq::opt::getQubitType(ctx)}, module);
 
+    // Add cz declaration as it may be referenced after peepholes run.
+    cudaq::opt::factory::createLLVMFunctionSymbol(
+        cudaq::opt::QIRCZ, LLVM::LLVMVoidType::get(ctx),
+        {cudaq::opt::getQubitType(ctx), cudaq::opt::getQubitType(ctx)}, module);
+
     // Add measure_body as it has a different signature than measure.
     cudaq::opt::factory::createLLVMFunctionSymbol(
         cudaq::opt::QIRMeasureBody, LLVM::LLVMVoidType::get(ctx),

diff --git a/runtime/nvqir/NVQIR.cpp b/runtime/nvqir/NVQIR.cpp
@@ -542,6 +542,14 @@ void __quantum__qis__cnot__body(Qubit *q, Qubit *r) {
   nvqir::getCircuitSimulatorInternal()->x(controls, rI);
 }
 
+void __quantum__qis__cz__body(Qubit *q, Qubit *r) {
+  auto qI = qubitToSizeT(q);
+  auto rI = qubitToSizeT(r);
+  ScopedTraceWithContext("NVQIR::cz", qI, rI);
+  std::vector<std::size_t> controls{qI};
+  nvqir::getCircuitSimulatorInternal()->z(controls, rI);
+}
+
 void __quantum__qis__reset(Qubit *q) {
   auto qI = qubitToSizeT(q);
   ScopedTraceWithContext("NVQIR::reset", qI);

diff --git a/targettests/execution/qspan_slices.cpp b/targettests/execution/qspan_slices.cpp
@@ -8,11 +8,12 @@
 
 // REQUIRES: c++20
 // clang-format off
-// RUN: nvq++ --target anyon                    --emulate %s -o %t && %t | FileCheck %s
-// RUN: nvq++ --target ionq                     --emulate %s -o %t && %t | FileCheck %s
-// RUN: nvq++ --target iqm --iqm-machine Adonis --emulate %s -o %t && %t | FileCheck %s
-// RUN: nvq++ --target oqc                      --emulate %s -o %t && %t | FileCheck %s
-// RUN: nvq++ --target quantinuum               --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target anyon                              --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target anyon --anyon-machine berkeley-25q --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target ionq                               --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target iqm --iqm-machine Adonis           --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target oqc                                --emulate %s -o %t && %t | FileCheck %s
+// RUN: nvq++ --target quantinuum                         --emulate %s -o %t && %t | FileCheck %s
 // Tests for --disable-qubit-mapping:
 // RUN: nvq++ -v %s -o %t --target oqc --emulate --disable-qubit-mapping && CUDAQ_MLIR_PRINT_EACH_PASS=1 %t |& FileCheck --check-prefix=DISABLE %s
 // RUN: nvq++ -v %s -o %t --target iqm --iqm-machine Adonis --emulate --disable-qubit-mapping && CUDAQ_MLIR_PRINT_EACH_PASS=1 %t |& FileCheck --check-prefix=DISABLE %s