From b7389a5cc01c92b9941f231e00269bc10e3be796 Mon Sep 17 00:00:00 2001
From: Pradnya Khalate <pkhalate@nvidia.com>
Date: Fri, 12 Jul 2024 11:45:24 -0700
Subject: [PATCH] * Addressing review comments:   -- No braces for single
 statements in if/else/for as per LLVM standards.   -- Additional test for
 remote execution of Python code   -- Check to emit error if custom operation
 is encountered in QIR

---
 lib/Frontend/nvqpp/ASTBridge.cpp               |  6 +++---
 lib/Optimizer/CodeGen/ConvertToQIRProfile.cpp  |  3 ++-
 lib/Optimizer/Transforms/GetConcreteMatrix.cpp |  3 +--
 python/tests/remote/test_remote_code_exec.py   | 18 ++++++++++++++++++
 runtime/nvqir/NVQIR.cpp                        | 13 +++++--------
 5 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/lib/Frontend/nvqpp/ASTBridge.cpp b/lib/Frontend/nvqpp/ASTBridge.cpp
index f853c55544..4a14c1cd7b 100644
--- a/lib/Frontend/nvqpp/ASTBridge.cpp
+++ b/lib/Frontend/nvqpp/ASTBridge.cpp
@@ -234,11 +234,11 @@ class QPUCodeFinder : public clang::RecursiveASTVisitor<QPUCodeFinder> {
               cudaq::details::getTagNameOfFunctionDecl(func, mangler);
           runChecks = true;
         }
-      if (cudaq::ASTBridgeAction::ASTBridgeConsumer::isQuantum(func)) {
+      if (cudaq::ASTBridgeAction::ASTBridgeConsumer::isQuantum(func))
         runChecks = true;
-      } else {
+      else
         quantumTypesNotAllowed = true;
-      }
+
       if (runChecks) {
         quantumTypesNotAllowed = false;
         // Run semantics checks on the kernel class.
diff --git a/lib/Optimizer/CodeGen/ConvertToQIRProfile.cpp b/lib/Optimizer/CodeGen/ConvertToQIRProfile.cpp
index 5b5e342c45..fcc719c947 100644
--- a/lib/Optimizer/CodeGen/ConvertToQIRProfile.cpp
+++ b/lib/Optimizer/CodeGen/ConvertToQIRProfile.cpp
@@ -547,7 +547,8 @@ struct VerifyQIRProfilePass
     func.walk([&](Operation *op) {
       if (auto call = dyn_cast<LLVM::CallOp>(op)) {
         auto funcName = call.getCalleeAttr().getValue();
-        if (!funcName.startswith("__quantum_")) {
+        if (!funcName.startswith("__quantum_") ||
+            funcName.equals(cudaq::opt::QIRCustomOp)) {
           call.emitOpError("unexpected call in QIR base profile");
           passFailed = true;
           return WalkResult::advance();
diff --git a/lib/Optimizer/Transforms/GetConcreteMatrix.cpp b/lib/Optimizer/Transforms/GetConcreteMatrix.cpp
index 0b2850a907..5cdaef9666 100644
--- a/lib/Optimizer/Transforms/GetConcreteMatrix.cpp
+++ b/lib/Optimizer/Transforms/GetConcreteMatrix.cpp
@@ -44,9 +44,8 @@ class CustomUnitaryPattern
 
     auto parentModule = customOp->getParentOfType<ModuleOp>();
     auto funcOp = parentModule.lookupSymbol<func::FuncOp>(generator);
-    if (!funcOp) {
+    if (!funcOp)
       return failure();
-    }
 
     // The generator function returns a concrete matrix. If prior passes have
     // run to constant fold and lift array values, the generator function will
diff --git a/python/tests/remote/test_remote_code_exec.py b/python/tests/remote/test_remote_code_exec.py
index e7a6c81e0d..00d2fea78f 100644
--- a/python/tests/remote/test_remote_code_exec.py
+++ b/python/tests/remote/test_remote_code_exec.py
@@ -12,6 +12,7 @@
 import subprocess
 import time
 import psutil
+import numpy as np
 
 import cudaq
 from cudaq import spin
@@ -351,6 +352,23 @@ def test_complex_vqe_named_lambda_sweep_grad(gradient):
     test_complex_vqe_named_lambda(cudaq.optimizers.Adam(), gradient)
 
 
+def test_arbitrary_unitary_synthesis():
+    cudaq.register_operation("custom_h",
+                             1. / np.sqrt(2.) * np.array([1, 1, 1, -1]))
+    cudaq.register_operation("custom_x", np.array([0, 1, 1, 0]))
+
+    @cudaq.kernel
+    def bell():
+        qubits = cudaq.qvector(2)
+        custom_h(qubits[0])
+        custom_x.ctrl(qubits[0], qubits[1])
+
+    counts = cudaq.sample(bell)
+    assert len(counts) == 2
+    assert "00" in counts
+    assert "11" in counts
+
+
 # leave for gdb debugging
 if __name__ == "__main__":
     loc = os.path.abspath(__file__)
diff --git a/runtime/nvqir/NVQIR.cpp b/runtime/nvqir/NVQIR.cpp
index ca65f113e0..ca1f19a6d9 100644
--- a/runtime/nvqir/NVQIR.cpp
+++ b/runtime/nvqir/NVQIR.cpp
@@ -597,20 +597,17 @@ void __quantum__qis__custom_unitary__adj(std::complex<double> *unitary,
   std::vector<std::vector<std::complex<double>>> unitaryConj2D;
   for (std::size_t r = 0; r < nToPowTwo; r++) {
     std::vector<std::complex<double>> row;
-    for (std::size_t c = 0; c < nToPowTwo; c++) {
+    for (std::size_t c = 0; c < nToPowTwo; c++)
       row.push_back(std::conj(unitary[r * nToPowTwo + c]));
-    }
     unitaryConj2D.push_back(row);
   }
-  for (std::size_t r = 0; r < nToPowTwo; r++) {
-    for (std::size_t c = 0; c < r; c++) {
+  for (std::size_t r = 0; r < nToPowTwo; r++)
+    for (std::size_t c = 0; c < r; c++)
       std::swap(unitaryConj2D[r][c], unitaryConj2D[c][r]);
-    }
-  }
   std::vector<std::complex<double>> unitaryFlattened;
-  for (auto const &row : unitaryConj2D) {
+  for (auto const &row : unitaryConj2D)
     unitaryFlattened.insert(unitaryFlattened.end(), row.begin(), row.end());
-  }
+
   nvqir::getCircuitSimulatorInternal()->applyCustomOperation(unitaryFlattened,
                                                              ctrlsVec, tgtsVec);
 }