MLIR: fix result positions to shadow (#2112)

EnzymeAD · Oct 8, 2024 · 152441a · 152441a
1 parent 2f92df3
commit 152441a
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.cpp b/enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.cpp
@@ -381,17 +381,20 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
   // table.
   SmallVector<Value> reps;
   size_t idx = 0;
-  for (Value r : op->getResults()) {
+  for (OpResult r : op->getResults()) {
     // TODO only if used
     reps.push_back(replacement->getResult(idx));
     idx++;
     if (!gutils->isConstantValue(r)) {
+      assert(resultPositionsToShadow.count(r.getResultNumber()));
       auto inverted = gutils->invertedPointers.lookupOrNull(r);
       assert(inverted);
       gutils->invertedPointers.map(r, replacement->getResult(idx));
       inverted.replaceAllUsesWith(replacement->getResult(idx));
       gutils->erase(inverted.getDefiningOp());
       idx++;
+    } else if (resultPositionsToShadow.count(r.getResultNumber())) {
+      idx++;
     }
   }