Add canonicalization pattern for merging nested CtrlOps

Author: denialhaag

mlir/include/mlir/Dialect/Flux/Builder/FluxProgramBuilder.h

@@ -380,7 +380,7 @@ class FluxProgramBuilder {
   Location loc;
 
 private:
-  bool inRegion = false;
+  int inRegion = 0;
 
   //===--------------------------------------------------------------------===//
   // Linear Type Tracking Helpers

mlir/include/mlir/Dialect/Flux/IR/FluxOps.td

@@ -433,9 +433,13 @@ def CtrlOp : FluxOp<"ctrl", traits = [UnitaryOpInterface, AttrSizedOperandSegmen
     let builders = [
         OpBuilder<(ins "ValueRange":$controls_in, "ValueRange":$targets_in), [{
             auto qubit_type = QubitType::get($_builder.getContext());
-            build($_builder, $_state, qubit_type, qubit_type, controls_in, targets_in);
+            SmallVector<Type> control_types(controls_in.size(), qubit_type);
+            SmallVector<Type> target_types(targets_in.size(), qubit_type);
+            build($_builder, $_state, TypeRange(control_types), TypeRange(target_types), controls_in, targets_in);
         }]>,
     ];
+
+    let hasCanonicalizer = 1;
 }
 
 #endif // FluxOPS

mlir/lib/Dialect/Flux/Builder/FluxProgramBuilder.cpp

@@ -178,7 +178,7 @@ Value FluxProgramBuilder::x(Value qubit) {
   const auto& qubitOut = xOp.getQubitOut();
 
   // Update tracking
-  if (!inRegion) {
+  if (inRegion == 0) {
     updateQubitTracking(qubit, qubitOut);
   }
 
@@ -232,21 +232,23 @@ std::pair<SmallVector<Value>, SmallVector<Value>> FluxProgramBuilder::ctrl(
   const mlir::OpBuilder::InsertionGuard guard(builder);
   builder.setInsertionPointToStart(&ctrlOp.getBody().emplaceBlock());
 
-  inRegion = true;
+  inRegion++;
   auto targetsYield = body(*this);
-  inRegion = false;
+  inRegion--;
 
   builder.create<YieldOp>(loc, targetsYield);
 
-  // Update tracking
   const auto& controlsOut = ctrlOp.getControlsOut();
-  for (const auto& [control, controlOut] : llvm::zip(controls, controlsOut)) {
-    updateQubitTracking(control, controlOut);
-  }
-
   const auto& targetsOut = ctrlOp.getTargetsOut();
-  for (const auto& [target, targetOut] : llvm::zip(targets, targetsOut)) {
-    updateQubitTracking(target, targetOut);
+
+  // Update tracking
+  if (inRegion == 0) {
+    for (const auto& [control, controlOut] : llvm::zip(controls, controlsOut)) {
+      updateQubitTracking(control, controlOut);
+    }
+    for (const auto& [target, targetOut] : llvm::zip(targets, targetsOut)) {
+      updateQubitTracking(target, targetOut);
+    }
   }
 
   return {controlsOut, targetsOut};

mlir/lib/Dialect/Flux/IR/FluxOps.cpp

@@ -556,6 +556,46 @@ struct MergeSubsequentRX final : OpRewritePattern<RXOp> {
   }
 };
 
+struct MergeNestedCtrl final : OpRewritePattern<CtrlOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(CtrlOp ctrlOp,
+                                PatternRewriter& rewriter) const override {
+    auto bodyUnitary = ctrlOp.getBodyUnitary();
+    auto bodyCtrlOp = llvm::dyn_cast<CtrlOp>(bodyUnitary.getOperation());
+
+    // Check if the body unitary is a CtrlOp
+    if (!bodyCtrlOp) {
+      return failure();
+    }
+
+    // Merge controls
+    SmallVector<Value> newControls;
+    newControls.append(ctrlOp.getControlsIn().begin(),
+                       ctrlOp.getControlsIn().end());
+    for (auto control : bodyCtrlOp.getControlsIn()) {
+      if (llvm::is_contained(newControls, control)) {
+        continue;
+      }
+      newControls.push_back(control);
+    }
+
+    // Create new CtrlOp
+    auto newCtrlOp = rewriter.create<CtrlOp>(ctrlOp.getLoc(), newControls,
+                                             bodyCtrlOp.getTargetsIn());
+
+    // Clone block
+    rewriter.cloneRegionBefore(bodyCtrlOp.getBody(), newCtrlOp.getBody(),
+                               newCtrlOp.getBody().end());
+
+    // Replace CtrlOps
+    rewriter.eraseOp(bodyCtrlOp);
+    rewriter.replaceOp(ctrlOp, newCtrlOp.getResults());
+
+    return success();
+  }
+};
+
 } // namespace
 
 void DeallocOp::getCanonicalizationPatterns(RewritePatternSet& results,
@@ -577,3 +617,8 @@ void RXOp::getCanonicalizationPatterns(RewritePatternSet& results,
                                        MLIRContext* context) {
   results.add<MergeSubsequentRX>(context);
 }
+
+void CtrlOp::getCanonicalizationPatterns(RewritePatternSet& results,
+                                         MLIRContext* context) {
+  results.add<MergeNestedCtrl>(context);
+}

mlir/unittests/pipeline/test_compiler_pipeline.cpp

@@ -1269,4 +1269,34 @@ TEST_F(SimpleConversionTest, CX) {
   EXPECT_TRUE(verify("Flux to Quartz", quartzIRConv, quartzExpected.get()));
 }
 
+TEST_F(SimpleConversionTest, CXMerge) {
+  const auto fluxInit = buildFluxIR([](flux::FluxProgramBuilder& b) {
+    auto reg = b.allocQubitRegister(3, "q");
+    auto q0a = reg[0];
+    auto q1a = reg[1];
+    auto q2a = reg[2];
+    b.ctrl({q0a}, {q1a, q2a}, [&](auto& b) {
+      auto q12b = b.ctrl({q1a}, {q2a}, [&](auto& b) {
+        auto q2b = b.x(q2a);
+        return SmallVector<Value>{q2b};
+      });
+      return SmallVector<Value>{q12b.first[0], q12b.second[0]};
+    });
+  });
+  const auto fluxInitIR = captureIR(fluxInit.get());
+
+  const auto fluxOpt = buildFluxIR([](flux::FluxProgramBuilder& b) {
+    auto reg = b.allocQubitRegister(3, "q");
+    auto q0a = reg[0];
+    auto q1a = reg[1];
+    auto q2a = reg[2];
+    b.ctrl({q0a, q1a}, {q2a}, [&](auto& b) {
+      auto q2b = b.x(q2a);
+      return SmallVector<Value>{q2b};
+    });
+  });
+
+  EXPECT_TRUE(verify("Flux Canonicalization", fluxInitIR, fluxOpt.get()));
+}
+
 } // namespace