[AutoDiff] Correctly propagate optional adjoint through switch_enum (#74985)

Fixes #74978

Author: asl

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -765,14 +765,15 @@ class PullbackCloner::Implementation final
                                         SILValue wrappedAdjoint,
                                         SILType optionalTy);
 
-  /// Accumulate optional buffer from `wrappedAdjoint`.
+  /// Accumulate adjoint of `wrappedAdjoint` into optionalBuffer.
   void accumulateAdjointForOptionalBuffer(SILBasicBlock *bb,
                                           SILValue optionalBuffer,
                                           SILValue wrappedAdjoint);
 
-  /// Set optional value from `wrappedAdjoint`.
-  void setAdjointValueForOptional(SILBasicBlock *bb, SILValue optionalValue,
-                                  SILValue wrappedAdjoint);
+  /// Accumulate adjoint of `wrappedAdjoint` into optionalValue.
+  void accumulateAdjointValueForOptional(SILBasicBlock *bb,
+                                         SILValue optionalValue,
+                                         SILValue wrappedAdjoint);
 
   //--------------------------------------------------------------------------//
   // Array literal initialization differentiation
@@ -2732,8 +2733,8 @@ void PullbackCloner::Implementation::accumulateAdjointForOptionalBuffer(
   builder.createDeallocStack(pbLoc, optTanAdjBuf);
 }
 
-// Set the adjoint value for the incoming `Optional` value.
-void PullbackCloner::Implementation::setAdjointValueForOptional(
+// Accumulate adjoint for the incoming `Optional` value.
+void PullbackCloner::Implementation::accumulateAdjointValueForOptional(
     SILBasicBlock *bb, SILValue optionalValue, SILValue wrappedAdjoint) {
   assert(getTangentValueCategory(optionalValue) == SILValueCategory::Object);
   auto pbLoc = getPullback().getLocation();
@@ -2745,10 +2746,11 @@ void PullbackCloner::Implementation::setAdjointValueForOptional(
 
   auto optTanAdjVal = builder.emitLoadValueOperation(
       pbLoc, optTanAdjBuf, LoadOwnershipQualifier::Take);
+
   recordTemporary(optTanAdjVal);
   builder.createDeallocStack(pbLoc, optTanAdjBuf);
 
-  setAdjointValue(bb, optionalValue, makeConcreteAdjointValue(optTanAdjVal));
+  addAdjointValue(bb, optionalValue, makeConcreteAdjointValue(optTanAdjVal), pbLoc);
 }
 
 SILBasicBlock *PullbackCloner::Implementation::buildPullbackSuccessor(
@@ -2959,12 +2961,12 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
           // Handle `switch_enum` on `Optional`.
           auto termInst = bbArg->getSingleTerminator();
           if (isSwitchEnumInstOnOptional(termInst)) {
-            setAdjointValueForOptional(bb, incomingValue, concreteBBArgAdjCopy);
+            accumulateAdjointValueForOptional(bb, incomingValue, concreteBBArgAdjCopy);
           } else {
             blockTemporaries[getPullbackBlock(predBB)].insert(
               concreteBBArgAdjCopy);
-            setAdjointValue(predBB, incomingValue,
-                            makeConcreteAdjointValue(concreteBBArgAdjCopy));
+            addAdjointValue(predBB, incomingValue,
+                            makeConcreteAdjointValue(concreteBBArgAdjCopy), pbLoc);
           }
         }
         break;

test/AutoDiff/validation-test/optional.swift

@@ -22,12 +22,30 @@ func optional_nil_coalescing(_ maybeX: Float?) -> Float {
 */
 
 OptionalTests.test("Active") {
+  @differentiable(reverse)
+  func id(y: Float) -> Float? {
+    return y
+  }
+
+  @differentiable(reverse)
+  func id2(y: Float?) -> Float {
+    return y!
+  }
+
   @differentiable(reverse)
   func square(y: Float) -> Float? {
     return y * y
   }
 
+  @differentiable(reverse)
+  func square2(y: Float?) -> Float {
+    return y! * y!
+  }
+
+  expectEqual(gradient(at: 10, of: {y in id(y:y)!}), .init(1.0))
+  expectEqual(gradient(at: 10, of: {y in id2(y:y)}), .init(1.0))
   expectEqual(gradient(at: 10, of: {y in square(y:y)!}), .init(20.0))
+  expectEqual(gradient(at: 10, of: {y in square2(y:y)}), .init(20.0))
 }
 
 OptionalTests.test("Let") {