[AutoDiff] Handle init_enum_data_addr and inject_enum_addr for Optional (swiftlang#68300)

Optional's `init_enum_data_addr` and `inject_enum_addr` instructions are generated in presence of non-loadable Optional values. The compiler used to treat these instructions as inactive, and this resulted in silent run-time
issues described in swiftlang#64223.

The patch marks `init_enum_data_addr` as "active" if its Optional operand is also active, and in PullbackCloner we differentiate through it and the related `inject_enum_addr`.

However, we only determine this relation in simple cases when both instructions are in the same block. There is no def-use relation between them (both take the same Optional operand), so if there is more than one set of instructions
operating on the same Optional, or there is some control flow, we currently bail out.

In PullbackCloner, we walk over instructions in reverse order and start from `inject_enum_addr` and its `Optional<Wrapped>.TangentVector` operand. Assuming that is is already initialized, we emit an `unchecked_take_enum_data_addr` and set it as the adjoint buffer of `init_enum_data_addr`. The Optional value is
invalidated, and we have to destroy the enum data address later when we reach `init_enum_data_addr`.

Author: asavonic

lib/SILOptimizer/Analysis/DifferentiableActivityAnalysis.cpp

@@ -408,7 +408,7 @@ void DifferentiableActivityInfo::propagateUsefulThroughAddress(
       SKIP_NODERIVATIVE(RefElementAddr)
 #undef SKIP_NODERIVATIVE
       if (Projection::isAddressProjection(res) || isa<BeginAccessInst>(res) ||
-          isa<BeginBorrowInst>(res))
+          isa<BeginBorrowInst>(res) || isa<InitEnumDataAddrInst>(res))
         propagateUsefulThroughAddress(res, dependentVariableIndex);
     }
   }

lib/SILOptimizer/Differentiation/LinearMapInfo.cpp

@@ -453,6 +453,19 @@ bool LinearMapInfo::shouldDifferentiateApplySite(FullApplySite applySite) {
   return hasActiveResults && hasActiveArguments;
 }
 
+static bool shouldDifferentiateInjectEnumAddr(
+    const InjectEnumAddrInst &inject,
+    const DifferentiableActivityInfo &activityInfo,
+    const AutoDiffConfig &config) {
+  SILValue en = inject.getOperand();
+  for (auto use : en->getUses()) {
+    auto *init = dyn_cast<InitEnumDataAddrInst>(use->getUser());
+    if (init && activityInfo.isActive(init, config))
+      return true;
+  }
+  return false;
+}
+
 /// Returns a flag indicating whether the instruction should be differentiated,
 /// given the differentiation indices of the instruction's parent function.
 /// Whether the instruction should be differentiated is determined sequentially
@@ -506,6 +519,13 @@ bool LinearMapInfo::shouldDifferentiateInstruction(SILInstruction *inst) {
         isa<DestroyValueInst>(inst) || isa<DestroyAddrInst>(inst))
       return true;
   }
+
+  // Should differentiate `inject_enum_addr` if the corresponding
+  // `init_enum_addr` has an active operand.
+  if (auto inject = dyn_cast<InjectEnumAddrInst>(inst))
+    if (shouldDifferentiateInjectEnumAddr(*inject, activityInfo, config))
+      return true;
+
   return false;
 }
 

lib/SILOptimizer/Differentiation/PullbackCloner.cpp

@@ -744,14 +744,23 @@ class PullbackCloner::Implementation final
   // Optional differentiation
   //--------------------------------------------------------------------------//
 
-  /// Given a `wrappedAdjoint` value of type `T.TangentVector`, creates an
-  /// `Optional<T>.TangentVector` value from it and adds it to the adjoint value
-  /// of `optionalValue`.
+  /// Given a `wrappedAdjoint` value of type `T.TangentVector` and `Optional<T>`
+  /// type, creates an `Optional<T>.TangentVector` buffer from it.
   ///
   /// `wrappedAdjoint` may be an object or address value, both cases are
   /// handled.
