[AutoDiff] Initial support for differentiation of throwing functions (#82653)

This adds initial support for differentiation of functions that may produce `Error` result. 

Essentially we wrap the pullback into `Optional` and emit a diamond-shape control flow pattern depending on whether the pullback value is available or not. VJP emission was modified to accommodate for this. In addition to this, some additional tricks are required as `try_apply` result is not available in the instruction parent block, it is available in normal successor basic block.

As a result we can now:
- differentiate an active `try_apply` result (that would be produced from `do ... try .. catch` constructions)
- `try_apply` when error result is unreachable (usually `try!` and similar source code constructs)
- Support (some) throwing functions with builtin differentiation operators. stdlib change will follow. Though we cannot support typed throws here (yet)
- Correctly propagate error types during currying around differentiable functions as well as type-checking for `@derivative(of:)` attribute, so we can register custom derivatives for functions producing error result
- Added custom derivative for `Optional.??` operator (note that support here is not yet complete as we cannot differentiate through autoclosures, so `x ?? y` works only if `y` is not active, e.g. a constant value).

Some fixes here and there

Author: asl

include/swift/SILOptimizer/Differentiation/Common.h

@@ -277,6 +277,14 @@ inline void createEntryArguments(SILFunction *f) {
       indResTy = indResTy.mapTypeOutOfContext();
     createFunctionArgument(f->mapTypeIntoContext(indResTy).getAddressType());
   }
+  if (auto indErrorResTy =
+          conv.getIndirectErrorResultType(f->getTypeExpansionContext())) {
+    if (indErrorResTy.hasArchetype())
+      indErrorResTy = indErrorResTy.mapTypeOutOfContext();
+    createFunctionArgument(
+        f->mapTypeIntoContext(indErrorResTy).getAddressType());
+  }
+
   for (auto paramTy : conv.getParameterSILTypes(f->getTypeExpansionContext())) {
     if (paramTy.hasArchetype())
       paramTy = paramTy.mapTypeOutOfContext();

lib/AST/Builtins.cpp

@@ -2920,6 +2920,7 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     if (!autodiff::getBuiltinApplyDerivativeConfig(
             OperationName, kind, arity, throws))
       return nullptr;
+    // TODO: Support somehow typed throws
     return getAutoDiffApplyDerivativeFunction(Context, Id, kind, arity,
                                               throws, /*thrownType=*/Type());
   }
@@ -2929,6 +2930,7 @@ ValueDecl *swift::getBuiltinValueDecl(ASTContext &Context, Identifier Id) {
     if (!autodiff::getBuiltinApplyTransposeConfig(
             OperationName, arity, throws))
       return nullptr;
+    // TODO: Support somehow typed throws
     return getAutoDiffApplyTransposeFunction(Context, Id, arity, throws,
                                              /*thrownType=*/Type());
   }

lib/SILGen/SILGenBuiltin.cpp

@@ -1190,12 +1190,9 @@ static ManagedValue emitBuiltinTypeTrait(SILGenFunction &SGF,
 }
 
 static ManagedValue emitBuiltinAutoDiffApplyDerivativeFunction(
-    AutoDiffDerivativeFunctionKind kind, unsigned arity,
-    bool throws, SILGenFunction &SGF, SILLocation loc,
-    SubstitutionMap substitutions, ArrayRef<ManagedValue> args, SGFContext C) {
-  // FIXME(https://github.com/apple/swift/issues/54259): Support throwing functions.
-  assert(!throws && "Throwing functions are not yet supported");
-
+    AutoDiffDerivativeFunctionKind kind, unsigned arity, SILGenFunction &SGF,
+    SILLocation loc, SubstitutionMap substitutions, ArrayRef<ManagedValue> args,
+    SGFContext C) {
   auto origFnVal = args[0];
   SmallVector<SILValue, 2> origFnArgVals;
   for (auto& arg : args.drop_front(1))
@@ -1213,7 +1210,8 @@ static ManagedValue emitBuiltinAutoDiffApplyDerivativeFunction(
   origFnVal = SGF.B.createBeginBorrow(loc, origFnVal);
   SILValue derivativeFn = SGF.B.createDifferentiableFunctionExtract(
       loc, kind, origFnVal.getValue());
-  auto derivativeFnType = derivativeFn->getType().castTo<SILFunctionType>();
+  SILType derivativeType = derivativeFn->getType();
+  auto derivativeFnType = derivativeType.castTo<SILFunctionType>();
   assert(derivativeFnType->getNumResults() == 2);
   assert(derivativeFnType->getNumParameters() == origFnArgVals.size());
 
@@ -1240,8 +1238,8 @@ static ManagedValue emitBuiltinAutoDiffApplyDerivativeFunction(
     applyArgs.push_back(SGF.B.createTupleElementAddr(loc, indResBuffer, 0));
     for (auto origFnArgVal : origFnArgVals)
       applyArgs.push_back(origFnArgVal);
-    auto differential = SGF.B.createApply(loc, derivativeFn, SubstitutionMap(),
-                                          applyArgs);
+    auto differential = SGF.emitApplyWithRethrow(
+        loc, derivativeFn, derivativeType, SubstitutionMap(), applyArgs);
 
     derivativeFn = SILValue();
 
@@ -1253,8 +1251,8 @@ static ManagedValue emitBuiltinAutoDiffApplyDerivativeFunction(
   }
 
   // Do the apply for the direct result case.
-  auto resultTuple = SGF.B.createApply(
-      loc, derivativeFn, SubstitutionMap(), origFnArgVals);
+  auto resultTuple = SGF.emitApplyWithRethrow(loc, derivativeFn, derivativeType,
+                                              SubstitutionMap(), origFnArgVals);
 
   derivativeFn = SILValue();
 
@@ -1323,8 +1321,8 @@ static ManagedValue emitBuiltinApplyDerivative(
   auto successfullyParsed = autodiff::getBuiltinApplyDerivativeConfig(
       builtinName, kind, arity, throws);
   assert(successfullyParsed);
-  return emitBuiltinAutoDiffApplyDerivativeFunction(
-      kind, arity, throws, SGF, loc, substitutions, args, C);
+  return emitBuiltinAutoDiffApplyDerivativeFunction(kind, arity, SGF, loc,
+                                                    substitutions, args, C);
 }
 
 static ManagedValue emitBuiltinApplyTranspose(

lib/SILOptimizer/Differentiation/LinearMapInfo.cpp

@@ -146,14 +146,15 @@ void LinearMapInfo::populateBranchingTraceDecl(SILBasicBlock *originalBB,
       heapAllocatedContext = true;
       decl->setInterfaceType(astCtx.TheRawPointerType);
     } else { // Otherwise the payload is the linear map tuple.
-      auto *linearMapStructTy = getLinearMapTupleType(predBB);
+      auto *linearMapTupleTy = getLinearMapTupleType(predBB);
       // Do not create entries for unreachable predecessors
-      if (!linearMapStructTy)
+      if (!linearMapTupleTy)
         continue;
-      auto canLinearMapStructTy = linearMapStructTy->getCanonicalType();
-      decl->setInterfaceType(
-          canLinearMapStructTy->hasArchetype()
-              ? canLinearMapStructTy->mapTypeOutOfContext() : canLinearMapStructTy);
+
+      auto canLinearMapTupleTy = linearMapTupleTy->getCanonicalType();
+      decl->setInterfaceType(canLinearMapTupleTy->hasArchetype()
+                                 ? canLinearMapTupleTy->mapTypeOutOfContext()
+                                 : canLinearMapTupleTy);
     }
     // Create enum element and enum case declarations.
     auto *paramList = ParameterList::create(astCtx, {decl});
@@ -185,6 +186,7 @@ Type LinearMapInfo::getLinearMapType(ADContext &context, FullApplySite fai) {
   auto hasActiveResults = llvm::any_of(allResults, [&](SILValue res) {
     return activityInfo.isActive(res, config);
   });
+
   bool hasActiveSemanticResultArgument = false;
   bool hasActiveArguments = false;
   auto numIndirectResults = fai.getNumIndirectSILResults();
@@ -313,10 +315,12 @@ Type LinearMapInfo::getLinearMapType(ADContext &context, FullApplySite fai) {
         params, silFnTy->getAllResultsInterfaceType().getASTType(), info);
   }
 
-  if (astFnTy->hasArchetype())
-    return astFnTy->mapTypeOutOfContext();
+  Type resultType =
+      astFnTy->hasArchetype() ? astFnTy->mapTypeOutOfContext() : astFnTy;
+  if (fai.getKind() == FullApplySiteKind::TryApplyInst)
+    resultType = resultType->wrapInOptionalType();
 
-  return astFnTy;
+  return resultType;
 }
 
 void LinearMapInfo::generateDifferentiationDataStructures(