From 82d4a6d968ee16915a8aaec75432b7d9e40afc59 Mon Sep 17 00:00:00 2001
From: Bruno Schmitt <bschmitt@nvidia.com>
Date: Wed, 28 Jun 2023 19:18:35 +0200
Subject: [PATCH] [bug] Fix multicontrol decomposition (#302)

---
 .../Transforms/MultiControlDecomposition.cpp   | 18 ++++++++++--------
 test/Transforms/multicontrol_decomposition.qke | 18 ++++++++++++++++++
 2 files changed, 28 insertions(+), 8 deletions(-)

diff --git a/lib/Optimizer/Transforms/MultiControlDecomposition.cpp b/lib/Optimizer/Transforms/MultiControlDecomposition.cpp
index 75ca22f51c..7570f0b4fb 100644
--- a/lib/Optimizer/Transforms/MultiControlDecomposition.cpp
+++ b/lib/Optimizer/Transforms/MultiControlDecomposition.cpp
@@ -63,11 +63,11 @@ class Decomposer {
                                 SmallVectorImpl<Value> &newControls,
                                 SmallVectorImpl<bool> &negatedControls);
 
-  void checkAndAddAncillas(Location loc, std::size_t numAncillas);
+  ArrayRef<Value> getAncillas(Location loc, std::size_t numAncillas);
 
   OpBuilder builder;
   Block *entryBlock;
-  SmallVector<Value> ancillas;
+  SmallVector<Value> allocatedAncillas;
 };
 
 } // namespace
@@ -95,11 +95,13 @@ Decomposer::extractControls(quake::OperatorInterface op,
   return success();
 }
 
-void Decomposer::checkAndAddAncillas(Location loc, std::size_t numAncillas) {
+ArrayRef<Value> Decomposer::getAncillas(Location loc, std::size_t numAncillas) {
   OpBuilder::InsertionGuard g(builder);
   builder.setInsertionPointToStart(entryBlock);
-  for (size_t i = ancillas.size(); i < numAncillas; ++i)
-    ancillas.push_back(builder.create<quake::AllocaOp>(loc));
+  // If we don't have enough ancillas, allocate new more.
+  for (size_t i = allocatedAncillas.size(); i < numAncillas; ++i)
+    allocatedAncillas.push_back(builder.create<quake::AllocaOp>(loc));
+  return {allocatedAncillas.begin(), allocatedAncillas.begin() + numAncillas};
 }
 
 LogicalResult Decomposer::v_decomposition(quake::OperatorInterface op) {
@@ -130,13 +132,13 @@ LogicalResult Decomposer::v_decomposition(quake::OperatorInterface op) {
   size_t requiredAncillas = isa<quake::XOp, quake::ZOp>(op)
                                 ? controls.size() - 2
                                 : controls.size() - 1;
-  checkAndAddAncillas(loc, requiredAncillas);
+  auto ancillas = getAncillas(loc, requiredAncillas);
 
   // Compute intermediate results
   SmallVector<Operation *> toCleanup;
   std::array<Value, 2> cs = {controls[0], controls[1]};
   toCleanup.push_back(builder.create<quake::XOp>(loc, cs, ancillas[0]));
-  if (!negatedControls.empty() && (negatedControls[0] || negatedControls[0]))
+  if (!negatedControls.empty() && (negatedControls[0] || negatedControls[1]))
     toCleanup.back()->setAttr("negated_qubit_controls",
                               builder.getDenseBoolArrayAttr(
                                   {negatedControls[0], negatedControls[1]}));
@@ -149,7 +151,7 @@ LogicalResult Decomposer::v_decomposition(quake::OperatorInterface op) {
   }
 
   // Compute output
-  if (isa<quake::XOp, quake::ZOp>(op)) {
+  if (!isa<quake::XOp, quake::ZOp>(op)) {
     createOperator(loc, name, parameters, ancillas.back(), targets, builder);
   } else {
     cs = {controls.back(), ancillas.back()};
diff --git a/test/Transforms/multicontrol_decomposition.qke b/test/Transforms/multicontrol_decomposition.qke
index a46eb6af68..5c4e0bc1f3 100644
--- a/test/Transforms/multicontrol_decomposition.qke
+++ b/test/Transforms/multicontrol_decomposition.qke
@@ -39,6 +39,13 @@ func.func @cccx_negated_controls(%c0: !quake.ref,  %c1: !quake.ref, %c2: !quake.
   return
 }
 
+// CHECK-LABEL: func.func @c4x_negated_controls
+func.func @c4x_negated_controls(%c0: !quake.ref,  %c1: !quake.ref, %c2: !quake.ref, %c3: !quake.ref, %t: !quake.ref) {
+  // CHECK-NOT: quake.x [%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}} neg [false, true, true, false]] %{{.*}}
+  quake.x [%c0, %c1, %c2, %c3 neg [false, true, true, false]] %t : (!quake.ref, !quake.ref, !quake.ref, !quake.ref, !quake.ref) -> ()
+  return
+}
+
 // CHECK-LABEL: func.func @cccx_veq
 func.func @cccx_veq(%c: !quake.veq<3>, %t: !quake.ref) {
   // CHECK-NOT: quake.x [%{{.*}}] %{{.*}} : (!quake.veq<3>, !quake.qref) -> ()
@@ -56,6 +63,17 @@ func.func @shared_ancilla(%c: !quake.veq<3>, %t: !quake.ref) {
   return
 }
 
+// CHECK-LABEL: func.func @shared_ancilla_2
+func.func @shared_ancilla_2(%c0: !quake.ref,  %c1: !quake.ref, %c2: !quake.ref, %c3: !quake.ref, %t: !quake.ref) {
+  // CHECK-NOT: quake.x [%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}] %{{.*}}
+  // CHECK-NOT: quake.z [%{{.*}}, %{{.*}}, %{{.*}}] %{{.*}}
+  // This will need two ancillas
+  quake.x [%c0, %c1, %c2, %c3] %t : (!quake.ref, !quake.ref, !quake.ref, !quake.ref, !quake.ref) -> ()
+  // This will need one ancilla
+  quake.z [%c0, %c1, %c2] %t : (!quake.ref, !quake.ref, !quake.ref, !quake.ref) -> ()
+  return
+}
+
 // CHECK-LABEL: func.func @cccz
 func.func @cccz(%c0: !quake.ref,  %c1: !quake.ref, %c2: !quake.ref, %t: !quake.ref) {
   // CHECK-NOT: quake.z [%{{.*}}, %{{.*}}, %{{.*}}] %{{.*}}