Keep the control when decomposing a controlled swap

cudaq/lib/Optimizer/Transforms/DecompositionPatterns.cpp

@@ -582,6 +582,12 @@ REGISTER_DECOMPOSITION_PATTERN(R1AdjToR1, {"r1<adj>", "r1"});
 // quake.cnot b, a;
 // quake.cnot a, b;
 // quake.cnot b, a;
+//
+// quake.swap [c] a, b
+// ───────────────────────────────────
+// quake.x [b] a;
+// quake.x [c, a] b;    // c is negated iff the swap control was negated
+// quake.x [b] a;
 struct SwapToCXType; // forward declare the pattern type, defined in the macro
                      // below
 struct SwapToCX
@@ -591,22 +597,55 @@ struct SwapToCX
 
   LogicalResult matchAndRewrite(cudaq::quake::SwapOp op,
                                 PatternRewriter &rewriter) const override {
+    auto numControls = cudaq::getKnownNumControls(op);
+    if (!numControls || *numControls > 1 || !isEnabled(*numControls))
+      return failure();
+
     // Op info
     Location loc = op->getLoc();
     Value a = op.getTarget(0);
     Value b = op.getTarget(1);
 
     QuakeOperatorCreator qRewriter(rewriter);
-    qRewriter.create<cudaq::quake::XOp>(loc, b, a);
-    qRewriter.create<cudaq::quake::XOp>(loc, a, b);
-    qRewriter.create<cudaq::quake::XOp>(loc, b, a);
+    if (*numControls == 1) {
+      // This is a Fredkin gate. This decomposition does not support
+      // `quake.control` types.
+      if (containsControlTypes(op))
+        return failure();
+      SmallVector<Value, 1> controls(1);
+      if (failed(checkAndExtractControls(op, controls, rewriter)))
+        return failure();
+      Value c = controls[0];
+
+      qRewriter.create<cudaq::quake::XOp>(loc, b, a);
+      SmallVector<Value, 2> ccxControls{c, a};
+      auto ccxOp = qRewriter.create<cudaq::quake::XOp>(loc, ccxControls, b);
+      // The outer cnots do not touch the control qubit, so a complemented
+      // swap control is equivalent to complementing the toffoli's first
+      // control. Let expand-control-negations materialize it downstream.
+      if (auto swapNegations = op.getNegatedQubitControls()) {
+        // One flag per control operand is required; pad for target a.
+        SmallVector<bool> flags{(*swapNegations)[0], false};
+        ccxOp.setNegatedQubitControls(
+            DenseBoolArrayAttr::get(rewriter.getContext(), flags));
+      }
+      qRewriter.create<cudaq::quake::XOp>(loc, b, ccxControls[1]);
 
-    qRewriter.selectWiresAndReplaceUses(op, ValueRange{a, b});
+      // The wires are ordered controls first, then targets.
+      qRewriter.selectWiresAndReplaceUses(
+          op, ValueRange{ccxControls[0], ccxControls[1], b});
+    } else {
+      qRewriter.create<cudaq::quake::XOp>(loc, b, a);
+      qRewriter.create<cudaq::quake::XOp>(loc, a, b);
+      qRewriter.create<cudaq::quake::XOp>(loc, b, a);
+      qRewriter.selectWiresAndReplaceUses(op, ValueRange{a, b});
+    }
     rewriter.eraseOp(op);
     return success();
   }
 };
-REGISTER_DECOMPOSITION_PATTERN(SwapToCX, {"swap", "x(1)"});
+REGISTER_DECOMPOSITION_PATTERN(SwapToCX, {"swap", "x(1)"},
+                               {"swap(1)", "x(1)", "x(2)"});
 
 // quake.h control, target
 // ───────────────────────────────────

cudaq/test/Transforms/BasisConversion/all_qir_gates.qke

