Fixing the Python bridge's loop and scope emission

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

@@ -1027,6 +1027,25 @@ 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` and before `cc-loop-normalize`, which rewrites
+    loop control and would otherwise tangle the dead value into it.
+  }];
+}
+
 def LoopPeeling : Pass<"cc-loop-peeling"> {
   let summary = "Peeling classical do-while loops.";
   let description = [{

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

cudaq/lib/Optimizer/Transforms/LoopPruneDeadArgs.cpp

@@ -0,0 +1,282 @@
+/*******************************************************************************
+ * 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<Operation *, unsigned>;
+
+/// 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<LoopCarriedSlot> forwardedSlot(OpOperand &use) {
+  Operation *user = use.getOwner();
+  unsigned pos = use.getOperandNumber();
+  if (auto loop = dyn_cast<cudaq::cc::LoopOp>(user))
+    return LoopCarriedSlot{loop.getOperation(), pos};
+  auto loop = dyn_cast_or_null<cudaq::cc::LoopOp>(user->getParentOp());
+  if (!loop)
+    return std::nullopt;
+  if (isa<cudaq::cc::ConditionOp>(user)) {
+    if (pos == 0)
+      return std::nullopt;
+    return LoopCarriedSlot{loop.getOperation(), pos - 1};
+  }
+  if (isa<cudaq::cc::ContinueOp, cudaq::cc::BreakOp>(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<Value> valuesInSlot(cudaq::cc::LoopOp loop, unsigned pos) {
+  SmallVector<Value> 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<cudaq::cc::ConditionOp, cudaq::cc::BreakOp>(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<cudaq::cc::LoopOp>(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<cudaq::cc::ConditionOp, cudaq::cc::ContinueOp,
+               cudaq::cc::BreakOp>(term))
+        continue;
+      unsigned offset = isa<cudaq::cc::ConditionOp>(term) ? 1 : 0;
+      SmallVector<Value> 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<unsigned>(pos) < entry.getNumArguments())
+        entry.eraseArgument(pos);
+  }
+
+  SmallVector<Value> newInitArgs;
+  SmallVector<Type> 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<cudaq::cc::LoopOp> 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<Value> candidates;
+  auto isPrunableComputation = [](Operation *op) {
+    return op->getNumRegions() == 0 && !op->hasTrait<OpTrait::IsTerminator>() &&
+           !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<Value> live;
+  auto slotIsLive = [&](LoopCarriedSlot slot) {
+    auto loop = cast<cudaq::cc::LoopOp>(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;
+      }
+  }
+
+  // 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<cudaq::cc::ConditionOp>(term) ? pos + 1 : pos;
+          if (isa<cudaq::cc::ConditionOp, cudaq::cc::ContinueOp,
+                  cudaq::cc::BreakOp>(term) &&
+              operandPos < term->getNumOperands())
+            term->setOperand(operandPos, initialArg);
+        }
+    }
+
+  // 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<cudaq::cc::LoopOp> 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;
+    }
+  }
+
+  // Delete the classical computations that just went dead.
+  for (bool erased = true; erased;) {
+    erased = false;
+    SmallVector<Operation *> deadOps;
+    func.walk([&](Operation *op) {
+      if (!cudaq::opt::hasQuantum(*op) && isOpTriviallyDead(op))
+        deadOps.push_back(op);
+    });
+    for (auto *op : deadOps) {
+      op->erase();
+      erased = true;
+    }
+  }
+}
+
+class LoopPruneDeadArgsPass
+    : public cudaq::opt::impl::LoopPruneDeadArgsBase<LoopPruneDeadArgsPass> {
+public:
+  using LoopPruneDeadArgsBase::LoopPruneDeadArgsBase;
+
+  void runOnOperation() override { pruneDeadLoopCarriedValues(getOperation()); }
+};
+} // namespace

cudaq/lib/Optimizer/Transforms/LowerUnwind.cpp

@@ -370,6 +370,12 @@ struct ScopeOpPattern : public OpRewritePattern<cudaq::cc::ScopeOp> {
     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<cudaq::cc::ScopeOp> {
         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<cudaq::cc::ScopeOp> {
         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 {

cudaq/lib/Optimizer/Transforms/MemToReg.cpp

@@ -826,10 +826,23 @@ class RegionDataFlow {
             getBinding(block, std::get<0>(info)) == std::get<1>(info))
           addBinding(block, std::get<0>(info), newReg);
         user->replaceUsesOfWith(std::get<1>(info), newReg);
+        // Other variables bound to this value follow it to the block
+        // argument, so `x = i` in a loop body gets this iteration's value.
+        updateBindingsOfValue(block, std::get<1>(info), newReg);
       }
     }
   }
 
+  /// Point any binding in block that still refers to oldVal at newVal.
+  void updateBindingsOfValue(Block *block, Value oldVal, Value newVal) {
+    auto iter = rMap.find(block);
+    if (iter == rMap.end())
+      return;
+    for (auto &binding : iter->second)
+      if (binding.second == oldVal)
+        binding.second = newVal;
+  }
+
   /// Track the memory reference \p mr as being live-out of the parent
   /// operation. (\p parent is passed for the assertion check only.)
   void addLiveOutOfParent(Operation *parent, MemRef mr) {

cudaq/lib/Optimizer/Transforms/Pipelines.cpp

@@ -111,6 +111,7 @@ static void createTargetPrepPipeline(OpPassManager &pm,
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createUnwindLowering());
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createClassicalMemToReg());
+  pm.addNestedPass<func::FuncOp>(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<func::FuncOp>(cudaq::opt::createClassicalMemToReg());
   pm.addNestedPass<func::FuncOp>(createCanonicalizerPass());
+  pm.addNestedPass<func::FuncOp>(cudaq::opt::createLoopPruneDeadArgs());
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createLoopNormalize());
   pm.addNestedPass<func::FuncOp>(cudaq::opt::createLoopInductionFusion());
   pm.addPass(cudaq::opt::createApplySpecialization());