Support batching properly

Author: vimarsh6739

enzyme/Enzyme/MLIR/Implementations/CoreDialectsAutoDiffImplementations.cpp

@@ -358,23 +358,42 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
     for (auto &&[i, regionValue, operand] :
          llvm::enumerate(targetValues, operandRange)) {
 
-      // check if all the predecessorValues are const too
-      SmallVector<Value> possibleActivePreds;
-      regionBranchOp.getPredecessorValues(successor, i, possibleActivePreds);
+      // check if we need to create a shadow for an inactive region value
+      if (gutils->isConstantValue(regionValue)) {
+
+        // if all the possible predecessors for this value are also const, then
+        // we can skip creating a shadow. Else we need to create a shadow for
+        // syntactic correctness
+
+        SmallVector<Value> possibleActivePreds;
+        SmallVector<RegionBranchPoint> predecessors;
+        regionBranchOp.getPredecessors(successor, predecessors);
+        for (RegionBranchPoint predecessor : predecessors) {
+          if (predecessor.isParent()) {
+            // if the predecessor is the parent itself, then it's just
+            // operand!
+            possibleActivePreds.push_back(operand);
+            continue;
+          }
+          auto terminator = predecessor.getTerminatorPredecessorOrNull();
+          auto predecessorOperands = terminator.getSuccessorOperands(successor);
+          if (i < predecessorOperands.size())
+            possibleActivePreds.push_back(predecessorOperands[i]);
+        }
 
-      bool skipOpShadow = true;
-      for (auto pv : possibleActivePreds) {
+        bool skipOpShadow = true;
+        for (auto pv : possibleActivePreds) {
+          if (!skipOpShadow)
+            break;
+          skipOpShadow = skipOpShadow && gutils->isConstantValue(pv);
+        };
+        // if there's any possible active predecessor, we create a shadow for it
         if (!skipOpShadow)
-          break;
-        skipOpShadow = skipOpShadow && gutils->isConstantValue(pv);
-      };
-
-      if (!skipOpShadow)
-        constOperandPositionsToShadow.insert(
-            operandRange.getBeginOperandIndex() + i);
-
-      if (skipOpShadow && gutils->isConstantValue(regionValue))
-        continue;
+          constOperandPositionsToShadow.insert(
+              operandRange.getBeginOperandIndex() + i);
+        if (skipOpShadow)
+          continue;
+      }
 
       operandPositionsToShadow.insert(operandRange.getBeginOperandIndex() + i);
       if (successor.isParent())
@@ -448,14 +467,11 @@ 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
+  // Re-fix block args for all successor regions
+  // Even though createWithShadows properly creates the differentiated control
+  // flow op(accounting for any const args which might have shadows),
+  // takeBody(...) replaces the successor regions entirely, including the block
+  // arguments. We fix the block arguments here for the entry successor regions.
   if (!constOperandPositionToShadow.empty()) {
     auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op);
     auto iface = dyn_cast<ControlFlowAutoDiffOpInterface>(op);
@@ -466,41 +482,55 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
         entrySuccessors);
 
     for (const RegionSuccessor &successor : entrySuccessors) {
+
       if (successor.isParent())
         continue;
 
-      OperandRange operandRange =
+      OperandRange oldRegionOperands =
           iface.getSuccessorOperands(regionBranchOp, successor);
-      ValueRange targetValues = regionBranchOp.getSuccessorInputs(successor);
+      ValueRange oldRegionInputs = regionBranchOp.getSuccessorInputs(successor);
+
+      // the new region corresponding to this successor(we want to modify the
+      // arguments of this region in-place)
+      auto &newRegion =
+          replacement->getRegion(successor.getSuccessor()->getRegionNumber());
 
-      for (int i = targetValues.size() - 1; i >= 0; --i) {
-        unsigned operandPosition = operandRange.getBeginOperandIndex() + i;
+      for (int i = oldRegionInputs.size() - 1; i >= 0; --i) {
+        unsigned operandPosition = oldRegionOperands.getBeginOperandIndex() + i;
         if (!constOperandPositionToShadow.contains(operandPosition))
           continue;
 
-        auto regionValue = dyn_cast<BlockArgument>(targetValues[i]);
-        if (!regionValue || gutils->invertedPointers.contains(regionValue))
+        auto oldRegionInput = dyn_cast<BlockArgument>(oldRegionInputs[i]);
+
+        if (!oldRegionInput ||
+            gutils->invertedPointers.contains(oldRegionInput))
           continue;
 
-        auto replacementArg =
-            cast<BlockArgument>(gutils->getNewFromOriginal(regionValue));
-        Block *replacementBlock = replacementArg.getOwner();
+        auto newRegionInput =
+            cast<BlockArgument>(gutils->getNewFromOriginal(oldRegionInput));
+
+        auto typeIface =
+            dyn_cast<AutoDiffTypeInterface>(oldRegionInput.getType());
+
+        if (!typeIface) {
+          op->emitError() << " AutoDiffTypeInterface not implemented for "
+                          << oldRegionInput.getType() << "\n";
+          return failure();
+        }
 
-        Value shadowArg;
-        if (replacementArg.getArgNumber() ==
-            replacementBlock->getNumArguments() - 1) {
-          shadowArg = replacementBlock->addArgument(
-              gutils->getShadowType(regionValue.getType()),
-              regionValue.getLoc());
+        Value newRegionShadow;
+        if (newRegionInput.getArgNumber() == newRegion.getNumArguments() - 1) {
+          newRegionShadow = newRegion.addArgument(
+              typeIface.getShadowType(gutils->width), newRegionInput.getLoc());
         } else {
-          shadowArg = replacementBlock->insertArgument(
-              replacementBlock->args_begin() + replacementArg.getArgNumber() +
-                  1,
-              gutils->getShadowType(regionValue.getType()),
-              regionValue.getLoc());
+          // insert at position i+1
+          newRegionShadow = newRegion.insertArgument(
+              newRegion.args_begin() + newRegionInput.getArgNumber() + 1,
+              typeIface.getShadowType(gutils->width), newRegionInput.getLoc());
         }
 
-        gutils->invertedPointers.map(regionValue, shadowArg);
+        // update the inverted pointer map
+        gutils->invertedPointers.map(oldRegionInput, newRegionShadow);
       }
     }
   }