@@ -315,7 +315,9 @@ module {
 // CHECK:           quake.x %[[VAL_73]] : (!quake.ref) -> ()
 // CHECK:           %[[VAL_74:.*]] = quake.extract_ref %[[VAL_71]][2] : (!quake.veq<3>) -> !quake.ref
 // CHECK:           quake.x {{\[}}%[[VAL_74]]] %[[VAL_73]] : (!quake.ref, !quake.ref) -> ()
-// CHECK:           quake.x {{\[}}%[[VAL_73]]] %[[VAL_74]] : (!quake.ref, !quake.ref) -> ()
+// The controlled swap lowers to cx b,a; ccx c,a,b; cx b,a; the toffoli in
+// the middle is further decomposed on this basis, so only pin its start.
+// CHECK:           quake.h %[[VAL_74]]
 // CHECK:           quake.x {{\[}}%[[VAL_74]]] %[[VAL_73]] : (!quake.ref, !quake.ref) -> ()
 // CHECK:           %[[VAL_75:.*]] = cc.alloca !cc.array<i8 x 3>
 // CHECK:           %[[VAL_76:.*]] = quake.mz %[[VAL_72]] : (!quake.ref) -> !quake.measure

cudaq/test/Transforms/DecompositionPatterns/SwapToCXControlled.qke

@@ -0,0 +1,51 @@
+// ========================================================================== //
+// Copyright (c) 2022 - 2026 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.   //
+// ========================================================================== //
+
+// CircuitCheck is deliberately not used here: its UnitaryBuilder has no
+// controlled multi-target path, so it miscomputes the unitary of a
+// controlled swap on the reference side.
+
+// RUN: cudaq-opt -pass-pipeline='builtin.module(decomposition{enable-patterns=SwapToCX})' %s | FileCheck %s
+
+func.func @test_controlled(%qc: !quake.ref, %qa : !quake.ref, %qb : !quake.ref) {
+  quake.swap [%qc] %qa, %qb : (!quake.ref, !quake.ref, !quake.ref) -> ()
+  return
+}
+
+// CHECK-LABEL:   func.func @test_controlled
+// CHECK:           quake.x {{\[}}%[[VAL_2:.*]]] %[[VAL_1:.*]] : (!quake.ref, !quake.ref) -> ()
+// CHECK:           quake.x {{\[}}%[[VAL_0:.*]], %[[VAL_1]]] %[[VAL_2]] : (!quake.ref, !quake.ref, !quake.ref) -> ()
+// CHECK:           quake.x {{\[}}%[[VAL_2]]] %[[VAL_1]] : (!quake.ref, !quake.ref) -> ()
+// CHECK:           return
+// CHECK:         }
+
+// A complemented control becomes a negated first control on the toffoli.
+func.func @test_complemented(%qc: !quake.ref, %qa : !quake.ref, %qb : !quake.ref) {
+  quake.swap [%qc neg [true]] %qa, %qb : (!quake.ref, !quake.ref, !quake.ref) -> ()
+  return
+}
+
+// CHECK-LABEL:   func.func @test_complemented(
+// CHECK:           quake.x {{\[}}%[[B_VAL:.*]]] %[[A_VAL:.*]] : (!quake.ref, !quake.ref) -> ()
+// CHECK-NEXT:      quake.x {{\[}}%[[C_VAL:.*]], %[[A_VAL]] neg {{\[}}true, false{{\]}}{{\]}} %[[B_VAL]] : (!quake.ref, !quake.ref, !quake.ref) -> ()
+// CHECK-NEXT:      quake.x {{\[}}%[[B_VAL]]] %[[A_VAL]] : (!quake.ref, !quake.ref) -> ()
+// CHECK:           return
+// CHECK:         }
+
+// The uncontrolled case must keep the plain three-CNOT lowering.
+func.func @test_bare(%qa: !quake.ref, %qb : !quake.ref) {
+  quake.swap %qa, %qb : (!quake.ref, !quake.ref) -> ()
+  return
+}
+
+// CHECK-LABEL:   func.func @test_bare
+// CHECK:           quake.x {{\[}}%[[VAL_1:.*]]] %[[VAL_0:.*]] : (!quake.ref, !quake.ref) -> ()
+// CHECK:           quake.x {{\[}}%[[VAL_0]]] %[[VAL_1]] : (!quake.ref, !quake.ref) -> ()
+// CHECK:           quake.x {{\[}}%[[VAL_1]]] %[[VAL_0]] : (!quake.ref, !quake.ref) -> ()
+// CHECK:           return
+// CHECK:         }

cudaq/unittests/Optimizer/DecompositionPatternSelectionTest.cpp

@@ -465,10 +465,25 @@ TEST_F(FullDecompositionPatternSelectionTest, DecomposeCCXToCZ) {
   std::vector<std::string> targetBasis{"h", "t", "z(1)"};
   auto selectedPatterns = selectPatterns(targetBasis);
 
-  std::vector<std::string> exp{"CCXToCCZ", "CCZToCX", "CXToCZ", "SwapToCX"};
+  // SwapToCX is selected once per registered source set: besides the plain
+  // swap entry, its single-control entry now also reaches this basis since
+  // x(2) decomposes through CCXToCCZ and CCZToCX.
+  std::vector<std::string> exp{"CCXToCCZ", "CCZToCX", "CXToCZ", "SwapToCX",
+                               "SwapToCX"};
   EXPECT_EQ(selectedPatterns, exp);
 }
 
+TEST_F(FullDecompositionPatternSelectionTest,
+       SwapToCXDisablesControlCountsCoveredByTheBasis) {
+  // Registering the controlled variant joins both source sets into an
+  // unbounded swap(n) source, so control counts already present in the
+  // basis are disabled and native swaps are preserved.
+  auto pattern = constructPattern({"swap", "x(1)"}, "SwapToCX");
+
+  std::vector<std::size_t> exp{0};
+  EXPECT_EQ(pattern->getDisabledControlCounts(), llvm::ArrayRef(exp));
+}
+
 // Regression: multi-hop chain where intermediate gates (t, z(2)) are not
 // in the basis but are reachable through further patterns.
 // Chain: x(2) -> CCXToCCZ -> {h,z(2)} -> CCZToCX -> {t,x(1)}

cudaq/unittests/Optimizer/DecompositionPatternsTest.cpp

@@ -85,7 +85,10 @@ ModuleOp createTestModule(MLIRContext *context, StringRef gateSpecStr) {
   numControls = std::min<size_t>(numControls, 2);
 
   size_t numQubits;
-  if (gateName == "swap" || gateName == "exp_pauli") {
+  if (gateName == "swap") {
+    // Swap needs two targets plus one qubit per control.
+    numQubits = numControls + 2;
+  } else if (gateName == "exp_pauli") {
     assert(numControls == 0);
     // exp_pauli can have any number of qubits, we hardcode to 2 for the test.
     numQubits = 2;
@@ -156,9 +159,10 @@ ModuleOp createTestModule(MLIRContext *context, StringRef gateSpecStr) {
         builder, loc, isAdj, ValueRange{{pi_2, pi_2}}, controls, target);
   } else if (gateName == "swap") {
     // Swap needs 2 targets
-    Value target = entry->getArgument(0);
-    Value target2 = entry->getArgument(1);
-    cudaq::quake::SwapOp::create(builder, loc, ValueRange{target, target2});
+    Value target = entry->getArgument(numControls);
+    Value target2 = entry->getArgument(numControls + 1);
+    cudaq::quake::SwapOp::create(builder, loc, controls,
+                                 ValueRange{target, target2});
   } else if (gateName == "exp_pauli") {
     Value target = entry->getArgument(0);
     Value target2 = entry->getArgument(1);
@@ -481,6 +485,177 @@ TEST_F(DecompositionPatternsTest, SAndTToR1AcceptDynamicControlsWhenNEnabled) {
   EXPECT_EQ(countOps<cudaq::quake::R1Op>(tModule), 1u);
 }
 
+TEST_F(DecompositionPatternsTest, SwapToCXKeepsSingleControl) {
+  // A controlled swap is a Fredkin gate: cx b,a; ccx c,a,b; cx b,a.
+  auto module = createTestModule(context.get(), "swap(1)");
+  ASSERT_TRUE(succeeded(applySinglePattern(module, "SwapToCX", {})));
+  EXPECT_EQ(countOps<cudaq::quake::SwapOp>(module), 0u);
+  EXPECT_EQ(countOps<cudaq::quake::XOp>(module), 3u);
+  auto gates = collectGateTypesInModule(module);
+  EXPECT_TRUE(gates.contains("x(1)"));
+  EXPECT_TRUE(gates.contains("x(2)"));
+
+  // The toffoli of the Fredkin decomposition must target the second swap
+  // target with controls {control qubit, first swap target}; targeting any
+  // other wire yields something that is not a cswap.
+  func::FuncOp testFunc;
+  SmallVector<Value> args;
+  module.walk([&](func::FuncOp op) {
+    if (!testFunc)
+      testFunc = op;
+  });
+  ASSERT_TRUE(static_cast<bool>(testFunc));
+  for (auto arg : testFunc.getArguments())
+    args.push_back(arg);
+  std::size_t toffolis = 0;
+  module.walk([&](cudaq::quake::XOp xop) {
+    if (xop.getControls().size() != 2)
+      return;
+    ++toffolis;
+    EXPECT_EQ(xop.getTarget(), args[2]);
+    EXPECT_TRUE(llvm::is_contained(xop.getControls(), args[0]));
+    EXPECT_TRUE(llvm::is_contained(xop.getControls(), args[1]));
+  });
+  EXPECT_EQ(toffolis, 1u);
+
+  // The uncontrolled case must remain a plain sequence of three CNOTs.
+  auto bareModule = createTestModule(context.get(), "swap");
+  ASSERT_TRUE(succeeded(applySinglePattern(bareModule, "SwapToCX", {})));
+  EXPECT_EQ(countOps<cudaq::quake::SwapOp>(bareModule), 0u);
+  EXPECT_EQ(countOps<cudaq::quake::XOp>(bareModule), 3u);
+  auto bareGates = collectGateTypesInModule(bareModule);
+  EXPECT_FALSE(bareGates.contains("x(2)"));
+}
+
+TEST_F(DecompositionPatternsTest, SwapToCXThreadsWireOperands) {
+  // Under wire semantics each rewritten gate must consume its predecessor's
+  // output wires, and the original op's results are replaced by the
+  // toffoli's control wire and the final wires of the two targets.
+  OpBuilder builder(context.get());
+  auto module = ModuleOp::create(builder, builder.getUnknownLoc());
+  builder.setInsertionPointToEnd(module.getBody());
+  auto wireTy = cudaq::quake::WireType::get(context.get());
+  SmallVector<Type> wireTypes{wireTy, wireTy, wireTy};
+  auto funcType = builder.getFunctionType(wireTypes, wireTypes);
+  auto func = func::FuncOp::create(builder, builder.getUnknownLoc(),
+                                   "test_func", funcType);
+  auto *entry = func.addEntryBlock();
+  builder.setInsertionPointToStart(entry);
+  Value c = entry->getArgument(0);
+  Value a = entry->getArgument(1);
+  Value b = entry->getArgument(2);
+  auto swapOp = builder.create<cudaq::quake::SwapOp>(
+      builder.getUnknownLoc(), wireTypes, /*is_adj=*/false, ValueRange{},
+      ValueRange{c}, ValueRange{a, b}, DenseBoolArrayAttr{});
+  builder.create<func::ReturnOp>(builder.getUnknownLoc(), swapOp.getResults());
+  ASSERT_TRUE(succeeded(applySinglePattern(module, "SwapToCX", {})));
+  EXPECT_EQ(countOps<cudaq::quake::SwapOp>(module), 0u);
+
+  SmallVector<cudaq::quake::XOp> xs;
+  func::ReturnOp ret;
+  module.walk([&](cudaq::quake::XOp xop) { xs.push_back(xop); });
+  module.walk([&](func::ReturnOp op) { ret = op; });
+  ASSERT_EQ(xs.size(), 3u);
+  auto &cx0 = xs[0];
+  auto &ccx = xs[1];
+  auto &cx1 = xs[2];
+  ASSERT_EQ(cx0.getControls().size(), 1u);
+  ASSERT_EQ(ccx.getControls().size(), 2u);
+  ASSERT_EQ(cx1.getControls().size(), 1u);
+  EXPECT_EQ(cx0.getControls()[0], b);
+  EXPECT_EQ(cx0.getTarget(), a);
+  EXPECT_EQ(ccx.getControls()[0], c);
+  EXPECT_EQ(ccx.getControls()[1], cx0.getResult(1));
+  EXPECT_EQ(ccx.getTarget(), cx0.getResult(0));
+  EXPECT_EQ(cx1.getControls()[0], ccx.getResult(2));
+  EXPECT_EQ(cx1.getTarget(), ccx.getResult(1));
+  ASSERT_TRUE(static_cast<bool>(ret));
+  EXPECT_EQ(ret.getOperands()[0], ccx.getResult(0));
+  EXPECT_EQ(ret.getOperands()[1], cx1.getResult(1));
+  EXPECT_EQ(ret.getOperands()[2], cx1.getResult(0));
+}
+
+TEST_F(DecompositionPatternsTest, SwapToCXForwardsComplementedControl) {
+  // A complemented swap control becomes a complemented first control on the
+  // toffoli of the Fredkin decomposition instead of an x pair around the
+  // whole sequence.
+  auto module = createTestModule(context.get(), "swap(1)");
+  cudaq::quake::SwapOp swapOp;
+  module.walk([&](cudaq::quake::SwapOp op) {
+    if (!swapOp)
+      swapOp = op;
+  });
+  ASSERT_TRUE(static_cast<bool>(swapOp.getOperation()));
+  swapOp.setNegatedQubitControls(
+      DenseBoolArrayAttr::get(context.get(), {true}));
+  ASSERT_TRUE(succeeded(applySinglePattern(module, "SwapToCX", {})));
+  EXPECT_EQ(countOps<cudaq::quake::SwapOp>(module), 0u);
+  EXPECT_EQ(countOps<cudaq::quake::XOp>(module), 3u);
+
+  func::FuncOp testFunc;
+  module.walk([&](func::FuncOp op) {
+    if (!testFunc)
+      testFunc = op;
+  });
+  ASSERT_TRUE(static_cast<bool>(testFunc));
+  Value ctrlArg = testFunc.getArgument(0);
+  std::size_t toffolis = 0;
+  module.walk([&](cudaq::quake::XOp xop) {
+    if (xop.getControls().size() != 2)
+      return;
+    ++toffolis;
+    auto negated = xop.getNegatedQubitControls();
+    ASSERT_TRUE(negated.has_value());
+    EXPECT_TRUE((*negated)[0]);
+    EXPECT_EQ(xop.getControls()[0], ctrlArg);
+  });
+  EXPECT_EQ(toffolis, 1u);
+}
+
+TEST_F(DecompositionPatternsTest, SwapToCXRejectsUnknownOrMultipleControls) {
+  // More than one control must be left untouched instead of silently
+  // dropping controls. The rewrite driver reports success either way, so
+  // assert on the surviving swap op.
+  auto twoCtrlModule = createTestModule(context.get(), "swap(2)");
+  ASSERT_TRUE(succeeded(applySinglePattern(twoCtrlModule, "SwapToCX", {})));
+  EXPECT_EQ(countOps<cudaq::quake::SwapOp>(twoCtrlModule), 1u);
+
+  // An unsized veq control has an unknown control count and must also be
+  // left untouched.
+  OpBuilder builder(context.get());
+  auto module = ModuleOp::create(builder, builder.getUnknownLoc());
+  builder.setInsertionPointToEnd(module.getBody());
+  auto refType = cudaq::quake::RefType::get(context.get());
+  SmallVector<Type> inputTypes{cudaq::quake::VeqType::getUnsized(context.get()),
+                               refType, refType};
+  auto funcType = builder.getFunctionType(inputTypes, {});
+  auto func = func::FuncOp::create(builder, builder.getUnknownLoc(),
+                                   "test_func", funcType);
+  auto *entry = func.addEntryBlock();
+  builder.setInsertionPointToStart(entry);
+  cudaq::quake::SwapOp::create(
+      builder, builder.getUnknownLoc(), ValueRange{entry->getArgument(0)},
+      ValueRange{entry->getArgument(1), entry->getArgument(2)});
+  ASSERT_TRUE(succeeded(applySinglePattern(module, "SwapToCX", {})));
+  EXPECT_EQ(countOps<cudaq::quake::SwapOp>(module), 1u);
+}
+
+TEST_F(DecompositionPatternsTest, SwapToCXRespectsDisabledControlCounts) {
+  // A basis that already provides a plain swap disables that control count,
+  // so bare swaps survive while the controlled variant still decomposes.
+  std::vector<std::size_t> disabled{0};
+
+  auto bareModule = createTestModule(context.get(), "swap");
+  ASSERT_TRUE(succeeded(applySinglePattern(bareModule, "SwapToCX", disabled)));
+  EXPECT_EQ(countOps<cudaq::quake::SwapOp>(bareModule), 1u);
+
+  auto ctrlModule = createTestModule(context.get(), "swap(1)");
+  ASSERT_TRUE(succeeded(applySinglePattern(ctrlModule, "SwapToCX", disabled)));
+  EXPECT_EQ(countOps<cudaq::quake::SwapOp>(ctrlModule), 0u);
+  auto gates = collectGateTypesInModule(ctrlModule);
+  EXPECT_TRUE(gates.contains("x(2)"));
+}
+
 // Test 4: Verify pattern decompositions produce only physical target gates
 TEST_F(DecompositionPatternsTest,
        DecompositionProducesOnlyPhysicalTargetGates) {