[Analysis] Create simple analysis pass to group unitary operations

cudaq/include/cudaq/Optimizer/Analysis/UnitaryOpGrouping.h

@@ -0,0 +1,49 @@
+/*******************************************************************************
+ * Copyright (c) 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.    *
+ ******************************************************************************/
+
+#pragma once
+
+#include "llvm/ADT/SmallVector.h"
+#include "mlir/Support/TypeID.h"
+
+namespace mlir {
+class Operation;
+class Block;
+class Region;
+} // namespace mlir
+
+namespace cudaq::quake::detail {
+struct UnitaryOpGroup {
+  UnitaryOpGroup() = default;
+
+  mlir::Block *block = nullptr;
+  mlir::Operation *firstOp = nullptr;
+  mlir::Operation *lastOp = nullptr;
+  llvm::SmallVector<mlir::Operation *> ops;
+};
+
+using UnitaryOpGroups = llvm::SmallVector<cudaq::quake::detail::UnitaryOpGroup>;
+
+/// Analysis to group unitary operations within a block together.
+struct UnitaryOpGroupingAnalysis {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnitaryOpGroupingAnalysis)
+
+  explicit UnitaryOpGroupingAnalysis(mlir::Operation *op) {
+    performAnalysis(op);
+  }
+
+  const UnitaryOpGroups &getGroups() const { return groups; }
+
+private:
+  void performAnalysis(mlir::Operation *operation);
+  void scanRegion(mlir::Region &region);
+  void scanBlock(mlir::Block &block);
+
+  UnitaryOpGroups groups;
+};
+} // namespace cudaq::quake::detail

cudaq/lib/Optimizer/Analysis/CMakeLists.txt

@@ -0,0 +1,14 @@
+# ============================================================================ #
+# Copyright (c) 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.     #
+# ============================================================================ #
+
+add_cudaq_library(OptAnalysis
+  UnitaryOpGrouping.cpp 
+
+  DEPENDS
+    QuakeDialect
+)

cudaq/lib/Optimizer/Analysis/UnitaryOpGrouping.cpp

@@ -0,0 +1,98 @@
+/*******************************************************************************
+ * Copyright (c) 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.    *
+ ******************************************************************************/
+
+#include "cudaq/Optimizer/Analysis/UnitaryOpGrouping.h"
+#include "cudaq/Optimizer/Dialect/Quake/QuakeOps.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/Debug.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Block.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Region.h"
+#include "mlir/Support/TypeID.h"
+
+#define DEBUG_TYPE "unitary-op-grouping-analysis"
+
+using namespace mlir;
+
+namespace {
+static bool isUnitaryOp(Operation *op) {
+  return op->hasTrait<cudaq::QuantumGate>() && !isa<cudaq::quake::ResetOp>(op);
+}
+
+/// If we've hit this function, that means we've reached an op that cannot be
+/// added to the current group of unitary ops. So, we do one of two things.
+/// 1) If the current group is empty, we don't do anything and just return
+/// 2) If the current group is non-empty, we have have a UnitaryOpsGroup to
+///    populate. We create the struct and then push it back to the
+///    UnitaryOpGroups vector.
+static void flushGroupIfNonEmpty(cudaq::quake::detail::UnitaryOpGroups &groups,
+                                 Block &block,
+                                 SmallVectorImpl<Operation *> &currUnitaryOps) {
+  if (currUnitaryOps.empty())
+    return;
+
+  cudaq::quake::detail::UnitaryOpGroup group;
+  group.block = &block;
+  group.firstOp = currUnitaryOps.front();
+  group.lastOp = currUnitaryOps.back();
+  group.ops.append(currUnitaryOps.begin(), currUnitaryOps.end());
+
+  LLVM_DEBUG(llvm::dbgs() << "Found unitary group with " << group.ops.size()
+                          << " op(s)\n");
+
+  groups.push_back(std::move(group));
+  currUnitaryOps.clear();
+}
+} // namespace
+
+void cudaq::quake::detail::UnitaryOpGroupingAnalysis::scanBlock(Block &block) {
+  SmallVector<Operation *> currUnitaryOps;
+
+  for (Operation &op : block) {
+    if (isUnitaryOp(&op)) {
+      currUnitaryOps.push_back(&op);
+      continue;
+    }
+
+    flushGroupIfNonEmpty(groups, block, currUnitaryOps);
+  }
+
+  flushGroupIfNonEmpty(groups, block, currUnitaryOps);
+}
+
+void cudaq::quake::detail::UnitaryOpGroupingAnalysis::scanRegion(
+    Region &region) {
+  for (Block &block : region) {
+    scanBlock(block);
+
+    // Nested regions, such as cc.if branches or cc.loop bodies, are scanned as
+    // independent blocks. MVP 1 never groups across the parent control-flow op.
+    for (Operation &op : block)
+      for (Region &nestedRegion : op.getRegions())
+        scanRegion(nestedRegion);
+  }
+}
+
+void cudaq::quake::detail::UnitaryOpGroupingAnalysis::performAnalysis(
+    Operation *operation) {
+  auto funcOp = dyn_cast<func::FuncOp>(operation);
+  if (!funcOp)
+    return;
+
+  LLVM_DEBUG(llvm::dbgs() << "Function to analyze: " << funcOp.getName()
+                          << '\n');
+
+  /// initially start with the func op region, then recursively scan nested
+  /// regions
+  for (Region &region : funcOp->getRegions())
+    scanRegion(region);
+
+  LLVM_DEBUG(llvm::dbgs() << "Found " << groups.size()
+                          << " unitary group(s)\n");
+}

