Fixing the Python bridge's loop and scope emission

cudaq/include/cudaq/Optimizer/Dialect/CC/CCOps.td

@@ -484,7 +484,8 @@ def cc_ConditionOp : CCOp<"condition",
 }
 
 def cc_ContinueOp : CCOp<"continue", [Pure, ReturnLike, Terminator,
-        ParentOneOf<["LoopOp", "ScopeOp", "IfOp"]>]> {
+        ParentOneOf<["LoopOp", "ScopeOp", "IfOp"]>,
+        DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>]> {
   let summary = "Continue branch.";
   let description = [{
     A ContinueOp is a generalized exit terminator for the dialect operations
@@ -521,7 +522,8 @@ def cc_ContinueOp : CCOp<"continue", [Pure, ReturnLike, Terminator,
 }
 
 def cc_BreakOp : CCOp<"break",
-        [Pure, ReturnLike, Terminator, ParentOneOf<["LoopOp"]>]> {
+        [Pure, ReturnLike, Terminator, ParentOneOf<["LoopOp"]>,
+         DeclareOpInterfaceMethods<RegionBranchTerminatorOpInterface>]> {
   let summary = "Break branch.";
   let description = [{
     A BreakOp can be used in a LoopOp's body region (only). A BreakOp is a

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

@@ -1027,6 +1027,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 = [{

cudaq/lib/Optimizer/Dialect/CC/CCOps.cpp

@@ -1983,7 +1983,6 @@ SmallVector<Region *> 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");
@@ -2050,7 +2049,20 @@ struct HoistLoopInvariantArgs : public OpRewritePattern<cudaq::cc::LoopOp> {
         if (block.hasNoSuccessors())
           terminators.push_back(block.getTerminator());
 
-    // 2. Determine if any arguments are invariant.
+    // 2. Everything below is indexed by slot number. Bail if the lists
+    // disagree on how many slots there are, rather than index out of bounds.
+    const auto numSlots = loop.getInitialArgs().size();
+    for (auto *term : terminators) {
+      if (term->getBlock()->getParent()->front().getNumArguments() != numSlots)
+        return failure();
+      auto forwarded = isa<cudaq::cc::ConditionOp>(term)
+                           ? term->getNumOperands() - 1
+                           : term->getNumOperands();
+      if (forwarded != numSlots)
+        return failure();
+    }
+
+    // 3. Determine if any arguments are invariant.
     SmallVector<bool> invariants;
     bool hasInvariants = false;
     for (auto iter : llvm::enumerate(loop.getInitialArgs())) {
@@ -2076,19 +2088,19 @@ struct HoistLoopInvariantArgs : public OpRewritePattern<cudaq::cc::LoopOp> {
       invariants.push_back(isInvar);
     }
 
-    // 3. For each invariant argument replace the uses with the original
+    // 4. For each invariant argument replace the uses with the original
     // invariant value throughout.
     if (hasInvariants) {
       for (auto iter : llvm::enumerate(invariants)) {
         if (iter.value()) {
           auto i = iter.index();
           Value initialVal = loop.getInitialArgs()[i];
-          loop.getResult(i).replaceAllUsesWith(initialVal);
+          rewriter.replaceAllUsesWith(loop.getResult(i), initialVal);
           for (auto *reg : loop.getRegions()) {
             if (reg->empty())
               continue;
             auto &entry = reg->front();
-            entry.getArgument(i).replaceAllUsesWith(initialVal);
+            rewriter.replaceAllUsesWith(entry.getArgument(i), initialVal);
           }
         }
       }

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/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;