diff --git a/cudaq/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.h b/cudaq/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.h index b99d59d8cd5..1d639fd6828 100644 --- a/cudaq/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.h +++ b/cudaq/include/cudaq/Optimizer/Dialect/Quake/QuakeOps.h @@ -183,6 +183,19 @@ constexpr bool isMeasure = std::is_same_v || std::is_same_v || std::is_same_v; +/// 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(op); +} + //===----------------------------------------------------------------------===// // Control and wire helpers. //===----------------------------------------------------------------------===// diff --git a/cudaq/include/cudaq/Optimizer/Transforms/Passes.td b/cudaq/include/cudaq/Optimizer/Transforms/Passes.td index 5eca8251ca6..2930a6362c9 100644 --- a/cudaq/include/cudaq/Optimizer/Transforms/Passes.td +++ b/cudaq/include/cudaq/Optimizer/Transforms/Passes.td @@ -433,6 +433,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 = [{ @@ -660,6 +689,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) -> () + ``` + into the `n` operations it stands for: + ```mlir + %arg0 = quake.extract_ref %veq[0] : (!quake.veq) -> !quake.ref + quake.* %arg0 : (!quake.ref) -> () + ... + %argn = quake.extract_ref %veq[n-1] : (!quake.veq) -> !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 = diff --git a/cudaq/lib/Optimizer/Transforms/CMakeLists.txt b/cudaq/lib/Optimizer/Transforms/CMakeLists.txt index 37ccb4b7e9f..e58c95dc3f4 100644 --- a/cudaq/lib/Optimizer/Transforms/CMakeLists.txt +++ b/cudaq/lib/Optimizer/Transforms/CMakeLists.txt @@ -24,6 +24,7 @@ add_cudaq_library(OptTransforms CombineMeasurements.cpp CombineQuantumAlloc.cpp CommutationAwareRewrite.cpp + ConsolidateBroadcasts.cpp ConstantPropagation.cpp DeadQuantumElimination.cpp DeadStoreRemoval.cpp @@ -36,6 +37,7 @@ add_cudaq_library(OptTransforms EraseNopCalls.cpp EraseQEC.cpp EraseVectorCopyCtor.cpp + ExpandBroadcasts.cpp ExpandControlNegations.cpp ExpandControlVeqs.cpp ExpandMeasurements.cpp diff --git a/cudaq/lib/Optimizer/Transforms/ConsolidateBroadcasts.cpp b/cudaq/lib/Optimizer/Transforms/ConsolidateBroadcasts.cpp new file mode 100644 index 00000000000..411961d2f01 --- /dev/null +++ b/cudaq/lib/Optimizer/Transforms/ConsolidateBroadcasts.cpp @@ -0,0 +1,131 @@ +/******************************************************************************* + * 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 "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 ®ion) { + return !region + .walk([](Operation *op) { + if (!isMemoryEffectFree(op) || + isa(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(block.front()); + if (!extract || extract.getIndex() != block.getArgument(0)) + return failure(); + + auto isBroadcastable = [&extract](Operation &op) { + auto gate = dyn_cast(op); + if (!gate) + return false; + if (!cudaq::quake::isBroadcastOperator(gate)) + return false; + if (!gate.getControls().empty() || gate.getTargets().size() != 1 || + gate.getTargets()[0] != extract.getRef()) + return false; + + return true; + }; + + SmallVector 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 loops; + getOperation().walk([&](cudaq::cc::LoopOp loop) { loops.push_back(loop); }); + for (auto loop : loops) + (void)rollLoop(loop); + } +}; +} // namespace diff --git a/cudaq/lib/Optimizer/Transforms/ExpandBroadcasts.cpp b/cudaq/lib/Optimizer/Transforms/ExpandBroadcasts.cpp new file mode 100644 index 00000000000..c85e611e5bc --- /dev/null +++ b/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 +class ExpandBroadcastPat : public OpRewritePattern { +public: + using OpRewritePattern::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(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 { + using ExpandBroadcastsBase::ExpandBroadcastsBase; + + void runOnOperation() override { + auto *ctx = &getContext(); + RewritePatternSet patterns(ctx); + patterns.insert< + ExpandBroadcastPat, + ExpandBroadcastPat, + ExpandBroadcastPat, ExpandBroadcastPat, + ExpandBroadcastPat, ExpandBroadcastPat, + ExpandBroadcastPat, ExpandBroadcastPat, + ExpandBroadcastPat, ExpandBroadcastPat, + ExpandBroadcastPat, ExpandBroadcastPat, + ExpandBroadcastPat>(ctx); + if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) + signalPassFailure(); + } +}; +} // namespace diff --git a/cudaq/test/Transforms/consolidate_broadcasts.qke b/cudaq/test/Transforms/consolidate_broadcasts.qke new file mode 100644 index 00000000000..30e72318241 --- /dev/null +++ b/cudaq/test/Transforms/consolidate_broadcasts.qke @@ -0,0 +1,506 @@ +// ========================================================================== // +// 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. // +// ========================================================================== // + +// RUN: cudaq-opt --consolidate-broadcasts --canonicalize %s | FileCheck %s + +// A counted loop covering the whole vector rolls into one operation. +func.func @loop() { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @loop() { +// CHECK: %[[VAL_0:.*]] = quake.alloca !quake.veq<4> +// CHECK: quake.h %[[VAL_0]] : (!quake.veq<4>) -> () +// CHECK-NOT: cc.loop +// CHECK: return +// CHECK: } + +// A parameter defined outside the loop is preserved. +func.func @parameters(%angle: f64) { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.rx (%angle) %e : (f64, !quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @parameters( +// CHECK-SAME: %[[VAL_0:.*]]: f64) { +// CHECK: %[[VAL_1:.*]] = quake.alloca !quake.veq<4> +// CHECK: quake.rx (%[[VAL_0]]) %[[VAL_1]] : (f64, !quake.veq<4>) -> () +// CHECK-NOT: cc.loop +// CHECK: return +// CHECK: } + +// Several operators on the same element all roll, in order. +func.func @several(%angle: f64) { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + quake.rx (%angle) %e : (f64, !quake.ref) -> () + quake.x %e : (!quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @several( +// CHECK-SAME: %[[VAL_0:.*]]: f64) { +// CHECK: %[[VAL_1:.*]] = quake.alloca !quake.veq<4> +// CHECK: quake.h %[[VAL_1]] : (!quake.veq<4>) -> () +// CHECK: quake.rx (%[[VAL_0]]) %[[VAL_1]] : (f64, !quake.veq<4>) -> () +// CHECK: quake.x %[[VAL_1]] : (!quake.veq<4>) -> () +// CHECK-NOT: cc.loop +// CHECK: return +// CHECK: } + +// An operator on some other qubit blocks the rewrite. +func.func @foreign_op(%q: !quake.ref) { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + quake.x %q : (!quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @foreign_op( +// CHECK-SAME: %[[VAL_0:.*]]: !quake.ref) { +// CHECK: %[[VAL_1:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_2:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_3:.*]] = arith.constant 4 : i64 +// CHECK: %[[VAL_4:.*]] = quake.alloca !quake.veq<4> +// CHECK: %[[VAL_5:.*]] = cc.loop while ((%[[VAL_6:.*]] = %[[VAL_1]]) -> (i64)) { +// CHECK: %[[VAL_7:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_3]] : i64 +// CHECK: cc.condition %[[VAL_7]](%[[VAL_6]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_8:.*]]: i64): +// CHECK: %[[VAL_9:.*]] = quake.extract_ref %[[VAL_4]][%[[VAL_8]]] : (!quake.veq<4>, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_9]] : (!quake.ref) -> () +// CHECK: quake.x %[[VAL_0]] : (!quake.ref) -> () +// CHECK: cc.continue %[[VAL_8]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_10:.*]]: i64): +// CHECK: %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_2]] : i64 +// CHECK: cc.continue %[[VAL_11]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } + +// A loop that does not cover the whole vector is left alone. +func.func @short_loop() { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c2 = arith.constant 2 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c2 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @short_loop() { +// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_2:.*]] = arith.constant 2 : i64 +// CHECK: %[[VAL_3:.*]] = quake.alloca !quake.veq<4> +// CHECK: %[[VAL_4:.*]] = cc.loop while ((%[[VAL_5:.*]] = %[[VAL_0]]) -> (i64)) { +// CHECK: %[[VAL_6:.*]] = arith.cmpi slt, %[[VAL_5]], %[[VAL_2]] : i64 +// CHECK: cc.condition %[[VAL_6]](%[[VAL_5]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_7:.*]]: i64): +// CHECK: %[[VAL_8:.*]] = quake.extract_ref %[[VAL_3]][%[[VAL_7]]] : (!quake.veq<4>, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_8]] : (!quake.ref) -> () +// CHECK: cc.continue %[[VAL_7]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_9:.*]]: i64): +// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_1]] : i64 +// CHECK: cc.continue %[[VAL_10]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } + +// A loop over a vector of unknown size is left alone. +func.func @dynamic(%vec: !quake.veq, %size: i64) { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %size : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %vec[%i] : (!quake.veq, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @dynamic( +// CHECK-SAME: %[[VAL_0:.*]]: !quake.veq, %[[VAL_1:.*]]: i64) { +// CHECK: %[[VAL_2:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_3:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_4:.*]] = cc.loop while ((%[[VAL_5:.*]] = %[[VAL_2]]) -> (i64)) { +// CHECK: %[[VAL_6:.*]] = arith.cmpi slt, %[[VAL_5]], %[[VAL_1]] : i64 +// CHECK: cc.condition %[[VAL_6]](%[[VAL_5]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_7:.*]]: i64): +// CHECK: %[[VAL_8:.*]] = quake.extract_ref %[[VAL_0]][%[[VAL_7]]] : (!quake.veq, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_8]] : (!quake.ref) -> () +// CHECK: cc.continue %[[VAL_7]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_9:.*]]: i64): +// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_3]] : i64 +// CHECK: cc.continue %[[VAL_10]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } + +// A rotation whose angle depends on the induction variable is left alone. +func.func @varying_angle() { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %f = arith.sitofp %i : i64 to f64 + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.rx (%f) %e : (f64, !quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @varying_angle() { +// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_2:.*]] = arith.constant 4 : i64 +// CHECK: %[[VAL_3:.*]] = quake.alloca !quake.veq<4> +// CHECK: %[[VAL_4:.*]] = cc.loop while ((%[[VAL_5:.*]] = %[[VAL_0]]) -> (i64)) { +// CHECK: %[[VAL_6:.*]] = arith.cmpi slt, %[[VAL_5]], %[[VAL_2]] : i64 +// CHECK: cc.condition %[[VAL_6]](%[[VAL_5]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_7:.*]]: i64): +// CHECK: %[[VAL_8:.*]] = arith.sitofp %[[VAL_7]] : i64 to f64 +// CHECK: %[[VAL_9:.*]] = quake.extract_ref %[[VAL_3]][%[[VAL_7]]] : (!quake.veq<4>, i64) -> !quake.ref +// CHECK: quake.rx (%[[VAL_8]]) %[[VAL_9]] : (f64, !quake.ref) -> () +// CHECK: cc.continue %[[VAL_7]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_10:.*]]: i64): +// CHECK: %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_1]] : i64 +// CHECK: cc.continue %[[VAL_11]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } + +// Only an uncontrolled operator broadcasts. +func.func @controlled(%ctrl: !quake.ref) { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.z [%ctrl] %e : (!quake.ref, !quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @controlled( +// CHECK-SAME: %[[VAL_0:.*]]: !quake.ref) { +// CHECK: %[[VAL_1:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_2:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_3:.*]] = arith.constant 4 : i64 +// CHECK: %[[VAL_4:.*]] = quake.alloca !quake.veq<4> +// CHECK: %[[VAL_5:.*]] = cc.loop while ((%[[VAL_6:.*]] = %[[VAL_1]]) -> (i64)) { +// CHECK: %[[VAL_7:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_3]] : i64 +// CHECK: cc.condition %[[VAL_7]](%[[VAL_6]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_8:.*]]: i64): +// CHECK: %[[VAL_9:.*]] = quake.extract_ref %[[VAL_4]][%[[VAL_8]]] : (!quake.veq<4>, i64) -> !quake.ref +// CHECK: quake.z [%[[VAL_0]]] %[[VAL_9]] : (!quake.ref, !quake.ref) -> () +// CHECK: cc.continue %[[VAL_8]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_10:.*]]: i64): +// CHECK: %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_2]] : i64 +// CHECK: cc.continue %[[VAL_11]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } + +func.func private @side_effect() -> i64 + +// The while region is dropped, so it must hold nothing but the loop control. +func.func @sneaky_while() { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %n = func.call @side_effect() : () -> i64 + %b = arith.cmpi slt, %i, %n : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @sneaky_while() { +// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_2:.*]] = quake.alloca !quake.veq<4> +// CHECK: %[[VAL_3:.*]] = cc.loop while ((%[[VAL_4:.*]] = %[[VAL_0]]) -> (i64)) { +// CHECK: %[[VAL_5:.*]] = func.call @side_effect() : () -> i64 +// CHECK: %[[VAL_6:.*]] = arith.cmpi slt, %[[VAL_4]], %[[VAL_5]] : i64 +// CHECK: cc.condition %[[VAL_6]](%[[VAL_4]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_7:.*]]: i64): +// CHECK: %[[VAL_8:.*]] = quake.extract_ref %[[VAL_2]][%[[VAL_7]]] : (!quake.veq<4>, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_8]] : (!quake.ref) -> () +// CHECK: cc.continue %[[VAL_7]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_9:.*]]: i64): +// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_1]] : i64 +// CHECK: cc.continue %[[VAL_10]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } + +// The step region is dropped as well. +func.func @sneaky_step(%q: !quake.ref) { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + quake.x %q : (!quake.ref) -> () + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @sneaky_step( +// CHECK-SAME: %[[VAL_0:.*]]: !quake.ref) { +// CHECK: %[[VAL_1:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_2:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_3:.*]] = arith.constant 4 : i64 +// CHECK: %[[VAL_4:.*]] = quake.alloca !quake.veq<4> +// CHECK: %[[VAL_5:.*]] = cc.loop while ((%[[VAL_6:.*]] = %[[VAL_1]]) -> (i64)) { +// CHECK: %[[VAL_7:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_3]] : i64 +// CHECK: cc.condition %[[VAL_7]](%[[VAL_6]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_8:.*]]: i64): +// CHECK: %[[VAL_9:.*]] = quake.extract_ref %[[VAL_4]][%[[VAL_8]]] : (!quake.veq<4>, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_9]] : (!quake.ref) -> () +// CHECK: cc.continue %[[VAL_8]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_10:.*]]: i64): +// CHECK: quake.x %[[VAL_0]] : (!quake.ref) -> () +// CHECK: %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_2]] : i64 +// CHECK: cc.continue %[[VAL_11]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } + +// An else region must not be dropped either. +func.func @with_else(%q: !quake.ref) { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + cc.continue %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } else { + ^bb0(%i: i64): + quake.x %q : (!quake.ref) -> () + cc.continue %i : i64 + } + return +} + +// CHECK-LABEL: func.func @with_else( +// CHECK-SAME: %[[VAL_0:.*]]: !quake.ref) { +// CHECK: %[[VAL_1:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_2:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_3:.*]] = arith.constant 4 : i64 +// CHECK: %[[VAL_4:.*]] = quake.alloca !quake.veq<4> +// CHECK: %[[VAL_5:.*]] = cc.loop while ((%[[VAL_6:.*]] = %[[VAL_1]]) -> (i64)) { +// CHECK: %[[VAL_7:.*]] = arith.cmpi slt, %[[VAL_6]], %[[VAL_3]] : i64 +// CHECK: cc.condition %[[VAL_7]](%[[VAL_6]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_8:.*]]: i64): +// CHECK: %[[VAL_9:.*]] = quake.extract_ref %[[VAL_4]][%[[VAL_8]]] : (!quake.veq<4>, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_9]] : (!quake.ref) -> () +// CHECK: cc.continue %[[VAL_8]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_10:.*]]: i64): +// CHECK: %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_2]] : i64 +// CHECK: cc.continue %[[VAL_11]] : i64 +// CHECK: } else { +// CHECK: ^bb0(%[[VAL_12:.*]]: i64): +// CHECK: quake.x %[[VAL_0]] : (!quake.ref) -> () +// CHECK: cc.continue %[[VAL_12]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } + +// A body that exits early does not visit every element. +func.func @early_exit() { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c4 = arith.constant 4 : i64 + %0 = quake.alloca !quake.veq<4> + %1 = cc.loop while ((%i = %c0) -> (i64)) { + %b = arith.cmpi slt, %i, %c4 : i64 + cc.condition %b(%i : i64) + } do { + ^bb0(%i: i64): + %e = quake.extract_ref %0[%i] : (!quake.veq<4>, i64) -> !quake.ref + quake.h %e : (!quake.ref) -> () + cc.break %i : i64 + } step { + ^bb0(%i: i64): + %n = arith.addi %i, %c1 : i64 + cc.continue %n : i64 + } + return +} + +// CHECK-LABEL: func.func @early_exit() { +// CHECK: %[[VAL_0:.*]] = arith.constant 0 : i64 +// CHECK: %[[VAL_1:.*]] = arith.constant 1 : i64 +// CHECK: %[[VAL_2:.*]] = arith.constant 4 : i64 +// CHECK: %[[VAL_3:.*]] = quake.alloca !quake.veq<4> +// CHECK: %[[VAL_4:.*]] = cc.loop while ((%[[VAL_5:.*]] = %[[VAL_0]]) -> (i64)) { +// CHECK: %[[VAL_6:.*]] = arith.cmpi slt, %[[VAL_5]], %[[VAL_2]] : i64 +// CHECK: cc.condition %[[VAL_6]](%[[VAL_5]] : i64) +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_7:.*]]: i64): +// CHECK: %[[VAL_8:.*]] = quake.extract_ref %[[VAL_3]][%[[VAL_7]]] : (!quake.veq<4>, i64) -> !quake.ref +// CHECK: quake.h %[[VAL_8]] : (!quake.ref) -> () +// CHECK: cc.break %[[VAL_7]] : i64 +// CHECK: } step { +// CHECK: ^bb0(%[[VAL_9:.*]]: i64): +// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_1]] : i64 +// CHECK: cc.continue %[[VAL_10]] : i64 +// CHECK: } +// CHECK: return +// CHECK: } diff --git a/cudaq/test/Transforms/expand_broadcasts.qke b/cudaq/test/Transforms/expand_broadcasts.qke new file mode 100644 index 00000000000..52c57a9601e --- /dev/null +++ b/cudaq/test/Transforms/expand_broadcasts.qke @@ -0,0 +1,84 @@ +// ========================================================================== // +// 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. // +// ========================================================================== // + +// RUN: cudaq-opt --expand-broadcasts %s | FileCheck %s + +func.func @broadcast() { + %0 = quake.alloca !quake.veq<3> + quake.h %0 : (!quake.veq<3>) -> () + return +} + +// CHECK-LABEL: func.func @broadcast() { +// CHECK: %[[VAL_0:.*]] = quake.alloca !quake.veq<3> +// CHECK: %[[VAL_1:.*]] = quake.extract_ref %[[VAL_0]][0] : (!quake.veq<3>) -> !quake.ref +// CHECK: quake.h %[[VAL_1]] : (!quake.ref) -> () +// CHECK: %[[VAL_2:.*]] = quake.extract_ref %[[VAL_0]][1] : (!quake.veq<3>) -> !quake.ref +// CHECK: quake.h %[[VAL_2]] : (!quake.ref) -> () +// CHECK: %[[VAL_3:.*]] = quake.extract_ref %[[VAL_0]][2] : (!quake.veq<3>) -> !quake.ref +// CHECK: quake.h %[[VAL_3]] : (!quake.ref) -> () +// CHECK: return +// CHECK: } + +// Parameters and adjointness are replicated on each operation. +func.func @parameters(%angle: f64) { + %0 = quake.alloca !quake.veq<2> + quake.rx (%angle) %0 : (f64, !quake.veq<2>) -> () + return +} + +// CHECK-LABEL: func.func @parameters( +// CHECK-SAME: %[[VAL_0:.*]]: f64) { +// CHECK: %[[VAL_1:.*]] = quake.alloca !quake.veq<2> +// CHECK: %[[VAL_2:.*]] = quake.extract_ref %[[VAL_1]][0] : (!quake.veq<2>) -> !quake.ref +// CHECK: quake.rx (%[[VAL_0]]) %[[VAL_2]] : (f64, !quake.ref) -> () +// CHECK: %[[VAL_3:.*]] = quake.extract_ref %[[VAL_1]][1] : (!quake.veq<2>) -> !quake.ref +// CHECK: quake.rx (%[[VAL_0]]) %[[VAL_3]] : (f64, !quake.ref) -> () +// CHECK: return +// CHECK: } + +// Only an uncontrolled operator broadcasts. +func.func @controlled(%ctrl: !quake.ref) { + %0 = quake.alloca !quake.veq<2> + quake.z [%ctrl] %0 : (!quake.ref, !quake.veq<2>) -> () + return +} + +// CHECK-LABEL: func.func @controlled( +// CHECK-SAME: %[[VAL_0:.*]]: !quake.ref) { +// CHECK: %[[VAL_1:.*]] = quake.alloca !quake.veq<2> +// CHECK: quake.z [%[[VAL_0]]] %[[VAL_1]] : (!quake.ref, !quake.veq<2>) -> () +// CHECK: return +// CHECK: } + +// A vector of unknown size cannot be expanded. +func.func @unsized(%vec: !quake.veq) { + quake.x %vec : (!quake.veq) -> () + return +} + +// CHECK-LABEL: func.func @unsized( +// CHECK-SAME: %[[VAL_0:.*]]: !quake.veq) { +// CHECK: quake.x %[[VAL_0]] : (!quake.veq) -> () +// CHECK: return +// CHECK: } + +// For a multi-qubit operator a veq target is an operand list, not a broadcast. +func.func @multi_qubit() { + %0 = quake.alloca !quake.veq<2> + %1 = quake.alloca !quake.veq<2> + quake.swap %0, %1 : (!quake.veq<2>, !quake.veq<2>) -> () + return +} + +// CHECK-LABEL: func.func @multi_qubit() { +// CHECK: %[[VAL_0:.*]] = quake.alloca !quake.veq<2> +// CHECK: %[[VAL_1:.*]] = quake.alloca !quake.veq<2> +// CHECK: quake.swap %[[VAL_0]], %[[VAL_1]] : (!quake.veq<2>, !quake.veq<2>) -> () +// CHECK: return +// CHECK: }