-  void accumulateAdjointForOptional(SILBasicBlock *bb, SILValue optionalValue,
-                                    SILValue wrappedAdjoint);
+  AllocStackInst *createOptionalAdjoint(SILBasicBlock *bb,
+                                        SILValue wrappedAdjoint,
+                                        SILType optionalTy);
+
+  /// Accumulate optional buffer from `wrappedAdjoint`.
+  void accumulateAdjointForOptionalBuffer(SILBasicBlock *bb,
+                                          SILValue optionalBuffer,
+                                          SILValue wrappedAdjoint);
+
+  /// Set optional value from `wrappedAdjoint`.
+  void setAdjointValueForOptional(SILBasicBlock *bb, SILValue optionalValue,
+                                  SILValue wrappedAdjoint);
 
   //--------------------------------------------------------------------------//
   // Array literal initialization differentiation
@@ -1687,6 +1696,104 @@ class PullbackCloner::Implementation final
     builder.emitZeroIntoBuffer(uccai->getLoc(), adjDest, IsInitialization);
   }
 
+  /// Handle a sequence of `init_enum_data_addr` and `inject_enum_addr`
+  /// instructions.
+  ///
+  /// Original: y = init_enum_data_addr x
+  ///           inject_enum_addr y
+  ///
+  ///  Adjoint: adj[x] += unchecked_take_enum_data_addr adj[y]
+  void visitInjectEnumAddrInst(InjectEnumAddrInst *inject) {
+    SILBasicBlock *bb = inject->getParent();
+    SILValue origEnum = inject->getOperand();
+
+    // Only `Optional`-typed operands are supported for now. Diagnose all other
+    // enum operand types.
+    auto *optionalEnumDecl = getASTContext().getOptionalDecl();
+    if (origEnum->getType().getEnumOrBoundGenericEnum() != optionalEnumDecl) {
+      LLVM_DEBUG(getADDebugStream()
+                 << "Unsupported enum type in PullbackCloner: " << *inject);
+      getContext().emitNondifferentiabilityError(
+          inject, getInvoker(),
+          diag::autodiff_expression_not_differentiable_note);
+      errorOccurred = true;
+      return;
+    }
+
+    InitEnumDataAddrInst *origData = nullptr;
+    for (auto use : origEnum->getUses()) {
+      if (auto *init = dyn_cast<InitEnumDataAddrInst>(use->getUser())) {
+        // We need a more complicated analysis when init_enum_data_addr and
+        // inject_enum_addr are in different blocks, or there is more than one
+        // such instruction. Bail out for now.
+        if (origData || init->getParent() != bb) {
+          LLVM_DEBUG(getADDebugStream()
+                     << "Could not find a matching init_enum_data_addr for: "
+                     << *inject);
+          getContext().emitNondifferentiabilityError(
+              inject, getInvoker(),
+              diag::autodiff_expression_not_differentiable_note);
+          errorOccurred = true;
+          return;
+        }
+
+        origData = init;
+      }
+    }
+
+    SILValue adjStruct = getAdjointBuffer(bb, origEnum);
+    StructDecl *adjStructDecl =
+        adjStruct->getType().getStructOrBoundGenericStruct();
+
+    VarDecl *adjOptVar = nullptr;
+    if (adjStructDecl) {
+      ArrayRef<VarDecl *> properties = adjStructDecl->getStoredProperties();
+      adjOptVar = properties.size() == 1 ? properties[0] : nullptr;
+    }
+
+    EnumDecl *adjOptDecl =
+        adjOptVar ? adjOptVar->getTypeInContext()->getEnumOrBoundGenericEnum()
+                  : nullptr;
+
+    // Optional<T>.TangentVector should be a struct with a single
+    // Optional<T.TangentVector> property. This is an implementation detail of
+    // OptionalDifferentiation.swift
+    if (!adjOptDecl || adjOptDecl != optionalEnumDecl)
+      llvm_unreachable("Unexpected type of Optional.TangentVector");
+
+    SILLocation loc = origData->getLoc();
+    StructElementAddrInst *adjOpt =
+        builder.createStructElementAddr(loc, adjStruct, adjOptVar);
+
+    // unchecked_take_enum_data_addr is destructive, so copy
+    // Optional<T.TangentVector> to a new alloca.
+    AllocStackInst *adjOptCopy =
+        createFunctionLocalAllocation(adjOpt->getType(), loc);
+    builder.createCopyAddr(loc, adjOpt, adjOptCopy, IsNotTake,
+                           IsInitialization);
+
+    EnumElementDecl *someElemDecl = getASTContext().getOptionalSomeDecl();
+    UncheckedTakeEnumDataAddrInst *adjData =
+        builder.createUncheckedTakeEnumDataAddr(loc, adjOptCopy, someElemDecl);
+
+    setAdjointBuffer(bb, origData, adjData);
+
+    // The Optional copy is invalidated, do not attempt to destroy it at the end
+    // of the pullback. The value returned from unchecked_take_enum_data_addr is
+    // destroyed in visitInitEnumDataAddrInst.
+    destroyedLocalAllocations.insert(adjOptCopy);
+  }
+
+  /// Handle `init_enum_data_addr` instruction.
+  /// Destroy the value returned from `unchecked_take_enum_data_addr`.
+  void visitInitEnumDataAddrInst(InitEnumDataAddrInst *init) {
+    auto bufIt = bufferMap.find({init->getParent(), SILValue(init)});
+    if (bufIt == bufferMap.end())
+      return;
+    SILValue adjData = bufIt->second;
+    builder.emitDestroyAddr(init->getLoc(), adjData);
+  }
+
   /// Handle `unchecked_ref_cast` instruction.
   ///   Original: y = unchecked_ref_cast x
   ///    Adjoint: adj[x] += adj[y]
