Initial support for differentiation of throwing functions

Author: asl

include/swift/SILOptimizer/Differentiation/Common.h

@@ -277,6 +277,12 @@ 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

@@ -2892,6 +2892,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());
   }
@@ -2901,6 +2902,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

@@ -1191,11 +1191,8 @@ 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");
-
+    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,10 @@ 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 +1253,10 @@ 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();
 
@@ -1324,7 +1326,7 @@ static ManagedValue emitBuiltinApplyDerivative(
       builtinName, kind, arity, throws);
   assert(successfullyParsed);
   return emitBuiltinAutoDiffApplyDerivativeFunction(
-      kind, arity, throws, SGF, loc, substitutions, args, C);
+      kind, arity, SGF, loc, substitutions, args, C);
 }
 
 static ManagedValue emitBuiltinApplyTranspose(

lib/SILOptimizer/Differentiation/LinearMapInfo.cpp

@@ -144,14 +144,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();
+
+      auto canLinearMapTupleTy = linearMapTupleTy->getCanonicalType();
       decl->setInterfaceType(
-          canLinearMapStructTy->hasArchetype()
-              ? canLinearMapStructTy->mapTypeOutOfContext() : canLinearMapStructTy);
+          canLinearMapTupleTy->hasArchetype()
+              ? canLinearMapTupleTy->mapTypeOutOfContext() : canLinearMapTupleTy);
     }
     // Create enum element and enum case declarations.
     auto *paramList = ParameterList::create(astCtx, {decl});
@@ -183,6 +184,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();
@@ -311,10 +313,13 @@ 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(