cudaq/lib/Optimizer/CMakeLists.txt

@@ -11,3 +11,4 @@ add_subdirectory(CodeGen)
 add_subdirectory(Dialect)
 add_subdirectory(Transforms)
 add_subdirectory(CAPI)
+add_subdirectory(Analysis)

cudaq/test/Analysis/unitary-op-grouping.qke

@@ -0,0 +1,48 @@
+// ========================================================================== //
+// Copyright (c) 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.   //
+// ========================================================================== //
+
+
+// RUN: cudaq-opt %s --test-unitary-op-grouping -o /dev/null 2>&1 | FileCheck %s
+
+module {
+  func.func @simple(%q0: !quake.ref, %q1: !quake.ref, %theta: f64) attributes {"cudaq-kernel"} {
+    quake.h %q0 : (!quake.ref) -> ()
+    quake.x %q1 : (!quake.ref) -> ()
+    %m = quake.mz %q0 : (!quake.ref) -> !cc.measure_handle
+    quake.z %q0 : (!quake.ref) -> ()
+    %c0 = arith.constant 0 : i64
+    quake.rx (%theta) %q1 : (f64, !quake.ref) -> ()
+    return
+  }
+
+  func.func @nested_if(%q0: !quake.ref, %q1: !quake.ref, %flag: i1) attributes {"cudaq-kernel"} {
+    cc.if(%flag) {
+      quake.h %q0 : (!quake.ref) -> ()
+      quake.x %q1 : (!quake.ref) -> ()
+    } else {
+      quake.z %q0 : (!quake.ref) -> ()
+    }
+    return
+  }
+}
+
+// CHECK-LABEL: func @simple
+// CHECK:       unitary group #0 -- found 2 ops
+// CHECK-NEXT:    quake.h
+// CHECK-NEXT:    quake.x
+// CHECK:       unitary group #1 -- found 1 ops
+// CHECK-NEXT:    quake.z
+// CHECK:       unitary group #2 -- found 1 ops
+// CHECK-NEXT:    quake.rx
+
+// CHECK-LABEL: func @nested_if
+// CHECK:       unitary group #0 -- found 2 ops
+// CHECK-NEXT:    quake.h
+// CHECK-NEXT:    quake.x
+// CHECK:       unitary group #1 -- found 1 ops
+// CHECK-NEXT:    quake.z

cudaq/test/CMakeLists.txt

@@ -21,6 +21,7 @@ if (CUDAQ_DISABLE_RUNTIME)
 endif()
 
 add_subdirectory(plugin)