@@ -1758,7 +1865,7 @@ class PullbackCloner::Implementation final
       errorOccurred = true;
       return;
     }
-    accumulateAdjointForOptional(bb, utedai->getOperand(), adjDest);
+    accumulateAdjointForOptionalBuffer(bb, utedai->getOperand(), adjDest);
     builder.emitZeroIntoBuffer(utedai->getLoc(), adjDest, IsNotInitialization);
   }
 
@@ -2342,12 +2449,11 @@ void PullbackCloner::Implementation::emitZeroDerivativesForNonvariedResult(
              << pullback);
 }
 
-void PullbackCloner::Implementation::accumulateAdjointForOptional(
-    SILBasicBlock *bb, SILValue optionalValue, SILValue wrappedAdjoint) {
+AllocStackInst *PullbackCloner::Implementation::createOptionalAdjoint(
+    SILBasicBlock *bb, SILValue wrappedAdjoint, SILType optionalTy) {
   auto pbLoc = getPullback().getLocation();
-  // Handle `switch_enum` on `Optional`.
   // `Optional<T>`
-  auto optionalTy = remapType(optionalValue->getType());
+  optionalTy = remapType(optionalTy);
   assert(optionalTy.getASTType()->isOptional());
   // `T`
   auto wrappedType = optionalTy.getOptionalObjectType();
@@ -2429,13 +2535,45 @@ void PullbackCloner::Implementation::accumulateAdjointForOptional(
   builder.createApply(pbLoc, initFnRef, subMap,
                       {optTanAdjBuf, optArgBuf, metatype});
   builder.createDeallocStack(pbLoc, optArgBuf);
+  return optTanAdjBuf;
+}
+
+// Accumulate adjoint for the incoming `Optional` buffer.
+void PullbackCloner::Implementation::accumulateAdjointForOptionalBuffer(
+    SILBasicBlock *bb, SILValue optionalBuffer, SILValue wrappedAdjoint) {
+  assert(getTangentValueCategory(optionalBuffer) == SILValueCategory::Address);
+  auto pbLoc = getPullback().getLocation();
 
-  // Accumulate adjoint for the incoming `Optional` value.
-  addToAdjointBuffer(bb, optionalValue, optTanAdjBuf, pbLoc);
+  // Allocate and initialize Optional<Wrapped>.TangentVector from
+  // Wrapped.TangentVector
+  AllocStackInst *optTanAdjBuf =
+      createOptionalAdjoint(bb, wrappedAdjoint, optionalBuffer->getType());
+
+  // Accumulate into optionalBuffer
+  addToAdjointBuffer(bb, optionalBuffer, optTanAdjBuf, pbLoc);
   builder.emitDestroyAddr(pbLoc, optTanAdjBuf);
   builder.createDeallocStack(pbLoc, optTanAdjBuf);
 }
 
+// Set the adjoint value for the incoming `Optional` value.
+void PullbackCloner::Implementation::setAdjointValueForOptional(
+    SILBasicBlock *bb, SILValue optionalValue, SILValue wrappedAdjoint) {
+  assert(getTangentValueCategory(optionalValue) == SILValueCategory::Object);
+  auto pbLoc = getPullback().getLocation();
+
+  // Allocate and initialize Optional<Wrapped>.TangentVector from
+  // Wrapped.TangentVector
+  AllocStackInst *optTanAdjBuf =
+      createOptionalAdjoint(bb, wrappedAdjoint, optionalValue->getType());
+
+  auto optTanAdjVal = builder.emitLoadValueOperation(
+      pbLoc, optTanAdjBuf, LoadOwnershipQualifier::Take);
+  recordTemporary(optTanAdjVal);
+  builder.createDeallocStack(pbLoc, optTanAdjBuf);
+
+  setAdjointValue(bb, optionalValue, makeConcreteAdjointValue(optTanAdjVal));
+}
+
 SILBasicBlock *PullbackCloner::Implementation::buildPullbackSuccessor(
     SILBasicBlock *origBB, SILBasicBlock *origPredBB,
     SmallDenseMap<SILValue, TrampolineBlockSet> &pullbackTrampolineBlockMap) {
@@ -2623,7 +2761,7 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
           // Handle `switch_enum` on `Optional`.
           auto termInst = bbArg->getSingleTerminator();
           if (isSwitchEnumInstOnOptional(termInst)) {
-            accumulateAdjointForOptional(bb, incomingValue, concreteBBArgAdjCopy);
+            setAdjointValueForOptional(bb, incomingValue, concreteBBArgAdjCopy);
           } else {
             blockTemporaries[getPullbackBlock(predBB)].insert(
               concreteBBArgAdjCopy);
@@ -2643,7 +2781,7 @@ void PullbackCloner::Implementation::visitSILBasicBlock(SILBasicBlock *bb) {
           // Handle `switch_enum` on `Optional`.
           auto termInst = bbArg->getSingleTerminator();
           if (isSwitchEnumInstOnOptional(termInst))
-            accumulateAdjointForOptional(bb, incomingValue, bbArgAdjBuf);
+            accumulateAdjointForOptionalBuffer(bb, incomingValue, bbArgAdjBuf);
           else
             addToAdjointBuffer(bb, incomingValue, bbArgAdjBuf, pbLoc);
         }

stdlib/public/Differentiation/OptionalDifferentiation.swift

@@ -13,6 +13,7 @@
 import Swift
 
 extension Optional: Differentiable where Wrapped: Differentiable {
+  @frozen
   public struct TangentVector: Differentiable, AdditiveArithmetic {
     public typealias TangentVector = Self
 

test/AutoDiff/SILOptimizer/optional_pullback.swift

@@ -0,0 +1,48 @@
+// RUN: %target-swift-emit-sil %s | %FileCheck %s
+
+import _Differentiation
+
+// CHECK: sil private{{.*}}@$s17optional_pullback23givesWrongTangentVector1xxSgx_t16_Differentiation14DifferentiableRzlFAeFRzlTJpSpSr
+// CHECK-SAME: $@convention(thin) <τ_0_0 where τ_0_0 : Differentiable>
+// CHECK-SAME: (@in_guaranteed Optional<τ_0_0>.TangentVector) -> @out τ_0_0.TangentVector
+//
+// CHECK: bb0(%[[RET_TAN:.+]] : $*τ_0_0.TangentVector, %[[OPT_TAN:.+]] : $*Optional<τ_0_0>.TangentVector):
+// CHECK: %[[RET_TAN_BUF:.+]] = alloc_stack $τ_0_0.TangentVector
+
+// CHECK: %[[ZERO1:.+]] = witness_method $τ_0_0.TangentVector, #AdditiveArithmetic.zero!getter
+// CHECK: apply %[[ZERO1]]<τ_0_0.TangentVector>(%[[RET_TAN_BUF]], %{{.*}})
+//
+// CHECK: %[[TAN_VAL_COPY:.+]] = alloc_stack $Optional<τ_0_0.TangentVector>
+// CHECK: %[[TAN_BUF:.+]] = alloc_stack $Optional<τ_0_0>.TangentVector
+
+// CHECK: copy_addr %[[OPT_TAN]] to [init] %[[TAN_BUF]] : $*Optional<τ_0_0>.TangentVector
+// CHECK: %[[TAN_VAL:.+]] = struct_element_addr %[[TAN_BUF]] : $*Optional<τ_0_0>.TangentVector, #Optional.TangentVector.value
+// CHECK: copy_addr %[[TAN_VAL]] to [init] %[[TAN_VAL_COPY]] : $*Optional<τ_0_0.TangentVector>
+//
+// CHECK: %[[TAN_DATA:.+]] = unchecked_take_enum_data_addr %[[TAN_VAL_COPY]] : $*Optional<τ_0_0.TangentVector>, #Optional.some!enumelt
+// CHECK: %[[PLUS_EQUAL:.+]] = witness_method $τ_0_0.TangentVector, #AdditiveArithmetic."+="
+// CHECK: apply %[[PLUS_EQUAL]]<τ_0_0.TangentVector>(%[[RET_TAN_BUF]], %[[TAN_DATA]], %{{.*}})
+//
+// CHECK: destroy_addr %[[TAN_DATA]] : $*τ_0_0.TangentVector
+// CHECK: %[[ZERO2:.+]] = witness_method $τ_0_0.TangentVector, #AdditiveArithmetic.zero!getter
+// CHECK: apply %[[ZERO2]]<τ_0_0.TangentVector>(%[[TAN_DATA]], %{{.*}})
+// CHECK: destroy_addr %[[TAN_DATA]] : $*τ_0_0.TangentVector
+//
+// CHECK: copy_addr [take] %[[RET_TAN_BUF:.+]] to [init] %[[RET_TAN:.+]]
+// CHECK: destroy_addr %[[TAN_BUF]] : $*Optional<τ_0_0>.TangentVector
+// CHECK: dealloc_stack %[[TAN_BUF]] : $*Optional<τ_0_0>.TangentVector
+// CHECK: dealloc_stack %[[TAN_VAL_COPY]] : $*Optional<τ_0_0.TangentVector>
+// CHECK: dealloc_stack %[[RET_TAN_BUF]] : $*τ_0_0.TangentVector
+
+@differentiable(reverse)
+func givesWrongTangentVector<Element>(x: Element) -> Element? where Element: Differentiable {
+    return x
+}
+
+@differentiable(reverse)
+func f(x: Double) -> Double {
+    let y = givesWrongTangentVector(x: x)
+    return y!
+}
+
+print(valueWithGradient(at: 0.0, of: f))