diff --git a/cudaq/lib/Optimizer/Transforms/LoopAnalysis.cpp b/cudaq/lib/Optimizer/Transforms/LoopAnalysis.cpp index 649940d9d7f..312d50c0816 100644 --- a/cudaq/lib/Optimizer/Transforms/LoopAnalysis.cpp +++ b/cudaq/lib/Optimizer/Transforms/LoopAnalysis.cpp @@ -847,6 +847,19 @@ cudaq::opt::getSecondaryInductions(cudaq::cc::LoopOp loop, if (i == primaryIdx) continue; + // The else region runs once, on normal loop exit, so a value it recomputes + // has no closed form in terms of the primary. Only fuse `i` if the else + // region passes it through untouched. + if (loop.hasPythonElse()) { + Block &elseEntry = loop.getElseRegion().front(); + if (i >= elseEntry.getNumArguments()) + continue; + LoopRegionSite elseSite{&loop.getElseRegion(), /*isWhile=*/false}; + Value carried = getCarriedValue(elseSite, i); + if (!carried || carried != elseEntry.getArgument(i)) + continue; + } + Value stepVal; bool isAdd = false; bool isPrimaryAlias = false; diff --git a/cudaq/lib/Optimizer/Transforms/LoopUnrollPatterns.inc b/cudaq/lib/Optimizer/Transforms/LoopUnrollPatterns.inc index 23e18f3fa95..dcfdba634c1 100644 --- a/cudaq/lib/Optimizer/Transforms/LoopUnrollPatterns.inc +++ b/cudaq/lib/Optimizer/Transforms/LoopUnrollPatterns.inc @@ -361,8 +361,23 @@ struct UnrollCountedLoop : public OpRewritePattern { iterationOpers[*components->induction] = nextIterCount; setIterationOpers(contBlock->getArguments()); } + // The else region runs once, when the loop ends without a break, so it + // goes between the last iteration and the exit block. Break edges branch + // to endBlock directly, which is what skips it. + Block *exitTarget = endBlock; + if (loop.hasPythonElse()) { + rewriter.cloneRegionBefore(loop.getElseRegion(), endBlock); + exitTarget = insBlock->getNextNode(); + for (Block *b = exitTarget; b != endBlock; b = b->getNextNode()) + if (auto cont = dyn_cast(b->getTerminator())) { + auto termOpers = cont.getOperands(); + rewriter.setInsertionPoint(cont); + rewriter.replaceOpWithNewOp(cont, endBlock, termOpers); + } + rewriter.setInsertionPointToEnd(insBlock); + } [[maybe_unused]] auto lastBranch = - cf::BranchOp::create(rewriter, loc, endBlock, iterationOpers); + cf::BranchOp::create(rewriter, loc, exitTarget, iterationOpers); rewriter.replaceOp(loop, endBlock->getArguments()); LLVM_DEBUG(llvm::dbgs() << "after unrolling a loop:\n"; diff --git a/cudaq/test/Transforms/python_for_else.qke b/cudaq/test/Transforms/python_for_else.qke new file mode 100644 index 00000000000..8f18f34c4b3 --- /dev/null +++ b/cudaq/test/Transforms/python_for_else.qke @@ -0,0 +1,57 @@ +// ========================================================================== // +// 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-unroll --canonicalize %s | FileCheck %s +// RUN: cudaq-opt --cc-loop-normalize --cc-loop-induction-fusion %s | \ +// RUN: FileCheck --check-prefix=FUSE %s + +// A for/else. The else block runs when the loop ends without a break, so this +// returns 3 + 10. + +func.func @for_else() -> i64 { + %c0 = arith.constant 0 : i64 + %c1 = arith.constant 1 : i64 + %c3 = arith.constant 3 : i64 + %c10 = arith.constant 10 : i64 + %0:2 = cc.loop while ((%acc = %c0, %i = %c0) -> (i64, i64)) { + %c = arith.cmpi slt, %i, %c3 : i64 + cc.condition %c(%acc, %i : i64, i64) + } do { + ^bb0(%acc: i64, %i: i64): + %a = arith.addi %acc, %c1 : i64 + cc.continue %a, %i : i64, i64 + } step { + ^bb0(%acc: i64, %i: i64): + %s = arith.addi %i, %c1 : i64 + cc.continue %acc, %s : i64, i64 + } else { + ^bb0(%acc: i64, %i: i64): + %e = arith.addi %acc, %c10 : i64 + cc.continue %e, %i : i64, i64 + } + return %0#0 : i64 +} + +// Unrolling must emit the else region. + +// CHECK-LABEL: func.func @for_else() -> i64 { +// CHECK: %[[VAL:.*]] = arith.constant 13 : i64 +// CHECK: return %[[VAL]] : i64 +// CHECK: } + +// `acc` steps by one per iteration like the induction variable, so it looks +// like a secondary induction. It must not be fused away, because the else +// region gives it a different value. + +// FUSE-LABEL: func.func @for_else() -> i64 { +// FUSE: cc.loop while ((%{{.*}} = %{{.*}}, %{{.*}} = %{{.*}}) -> (i64, i64)) { +// FUSE: } else { +// FUSE: ^bb0(%[[ACC:.*]]: i64, %{{.*}}: i64): +// FUSE: %[[SUM:.*]] = arith.addi %[[ACC]], %{{.*}} : i64 +// FUSE: cc.continue %[[SUM]], %{{.*}} : i64, i64 +// FUSE: } diff --git a/python/tests/regression/test_loop_else.py b/python/tests/regression/test_loop_else.py new file mode 100644 index 00000000000..3384d99efaa --- /dev/null +++ b/python/tests/regression/test_loop_else.py @@ -0,0 +1,106 @@ +# ============================================================================ # +# 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. # +# ============================================================================ # + +# A loop's `else` block runs when the loop ends without a `break`. Two passes +# used to drop it. `cc-loop-unroll` never emitted the else region, and +# `cc-loop-induction-fusion` fused away a variable the else block updates. + +import cudaq + + +def test_for_else(): + + @cudaq.kernel + def kernel() -> int: + acc = 0 + for i in range(3): + acc = acc + 1 + else: + acc = acc + 10 + return acc + + assert kernel() == 13 + + +def test_while_else(): + + @cudaq.kernel + def kernel() -> int: + acc = 0 + i = 0 + while i < 3: + acc = acc + 1 + i = i + 1 + else: + acc = acc + 10 + return acc + + assert kernel() == 13 + + +def test_break_skips_else(): + + @cudaq.kernel + def kernel() -> int: + acc = 0 + for i in range(3): + acc = acc + 1 + if i == 1: + break + else: + acc = acc + 10 + return acc + + assert kernel() == 2 + + +def test_zero_trip_still_runs_else(): + + @cudaq.kernel + def kernel() -> int: + acc = 0 + for i in range(0): + acc = acc + 1 + else: + acc = acc + 10 + return acc + + assert kernel() == 10 + + +def test_nested_else(): + + @cudaq.kernel + def kernel() -> int: + acc = 0 + for i in range(2): + for j in range(2): + acc = acc + 1 + else: + acc = acc + 10 + else: + acc = acc + 100 + return acc + + assert kernel() == 124 + + +def test_else_with_quantum_op(): + + @cudaq.kernel + def kernel() -> int: + q = cudaq.qvector(2) + acc = 0 + for i in range(2): + acc = acc + 1 + else: + x(q[0]) + acc = acc + 10 + return acc + + assert kernel() == 12