diff --git a/cudaq/include/cudaq/Optimizer/Transforms/Passes.td b/cudaq/include/cudaq/Optimizer/Transforms/Passes.td index 4393d8f78a0..2b16649713a 100644 --- a/cudaq/include/cudaq/Optimizer/Transforms/Passes.td +++ b/cudaq/include/cudaq/Optimizer/Transforms/Passes.td @@ -1028,6 +1028,31 @@ def LoopNormalize : Pass<"cc-loop-normalize"> { ]; } +def LoopPruneDeadArgs : Pass<"cc-loop-prune-dead-args", "mlir::func::FuncOp"> { + let summary = "Drop loop-carried values that nothing reads."; + let description = [{ + A Python local lives for the whole function, so its storage is in the entry + block and `memtoreg` promotes it to a value carried by every loop around + it. When nothing reads the value, the loop is left carrying it along with + the arithmetic that computes it. + + That residue reads like a real dependence on an earlier loop's result, and + adjoint generation cannot reverse a block where a value produced by an op + it moves is read by one that stays. Where a carried value is dead, thread + the loop's own initial value around the loop instead and delete the + computations that go dead as a result. + + Run this after `memtoreg`, which is what creates the carried value, and + before `apply-op-specialization`. Order relative to `cc-loop-normalize` + does not matter. + + MLIR's generic `RegionBranchOpInterface` canonicalization cannot do this + job. It erases one dead value at a time, so it drops a slot from some of + the loop's regions and not the rest, and the `cc.loop` verifier rejects + that. A slot has to go from all four regions and the results at once. + }]; +} + def LoopPeeling : Pass<"cc-loop-peeling"> { let summary = "Peeling classical do-while loops."; let description = [{ diff --git a/cudaq/lib/Optimizer/Dialect/CC/CCOps.cpp b/cudaq/lib/Optimizer/Dialect/CC/CCOps.cpp index 654c1a0673e..7ac7e79f59c 100644 --- a/cudaq/lib/Optimizer/Dialect/CC/CCOps.cpp +++ b/cudaq/lib/Optimizer/Dialect/CC/CCOps.cpp @@ -1983,7 +1983,6 @@ SmallVector cudaq::cc::LoopOp::getLoopRegions() { OperandRange cudaq::cc::LoopOp::getEntrySuccessorOperands(RegionBranchPoint point) { - llvm::errs() << "getEntrySuccessorOperands: " << point << "\n"; assert(!point.isParent() && "invalid index region"); Operation *pred = point.getTerminatorPredecessorOrNull(); assert(pred && "must have a terminator"); diff --git a/cudaq/lib/Optimizer/Transforms/CMakeLists.txt b/cudaq/lib/Optimizer/Transforms/CMakeLists.txt index 087030d85ac..50100b3f22f 100644 --- a/cudaq/lib/Optimizer/Transforms/CMakeLists.txt +++ b/cudaq/lib/Optimizer/Transforms/CMakeLists.txt @@ -51,6 +51,7 @@ add_cudaq_library(OptTransforms LoopAnalysis.cpp LoopInductionFusion.cpp LoopNormalize.cpp + LoopPruneDeadArgs.cpp LoopPeeling.cpp LoopUnroll.cpp LowerPhase.cpp diff --git a/cudaq/lib/Optimizer/Transforms/LoopPruneDeadArgs.cpp b/cudaq/lib/Optimizer/Transforms/LoopPruneDeadArgs.cpp new file mode 100644 index 00000000000..03b1127d158 --- /dev/null +++ b/cudaq/lib/Optimizer/Transforms/LoopPruneDeadArgs.cpp @@ -0,0 +1,290 @@ +/******************************************************************************* + * 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 "PassDetails.h" +#include "cudaq/Optimizer/Dialect/CC/CCOps.h" +#include "cudaq/Optimizer/Dialect/Characteristics.h" +#include "cudaq/Optimizer/Transforms/Passes.h" +#include "mlir/Transforms/Passes.h" + +namespace cudaq::opt { +#define GEN_PASS_DEF_LOOPPRUNEDEADARGS +#include "cudaq/Optimizer/Transforms/Passes.h.inc" +} // namespace cudaq::opt + +#define DEBUG_TYPE "cc-loop-prune-dead-args" + +using namespace mlir; + +namespace { + +using LoopCarriedSlot = std::pair; + +/// Return the slot a use passes its value into, if passing it along is all the +/// use does (a loop's initial argument, or a region terminator forwarding it +/// around the loop). +std::optional forwardedSlot(OpOperand &use) { + Operation *user = use.getOwner(); + unsigned pos = use.getOperandNumber(); + if (auto loop = dyn_cast(user)) + return LoopCarriedSlot{loop.getOperation(), pos}; + auto loop = dyn_cast_or_null(user->getParentOp()); + if (!loop) + return std::nullopt; + if (isa(user)) { + if (pos == 0) + return std::nullopt; + return LoopCarriedSlot{loop.getOperation(), pos - 1}; + } + if (isa(user)) + return LoopCarriedSlot{loop.getOperation(), pos}; + return std::nullopt; +} + +/// Collect every value occupying slot pos of loop (the loop's result and the +/// matching entry block argument of each of its regions). +SmallVector valuesInSlot(cudaq::cc::LoopOp loop, unsigned pos) { + SmallVector values; + if (pos < loop.getNumResults()) + values.push_back(loop.getResult(pos)); + for (auto *region : loop.getRegions()) { + if (region->empty()) + continue; + Block &entry = region->front(); + if (pos < entry.getNumArguments()) + values.push_back(entry.getArgument(pos)); + } + return values; +} + +/// The slots in dead can be removed from loop's signature only if every +/// terminator forwarding values around the loop sits directly in one of its +/// regions, and nothing reads the block arguments being dropped. +bool canEraseDeadSlots(cudaq::cc::LoopOp loop, + const llvm::SmallBitVector &dead) { + bool ok = true; + loop.getOperation()->walk([&](Operation *op) { + if (!isa(op) || + op->getParentOp() == loop.getOperation() || !op->getNumOperands()) + return; + // A break nested in a scope still exits this loop and carries its values. + for (auto *p = op->getParentOp(); p; p = p->getParentOp()) { + if (p == loop.getOperation()) { + ok = false; + return; + } + if (isa(p)) + return; + } + }); + if (!ok) + return false; + for (auto *region : loop.getRegions()) { + if (region->empty()) + continue; + Block &entry = region->front(); + for (unsigned pos = 0, end = dead.size(); pos != end; ++pos) + if (dead[pos] && pos < entry.getNumArguments() && + !entry.getArgument(pos).use_empty()) + return false; + } + return true; +} + +/// Rebuild loop without the slots marked in dead. +void eraseDeadSlots(cudaq::cc::LoopOp loop, const llvm::SmallBitVector &dead) { + // Trim the operands each region terminator forwards. + for (auto *region : loop.getRegions()) + for (Block &block : *region) { + if (!block.hasNoSuccessors()) + continue; + Operation *term = block.getTerminator(); + if (term->getParentOp() != loop.getOperation() || + !isa(term)) + continue; + unsigned offset = isa(term) ? 1 : 0; + SmallVector keep; + for (unsigned i = 0, n = term->getNumOperands(); i != n; ++i) + if (i < offset || dead.size() <= i - offset || !dead[i - offset]) + keep.push_back(term->getOperand(i)); + term->setOperands(keep); + } + + // Drop the matching block arguments. + for (auto *region : loop.getRegions()) { + if (region->empty()) + continue; + Block &entry = region->front(); + for (int pos = dead.size() - 1; pos >= 0; --pos) + if (dead[pos] && static_cast(pos) < entry.getNumArguments()) + entry.eraseArgument(pos); + } + + SmallVector newInitArgs; + SmallVector newResultTypes; + for (unsigned pos = 0, end = loop.getInitialArgs().size(); pos != end; ++pos) + if (!dead[pos]) { + newInitArgs.push_back(loop.getInitialArgs()[pos]); + if (pos < loop.getNumResults()) + newResultTypes.push_back(loop.getResultTypes()[pos]); + } + + OpBuilder builder(loop); + auto newLoop = cudaq::cc::LoopOp::create( + builder, loop.getLoc(), newResultTypes, newInitArgs, + loop.isPostConditional(), [](OpBuilder &, Location, Region &) {}, + [](OpBuilder &, Location, Region &) {}, + /*stepBuilder=*/nullptr); + newLoop->setDiscardableAttrs(loop->getDiscardableAttrDictionary()); + newLoop.getWhileRegion().takeBody(loop.getWhileRegion()); + newLoop.getBodyRegion().takeBody(loop.getBodyRegion()); + newLoop.getStepRegion().takeBody(loop.getStepRegion()); + newLoop.getElseRegion().takeBody(loop.getElseRegion()); + + unsigned newPos = 0; + for (unsigned pos = 0, end = loop.getNumResults(); pos != end; ++pos) + if (!dead[pos]) + loop.getResult(pos).replaceAllUsesWith(newLoop.getResult(newPos++)); + loop.erase(); +} + +void pruneDeadLoopCarriedValues(func::FuncOp func) { + SmallVector loops; + func.walk([&](cudaq::cc::LoopOp loop) { loops.push_back(loop); }); + if (loops.empty()) + return; + + // Reaching the loop-carried values means going through the arithmetic that + // computes them, and that arithmetic reads the loop's own block arguments. + SmallVector candidates; + auto isPrunableComputation = [](Operation *op) { + return op->getNumRegions() == 0 && !op->hasTrait() && + !cudaq::opt::hasQuantum(*op) && isMemoryEffectFree(op); + }; + for (auto loop : loops) + for (unsigned pos = 0, end = loop.getInitialArgs().size(); pos != end; + ++pos) + llvm::append_range(candidates, valuesInSlot(loop, pos)); + func.walk([&](Operation *op) { + if (isPrunableComputation(op)) + llvm::append_range(candidates, op->getResults()); + }); + + // Start from the assumption that every candidate is dead and propagate + // liveness until it stops spreading. A value is live if + // 1. it has an opaque use + // 2. it is forwarded into a slot already known to be live + // 3. it feeds a computation whose own result is already known to be live + // Liveness only ever grows and the candidate set is finite, so this + // terminates. + DenseSet live; + auto slotIsLive = [&](LoopCarriedSlot slot) { + auto loop = cast(slot.first); + return llvm::any_of(valuesInSlot(loop, slot.second), + [&](Value val) { return live.count(val); }); + }; + auto hasLiveUse = [&](Value val) { + for (OpOperand &use : val.getUses()) { + if (auto forwarded = forwardedSlot(use)) { + if (slotIsLive(*forwarded)) + return true; + continue; + } + Operation *user = use.getOwner(); + if (!isPrunableComputation(user)) + return true; + if (llvm::any_of(user->getResults(), + [&](Value res) { return live.count(res); })) + return true; + } + return false; + }; + for (bool changed = true; changed;) { + changed = false; + for (Value val : candidates) + if (!live.count(val) && hasLiveUse(val)) { + live.insert(val); + changed = true; + } + } + + auto deleteDeadComputations = [&]() { + for (bool erased = true; erased;) { + erased = false; + SmallVector deadOps; + func.walk([&](Operation *op) { + if (!cudaq::opt::hasQuantum(*op) && isOpTriviallyDead(op)) + deadOps.push_back(op); + }); + for (auto *op : deadOps) { + op->erase(); + erased = true; + } + } + }; + + // Short-circuit each dead slot with the loop's initial value for that slot. + // The initial value is an operand of the loop, so it dominates every use we + // rewrite. + for (auto loop : loops) + for (auto [pos, initialArg] : llvm::enumerate(loop.getInitialArgs())) { + if (slotIsLive(LoopCarriedSlot{loop.getOperation(), pos})) + continue; + for (auto *region : loop.getRegions()) + for (Block &block : *region) { + if (!block.hasNoSuccessors()) + continue; + Operation *term = block.getTerminator(); + if (term->getParentOp() != loop.getOperation()) + continue; + unsigned operandPos = + isa(term) ? pos + 1 : pos; + if (isa(term) && + operandPos < term->getNumOperands()) + term->setOperand(operandPos, initialArg); + } + } + + // Delete the computations the short-circuiting just made dead. This has to + // happen before the slots are erased. While a dead slot's block argument + // still feeds arithmetic, `canEraseDeadSlots` refuses to drop it. + deleteDeadComputations(); + + // Drop each dead slot from the loop's signature. Erasing an inner loop's + // slots is what frees up the enclosing loop's, so keep going until nothing + // more can be dropped. + for (bool changed = true; changed;) { + changed = false; + SmallVector current; + func.walk([&](cudaq::cc::LoopOp loop) { current.push_back(loop); }); + for (auto loop : llvm::reverse(current)) { + llvm::SmallBitVector dead(loop.getInitialArgs().size()); + for (unsigned pos = 0, end = dead.size(); pos != end; ++pos) + if (!slotIsLive(LoopCarriedSlot{loop.getOperation(), pos})) + dead.set(pos); + if (dead.none() || !canEraseDeadSlots(loop, dead)) + continue; + eraseDeadSlots(loop, dead); + changed = true; + break; + } + } + + deleteDeadComputations(); +} + +class LoopPruneDeadArgsPass + : public cudaq::opt::impl::LoopPruneDeadArgsBase { +public: + using LoopPruneDeadArgsBase::LoopPruneDeadArgsBase; + + void runOnOperation() override { pruneDeadLoopCarriedValues(getOperation()); } +}; +} // namespace diff --git a/cudaq/lib/Optimizer/Transforms/LowerUnwind.cpp b/cudaq/lib/Optimizer/Transforms/LowerUnwind.cpp index b07ed2700d4..dfdc7d930db 100644 --- a/cudaq/lib/Optimizer/Transforms/LowerUnwind.cpp +++ b/cudaq/lib/Optimizer/Transforms/LowerUnwind.cpp @@ -370,6 +370,12 @@ struct ScopeOpPattern : public OpRewritePattern { assert(iter != infoMap.opParentMap.end()); bool asPrimitive = anyPrimitiveAncestor(infoMap.opParentMap, scope.getOperation()); + // A break or continue leaving this scope may still have to pass through + // the landing pad of an enclosing scope, which owns allocations of its own + // to deallocate. Ending the chain here with a `cc.break` or `cc.continue` + // is correct only when nothing but the loop itself lies above. + bool toLandingPad = asPrimitive || anyScopeAncestor(infoMap.opParentMap, + scope.getOperation()); LLVM_DEBUG(llvm::dbgs() << "replacing scope @" << scope.getLoc() << '\n'); auto loc = scope.getLoc(); auto *initBlock = rewriter.getInsertionBlock(); @@ -415,7 +421,7 @@ struct ScopeOpPattern : public OpRewritePattern { for (auto a : llvm::reverse(qallocas)) cudaq::quake::DeallocOp::create(rewriter, a->getLoc(), adjustedDeallocArg(a)); - if (asPrimitive) { + if (toLandingPad) { Block *landingPad = getLandingPad(infoMap, scope).continueBlock; cf::BranchOp::create(rewriter, loc, landingPad, blk->getArguments()); } else { @@ -429,7 +435,7 @@ struct ScopeOpPattern : public OpRewritePattern { for (auto a : llvm::reverse(qallocas)) cudaq::quake::DeallocOp::create(rewriter, a->getLoc(), adjustedDeallocArg(a)); - if (asPrimitive) { + if (toLandingPad) { Block *landingPad = getLandingPad(infoMap, scope).breakBlock; cf::BranchOp::create(rewriter, loc, landingPad, blk->getArguments()); } else { diff --git a/cudaq/lib/Optimizer/Transforms/MemToReg.cpp b/cudaq/lib/Optimizer/Transforms/MemToReg.cpp index 8a8f8410be2..552fe86c60d 100644 --- a/cudaq/lib/Optimizer/Transforms/MemToReg.cpp +++ b/cudaq/lib/Optimizer/Transforms/MemToReg.cpp @@ -994,11 +994,25 @@ class RegionDataFlow { // Phase 3: Update bindings and replace uses with the new block args. for (auto &info : defInfos) { for (auto [user, block] : std::get(info)) { - Value newReg = liveInMap[block][std::get<0>(info)]; - if (!hasBinding(block, std::get<0>(info)) || - getBinding(block, std::get<0>(info)) == std::get<1>(info)) - addBinding(block, std::get<0>(info), newReg); - user->replaceUsesOfWith(std::get<1>(info), newReg); + auto memref = std::get<0>(info); + auto oldVal = std::get<1>(info); + Value newReg = liveInMap[block][memref]; + if (!hasBinding(block, memref) || getBinding(block, memref) == oldVal) + addBinding(block, memref, newReg); + // A copy such as `x = i` stores this def's value into another + // variable, binding that variable to the same value. Note the copy's + // target before rewriting the store. + Value copyTarget; + if (auto store = dyn_cast(user)) + if (store.getValue() == oldVal) + copyTarget = store.getPtrvalue(); + user->replaceUsesOfWith(oldVal, newReg); + // The copy's target must follow the value it holds to the block + // argument, or `x` reads the value the def had before the loop + // instead of this iteration's. + if (copyTarget && hasBinding(block, copyTarget) && + getBinding(block, copyTarget) == oldVal) + addBinding(block, copyTarget, newReg); } } } diff --git a/cudaq/lib/Optimizer/Transforms/Pipelines.cpp b/cudaq/lib/Optimizer/Transforms/Pipelines.cpp index 9c17c3f5175..d96b72bd547 100644 --- a/cudaq/lib/Optimizer/Transforms/Pipelines.cpp +++ b/cudaq/lib/Optimizer/Transforms/Pipelines.cpp @@ -111,6 +111,7 @@ static void createTargetPrepPipeline(OpPassManager &pm, pm.addNestedPass(cudaq::opt::createUnwindLowering()); pm.addNestedPass(createCanonicalizerPass()); pm.addNestedPass(cudaq::opt::createClassicalMemToReg()); + pm.addNestedPass(cudaq::opt::createLoopPruneDeadArgs()); cudaq::opt::createClassicalOptimizationPipeline( pm, std::nullopt, {options.allowEarlyExit}, std::nullopt, {options.disableLoopUnrolling}); @@ -305,6 +306,7 @@ static void createPythonAOTPipeline(OpPassManager &pm, pm.addPass(cudaq::opt::createLambdaLifting()); pm.addNestedPass(cudaq::opt::createClassicalMemToReg()); pm.addNestedPass(createCanonicalizerPass()); + pm.addNestedPass(cudaq::opt::createLoopPruneDeadArgs()); pm.addNestedPass(cudaq::opt::createLoopNormalize()); pm.addNestedPass(cudaq::opt::createLoopInductionFusion()); pm.addPass(cudaq::opt::createApplySpecialization()); diff --git a/cudaq/lib/Optimizer/Transforms/VariableCoalesce.cpp b/cudaq/lib/Optimizer/Transforms/VariableCoalesce.cpp index 3b29729cfce..74d12ba01dd 100644 --- a/cudaq/lib/Optimizer/Transforms/VariableCoalesce.cpp +++ b/cudaq/lib/Optimizer/Transforms/VariableCoalesce.cpp @@ -23,22 +23,6 @@ namespace cudaq::opt { using namespace mlir; namespace { -/// A variable allocated inside a loop gets fresh storage on each iteration. -/// Raising it to the entry block replaces that with a single slot that -/// persists across iterations, which forces `memtoreg` to thread the variable -/// around the loop as a loop-carried value whether or not anything reads it. -/// Dead values like those defeat `cc.loop` reversal in the -/// apply-op-specialization pass, so such a variable is left where it is. -static bool isNestedInLoop(Operation *op) { - for (auto *p = op->getParentOp(); p; p = p->getParentOp()) { - if (isa(p)) - break; - if (isa(p)) - return true; - } - return false; -} - struct AllocationAnalysis { explicit AllocationAnalysis(Operation *op, bool hoistOnly) : hoistOnly(hoistOnly) { @@ -89,8 +73,6 @@ struct AllocationAnalysis { auto *parent = alloc->getParentOp(); if (isa(parent)) return WalkResult::advance(); - if (isNestedInLoop(alloc)) - return WalkResult::advance(); if (auto scope = dyn_cast(parent)) { varsToMove.insert(alloc); scopeMap[alloc] = scope; diff --git a/cudaq/test/Transforms/apply-8.qke b/cudaq/test/Transforms/apply-8.qke index b4650e910a1..3524c730933 100644 --- a/cudaq/test/Transforms/apply-8.qke +++ b/cudaq/test/Transforms/apply-8.qke @@ -6,7 +6,7 @@ // the terms of the Apache License 2.0 which accompanies this distribution. // // ========================================================================== // -// RUN: cudaq-opt --pass-pipeline='builtin.module(func.func(variable-coalesce,canonicalize,memtoreg{classical=1 quantum=0},canonicalize,cc-loop-normalize,cc-loop-induction-fusion),apply-op-specialization)' %s | FileCheck %s +// RUN: cudaq-opt --pass-pipeline='builtin.module(func.func(variable-coalesce,canonicalize,memtoreg{classical=1 quantum=0},canonicalize,cc-loop-prune-dead-args,cc-loop-normalize,cc-loop-induction-fusion),apply-op-specialization)' %s | FileCheck %s // Regression test for issue #5223. Taking the adjoint of a kernel with a nested // loop. The inner loop's induction variable `j` is local to the outer loop's @@ -90,10 +90,10 @@ func.func @entry() attributes {"cudaq-entrypoint", "cudaq-kernel"} { return } -// The reversed body runs `r1` before `h`, and the forward loops each carry a -// single iteration argument (the induction variable). A `j` raised to the -// entry block would show up here as a second, unread iteration argument on both -// loops, and reversal would fail outright. +// The reversed body runs `r1` before `h`. `j` and `angle` are allocated in the +// entry block, so both loops carry them as iteration arguments that nothing +// reads. `cc-loop-prune-dead-args` drops those slots, leaving each loop with +// just its induction variable, and that is what makes the loops reversible. // CHECK-LABEL: func.func private @qft.adj( // CHECK: cc.loop while @@ -102,6 +102,7 @@ func.func @entry() attributes {"cudaq-entrypoint", "cudaq-kernel"} { // CHECK: quake.h // CHECK-LABEL: func.func @qft( +// CHECK-NOT: cc.undef // CHECK: cc.loop while ((%[[I:.*]] = %{{.*}}) -> (i32)) { // CHECK: quake.h // CHECK: cc.loop while ((%[[J:.*]] = %{{.*}}) -> (i32)) { diff --git a/cudaq/test/Transforms/coalesce.qke b/cudaq/test/Transforms/coalesce.qke index 457959512db..f604f85510b 100644 --- a/cudaq/test/Transforms/coalesce.qke +++ b/cudaq/test/Transforms/coalesce.qke @@ -161,10 +161,6 @@ func.func @f2() { return } -// Variables local to a loop body keep their allocas in the body. Raising them -// to the entry block would turn them into loop-carried values that nothing -// reads, which blocks `cc.loop` reversal in apply-op-specialization. - // CHECK-LABEL: func.func @f2() { // CHECK-DAG: %[[VAL_11:.*]] = cc.alloca f64 // CHECK-DAG: %[[VAL_12:.*]] = cc.alloca i32 @@ -173,26 +169,5 @@ func.func @f2() { // CHECK-DAG: %[[VAL_15:.*]] = cc.alloca i32 // CHECK-DAG: %[[VAL_16:.*]] = cc.alloca f64 // CHECK-NOT: cc.alloca -// CHECK: cc.scope { -// CHECK: cc.loop while { -// CHECK: } do { -// CHECK: cc.scope { -// CHECK: %[[VAL_17:.*]] = cc.alloca i32 -// CHECK: %[[VAL_18:.*]] = cc.alloca f64 -// CHECK: func.call @test(%[[VAL_17]], %[[VAL_18]]) -// CHECK: } -// CHECK: cc.scope { -// CHECK: %[[VAL_19:.*]] = cc.alloca i32 -// CHECK: %[[VAL_20:.*]] = cc.alloca f64 -// CHECK: func.call @test(%[[VAL_19]], %[[VAL_20]]) -// CHECK: } -// CHECK: cc.continue -// CHECK: } step { -// CHECK: } -// CHECK: cc.scope { -// CHECK-NOT: cc.alloca -// CHECK: func.call @test(%[[VAL_12]], %[[VAL_11]]) -// CHECK: } -// CHECK: } -// CHECK: return -// CHECK: } +// CHECK: cc.scope +// CHECK: cc.loop while diff --git a/cudaq/test/Transforms/loop-prune-dead-args-wires.qke b/cudaq/test/Transforms/loop-prune-dead-args-wires.qke new file mode 100644 index 00000000000..f611a284674 --- /dev/null +++ b/cudaq/test/Transforms/loop-prune-dead-args-wires.qke @@ -0,0 +1,76 @@ +// ========================================================================== // +// 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 --cc-loop-prune-dead-args %s | FileCheck %s + +// The pass must never drop a loop slot that carries a wire. Every borrowed +// wire is returned, so a wire slot is always consumed and always live. + +quake.wire_set @wires[2147483647] + +// A wire gated inside the body and consumed after the loop. +func.func @wire_live() { + %cond = cc.undef i1 + %0 = quake.borrow_wire @wires[0] : !quake.wire + %loop = cc.loop while ((%w = %0) -> !quake.wire) { + cc.condition %cond (%w : !quake.wire) + } do { + ^bb0(%bw : !quake.wire): + %out = quake.h %bw : (!quake.wire) -> !quake.wire + cc.continue %out : !quake.wire + } + quake.return_wire %loop : !quake.wire + return +} + +// CHECK-LABEL: func.func @wire_live() +// CHECK: cc.loop while ((%{{.*}} = %{{.*}}) -> (!quake.wire)) { +// CHECK: quake.return_wire + +// Two wires threaded through one loop, both gated in the body, both returned. +func.func @wire_mixed() { + %cond = cc.undef i1 + %0 = quake.borrow_wire @wires[0] : !quake.wire + %1 = quake.borrow_wire @wires[1] : !quake.wire + %loop:2 = cc.loop while ((%a = %0, %b = %1) -> (!quake.wire, !quake.wire)) { + cc.condition %cond (%a, %b : !quake.wire, !quake.wire) + } do { + ^bb0(%ba : !quake.wire, %bb : !quake.wire): + %o:2 = quake.x [%ba] %bb : (!quake.wire, !quake.wire) -> (!quake.wire, !quake.wire) + cc.continue %o#0, %o#1 : !quake.wire, !quake.wire + } + quake.return_wire %loop#0 : !quake.wire + quake.return_wire %loop#1 : !quake.wire + return +} + +// CHECK-LABEL: func.func @wire_mixed() +// CHECK: cc.loop while ((%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) -> (!quake.wire, !quake.wire)) { + +// %1 is threaded through untouched. It is invariant but not dead, because the +// loop result is returned, so the slot must survive. +func.func @wire_invariant_but_returned() { + %cond = cc.undef i1 + %0 = quake.borrow_wire @wires[0] : !quake.wire + %1 = quake.borrow_wire @wires[1] : !quake.wire + %loop:2 = cc.loop while ((%a = %0, %b = %1) -> (!quake.wire, !quake.wire)) { + cc.condition %cond (%a, %b : !quake.wire, !quake.wire) + } do { + ^bb0(%ba : !quake.wire, %bb : !quake.wire): + %h = quake.h %ba : (!quake.wire) -> !quake.wire + cc.continue %h, %bb : !quake.wire, !quake.wire + } + quake.return_wire %loop#0 : !quake.wire + quake.return_wire %loop#1 : !quake.wire + return +} + +// CHECK-LABEL: func.func @wire_invariant_but_returned() +// CHECK-COUNT-2: quake.borrow_wire +// CHECK: cc.loop while +// CHECK-COUNT-2: quake.return_wire diff --git a/cudaq/test/Transforms/loop-prune-dead-args.qke b/cudaq/test/Transforms/loop-prune-dead-args.qke new file mode 100644 index 00000000000..8af9aaee18b --- /dev/null +++ b/cudaq/test/Transforms/loop-prune-dead-args.qke @@ -0,0 +1,103 @@ +// ========================================================================== // +// 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 --cc-loop-prune-dead-args %s | FileCheck %s + +// Slot 0 is carried around the loop and never read. The body computes it from +// the loop's own induction value, which is what makes the loop unreversible. +// The pass drops the slot and the arithmetic that fed it. + +func.func @dead_slot(%n: i32, %q: !quake.veq) { + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %u = cc.undef i32 + %0:2 = cc.loop while ((%arg0 = %u, %arg1 = %c0) -> (i32, i32)) { + %c = arith.cmpi slt, %arg1, %n : i32 + cc.condition %c(%arg0, %arg1 : i32, i32) + } do { + ^bb0(%arg0: i32, %arg1: i32): + %i = cc.cast signed %arg1 : (i32) -> i64 + %r = quake.extract_ref %q[%i] : (!quake.veq, i64) -> !quake.ref + quake.h %r : (!quake.ref) -> () + %v = arith.addi %arg1, %c1 : i32 + cc.continue %v, %arg1 : i32, i32 + } step { + ^bb0(%arg0: i32, %arg1: i32): + %s = arith.addi %arg1, %c1 : i32 + cc.continue %arg0, %s : i32, i32 + } + return +} + +// CHECK-LABEL: func.func @dead_slot( +// CHECK-NOT: cc.undef +// CHECK: cc.loop while ((%[[IV:.*]] = %{{.*}}) -> (i32)) { +// CHECK: cc.condition %{{.*}}(%[[IV]] : i32) +// CHECK: } do { +// CHECK: quake.h +// CHECK: cc.continue %{{.*}} : i32 +// CHECK: } step { +// CHECK: } + +// A slot the code below the loop reads is left alone. + +func.func @live_slot(%n: i32, %q: !quake.veq) -> i32 { + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %u = cc.undef i32 + %0:2 = cc.loop while ((%arg0 = %u, %arg1 = %c0) -> (i32, i32)) { + %c = arith.cmpi slt, %arg1, %n : i32 + cc.condition %c(%arg0, %arg1 : i32, i32) + } do { + ^bb0(%arg0: i32, %arg1: i32): + %v = arith.addi %arg1, %c1 : i32 + cc.continue %v, %arg1 : i32, i32 + } step { + ^bb0(%arg0: i32, %arg1: i32): + %s = arith.addi %arg1, %c1 : i32 + cc.continue %arg0, %s : i32, i32 + } + return %0#0 : i32 +} + +// CHECK-LABEL: func.func @live_slot( +// CHECK: %[[ADDI:.*]] = arith.addi +// CHECK: cc.continue %[[ADDI]], %{{.*}} : i32, i32 + +// A dead accumulator slot. Its block argument feeds arithmetic in the body, so +// the slot can only be erased after that arithmetic is gone. One run of the +// pass must do both. +func.func @dead_accumulator(%n: i32, %q: !quake.veq) { + %c0 = arith.constant 0 : i32 + %c1 = arith.constant 1 : i32 + %0:2 = cc.loop while ((%acc = %c0, %i = %c0) -> (i32, i32)) { + %c = arith.cmpi slt, %i, %n : i32 + cc.condition %c(%acc, %i : i32, i32) + } do { + ^bb0(%acc: i32, %i: i32): + %na = arith.addi %acc, %i : i32 + %x = cc.cast signed %i : (i32) -> i64 + %r = quake.extract_ref %q[%x] : (!quake.veq, i64) -> !quake.ref + quake.h %r : (!quake.ref) -> () + cc.continue %na, %i : i32, i32 + } step { + ^bb0(%acc: i32, %i: i32): + %s = arith.addi %i, %c1 : i32 + cc.continue %acc, %s : i32, i32 + } + return +} + +// CHECK-LABEL: func.func @dead_accumulator( +// CHECK: cc.loop while ((%[[IV:.*]] = %{{.*}}) -> (i32)) { +// CHECK: cc.condition %{{.*}}(%[[IV]] : i32) +// CHECK: } do { +// CHECK: ^bb0(%[[BIV:.*]]: i32): +// CHECK-NOT: arith.addi %{{.*}}, %[[BIV]] : i32 +// CHECK: cc.continue %[[BIV]] : i32 +// CHECK: } step { diff --git a/cudaq/test/Transforms/memtoreg-11.qke b/cudaq/test/Transforms/memtoreg-11.qke new file mode 100644 index 00000000000..e22e47e36e2 --- /dev/null +++ b/cudaq/test/Transforms/memtoreg-11.qke @@ -0,0 +1,47 @@ +// ========================================================================== // +// 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 --memtoreg=quantum=0 %s | FileCheck %s + +// A loop body that copies the counter into another variable. `v` must get the +// counter's value for this iteration, not the one it had before the loop. + +func.func @copy_counter(%n: i64) -> i64 { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %ctr = cc.alloca i64 + %v = cc.alloca i64 + cc.store %c0, %ctr : !cc.ptr + cc.store %c0, %v : !cc.ptr + cc.loop while { + %a = cc.load %ctr : !cc.ptr + %c = arith.cmpi slt, %a, %n : i64 + cc.condition %c + } do { + %a = cc.load %ctr : !cc.ptr + cc.store %a, %v : !cc.ptr + cc.continue + } step { + %a = cc.load %ctr : !cc.ptr + %b = arith.addi %a, %c1 : i64 + cc.store %b, %ctr : !cc.ptr + } + %r = cc.load %v : !cc.ptr + return %r : i64 +} + +// CHECK-LABEL: func.func @copy_counter( +// CHECK-SAME: %[[VAL_0:.*]]: i64) -> i64 { +// CHECK: %[[LOOP:.*]]:2 = cc.loop while ((%[[VAL_1:.*]] = %[[VAL_2:.*]], %[[VAL_3:.*]] = %[[VAL_4:.*]]) -> (i64, i64)) { +// CHECK: } do { +// CHECK: ^bb0(%[[VAL_5:.*]]: i64, %[[VAL_6:.*]]: i64): +// CHECK: cc.continue %[[VAL_5]], %[[VAL_5]] : i64, i64 +// CHECK: } step { +// CHECK: } +// CHECK: return %[[LOOP]]#1 : i64 +// CHECK: } diff --git a/cudaq/test/Transforms/unwind_lowering.qke b/cudaq/test/Transforms/unwind_lowering.qke index fd1c1cce169..6b08a5ee898 100644 --- a/cudaq/test/Transforms/unwind_lowering.qke +++ b/cudaq/test/Transforms/unwind_lowering.qke @@ -230,3 +230,126 @@ func.func @__nvqpp__mlirgen__3(%arg0: i32) { // CHECK: return // CHECK: } + +// The unwind sits in a scope of its own, nested inside the `if`, inside the +// scope that owns the allocation. The jump must thread through the landing pad +// of every scope on the way out, not just the innermost one, or the allocation +// is never deallocated on the taken branch. + +func.func @__nvqpp__mlirgen__4(%arg0: i32) { + %c1_i32 = arith.constant 1 : i32 + %c0_i32 = arith.constant 0 : i32 + %0 = cc.alloca i32 + cc.store %arg0, %0 : !cc.ptr + %1 = quake.alloca !quake.ref + cc.scope { + %2 = cc.alloca i32 + cc.store %c0_i32, %2 : !cc.ptr + cc.loop while { + %3 = cc.load %2 : !cc.ptr + %4 = cc.load %0 : !cc.ptr + %5 = arith.cmpi slt, %3, %4 : i32 + cc.condition %5 + } do { + cc.scope { + %3 = quake.alloca !quake.ref + quake.x %3 : (!quake.ref) -> () + %4 = cc.load %2 : !cc.ptr + %5 = arith.cmpi sgt, %4, %c0_i32 : i32 + cc.if(%5) { + cc.scope { + cc.unwind_break + } + } + quake.h %1 : (!quake.ref) -> () + } + cc.continue + } step { + %3 = cc.load %2 : !cc.ptr + %4 = arith.addi %3, %c1_i32 : i32 + cc.store %4, %2 : !cc.ptr + } + } + %measOut = quake.mz %1 : (!quake.ref) -> !cc.measure_handle + return +} + +// CHECK-LABEL: func.func @__nvqpp__mlirgen__4( +// CHECK: %[[ALLOCA_1:.*]] = quake.alloca !quake.ref +// CHECK: cc.scope { +// CHECK: cc.loop while { +// CHECK: } do { +// CHECK: %[[ALLOCA_2:.*]] = quake.alloca !quake.ref +// CHECK: quake.x %[[ALLOCA_2]] : (!quake.ref) -> () +// CHECK: cf.cond_br %{{.*}}, ^bb2, ^bb1 +// CHECK: ^bb1: +// CHECK: quake.h %[[ALLOCA_1]] : (!quake.ref) -> () +// CHECK: quake.dealloc %[[ALLOCA_2]] : !quake.ref +// CHECK: cc.continue +// CHECK: ^bb2: +// CHECK: quake.dealloc %[[ALLOCA_2]] : !quake.ref +// CHECK: cc.break +// CHECK: } step { +// CHECK: } +// CHECK: } +// CHECK: return +// CHECK: } + +func.func @__nvqpp__mlirgen__5(%arg0: i32) { + %c1_i32 = arith.constant 1 : i32 + %c0_i32 = arith.constant 0 : i32 + %0 = cc.alloca i32 + cc.store %arg0, %0 : !cc.ptr + %1 = quake.alloca !quake.ref + cc.scope { + %2 = cc.alloca i32 + cc.store %c0_i32, %2 : !cc.ptr + cc.loop while { + %3 = cc.load %2 : !cc.ptr + %4 = cc.load %0 : !cc.ptr + %5 = arith.cmpi slt, %3, %4 : i32 + cc.condition %5 + } do { + cc.scope { + %3 = quake.alloca !quake.ref + quake.x %3 : (!quake.ref) -> () + %4 = cc.load %2 : !cc.ptr + %5 = arith.cmpi sgt, %4, %c0_i32 : i32 + cc.if(%5) { + cc.scope { + cc.unwind_continue + } + } + quake.h %1 : (!quake.ref) -> () + } + cc.continue + } step { + %3 = cc.load %2 : !cc.ptr + %4 = arith.addi %3, %c1_i32 : i32 + cc.store %4, %2 : !cc.ptr + } + } + %measOut = quake.mz %1 : (!quake.ref) -> !cc.measure_handle + return +} + +// CHECK-LABEL: func.func @__nvqpp__mlirgen__5( +// CHECK: %[[ALLOCA_1:.*]] = quake.alloca !quake.ref +// CHECK: cc.scope { +// CHECK: cc.loop while { +// CHECK: } do { +// CHECK: %[[ALLOCA_2:.*]] = quake.alloca !quake.ref +// CHECK: quake.x %[[ALLOCA_2]] : (!quake.ref) -> () +// CHECK: cf.cond_br %{{.*}}, ^bb2, ^bb1 +// CHECK: ^bb1: +// CHECK: quake.h %[[ALLOCA_1]] : (!quake.ref) -> () +// CHECK: quake.dealloc %[[ALLOCA_2]] : !quake.ref +// CHECK: cc.continue +// CHECK: ^bb2: +// CHECK: quake.dealloc %[[ALLOCA_2]] : !quake.ref +// CHECK: cc.continue +// CHECK: } step { +// CHECK: } +// CHECK: } +// CHECK: return +// CHECK: } diff --git a/python/cudaq/kernel/ast_bridge.py b/python/cudaq/kernel/ast_bridge.py index ecee44356b6..fb6915adb97 100644 --- a/python/cudaq/kernel/ast_bridge.py +++ b/python/cudaq/kernel/ast_bridge.py @@ -432,54 +432,6 @@ def node_error(msg): self.isSubscriptRoot = False self.verbose = verbose self.currentNode = None - # `for` loop targets that are used nowhere outside their loop, keyed on - # `id()` - self.loopLocalTargets = {} - self.sinkAllocaNames = set() - - def __analyzeLoopLocalTargets(self, statements, argNames): - """Record, for each `for` loop in `statements`, which of its target - variables never occur outside that loop. - - Python keeps a loop variable alive after its loop, so by default the - storage for one is allocated in the function's entry block. That is - needed only when something below the loop can still read it; a variable - that no code outside the loop mentions can live in the loop body - instead. Keeping it there matters because `memtoreg` promotes an - entry-block slot into a value carried by every enclosing loop, whether - or not anything reads it, and those dead loop-carried values defeat - `cc.loop` reversal in the apply-op-specialization pass. - """ - forNodes = [ - n for stmt in statements for n in ast.walk(stmt) - if isinstance(n, ast.For) - ] - if not forNodes: - return - allNames = [ - n for stmt in statements for n in ast.walk(stmt) - if isinstance(n, ast.Name) - ] - for forNode in forNodes: - if forNode.orelse: - continue - targets = { - t.id - for t in ast.walk(forNode.target) - if isinstance(t, ast.Name) - } - targets -= set(argNames) - if not targets: - continue - insideLoop = {id(n) for n in ast.walk(forNode)} - usedOutside = { - n.id - for n in allNames - if n.id in targets and id(n) not in insideLoop - } - local = targets - usedOutside - if local: - self.loopLocalTargets[id(forNode)] = local def isCudaqName(self, name): """Return True if `name` is 'cudaq' or a known alias for the cudaq @@ -1567,6 +1519,52 @@ def createMonotonicForLoop(self, None if orElseBuilder is None else (lambda args: orElseBuilder(args[0]))) + def createMonotonicForLoopInMemory(self, + bodyBuilder, + startVal, + stepVal, + endVal, + isDecrementing=False, + orElseBuilder=None): + """Create a `for` loop whose counter lives in memory. + + The loop loads the counter to test it and loads/increments/stores it + to step. `memtoreg` promotes it to a value later. + + The counter is separate from the Python loop variable, which takes a + copy of it each iteration, the way Python takes each value from the + iterator. + """ + iTy = self.getIntegerType() + condPred = IntegerAttr.get( + iTy, 4) if isDecrementing else IntegerAttr.get(iTy, 2) + + # The loop gets its own scope for the counter, so the slot dies with + # the loop. + scope = cc.ScopeOp([]) + scopeBlock = Block.create_at_start(scope.initRegion, []) + with InsertionPoint(scopeBlock): + counter = cc.AllocaOp(cc.PointerType.get(iTy), + TypeAttr.get(iTy)).result + cc.StoreOp(startVal, counter) + + def evalCond(_): + return arith.CmpIOp(condPred, + cc.LoadOp(counter).result, endVal).result + + def evalStep(_): + next = arith.AddIOp(cc.LoadOp(counter).result, stepVal).result + cc.StoreOp(next, counter) + return [] + + loop = self.createForLoop( + [], lambda _: bodyBuilder(cc.LoadOp(counter).result), [], + evalCond, evalStep, None if orElseBuilder is None else + (lambda _: orElseBuilder(cc.LoadOp(counter).result))) + if not self.hasTerminator(scopeBlock): + cc.ContinueOp([]) + return loop + def createInvariantForLoop(self, bodyBuilder, endVal): """Create an invariant loop using the CC dialect.""" @@ -2115,8 +2113,6 @@ def visit_FunctionDef(self, node): # errors on assignments that may lead to unexpected behavior # (i.e. behavior not following expected Python behavior). self.buildingFunctionBody = True - self.__analyzeLoopLocalTargets( - node.body, [arg.arg for arg in node.args.args]) with trace.span("ast_bridge.visit_function_body", statement_count=len(node.body)): for n in node.body: @@ -2482,12 +2478,9 @@ def update_in_parent_block(destination, value): if storeAsVal or cc.PointerType.isinstance(value.type): return target, value - # A variable that outlives the block it is assigned in needs - # its storage in the function's entry block. - allocaBlock = (InsertionPoint.current.block - if target.id in self.sinkAllocaNames else - self.symbolTable.scopeRoot) - with InsertionPoint.at_block_begin(allocaBlock): + # A Python local lives for the whole function, so its + # storage goes in the entry block. + with InsertionPoint.at_block_begin(self.symbolTable.scopeRoot): address = cc.AllocaOp(cc.PointerType.get(value.type), TypeAttr.get(value.type)).result cc.StoreOp(value, address) @@ -5344,8 +5337,6 @@ def loadElement(iterVar): else: self.emitFatalError('{} iterable type not supported.', node) - loopLocal = self.loopLocalTargets.get(id(node), set()) - def blockBuilder(iterVar, stmts): self.symbolTable.beginBlock() values = getValues(iterVar) @@ -5353,25 +5344,21 @@ def blockBuilder(iterVar, stmts): # iteration variable(s) to have consistent behavior. assignNode = ast.Assign(targets=[node.target], value=values) assignNode.lineno = node.lineno - outerSink = self.sinkAllocaNames - self.sinkAllocaNames = { - name for name in loopLocal if name not in self.symbolTable - } - try: - self.visit(assignNode) - finally: - self.sinkAllocaNames = outerSink + self.visit(assignNode) self.buildScopedBlock(stmts) self.symbolTable.endBlock() - self.createMonotonicForLoop( - lambda iterVar: blockBuilder(iterVar, node.body), - startVal=startVal, - stepVal=stepVal, - endVal=endVal, - isDecrementing=isDecrementing, - orElseBuilder=None if not node.orelse else - lambda iterVar: blockBuilder(iterVar, node.orelse)) + # `range()` counts on its own; other `iterables` use the counter as an + # index. + createLoop = (self.createMonotonicForLoop + if iterable else self.createMonotonicForLoopInMemory) + createLoop(lambda iterVar: blockBuilder(iterVar, node.body), + startVal=startVal, + stepVal=stepVal, + endVal=endVal, + isDecrementing=isDecrementing, + orElseBuilder=None if not node.orelse else + lambda iterVar: blockBuilder(iterVar, node.orelse)) def visit_While(self, node): """Convert Python while statements into the equivalent CC `LoopOp`.""" diff --git a/python/tests/mlir/ast_break.py b/python/tests/mlir/ast_break.py index 421ec9a549c..90c6c7e498d 100644 --- a/python/tests/mlir/ast_break.py +++ b/python/tests/mlir/ast_break.py @@ -43,14 +43,14 @@ def kernel(x: float): # CHECK: %[[VAL_17:.*]] = math.fpowi %[[VAL_16]], %[[VAL_2]] : f64, i64 # CHECK: %[[VAL_18:.*]] = arith.addf %[[VAL_16]], %[[VAL_17]] : f64 # CHECK: %[[VAL_19:.*]] = arith.cmpf ogt, %[[VAL_18]], %[[VAL_1]] : f64 -# CHECK: cf.cond_br %[[VAL_19]], ^bb1, ^bb2 +# CHECK: cf.cond_br %[[VAL_19]], ^bb2, ^bb1 # CHECK: ^bb1: -# CHECK: cc.break %[[VAL_14]], %[[VAL_18]] : i64, f64 -# CHECK: ^bb2: # CHECK: %[[VAL_24:.*]] = arith.remui %[[VAL_14]], %[[VAL_6]] : i64 # CHECK: %[[VAL_25:.*]] = quake.extract_ref %[[VAL_8]]{{\[}}%[[VAL_24]]] : (!quake.veq<4>, i64) -> !quake.ref # CHECK: quake.ry (%[[VAL_18]]) %[[VAL_25]] : (f64, !quake.ref) -> () # CHECK: cc.continue %[[VAL_14]], %[[VAL_18]] : i64, f64 +# CHECK: ^bb2: +# CHECK: cc.break %[[VAL_14]], %[[VAL_18]] : i64, f64 # CHECK: } step { # CHECK: ^bb0(%[[VAL_26:.*]]: i64, %[[VAL_28:.*]]: f64): # CHECK: %[[VAL_29:.*]] = arith.addi %[[VAL_26]], %[[VAL_4]] : i64 diff --git a/python/tests/mlir/measure_handle.py b/python/tests/mlir/measure_handle.py index 63ee78b5f6c..fdb2e201c5d 100644 --- a/python/tests/mlir/measure_handle.py +++ b/python/tests/mlir/measure_handle.py @@ -470,7 +470,7 @@ def kernel_handle_vec_cross_round(): # CHECK: %[[MVEC:.*]] = quake.mz %[[VEQ]] name "mvec" : (!quake.veq<3>) -> !cc.sequence # CHECK: cc.loop while # CHECK: %[[MNEW:.*]] = quake.mz %[[VEQ]] name "m_new" : (!quake.veq<3>) -> !cc.sequence -# CHECK: cc.continue {{.*}}%[[MNEW]], %[[MNEW]] +# CHECK: cc.continue # CHECK-NOT: quake.discriminate # CHECK: return diff --git a/python/tests/mlir/qalloc_scope_loop.py b/python/tests/mlir/qalloc_scope_loop.py index f7756aad8a4..d2dc6078cd8 100644 --- a/python/tests/mlir/qalloc_scope_loop.py +++ b/python/tests/mlir/qalloc_scope_loop.py @@ -80,6 +80,40 @@ def for_else_kernel(n: int): print(for_else_kernel) +def test_qalloc_freed_on_break(): + """A qubit allocated in a `for` body is freed on the `break` path.""" + + @cudaq.kernel + def break_kernel(n: int): + r = cudaq.qubit() + for i in range(n): + q = cudaq.qvector(2) + h(q[0]) + x.ctrl(q[0], r) + if i == 1: + break + mz(r) + + print(break_kernel) + + +def test_qalloc_freed_on_continue(): + """A qubit allocated in a `for` body is freed on the `continue` path.""" + + @cudaq.kernel + def continue_kernel(n: int): + r = cudaq.qubit() + for i in range(n): + q = cudaq.qvector(2) + h(q[0]) + if i == 1: + continue + x.ctrl(q[0], r) + mz(r) + + print(continue_kernel) + + # CHECK-LABEL: func.func @__nvqpp__mlirgen__for_kernel.. # CHECK: %[[VAL_0:.*]] = quake.alloca !quake.ref # CHECK: cc.loop while @@ -123,3 +157,34 @@ def for_else_kernel(n: int): # CHECK: } # CHECK: } # CHECK: quake.dealloc %[[VAL_0]] : !quake.ref + +# CHECK-LABEL: func.func @__nvqpp__mlirgen__break_kernel.. +# CHECK: %[[VAL_0:.*]] = quake.alloca !quake.ref +# CHECK: cc.loop while +# CHECK: } do { +# CHECK: %[[VAL_1:.*]] = quake.alloca !quake.veq<2> +# CHECK: cf.cond_br %{{.*}}, ^bb2, ^bb1 +# CHECK: ^bb1: +# CHECK: quake.dealloc %[[VAL_1]] : !quake.veq<2> +# CHECK: cc.continue +# CHECK: ^bb2: +# CHECK: quake.dealloc %[[VAL_1]] : !quake.veq<2> +# CHECK: cc.break +# CHECK: } step { +# CHECK: quake.dealloc %[[VAL_0]] : !quake.ref + +# CHECK-LABEL: func.func @__nvqpp__mlirgen__continue_kernel.. +# CHECK: %[[VAL_0:.*]] = quake.alloca !quake.ref +# CHECK: cc.loop while +# CHECK: } do { +# CHECK: %[[VAL_1:.*]] = quake.alloca !quake.veq<2> +# CHECK: cf.cond_br %{{.*}}, ^bb2, ^bb1 +# CHECK: ^bb1: +# CHECK: quake.x [%{{.*}}] %[[VAL_0]] : (!quake.ref, !quake.ref) -> () +# CHECK: quake.dealloc %[[VAL_1]] : !quake.veq<2> +# CHECK: cc.continue +# CHECK: ^bb2: +# CHECK: quake.dealloc %[[VAL_1]] : !quake.veq<2> +# CHECK: cc.continue +# CHECK: } step { +# CHECK: quake.dealloc %[[VAL_0]] : !quake.ref diff --git a/python/tests/mlir/qec_ops.py b/python/tests/mlir/qec_ops.py index 276985bdf42..720a74348f3 100644 --- a/python/tests/mlir/qec_ops.py +++ b/python/tests/mlir/qec_ops.py @@ -278,18 +278,16 @@ def kernel_rep_code_d3(n_rounds: int): # CHECK-SAME: %[[ARG0:[0-9]+|[a-zA-Z$._-][a-zA-Z0-9$._-]*]]: i64) attributes {"cudaq-entrypoint", "cudaq-kernel"} { # CHECK-DAG: %[[CONSTANT_0:.*]] = arith.constant 1 : i64 # CHECK-DAG: %[[CONSTANT_1:.*]] = arith.constant 0 : i64 -# CHECK: %[[UNDEF_0:.*]] = cc.undef !cc.measure_handle -# CHECK: %[[UNDEF_1:.*]] = cc.undef !cc.measure_handle # CHECK: %[[ALLOCA_0:.*]] = quake.alloca !quake.veq<3> # CHECK: %[[ALLOCA_1:.*]] = quake.alloca !quake.ref # CHECK: %[[ALLOCA_2:.*]] = quake.alloca !quake.ref # CHECK: %[[UNDEF_2:.*]] = cc.undef !cc.measure_handle # CHECK: %[[UNDEF_3:.*]] = cc.undef !cc.measure_handle -# CHECK: %[[LOOP_0:.*]]:5 = cc.loop while ((%[[VAL_0:.*]] = %[[CONSTANT_1]], %[[VAL_2:.*]] = %[[UNDEF_2]], %[[VAL_3:.*]] = %[[UNDEF_3]], %[[VAL_4:.*]] = %[[UNDEF_1]], %[[VAL_5:.*]] = %[[UNDEF_0]]) -> (i64, !cc.measure_handle, !cc.measure_handle, !cc.measure_handle, !cc.measure_handle)) { +# CHECK: %[[LOOP_0:.*]]:3 = cc.loop while ((%[[VAL_0:.*]] = %[[CONSTANT_1]], %[[VAL_4:.*]] = %[[UNDEF_2]], %[[VAL_5:.*]] = %[[UNDEF_3]]) -> (i64, !cc.measure_handle, !cc.measure_handle)) { # CHECK: %[[CMPI_0:.*]] = arith.cmpi slt, %[[VAL_0]], %[[ARG0]] : i64 -# CHECK: cc.condition %[[CMPI_0]](%[[VAL_0]], %[[VAL_2]], %[[VAL_3]], %[[VAL_4]], %[[VAL_5]] : i64, !cc.measure_handle, !cc.measure_handle, !cc.measure_handle, !cc.measure_handle) +# CHECK: cc.condition %[[CMPI_0]](%[[VAL_0]], %[[VAL_4]], %[[VAL_5]] : i64, !cc.measure_handle, !cc.measure_handle) # CHECK: } do { -# CHECK: ^bb0(%[[VAL_6:.*]]: i64, %[[VAL_8:.*]]: !cc.measure_handle, %[[VAL_9:.*]]: !cc.measure_handle, %[[VAL_10:.*]]: !cc.measure_handle, %[[VAL_11:.*]]: !cc.measure_handle): +# CHECK: ^bb0(%[[VAL_6:.*]]: i64, %[[VAL_10:.*]]: !cc.measure_handle, %[[VAL_11:.*]]: !cc.measure_handle): # CHECK: %[[EXTRACT_REF_0:.*]] = quake.extract_ref %[[ALLOCA_0]][0] : (!quake.veq<3>) -> !quake.ref # CHECK: quake.x {{\[}}%[[EXTRACT_REF_0]]] %[[ALLOCA_1]] : (!quake.ref, !quake.ref) -> () # CHECK: %[[EXTRACT_REF_1:.*]] = quake.extract_ref %[[ALLOCA_0]][1] : (!quake.veq<3>) -> !quake.ref @@ -303,15 +301,15 @@ def kernel_rep_code_d3(n_rounds: int): # CHECK: quake.reset %[[ALLOCA_2]] : (!quake.ref) -> () # CHECK: %[[CMPI_1:.*]] = arith.cmpi sgt, %[[VAL_6]], %[[CONSTANT_1]] : i64 # CHECK: cc.if(%[[CMPI_1]]) { -# CHECK: qec.detector %[[VAL_8]], %[[MZ_0]] : !cc.measure_handle, !cc.measure_handle -# CHECK: qec.detector %[[VAL_9]], %[[MZ_1]] : !cc.measure_handle, !cc.measure_handle +# CHECK: qec.detector %[[VAL_10]], %[[MZ_0]] : !cc.measure_handle, !cc.measure_handle +# CHECK: qec.detector %[[VAL_11]], %[[MZ_1]] : !cc.measure_handle, !cc.measure_handle # CHECK: } else { # CHECK: } -# CHECK: cc.continue %[[VAL_6]], %[[MZ_0]], %[[MZ_1]], %[[MZ_0]], %[[MZ_1]] : i64, !cc.measure_handle, !cc.measure_handle, !cc.measure_handle, !cc.measure_handle +# CHECK: cc.continue %[[VAL_6]], %[[MZ_0]], %[[MZ_1]] : i64, !cc.measure_handle, !cc.measure_handle # CHECK: } step { -# CHECK: ^bb0(%[[VAL_12:.*]]: i64, %[[VAL_14:.*]]: !cc.measure_handle, %[[VAL_15:.*]]: !cc.measure_handle, %[[VAL_16:.*]]: !cc.measure_handle, %[[VAL_17:.*]]: !cc.measure_handle): +# CHECK: ^bb0(%[[VAL_12:.*]]: i64, %[[VAL_16:.*]]: !cc.measure_handle, %[[VAL_17:.*]]: !cc.measure_handle): # CHECK: %[[ADDI_0:.*]] = arith.addi %[[VAL_12]], %[[CONSTANT_0]] : i64 -# CHECK: cc.continue %[[ADDI_0]], %[[VAL_14]], %[[VAL_15]], %[[VAL_16]], %[[VAL_17]] : i64, !cc.measure_handle, !cc.measure_handle, !cc.measure_handle, !cc.measure_handle +# CHECK: cc.continue %[[ADDI_0]], %[[VAL_16]], %[[VAL_17]] : i64, !cc.measure_handle, !cc.measure_handle # CHECK: } # CHECK: %[[MZ_2:.*]] = quake.mz %[[ALLOCA_0]] name "readout" : (!quake.veq<3>) -> !cc.sequence # CHECK: qec.observable %[[MZ_2]] : !cc.sequence diff --git a/python/tests/mlir/qft.py b/python/tests/mlir/qft.py index 503b5d63f87..d8946c2d2f4 100644 --- a/python/tests/mlir/qft.py +++ b/python/tests/mlir/qft.py @@ -39,7 +39,6 @@ def iqft(qubits: cudaq.qview): # CHECK-DAG: %[[CONSTANT_2:.*]] = arith.constant 2 : i64 # CHECK-DAG: %[[CONSTANT_3:.*]] = arith.constant 1 : i64 # CHECK-DAG: %[[CONSTANT_4:.*]] = arith.constant 0 : i64 -# CHECK-DAG: %[[UNDEF_1:.*]] = cc.undef i64 # CHECK: %[[VEQ_SIZE_0:.*]] = quake.veq_size %[[ARG0]] : (!quake.veq) -> i64 # CHECK: %[[FLOORDIVSI_0:.*]] = arith.floordivsi %[[VEQ_SIZE_0]], %[[CONSTANT_2]] : i64 # CHECK: %[[LOOP_0:.*]] = cc.loop while ((%[[VAL_0:.*]] = %[[CONSTANT_4]]) -> (i64)) { @@ -57,13 +56,13 @@ def iqft(qubits: cudaq.qview): # CHECK: ^bb0(%[[VAL_2:.*]]: i64): # CHECK: %[[ADDI_0:.*]] = arith.addi %[[VAL_2]], %[[CONSTANT_3]] : i64 # CHECK: cc.continue %[[ADDI_0]] : i64 -# CHECK: } {normalized} +# CHECK: } # CHECK: %[[SUBI_2:.*]] = arith.subi %[[VEQ_SIZE_0]], %[[CONSTANT_3]] : i64 -# CHECK: %[[LOOP_1:.*]]:2 = cc.loop while ((%[[VAL_3:.*]] = %[[CONSTANT_4]], %[[VAL_4:.*]] = %[[UNDEF_1]]) -> (i64, i64)) { +# CHECK: %[[LOOP_1:.*]] = cc.loop while ((%[[VAL_3:.*]] = %[[CONSTANT_4]]) -> (i64)) { # CHECK: %[[CMPI_1:.*]] = arith.cmpi slt, %[[VAL_3]], %[[SUBI_2]] : i64 -# CHECK: cc.condition %[[CMPI_1]](%[[VAL_3]], %[[VAL_4]] : i64, i64) +# CHECK: cc.condition %[[CMPI_1]](%[[VAL_3]] : i64) # CHECK: } do { -# CHECK: ^bb0(%[[VAL_6:.*]]: i64, %[[VAL_7:.*]]: i64): +# CHECK: ^bb0(%[[VAL_6:.*]]: i64): # CHECK: %[[EXTRACT_REF_2:.*]] = quake.extract_ref %[[ARG0]]{{\[}}%[[VAL_6]]] : (!quake.veq, i64) -> !quake.ref # CHECK: quake.h %[[EXTRACT_REF_2]] : (!quake.ref) -> () # CHECK: %[[ADDI_1:.*]] = arith.addi %[[VAL_6]], %[[CONSTANT_3]] : i64 @@ -90,12 +89,12 @@ def iqft(qubits: cudaq.qview): # CHECK: %[[ADDI_2:.*]] = arith.addi %[[VAL_11]], %[[CONSTANT_3]] : i64 # CHECK: cc.continue %[[ADDI_2]] : i64 # CHECK: } {normalized} -# CHECK: cc.continue %[[VAL_6]], %[[ADDI_1]] : i64, i64 +# CHECK: cc.continue %[[VAL_6]] : i64 # CHECK: } step { -# CHECK: ^bb0(%[[VAL_12:.*]]: i64, %[[VAL_13:.*]]: i64): +# CHECK: ^bb0(%[[VAL_12:.*]]: i64): # CHECK: %[[ADDI_3:.*]] = arith.addi %[[VAL_12]], %[[CONSTANT_3]] : i64 -# CHECK: cc.continue %[[ADDI_3]], %[[VAL_13]] : i64, i64 -# CHECK: } {normalized} +# CHECK: cc.continue %[[ADDI_3]] : i64 +# CHECK: } # CHECK: %[[SUBI_7:.*]] = arith.subi %[[VEQ_SIZE_0]], %[[CONSTANT_3]] : i64 # CHECK: %[[EXTRACT_REF_5:.*]] = quake.extract_ref %[[ARG0]]{{\[}}%[[SUBI_7]]] : (!quake.veq, i64) -> !quake.ref # CHECK: quake.h %[[EXTRACT_REF_5]] : (!quake.ref) -> ()