✨ Add fuse-two-qubit-unitary-runs pass

CHANGELOG.md

@@ -20,6 +20,9 @@ releases may include breaking changes.
   [**@li-mingbao**], [**@DRovara**], [**@MatthiasReumann**])
 - ✨ Add decision diagram-based construction and simulation of static unitary
   QCO functions ([#1915]) ([**@simon1hofmann**])
+- ✨ Add a `fuse-two-qubit-unitary-runs` pass for fusing compile-time two-qubit
+  unitary windows via Weyl/KAK resynthesis ([#1865]) ([**@simon1hofmann**],
+  [**@burgholzer**])
 - ✨ Add native relative-phase CCX (`rccx`) support across the IR, DD package,
   ZX diagrams, OpenQASM import/export, and Python/Qiskit bindings ([#1886])
   ([**@simon1hofmann**])
@@ -36,7 +39,7 @@ releases may include breaking changes.
 - ✨ Add support for IQM's `move` gate in the QDMI Qiskit backend converter
   ([#1844], [#1848]) ([**@burgholzer**], [**@marcelwa**])
 - 🚸 Add `const` version of the `CompoundOperation`'s `getOps()` function
-  ([#1826]) ([**@ystade])
+  ([#1826]) ([**@ystade**])
 - 🐳 Add dev container configuration for consistent local development
   environment ([#1786]) ([**@denialhaag**])
 - ✨ Add two-qubit Weyl (KAK) decomposition and native-gateset synthesis support
@@ -667,6 +670,7 @@ changelogs._
 [#1872]: https://github.com/munich-quantum-toolkit/core/pull/1872
 [#1870]: https://github.com/munich-quantum-toolkit/core/pull/1870
 [#1869]: https://github.com/munich-quantum-toolkit/core/pull/1869
+[#1865]: https://github.com/munich-quantum-toolkit/core/pull/1865
 [#1850]: https://github.com/munich-quantum-toolkit/core/pull/1850
 [#1849]: https://github.com/munich-quantum-toolkit/core/pull/1849
 [#1848]: https://github.com/munich-quantum-toolkit/core/pull/1848

bindings/mlir/register_mlir.cpp

@@ -387,6 +387,11 @@ operations.)pb");
               &mlir::QCOProgram::fuseSingleQubitUnitaryRuns>::call,
           nb::kw_only(), "basis"_a = "zyz",
           "Fuse single-qubit unitary runs into the chosen decomposition basis.")
+      .def("fuse_two_qubit_unitary_runs",
+           &BooleanMemberAdapter<
+               &mlir::QCOProgram::fuseTwoQubitUnitaryRuns>::call,
+           nb::kw_only(), "native_gates"_a = "",
+           "Lower unitaries to a native gate menu via two-qubit run fusion.")
       .def("unroll_quantum_loops",
            &BooleanMemberAdapter<&mlir::QCOProgram::unrollQuantumLoops>::call,
            nb::kw_only(), "unroll_factor"_a = -1,

mlir/include/mlir/Compiler/Programs.h

@@ -181,6 +181,9 @@ class QCOProgram final : public Program {
   /// Fuse single-qubit unitary runs into the selected Euler basis.
   [[nodiscard]] bool fuseSingleQubitUnitaryRuns(std::string_view basis = "zyz");
 
+  /// Lower unitaries to a native gate menu via two-qubit run fusion.
+  [[nodiscard]] bool fuseTwoQubitUnitaryRuns(std::string_view nativeGates = "");
+
   /// Unroll loops containing quantum operations.
   [[nodiscard]] bool unrollQuantumLoops(int64_t factor = -1);
 

mlir/include/mlir/Dialect/QCO/Transforms/Decomposition/Euler.h

@@ -20,6 +20,10 @@
 #include <cstdint>
 #include <optional>
 
+namespace mlir {
+class RewritePatternSet;
+} // namespace mlir
+
 namespace mlir::qco::decomposition {
 
 /**
@@ -99,4 +103,15 @@ synthesizeUnitary1QEuler(OpBuilder& builder, Location loc, Value qubit,
  */
 void emitGPhaseIfNeeded(OpBuilder& builder, Location loc, double phase);
 
+/**
+ * @brief Populates @p patterns with the single-qubit run fusion rewrite for
+ * @p basis (the reusable core of `fuse-single-qubit-unitary-runs`).
+ *
+ * @param skipControlledBodies When set, single-qubit gates nested in `qco.ctrl`
+ * bodies are left untouched.
+ */
+void populateFuseSingleQubitUnitaryRunsPatterns(
+    RewritePatternSet& patterns, EulerBasis basis,
+    bool skipControlledBodies = false);
+
 } // namespace mlir::qco::decomposition

mlir/include/mlir/Dialect/QCO/Transforms/Decomposition/NativeGateset.h

@@ -18,6 +18,10 @@
 #include <cstdint>
 #include <optional>
 
+namespace mlir {
+class Operation;
+} // namespace mlir
+
 namespace mlir::qco::decomposition {
 
 /**
@@ -63,6 +67,15 @@ struct NativeGateset {
    */
   [[nodiscard]] std::optional<TwoQubitNativeDecomposition>
   decomposeTarget(const Matrix4x4& target) const;
+
+  /**
+   * @brief Whether @p op is already on this native gateset.
+   *
+   * `qco.barrier` and `qco.gphase` are always allowed. Single-qubit primitives
+   * and single-target `qco.ctrl` shells (with an `X`/`Z` body) are checked
+   * against @p gates. All other ops are rejected.
+   */
+  [[nodiscard]] bool allowsOp(Operation* op) const;
 };
 
 } // namespace mlir::qco::decomposition

mlir/include/mlir/Dialect/QCO/Transforms/Passes.td

@@ -67,6 +67,57 @@ def FuseSingleQubitUnitaryRuns
       "Target Euler basis (zyz, zxz, xzx, xyx, u, zsxx, r).">];
 }
 
+def FuseTwoQubitUnitaryRuns
+    : Pass<"fuse-two-qubit-unitary-runs", "mlir::ModuleOp"> {
+  let dependentDialects = ["mlir::qco::QCODialect"];
+  let summary = "Lower QCO unitary gates to a user-specified native gate menu.";
+  let description = [{
+    Rewrites unitaries to match the comma-separated `native-gates` gateset.
+    `qco.barrier` and `qco.gphase` are preserved. A `qco.ctrl` is treated as a
+    two-qubit gate only when it has a single control and a single target; gates
+    acting on more than two qubits (e.g. multi-controlled `qco.ctrl`) are out of
+    scope and left untouched for a dedicated multi-controlled synthesis pass.
+
+    The gateset is a comma-separated list of gate tokens (order not
+    significant) from which the pass resolves a single-qubit Euler basis
+    (generic `qco.u` when `u` is present; IBM-style surface gates when all of
+    `x`, `sx`, and `rz` are present; IQM-style `qco.r` when `r` is present;
+    or a supported rotation pair chosen from `rx`, `ry`, `rz`) plus the
+    two-qubit entangler `cx` or `cz`.
+
+    Recognised tokens: `u`, `x`, `sx`, `rz`, `rx`, `ry`, `r`, `cx`,
+    `cz`. An empty or whitespace-only gateset is a no-op, which is the intended
+    pipeline default when synthesis is not needed. An unrecognised token causes
+    the pass to fail.
+
+    Example gatesets (each line is one illustrative gateset; pick either `cx`
+    or `cz` as the entangler, or list both if both are native):
+    - IBM: `x,sx,rz,cx` or `x,sx,rz,cz`
+    - Generic single-qubit U: `u,cx` or `u,cz`
+    - IQM default: `r,cz` (or `r,cx` if CX is the native entangler)
+    - Rotation pair + entangler: `rx,rz,cx`, `rx,ry,cz`, `ry,rz,cx`, etc.
+      Supported pairs are exactly `rx`+`rz`, `rx`+`ry`, and `ry`+`rz`.
+
+    Stages: fuse single-qubit runs (reusing `fuse-single-qubit-unitary-runs`,
+    which also lowers lone off-gateset single-qubit gates); fuse two-qubit runs
+    and lower any remaining off-gateset two-qubit ops via Weyl synthesis; fuse
+    the single-qubit seams introduced by that synthesis. The pass fails if any
+    single- or two-qubit op remains off the native gateset afterwards; gates on
+    more than two qubits are left untouched and do not cause a failure.
+
+    Lowering is deterministic: `cx` is preferred over `cz`, single-qubit
+    factors use the resolved Euler basis, and two-qubit run replacement uses
+    the minimal entangler count from the synthesizer.
+  }];
+  let options = [Option<
+      "nativeGates", "native-gates", "std::string", "\"\"",
+      "Comma-separated native gateset. Empty or whitespace-only is "
+      "a no-op. Tokens: u, x, sx, rz, rx, ry, r, cx, cz. "
+      "Examples: x,sx,rz,cx; u,cx; r,cz; rx,rz,cx.">];
+}
+
+//===----------------------------------------------------------------------===//
+
 def QuantumLoopUnroll
     : InterfacePass<"quantum-loop-unroll", "FunctionOpInterface"> {
   let dependentDialects = ["mlir::qco::QCODialect", "mlir::scf::SCFDialect"];

mlir/lib/Compiler/Programs.cpp

@@ -349,6 +349,17 @@ bool QCOProgram::fuseSingleQubitUnitaryRuns(const std::string_view basis) {
       "failed to fuse single-qubit unitary runs"));
 }
 
+bool QCOProgram::fuseTwoQubitUnitaryRuns(const std::string_view nativeGates) {
+  qco::FuseTwoQubitUnitaryRunsOptions options;
+  options.nativeGates = nativeGates;
+  return succeeded(runPasses(
+      mod(),
+      [&options](OpPassManager& pm) {
+        pm.addPass(qco::createFuseTwoQubitUnitaryRuns(options));
+      },
+      "failed to fuse two-qubit unitary runs"));
+}
+
 bool QCOProgram::unrollQuantumLoops(const int64_t factor) {
   qco::QuantumLoopUnrollOptions options;
   options.unrollFactor = factor;

mlir/lib/Dialect/QCO/Transforms/Decomposition/NativeGateset.cpp

@@ -10,12 +10,16 @@
 
 #include "mlir/Dialect/QCO/Transforms/Decomposition/NativeGateset.h"
 
+#include "mlir/Dialect/QCO/IR/QCOInterfaces.h"
+#include "mlir/Dialect/QCO/IR/QCOOps.h"
 #include "mlir/Dialect/QCO/Transforms/Decomposition/Euler.h"
 #include "mlir/Dialect/QCO/Transforms/Decomposition/Weyl.h"
 #include "mlir/Dialect/QCO/Utils/Matrix.h"
 
 #include <llvm/ADT/StringSwitch.h>
+#include <llvm/ADT/TypeSwitch.h>
 #include <llvm/Support/ErrorHandling.h>
+#include <mlir/IR/Operation.h>
 #include <mlir/Support/LLVM.h>
 
 #include <optional>
@@ -141,6 +145,48 @@ NativeGateset::decomposeTarget(const Matrix4x4& target) const {
   return cachedNativeBasisDecomposer(*entangler).decomposeTarget(target);
 }
 
+static std::optional<NativeGateKind> gateKindFor(UnitaryOpInterface op) {
+  return TypeSwitch<Operation*, std::optional<NativeGateKind>>(
+             op.getOperation())
+      .Case<UOp>([](UOp) { return NativeGateKind::U; })
+      .Case<XOp>([](XOp) { return NativeGateKind::X; })
+      .Case<SXOp>([](SXOp) { return NativeGateKind::SX; })
+      .Case<RZOp>([](RZOp) { return NativeGateKind::RZ; })
+      .Case<RXOp>([](RXOp) { return NativeGateKind::RX; })
+      .Case<RYOp>([](RYOp) { return NativeGateKind::RY; })
+      .Case<ROp>([](ROp) { return NativeGateKind::R; })
+      .Default([](Operation*) { return std::nullopt; });
+}
+
+static std::optional<NativeGateKind> entanglerKindFor(CtrlOp ctrl) {
+  if (ctrl.getNumControls() != 1 || ctrl.getNumTargets() != 1 ||
+      ctrl.getNumBodyUnitaries() != 1) {
+    return std::nullopt;
+  }
+  return TypeSwitch<Operation*, std::optional<NativeGateKind>>(
+             ctrl.getBodyUnitary(0).getOperation())
+      .Case<XOp>([](XOp) { return NativeGateKind::CX; })
+      .Case<ZOp>([](ZOp) { return NativeGateKind::CZ; })
+      .Default([](Operation*) { return std::nullopt; });
+}
+
+bool NativeGateset::allowsOp(Operation* op) const {
+  return TypeSwitch<Operation*, bool>(op)
+      .Case<BarrierOp, GPhaseOp>([](auto) { return true; })
+      .Case<CtrlOp>([&](CtrlOp ctrl) {
+        const auto kind = entanglerKindFor(ctrl);
+        return kind && gates.contains(*kind);
+      })
+      .Case<UnitaryOpInterface>([&](UnitaryOpInterface unitary) {
+        if (!unitary.isSingleQubit()) {
+          return false;
+        }
+        const auto gate = gateKindFor(unitary);
+        return gate && gates.contains(*gate);
+      })
+      .Default([](Operation*) { return false; });
+}
+
 std::optional<NativeGateset> NativeGateset::parse(StringRef nativeGates) {
   auto gates = parseGateSet(nativeGates);
   if (!gates) {

mlir/lib/Dialect/QCO/Transforms/NativeSynthesis/FuseSingleQubitUnitaryRuns.cpp

@@ -150,10 +150,13 @@ namespace {
 struct FuseSingleQubitUnitaryRunsPattern final
     : OpInterfaceRewritePattern<UnitaryOpInterface> {
   FuseSingleQubitUnitaryRunsPattern(MLIRContext* context,
-                                    const decomposition::EulerBasis basis)
-      : OpInterfaceRewritePattern(context), basis(basis) {}
+                                    const decomposition::EulerBasis basis,
+                                    const bool skipControlledBodies)
+      : OpInterfaceRewritePattern(context), basis(basis),
+        skipControlledBodies(skipControlledBodies) {}
 
   decomposition::EulerBasis basis;
+  bool skipControlledBodies;
 
   /**
    * @brief Fuses the run anchored at `op` when beneficial.
@@ -167,6 +170,10 @@ struct FuseSingleQubitUnitaryRunsPattern final
    */
   LogicalResult matchAndRewrite(UnitaryOpInterface op,
                                 PatternRewriter& rewriter) const override {
+    if (skipControlledBodies &&
+        (op.getOperation()->getParentOfType<CtrlOp>() != nullptr)) {
+      return failure();
+    }
     if (!isRunMemberCandidate(op)) {
       return failure();
     }
@@ -219,8 +226,8 @@ struct FuseSingleQubitUnitaryRunsPass final
     }
 
     RewritePatternSet patterns(&getContext());
-    patterns.add<FuseSingleQubitUnitaryRunsPattern>(patterns.getContext(),
-                                                    *parsed);
+    decomposition::populateFuseSingleQubitUnitaryRunsPatterns(
+        patterns, *parsed, /*skipControlledBodies=*/false);
 
     if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
       signalPassFailure();
@@ -231,3 +238,14 @@ struct FuseSingleQubitUnitaryRunsPass final
 } // namespace
 
 } // namespace mlir::qco
+
+namespace mlir::qco::decomposition {
+
+void populateFuseSingleQubitUnitaryRunsPatterns(
+    RewritePatternSet& patterns, const EulerBasis basis,
+    const bool skipControlledBodies) {
+  patterns.add<FuseSingleQubitUnitaryRunsPattern>(patterns.getContext(), basis,
+                                                  skipControlledBodies);
+}
+
+} // namespace mlir::qco::decomposition