Add broadcast operation transformation passes

cudaq/include/cudaq/Optimizer/Transforms/Passes.td

@@ -371,6 +371,35 @@ def CombineQuantumAllocations :
     "cudaq::quake::QuakeDialect"];
 }
 
+def ConsolidateBroadcasts :
+    Pass<"consolidate-broadcasts", "mlir::func::FuncOp"> {
+  let summary = "Roll loops over a veq back into broadcast form.";
+  let description = [{
+    A counted loop that applies one-target operators to every element of a
+    `veq` in order is replaced by those operators applied to the `veq` itself.
+    ```mlir
+    %0 = quake.alloca !quake.veq<2>
+    cc.loop while ((%i = %c0) -> (i64)) {
+      ...
+    } do {
+    ^bb0(%i: i64):
+      %r = quake.extract_ref %0[%i] : (!quake.veq<2>, i64) -> !quake.ref
+      quake.h %r : (!quake.ref) -> ()
+      ...
+    }
+    ────────────────────────────────────────────────────────────────────────
+    %0 = quake.alloca !quake.veq<2>
+    quake.h %0 : (!quake.veq<2>) -> ()
+    ```
+    The loop is only rolled if it covers the entire vector and its body does
+    nothing but extract the element and operate on it. Only an uncontrolled
+    operator broadcasts: given a control, the last qubit is the sole target of
+    a single operation.
+  }];
+
+  let dependentDialects = ["cudaq::cc::CCDialect"];
+}
+
 def ConstantPropagation : Pass<"constant-propagation", "mlir::func::FuncOp"> {
   let summary = "Propagate constants to their uses.";
   let description = [{
@@ -598,6 +627,31 @@ def EraseVectorCopyCtor : Pass<"erase-vector-copy-ctor"> {
   }];
 }
 
+def ExpandBroadcasts : Pass<"expand-broadcasts", "mlir::func::FuncOp"> {
+  let summary = "Expands veqs used as targets into individual qubits.";
+  let description = [{
+    A one-target operator applied to a `veq` broadcasts that operator over
+    every element of the vector. Given a vector of constant size `n`, this
+    pass rewrites
+    ```mlir
+    quake.* %veq : (!quake.veq<n>) -> ()
+    ```
+    into the `n` operations it stands for:
+    ```mlir
+    %arg0 = quake.extract_ref %veq[0] : (!quake.veq<n>) -> !quake.ref
+    quake.* %arg0 : (!quake.ref) -> ()
+    ...
+    %argn = quake.extract_ref %veq[n-1] : (!quake.veq<n>) -> !quake.ref
+    quake.* %argn : (!quake.ref) -> ()
+    ```
+    Parameters and attributes are replicated on each operation. Multi-qubit
+    operators (`swap`, `exp_pauli`, custom unitaries) are left alone: for those
+    a `veq` target is the operand list of a single N-qubit gate rather than a
+    broadcast. Controlled operators are left alone as well: only an
+    uncontrolled operator broadcasts.
+  }];
+}
+
 def ExpandControlNegations :
     Pass<"expand-control-negations", "mlir::func::FuncOp"> {
   let summary =

cudaq/lib/Optimizer/Transforms/CMakeLists.txt

@@ -23,6 +23,7 @@ add_cudaq_library(OptTransforms
   CliffordTSynthesis.cpp
   CombineMeasurements.cpp
   CombineQuantumAlloc.cpp
+  ConsolidateBroadcasts.cpp
   ConstantPropagation.cpp
   DeadQuantumElimination.cpp
   DeadStoreRemoval.cpp
@@ -35,6 +36,7 @@ add_cudaq_library(OptTransforms
   EraseNopCalls.cpp
   EraseQEC.cpp
   EraseVectorCopyCtor.cpp
+  ExpandBroadcasts.cpp
   ExpandControlNegations.cpp
   ExpandControlVeqs.cpp
   ExpandMeasurements.cpp

cudaq/lib/Optimizer/Transforms/ConsolidateBroadcasts.cpp

@@ -0,0 +1,132 @@
+/*******************************************************************************
+ * 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.    *
+ ******************************************************************************/
+
+#include "LoopAnalysis.h"
+#include "PassDetails.h"
+#include "QuakeOperatorUtilities.h"
+#include "cudaq/Optimizer/Transforms/Passes.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+namespace cudaq::opt {
+#define GEN_PASS_DEF_CONSOLIDATEBROADCASTS
+#include "cudaq/Optimizer/Transforms/Passes.h.inc"
+} // namespace cudaq::opt
+
+#define DEBUG_TYPE "consolidate-broadcasts"
+
+using namespace mlir;
+
+namespace {
+
+/// Does \p region hold nothing but the loop's own control: no side effects,
+/// and no way out of the loop?
+bool isControlOnly(Region &region) {
+  return !region
+              .walk([](Operation *op) {
+                if (!isMemoryEffectFree(op) ||
+                    isa<cudaq::cc::BreakOp, cudaq::cc::UnwindBreakOp>(op))
+                  return WalkResult::interrupt();
+                return WalkResult::advance();
+              })
+              .wasInterrupted();
+}
+
+/// Match `for (i = 0; i < N; ++i) { op(v[i]); ... }`, where `v` is a veq of
+/// size `N`, and replace the entire loop with `op(v); ...`.
+LogicalResult rollLoop(cudaq::cc::LoopOp loop) {
+  // A counted loop runs a constant number of iterations from 0, stepping by 1,
+  // with no early exit and no `do while` form. It must also leave nothing
+  // behind.
+  if (!cudaq::opt::isaCountedLoop(loop) ||
+      !llvm::all_of(loop->getResults(), [](Value v) { return v.use_empty(); }))
+    return failure();
+  auto components = cudaq::opt::getLoopComponents(loop);
+  assert(components && "counted loop must have components");
+  auto iterations = components->getIterationsConstant();
+  if (!iterations)
+    return failure();
+
+  // An `else` region must not be dropped, and the while and step regions,
+  // which the rewrite also drops, must hold nothing but the loop's control.
+  if (loop.hasPythonElse() || !isControlOnly(loop.getWhileRegion()) ||
+      !isControlOnly(loop.getStepRegion()))
+    return failure();
+
+  Region &body = loop.getBodyRegion();
+  if (!body.hasOneBlock())
+    return failure();
+  Block &block = body.front();
+  // Only one argument, the induction variable
+  if (block.getNumArguments() != 1)
+    return failure();
+
+  // In the loop body, we're matching on the form
+  // %var = quake.extract_ref [%induction_var] %veq
+  // quake.op %var
+  // ...
+  // cc.continue
+  // Where each op is a broadcast operator
+  auto extract = dyn_cast<cudaq::quake::ExtractRefOp>(block.front());
+  if (!extract || extract.getIndex() != block.getArgument(0))
+    return failure();
+
+  auto isBroadcastable = [&extract](Operation &op) {
+    auto gate = dyn_cast<cudaq::quake::OperatorInterface>(op);
+    if (!gate)
+      return false;
+    if (!cudaq::opt::isBroadcastOperator(gate))
+      return false;
+    if (!gate.getControls().empty() || gate.getTargets().size() != 1 ||
+        gate.getTargets()[0] != extract.getRef())
+      return false;
+
+    return true;
+  };
+
+  SmallVector<cudaq::quake::OperatorInterface> broadcastable;
+
+  for (Operation &op : block.without_terminator()) {
+    if (extract == &op)
+      continue;
+    if (!isBroadcastable(op))
+      return failure();
+    broadcastable.emplace_back(&op);
+  }
+  if (broadcastable.empty())
+    return failure();
+
+  // The loop must walk the whole vector.
+  Value veq = extract.getVeq();
+  if (cudaq::quake::getVeqSize(veq) != iterations)
+    return failure();
+
+  // Given the body above, the operators' parameters are all defined outside
+  // the loop, so the clones are well-formed there. The sole target is the last
+  // operand.
+  OpBuilder builder(loop);
+  for (auto gate : broadcastable) {
+    Operation *broadcast = builder.clone(*gate.getOperation());
+    broadcast->setOperand(broadcast->getNumOperands() - 1, veq);
+  }
+  loop.erase();
+  return success();
+}
+
+struct ConsolidateBroadcastsPass
+    : public cudaq::opt::impl::ConsolidateBroadcastsBase<
+          ConsolidateBroadcastsPass> {
+  using ConsolidateBroadcastsBase::ConsolidateBroadcastsBase;
+
+  void runOnOperation() override {
+    SmallVector<cudaq::cc::LoopOp> loops;
+    getOperation().walk([&](cudaq::cc::LoopOp loop) { loops.push_back(loop); });
+    for (auto loop : loops)
+      (void)rollLoop(loop);
+  }
+};
+} // namespace

cudaq/lib/Optimizer/Transforms/ExpandBroadcasts.cpp

@@ -0,0 +1,78 @@
+/*******************************************************************************
+ * 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.    *
+ ******************************************************************************/
+
+#include "PassDetails.h"
+#include "cudaq/Optimizer/Transforms/Passes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+namespace cudaq::opt {
+#define GEN_PASS_DEF_EXPANDBROADCASTS
+#include "cudaq/Optimizer/Transforms/Passes.h.inc"
+} // namespace cudaq::opt
+
+#define DEBUG_TYPE "expand-broadcasts"
+
+using namespace mlir;
+
+namespace {
+/// Replace a single-qubit operator whose target is a constant sized veq of
+/// size \e N with \e N copies of that operator, one per element of the veq.
+/// The controls, parameters, and attributes of the original operator are
+/// replicated verbatim on each copy.
+template <typename OP>
+class ExpandBroadcastPat : public OpRewritePattern<OP> {
+public:
+  using OpRewritePattern<OP>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OP op,
+                                PatternRewriter &rewriter) const override {
+    // Only an uncontrolled operator broadcasts
+    if (op.getTargets().size() != 1 || !op.getControls().empty())
+      return failure();
+    Value target = op.getTargets()[0];
+    if (!isa<cudaq::quake::VeqType>(target.getType()))
+      return failure();
+    auto size = cudaq::quake::getVeqSize(target);
+    if (!size)
+      return failure();
+
+    auto loc = op.getLoc();
+    // The sole target is the last operand (skip angles for rotations)
+    unsigned targetPos = op->getNumOperands() - 1;
+    for (std::size_t i = 0; i < *size; ++i) {
+      Value ref = cudaq::quake::ExtractRefOp::create(rewriter, loc, target, i);
+      Operation *clone = rewriter.clone(*op.getOperation());
+      clone->setOperand(targetPos, ref);
+    }
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+struct ExpandBroadcastsPass
+    : public cudaq::opt::impl::ExpandBroadcastsBase<ExpandBroadcastsPass> {
+  using ExpandBroadcastsBase::ExpandBroadcastsBase;
+
+  void runOnOperation() override {
+    auto *ctx = &getContext();
+    RewritePatternSet patterns(ctx);
+    patterns.insert<
+        ExpandBroadcastPat<cudaq::quake::HOp>,
+        ExpandBroadcastPat<cudaq::quake::PhasedRxOp>,
+        ExpandBroadcastPat<cudaq::quake::R1Op>, ExpandBroadcastPat<cudaq::quake::RxOp>,
+        ExpandBroadcastPat<cudaq::quake::RyOp>, ExpandBroadcastPat<cudaq::quake::RzOp>,
+        ExpandBroadcastPat<cudaq::quake::SOp>, ExpandBroadcastPat<cudaq::quake::TOp>,
+        ExpandBroadcastPat<cudaq::quake::U2Op>, ExpandBroadcastPat<cudaq::quake::U3Op>,
+        ExpandBroadcastPat<cudaq::quake::XOp>, ExpandBroadcastPat<cudaq::quake::YOp>,
+        ExpandBroadcastPat<cudaq::quake::ZOp>>(ctx);
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
+      signalPassFailure();
+  }
+};
+} // namespace

cudaq/lib/Optimizer/Transforms/QuakeOperatorUtilities.h

@@ -19,6 +19,19 @@
 
 namespace cudaq::opt {
 
+/// Return true when \p op is a one-target operator for which a `veq` operand
+/// in the target position means "apply this operator to every element of the
+/// vector". Multi-qubit operators (`swap`, `exp_pauli`, custom unitaries) are
+/// excluded: for those a `veq` target is the operand list of a single N-qubit
+/// gate, not a broadcast.
+inline bool isBroadcastOperator(mlir::Operation *op) {
+  return mlir::isa<cudaq::quake::HOp, cudaq::quake::PhasedRxOp,
+                   cudaq::quake::R1Op, cudaq::quake::RxOp, cudaq::quake::RyOp,
+                   cudaq::quake::RzOp, cudaq::quake::SOp, cudaq::quake::TOp,
+                   cudaq::quake::U2Op, cudaq::quake::U3Op, cudaq::quake::XOp,
+                   cudaq::quake::YOp, cudaq::quake::ZOp>(op);
+}
+
 /// The controls and polarities resulting from expanding statically sized
 /// vector controls. Controls with unresolved vector sizes remain intact for
 /// callers that can lower them without making the predicate scalar.