+add_subdirectory(lib)
 
 configure_lit_site_cfg(
   ${CMAKE_CURRENT_SOURCE_DIR}/lit.site.cfg.py.in

cudaq/test/lib/Analysis/CMakeLists.txt

@@ -0,0 +1,28 @@
+# ============================================================================ #
+# Copyright (c) 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.     #
+# ============================================================================ #
+
+add_mlir_library(CUDAQTestAnalysis
+  TestUnitaryOpGrouping.cpp
+
+  DISABLE_INSTALL
+
+  DEPENDS
+    QuakeDialect
+    OptAnalysis
+
+  LINK_LIBS PUBLIC
+    MLIRIR
+    MLIRPass
+    MLIRFuncDialect
+    OptAnalysis
+)
+
+if(TARGET cudaq-opt)
+  target_link_libraries(cudaq-opt PRIVATE CUDAQTestAnalysis)
+  target_compile_definitions(cudaq-opt PRIVATE CUDAQ_INCLUDE_TESTS)
+endif()

cudaq/test/lib/Analysis/TestUnitaryOpGrouping.cpp

@@ -0,0 +1,59 @@
+/*******************************************************************************
+ * Copyright (c) 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.    *
+ ******************************************************************************/
+
+#include "cudaq/Optimizer/Analysis/UnitaryOpGrouping.h"
+#include "llvm/Support/raw_ostream.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Operation.h"
+#include "mlir/Pass/Pass.h"
+
+using namespace mlir;
+
+namespace {
+class TestUnitaryOpGroupingPass
+    : public PassWrapper<TestUnitaryOpGroupingPass,
+                         OperationPass<ModuleOp>> {
+public:
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TestUnitaryOpGroupingPass)
+
+  StringRef getArgument() const final { return "test-unitary-op-grouping"; }
+
+  StringRef getDescription() const final {
+    return "Print found groups of unitary operations";
+  }
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+
+    for (auto funcOp : module.getOps<func::FuncOp>()) {
+      auto &analysis = getChildAnalysis<
+          cudaq::quake::detail::UnitaryOpGroupingAnalysis>(
+          funcOp.getOperation());
+
+      llvm::errs() << "func @" << funcOp.getName() << '\n';
+
+      unsigned groupIndex = 0;
+      for (const auto &group : analysis.getGroups()) {
+        llvm::errs() << "  unitary group #" << groupIndex++ << " -- found "
+                     << group.ops.size() << " ops\n";
+        for (Operation *op : group.ops)
+          llvm::errs() << "    " << op->getName().getStringRef() << '\n';
+      }
+    }
+
+    markAllAnalysesPreserved();
+  }
+};
+} // namespace
+
+namespace cudaq::test {
+void registerTestUnitaryOpGroupingPass() {
+  PassRegistration<TestUnitaryOpGroupingPass>();
+}
+} // namespace cudaq::test

cudaq/test/lib/CMakeLists.txt

@@ -0,0 +1,9 @@
+# ============================================================================ #
+# Copyright (c) 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.     #
+# ============================================================================ #
+
+add_subdirectory(Analysis)

cudaq/test/lit.cfg.py

@@ -49,7 +49,7 @@
 
 # Exclude a list of directories from the test suite:
 #   - 'Inputs' contain auxiliary inputs for various tests.
-local_excludes = ['Inputs', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt', 'plugin']
+local_excludes = ['Inputs', 'CMakeLists.txt', 'README.txt', 'LICENSE.txt', 'plugin', 'lib']
 config.excludes = [exclude for exclude in config.excludes] + local_excludes
 
 # The root path where tests are located.

cudaq/tools/cudaq-opt/cudaq-opt.cpp

@@ -31,6 +31,12 @@
 
 using namespace llvm;
 
+#ifdef CUDAQ_INCLUDE_TESTS
+namespace cudaq::test {
+void registerTestUnitaryOpGroupingPass();
+} // namespace cudaq::test
+#endif
+
 /// Dialect extension to allow inlining of the MLIR defined LLVM-IR dialects
 /// which lacks inlining support out of the box.
 class InlinerExtension
@@ -60,6 +66,10 @@ int main(int argc, char **argv) {
   cudaq::registerAllCLOptions();
   cudaq::registerAllPasses();
 
+#ifdef CUDAQ_INCLUDE_TESTS
+  cudaq::test::registerTestUnitaryOpGroupingPass();
+#endif
+
   // See if we have been asked to load a pass plugin,
   // if so load it.
   std::vector<std::string> args(&argv[0], &argv[0] + argc);