@@ -318,15 +318,14 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
318318 // add the shadow as operand.
319319 auto regionBranchOp = dyn_cast<RegionBranchOpInterface>(op);
320320 if (!regionBranchOp) {
321- op->emitError () << " RegionBranchOpInterface not implemented for "
322- << " \n "
323- return failure ();
321+ return op->emitError () << " RegionBranchOpInterface not implemented for "
322+ << *op << " \n "
324323 }
325324 auto iface = dyn_cast<ControlFlowAutoDiffOpInterface>(op);
326325 if (!iface) {
327- op->emitError () << " ControlFlowAutoDiffOpInterface not implemented for "
328- << *op << " \n " ;
329- return failure () ;
326+ return op->emitError ()
327+ << " ControlFlowAutoDiffOpInterface not implemented for " << *op
328+ << " \n "
330329 }
331330
332331 // TODO: we may need to record, for every successor, which of its inputs
@@ -352,8 +351,37 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
352351 // operands.
353352 for (auto &&[i, regionValue, operand] :
354353 llvm::enumerate (targetValues, operandRange)) {
355- if (gutils->isConstantValue (regionValue))
356- continue ;
354+
355+ // if all the possible predecessors for this value are also const, then
356+ // we can skip creating a shadow. Else we need to create a shadow for
357+ // activity correctness
358+ if (gutils->isConstantValue (regionValue)) {
359+ SmallVector<Value> possibleActivePreds;
360+ SmallVector<RegionBranchPoint> predecessors;
361+ regionBranchOp.getPredecessors (successor, predecessors);
362+ for (RegionBranchPoint predecessor : predecessors) {
363+ if (predecessor.isParent ()) {
364+ // if the predecessor is the parent itself, then it's just
365+ // `operand`
366+ possibleActivePreds.push_back (operand);
367+ continue ;
368+ }
369+ auto terminator = predecessor.getTerminatorPredecessorOrNull ();
370+ auto predecessorOperands = terminator.getSuccessorOperands (successor);
371+ if (i < predecessorOperands.size ())
372+ possibleActivePreds.push_back (predecessorOperands[i]);
373+ }
374+
375+ bool skipOpShadow = true ;
376+ for (auto pv : possibleActivePreds) {
377+ if (!skipOpShadow)
378+ break ;
379+ skipOpShadow = skipOpShadow && gutils->isConstantValue (pv);
380+ };
381+ if (skipOpShadow)
382+ continue ;
383+ // if there's any possible active predecessor, we create a shadow for it
384+ }
357385 operandPositionsToShadow.insert (operandRange.getBeginOperandIndex () + i);
358386 if (successor.isParent ())
359387 resultPositionsToShadow.insert (i);
@@ -388,33 +416,87 @@ LogicalResult mlir::enzyme::detail::controlFlowForwardHandler(
388416 continue ;
389417 auto typeIface = dyn_cast<AutoDiffTypeInterface>(result.
getType ());
390418 if (!typeIface) {
391- op->emitError () << " AutoDiffTypeInterface not implemented for "
392- << result.getType () << " \n "
393- return failure ();
419+ return op->emitError () << " AutoDiffTypeInterface not implemented for "
420+ << result.getType () << " \n "
394421 }
395422 newOpResultTypes.push_back (typeIface.getShadowType (gutils->width ));
396423 }
397424
398425 SmallVector<Value> newOperands;
399426 newOperands.reserve (op->getNumOperands () + operandPositionsToShadow.size ());
427+
428+ auto iface = dyn_cast<ControlFlowAutoDiffOpInterface>(op);
429+ if (!iface) {
430+ return op->emitError ()
431+ << " ControlFlowAutoDiffOpInterface not implemented for "
432+ << " \n "
433+ }
434+
435+ auto regionBranchOp = cast<RegionBranchOpInterface>(op);
436+ SmallVector<RegionSuccessor> entrySuccessors;
437+ regionBranchOp.getEntrySuccessorRegions (
438+ SmallVector<Attribute>(op->getNumOperands (), Attribute ()),
439+ entrySuccessors);
400440 for (OpOperand &operand : op->getOpOperands ()) {
401441 newOperands.push_back (gutils->getNewFromOriginal (operand.get ()));
402- if (operandPositionsToShadow.contains (operand.getOperandNumber ()))
403- newOperands.push_back (gutils->invertPointerM (operand.get (), builder));
442+ if (operandPositionsToShadow.contains (operand.getOperandNumber ())) {
443+ Value shadowValue = nullptr ;
444+ if (!gutils->isConstantValue (operand.get ()))
445+ shadowValue = gutils->invertPointerM (operand.get (), builder);
446+ else {
447+ auto Ty = operand.get ().getType ();
448+ auto shadowType =
449+ cast<AutoDiffTypeInterface>(Ty).getShadowType (gutils->width );
450+ shadowValue = cast<AutoDiffTypeInterface>(shadowType)
451+ .createNullValue (builder, operand.get ().getLoc ());
452+
453+ // modify block arguments for entry successors to newOp, since
454+ // forceAugmentedReturns will not shadow const operands. No need to add
455+ // to the invertPointers map since `operand` is const (the shadow will
456+ // be unused)
457+ for (const RegionSuccessor &successor : entrySuccessors) {
458+ if (successor.isParent ())
459+ continue ;
460+ auto &newOpRegion =
461+ newOp->getRegion (successor.getSuccessor ()->getRegionNumber ());
462+ OperandRange succOperands =
463+ iface.getSuccessorOperands (regionBranchOp, successor);
464+ ValueRange succInputs = regionBranchOp.getSuccessorInputs (successor);
465+
466+ if (succOperands.empty ())
467+ continue ;
468+
469+ auto succInputPos =
470+ operand.getOperandNumber () - succOperands.getBeginOperandIndex ();
471+
472+ if (succInputPos >= 0 && succInputPos < succInputs.size ()) {
473+ auto oldRegionInput =
474+ dyn_cast<BlockArgument>(succInputs[succInputPos]);
475+ if (!oldRegionInput)
476+ continue ;
477+ if (gutils->invertedPointers .contains (oldRegionInput))
478+ continue ;
479+ auto newOpBlockVal =
480+ cast<BlockArgument>(gutils->getNewFromOriginal (oldRegionInput));
481+ auto i = newOpBlockVal.getArgNumber ();
482+ if (i == newOpRegion.getNumArguments () - 1 ) {
483+ newOpRegion.addArgument (shadowType, newOpBlockVal.getLoc ());
484+ } else {
485+ newOpRegion.insertArgument (newOpRegion.args_begin () + i + 1 ,
486+ shadowType, newOpBlockVal.getLoc ());
487+ }
488+ }
489+ }
490+ }
491+ newOperands.push_back (shadowValue);
492+ }
404493 }
405494 // We are assuming the op can forward additional operands, listed
406495 // immediately after the original operands, to the same regions.
407496 // ^^
408497 // Our interface guarantees this.
409498 // We also assume that the region-holding op returns all of the values
410499 // yielded by terminators, and only those values.
411-
412- auto iface = dyn_cast<ControlFlowAutoDiffOpInterface>(op);
413- if (!iface) {
414- op->emitError () << " ControlFlowAutoDiffOpInterface not implemented for "
415- << *op << " \n "
416- return failure ();
417- }
418500 Operation *replacement = iface.createWithShadows (
419501 builder, gutils, op, newOperands, newOpResultTypes);
420502 assert (replacement->getNumResults () == newOpResultTypes.size ());
0 commit comments