Canonicalize const ifs

include/cudaq/Optimizer/Dialect/CC/CCOps.td

@@ -393,6 +393,7 @@ def cc_IfOp : CCOp<"if",
 
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
+  let hasCanonicalizer = 1;
 
   let extraClassDeclaration = [{
     using RegionBuilderFn = llvm::function_ref<void(mlir::OpBuilder &,

include/cudaq/Optimizer/Dialect/CC/LowerToCFGPatterns.inc

@@ -0,0 +1,91 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2025 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.    *
+ ******************************************************************************/
+
+// These canonicalization patterns are used by the canonicalize pass and
+// lower-to-cfg pass.
+
+// This file must be included after a `using namespace mlir;` as it uses bare
+// identifiers from that namespace.
+
+namespace {
+
+class RewriteIf : public OpRewritePattern<cudaq::cc::IfOp> {
+public:
+  using OpRewritePattern::OpRewritePattern;
+
+  /// Rewrites an if construct like
+  /// ```mlir
+  /// (0)
+  /// quake.if %cond {
+  ///   (1)
+  /// } else {
+  ///   (2)
+  /// }
+  /// (3)
+  /// ```
+  /// to a CFG like
+  /// ```mlir
+  ///   (0)
+  ///   cf.cond_br %cond, ^bb1, ^bb2
+  /// ^bb1:
+  ///   (1)
+  ///   cf.br ^bb3
+  /// ^bb2:
+  ///   (2)
+  ///   cf.br ^bb3
+  /// ^bb3:
+  ///   (3)
+  /// ```
+  LogicalResult matchAndRewrite(cudaq::cc::IfOp ifOp,
+                                PatternRewriter &rewriter) const override {
+    auto loc = ifOp.getLoc();
+    auto *initBlock = rewriter.getInsertionBlock();
+    auto initPos = rewriter.getInsertionPoint();
+    auto *endBlock = rewriter.splitBlock(initBlock, initPos);
+    if (ifOp.getNumResults() != 0) {
+      Block *continueBlock = rewriter.createBlock(
+          endBlock, ifOp.getResultTypes(),
+          SmallVector<Location>(ifOp.getNumResults(), loc));
+      rewriter.create<cf::BranchOp>(loc, endBlock);
+      endBlock = continueBlock;
+    }
+    auto *thenBlock = &ifOp.getThenRegion().front();
+    bool hasElse = !ifOp.getElseRegion().empty();
+    auto *elseBlock = hasElse ? &ifOp.getElseRegion().front() : endBlock;
+    updateBodyBranches(&ifOp.getThenRegion(), rewriter, endBlock);
+    updateBodyBranches(&ifOp.getElseRegion(), rewriter, endBlock);
+    rewriter.inlineRegionBefore(ifOp.getThenRegion(), endBlock);
+    if (hasElse)
+      rewriter.inlineRegionBefore(ifOp.getElseRegion(), endBlock);
+    rewriter.setInsertionPointToEnd(initBlock);
+    rewriter.create<cf::CondBranchOp>(loc, ifOp.getCondition(), thenBlock,
+                                      ifOp.getLinearArgs(), elseBlock,
+                                      ifOp.getLinearArgs());
+    rewriter.replaceOp(ifOp, endBlock->getArguments());
+    return success();
+  }
+
+  // Replace all the ContinueOp in the body region with branches to the correct
+  // basic blocks.
+  void updateBodyBranches(Region *bodyRegion, PatternRewriter &rewriter,
+                          Block *continueBlock) const {
+    // Walk body region and replace all continue and break ops.
+    for (Block &block : *bodyRegion) {
+      auto *terminator = block.getTerminator();
+      if (auto cont = dyn_cast<cudaq::cc::ContinueOp>(terminator)) {
+        rewriter.setInsertionPointToEnd(&block);
+        LLVM_DEBUG(llvm::dbgs() << "replacing " << *terminator << '\n');
+        rewriter.replaceOpWithNewOp<cf::BranchOp>(cont, continueBlock,
+                                                  cont.getOperands());
+      }
+      // Other ad-hoc control flow in the region need not be rewritten.
+    }
+  }
+};
+
+} // namespace

lib/Optimizer/Dialect/CC/CCOps.cpp

@@ -11,6 +11,7 @@
 #include "cudaq/Optimizer/Dialect/CC/CCDialect.h"
 #include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Debug.h"
 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlowOps.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
@@ -21,8 +22,12 @@
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 
+#define DEBUG_TYPE "canonicalize"
+
 using namespace mlir;
 
+#include "cudaq/Optimizer/Dialect/CC/LowerToCFGPatterns.inc"
+
 template <typename R>
 R getParentOfType(Operation *op) {
   do {
@@ -2038,6 +2043,14 @@ MutableOperandRange cudaq::cc::ConditionOp::getMutableSuccessorOperands(
   return getResultsMutable();
 }
 
+//===----------------------------------------------------------------------===//
+// IfOp
+//===----------------------------------------------------------------------===//
+
+void cudaq::cc::IfOp::getCanonicalizationPatterns(RewritePatternSet &patterns,
+                                                  MLIRContext *context) {
+  patterns.add<RewriteIf>(context);
+}
 //===----------------------------------------------------------------------===//
 // OffsetOfOp
 //===----------------------------------------------------------------------===//

lib/Optimizer/Transforms/LowerToCFG.cpp

@@ -12,6 +12,7 @@
 #include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
 #include "cudaq/Optimizer/Transforms/Passes.h"
 #include "cudaq/Todo.h"
+#include "llvm/Support/Debug.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
@@ -20,6 +21,8 @@
 
 using namespace mlir;
 
+#include "cudaq/Optimizer/Dialect/CC/LowerToCFGPatterns.inc"
+
 namespace {
 class RewriteScope : public OpRewritePattern<cudaq::cc::ScopeOp> {
 public:
@@ -278,80 +281,6 @@ class RewriteLoop : public OpRewritePattern<cudaq::cc::LoopOp> {
   }
 };
 
-class RewriteIf : public OpRewritePattern<cudaq::cc::IfOp> {
-public:
-  using OpRewritePattern::OpRewritePattern;
-
-  /// Rewrites an if construct like
-  /// ```mlir
-  /// (0)
-  /// quake.if %cond {
-  ///   (1)
-  /// } else {
-  ///   (2)
-  /// }
-  /// (3)
-  /// ```
-  /// to a CFG like
-  /// ```mlir
-  ///   (0)
-  ///   cf.cond_br %cond, ^bb1, ^bb2
-  /// ^bb1:
-  ///   (1)
-  ///   cf.br ^bb3
-  /// ^bb2:
-  ///   (2)
-  ///   cf.br ^bb3
-  /// ^bb3:
-  ///   (3)
-  /// ```
-  LogicalResult matchAndRewrite(cudaq::cc::IfOp ifOp,
-                                PatternRewriter &rewriter) const override {
-    auto loc = ifOp.getLoc();
-    auto *initBlock = rewriter.getInsertionBlock();
-    auto initPos = rewriter.getInsertionPoint();
-    auto *endBlock = rewriter.splitBlock(initBlock, initPos);
-    if (ifOp.getNumResults() != 0) {
-      Block *continueBlock = rewriter.createBlock(
-          endBlock, ifOp.getResultTypes(),
-          SmallVector<Location>(ifOp.getNumResults(), loc));
-      rewriter.create<cf::BranchOp>(loc, endBlock);
-      endBlock = continueBlock;
-    }
-    auto *thenBlock = &ifOp.getThenRegion().front();
-    bool hasElse = !ifOp.getElseRegion().empty();
-    auto *elseBlock = hasElse ? &ifOp.getElseRegion().front() : endBlock;
-    updateBodyBranches(&ifOp.getThenRegion(), rewriter, endBlock);
-    updateBodyBranches(&ifOp.getElseRegion(), rewriter, endBlock);
-    rewriter.inlineRegionBefore(ifOp.getThenRegion(), endBlock);
-    if (hasElse)
-      rewriter.inlineRegionBefore(ifOp.getElseRegion(), endBlock);
-    rewriter.setInsertionPointToEnd(initBlock);
-    rewriter.create<cf::CondBranchOp>(loc, ifOp.getCondition(), thenBlock,
-                                      ifOp.getLinearArgs(), elseBlock,
-                                      ifOp.getLinearArgs());
-    rewriter.replaceOp(ifOp, endBlock->getArguments());
-    return success();
-  }
-
-  // Replace all the ContinueOp in the body region with branches to the correct
-  // basic blocks.
-  void updateBodyBranches(Region *bodyRegion, PatternRewriter &rewriter,
-                          Block *continueBlock) const {
-    // Walk body region and replace all continue and break ops.
-    for (Block &block : *bodyRegion) {
-      auto *terminator = block.getTerminator();
-      if (auto cont = dyn_cast<cudaq::cc::ContinueOp>(terminator)) {
-        rewriter.setInsertionPointToEnd(&block);
-        LLVM_DEBUG(llvm::dbgs() << "replacing " << *terminator << '\n');
-        rewriter.replaceOpWithNewOp<cf::BranchOp>(cont, continueBlock,
-                                                  cont.getOperands());
-      }
-      // Other ad-hoc control flow in the region need not be rewritten.
-    }
-  }
-};
-
 class RewriteReturn : public OpRewritePattern<cudaq::cc::ReturnOp> {
 public:
   using OpRewritePattern::OpRewritePattern;