Fix forward mode AD for for-like ops

Author: vimarsh6739

enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.cpp

@@ -334,6 +334,11 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
   llvm::SmallDenseSet<unsigned> operandPositionsToShadow;
   llvm::SmallDenseSet<unsigned> resultPositionsToShadow;
 
+  // while these operands are inactive in the op region(s), we may still need to
+  // create placeholder shadows for them to ensure syntactic correctness for the
+  // IR
+  llvm::SmallDenseSet<unsigned> constOperandPositionsToShadow;
+
   SmallVector<RegionSuccessor> entrySuccessors;
   regionBranchOp.getEntrySuccessorRegions(
       SmallVector<Attribute>(op->getNumOperands(), Attribute()),
@@ -352,8 +357,25 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
     // operands.
     for (auto &&[i, regionValue, operand] :
          llvm::enumerate(targetValues, operandRange)) {
-      if (gutils->isConstantValue(regionValue))
+
+      // check if all the predecessorValues are const too
+      SmallVector<Value> possibleActivePreds;
+      regionBranchOp.getPredecessorValues(successor, i, possibleActivePreds);
+
+      bool skipOpShadow = true;
+      for (auto pv : possibleActivePreds) {
+        if (!skipOpShadow)
+          break;
+        skipOpShadow = skipOpShadow && gutils->isConstantValue(pv);
+      };
+
+      if (!skipOpShadow)
+        constOperandPositionsToShadow.insert(
+            operandRange.getBeginOperandIndex() + i);
+
+      if (skipOpShadow && gutils->isConstantValue(regionValue))
         continue;
+
       operandPositionsToShadow.insert(operandRange.getBeginOperandIndex() + i);
       if (successor.isParent())
         resultPositionsToShadow.insert(i);
@@ -365,13 +387,16 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
       resultPositionsToShadow.insert(res.getResultNumber());
 
   return controlFlowForwardHandler(
-      op, builder, gutils, operandPositionsToShadow, resultPositionsToShadow);
+      op, builder, gutils, operandPositionsToShadow, resultPositionsToShadow,
+      constOperandPositionsToShadow);
 }
 
 LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
     Operation *op, OpBuilder &builder, MGradientUtils *gutils,
     const llvm::SmallDenseSet<unsigned> &operandPositionsToShadow,
-    const llvm::SmallDenseSet<unsigned> &resultPositionsToShadow) {
+    const llvm::SmallDenseSet<unsigned> &resultPositionsToShadow,
+    const llvm::SmallDenseSet<unsigned> &constOperandPositionToShadow) {
+
   // For all active results, add shadow types.
   // For now, assuming all results are relevant.
   Operation *newOp = gutils->getNewFromOriginal(op);
@@ -423,6 +448,63 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
     replacementRegion.takeBody(region);
   }
 
+  // Re-fix block args for entry -> blk-successors
+  // if constOperandPositionsToShadow is non-empty, the takeBody(...) from
+  // earlier replaces the body for replacement(which has well formed
+  // successor-args) with newOp's successor regions
+  //
+  // We fix it by modifying the blockarguments for all the entry successors,
+  // adding a newblockarg for every entry in constOperandPositionsToShadow and
+  // updating the invertedPointerMap
+  if (!constOperandPositionToShadow.empty()) {
+    auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op);
+    auto iface = dyn_cast<ControlFlowAutoDiffOpInterface>(op);
+
+    SmallVector<RegionSuccessor> entrySuccessors;
+    regionBranchOp.getEntrySuccessorRegions(
+        SmallVector<Attribute>(op->getNumOperands(), Attribute()),
+        entrySuccessors);
+
+    for (const RegionSuccessor &successor : entrySuccessors) {
+      if (successor.isParent())
+        continue;
+
+      OperandRange operandRange =
+          iface.getSuccessorOperands(regionBranchOp, successor);
+      ValueRange targetValues = regionBranchOp.getSuccessorInputs(successor);
+
+      for (int i = targetValues.size() - 1; i >= 0; --i) {
+        unsigned operandPosition = operandRange.getBeginOperandIndex() + i;
+        if (!constOperandPositionToShadow.contains(operandPosition))
+          continue;
+
+        auto regionValue = dyn_cast<BlockArgument>(targetValues[i]);
+        if (!regionValue || gutils->invertedPointers.contains(regionValue))
+          continue;
+
+        auto replacementArg =
+            cast<BlockArgument>(gutils->getNewFromOriginal(regionValue));
+        Block *replacementBlock = replacementArg.getOwner();
+
+        Value shadowArg;
+        if (replacementArg.getArgNumber() ==
+            replacementBlock->getNumArguments() - 1) {
+          shadowArg = replacementBlock->addArgument(
+              gutils->getShadowType(regionValue.getType()),
+              regionValue.getLoc());
+        } else {
+          shadowArg = replacementBlock->insertArgument(
+              replacementBlock->args_begin() + replacementArg.getArgNumber() +
+                  1,
+              gutils->getShadowType(regionValue.getType()),
+              regionValue.getLoc());
+        }
+
+        gutils->invertedPointers.map(regionValue, shadowArg);
+      }
+    }
+  }
+
   // Inject the mapping for the new results into GradientUtil's shadow
   // table.
   SmallVector<Value> reps;

enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.h

@@ -40,7 +40,8 @@ LogicalResult controlFlowForwardHandler(Operation *op, OpBuilder &builder,
 LogicalResult controlFlowForwardHandler(
     Operation *op, OpBuilder &builder, MGradientUtils *gutils,
     const llvm::SmallDenseSet<unsigned> &operandPositionsToShadow,
-    const llvm::SmallDenseSet<unsigned> &resultPositionsToShadow);
+    const llvm::SmallDenseSet<unsigned> &resultPositionsToShadow,
+    const llvm::SmallDenseSet<unsigned> &constOperandPositionToShadow);
 
 // Implements forward-mode differentiation of branching operations.
 // Assumes that successive shadows are legal

enzyme/test/MLIR/ForwardMode/for3.mlir

@@ -0,0 +1,51 @@
+// RUN: %eopt --enzyme %s | FileCheck %s
+
+module {
+  func.func @carry_mismatch_scf(%x : f64) -> f64 {
+    %zero = arith.constant 0.0 : f64
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c10 = arith.constant 10 : index
+    %r = scf.for %i = %c0 to %c10 step %c1 iter_args(%acc = %zero) -> (f64) {
+      scf.yield %x : f64
+    }
+    return %r : f64
+  }
+
+  func.func @dcarry_mismatch_scf(%x : f64, %dx : f64) -> f64 {
+    %r = enzyme.fwddiff @carry_mismatch_scf(%x, %dx) { activity=[#enzyme<activity enzyme_dup>], ret_activity=[#enzyme<activity enzyme_dupnoneed>] } : (f64, f64) -> (f64)
+    return %r : f64
+  }
+
+  func.func @carry_mismatch_affine(%x : f64) -> f64 {
+    %zero = arith.constant 0.0 : f64
+    %r = affine.for %i = 0 to 10 iter_args(%acc = %zero) -> (f64) {
+      affine.yield %x : f64
+    }
+    return %r : f64
+  }
+
+  func.func @dcarry_mismatch_affine(%x : f64, %dx : f64) -> f64 {
+    %r = enzyme.fwddiff @carry_mismatch_affine(%x, %dx) { activity=[#enzyme<activity enzyme_dup>], ret_activity=[#enzyme<activity enzyme_dupnoneed>] } : (f64, f64) -> (f64)
+    return %r : f64
+  }
+}
+
+// CHECK-LABEL: func.func private @fwddiffecarry_mismatch_scf(
+// CHECK-DAG: %[[ZERO0:.+]] = arith.constant 0.000000e+00 : f64
+// CHECK-DAG: %[[ZERO1:.+]] = arith.constant 0.000000e+00 : f64
+// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index
+// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index
+// CHECK-DAG: %[[C10:.+]] = arith.constant 10 : index
+// CHECK: %[[LOOP:.+]]:2 = scf.for %[[IV:.+]] = %[[C0]] to %[[C10]] step %[[C1]] iter_args(%[[DACC:.+]] = %[[ZERO1]], %[[ACC:.+]] = %[[ZERO0]]) -> (f64, f64) {
+// CHECK-NEXT:   scf.yield %[[ARG0:.+]], %[[ARG1:.+]] : f64, f64
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[LOOP]]#1 : f64
+
+// CHECK-LABEL: func.func private @fwddiffecarry_mismatch_affine(
+// CHECK-DAG: %[[AZERO0:.+]] = arith.constant 0.000000e+00 : f64
+// CHECK-DAG: %[[AZERO1:.+]] = arith.constant 0.000000e+00 : f64
+// CHECK: %[[ALOOP:.+]]:2 = affine.for %[[AIV:.+]] = 0 to 10 iter_args(%[[ADACC:.+]] = %[[AZERO1]], %[[AACC:.+]] = %[[AZERO0]]) -> (f64, f64) {
+// CHECK-NEXT:   affine.yield %[[ARG0:.+]], %[[ARG1:.+]] : f64, f64
+// CHECK-NEXT: }
+// CHECK-NEXT: return %[[ALOOP]]#1 : f64