[AutoDiff] Fix custom derivative thunk for Optional (#71721)

Enable the nil coalescing operator (aka `??`) for Optional type.

Fixes #55882

Co-authored-by: Anton Korobeynikov <[email protected]>

Author: asavonic

lib/SILGen/SILGenPoly.cpp

@@ -6305,10 +6305,13 @@ SILFunction *SILGenModule::getOrCreateCustomDerivativeThunk(
     arguments.push_back(indErrorRes.getLValueAddress());
   forwardFunctionArguments(thunkSGF, loc, fnRefType, params, arguments);
 
+  SubstitutionMap subs = thunk->getForwardingSubstitutionMap();
+  SILType substFnType = fnRef->getType().substGenericArgs(
+      M, subs, thunk->getTypeExpansionContext());
+
   // Apply function argument.
-  auto apply = thunkSGF.emitApplyWithRethrow(
-      loc, fnRef, /*substFnType*/ fnRef->getType(),
-      thunk->getForwardingSubstitutionMap(), arguments);
+  auto apply =
+      thunkSGF.emitApplyWithRethrow(loc, fnRef, substFnType, subs, arguments);
 
   // Self reordering thunk is necessary if wrt at least two parameters,
   // including self.

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -906,6 +906,15 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
       traceMessage.c_str(), witness->getOriginalFunction());
 
   assert(witness->isDefinition());
+  SILFunction *orig = witness->getOriginalFunction();
+
+  // We can generate empty JVP / VJP for functions available externally. These
+  // functions have the same linkage as the original ones sans `external`
+  // flag. Important exception here hidden_external functions as they are
+  // serializable but corresponding hidden ones would be not and the SIL
+  // verifier will fail. Patch `serializeFunctions` for this case.
+  if (orig->getLinkage() == SILLinkage::HiddenExternal)
+    serializeFunctions = IsNotSerialized;
 
   // If the JVP doesn't exist, need to synthesize it.
   if (!witness->getJVP()) {
@@ -914,9 +923,8 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
     // - Functions with unsupported control flow.
     if (context.getASTContext()
             .LangOpts.hasFeature(Feature::ForwardModeDifferentiation) &&
-        (diagnoseNoReturn(context, witness->getOriginalFunction(), invoker) ||
-         diagnoseUnsupportedControlFlow(
-             context, witness->getOriginalFunction(), invoker)))
+        (diagnoseNoReturn(context, orig, invoker) ||
+         diagnoseUnsupportedControlFlow(context, orig, invoker)))
       return true;
 
     // Create empty JVP.
@@ -933,10 +941,10 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
         !witness->getVJP()) {
       // JVP and differential generation do not currently support functions with
       // multiple basic blocks.
-      if (witness->getOriginalFunction()->size() > 1) {
-        context.emitNondifferentiabilityError(
-            witness->getOriginalFunction()->getLocation().getSourceLoc(),
-            invoker, diag::autodiff_jvp_control_flow_not_supported);
+      if (orig->size() > 1) {
+        context.emitNondifferentiabilityError(orig->getLocation().getSourceLoc(),
+                                              invoker,
+                                              diag::autodiff_jvp_control_flow_not_supported);
         return true;
       }
       // Emit JVP function.
@@ -950,7 +958,7 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
                      "_fatalErrorForwardModeDifferentiationDisabled");
       LLVM_DEBUG(getADDebugStream()
                  << "Generated empty JVP for "
-                 << witness->getOriginalFunction()->getName() << ":\n"
+                 << orig->getName() << ":\n"
                  << *jvp);
     }
   }
@@ -960,9 +968,8 @@ bool DifferentiationTransformer::canonicalizeDifferentiabilityWitness(
     // Diagnose:
     // - Functions with no return.
     // - Functions with unsupported control flow.
-    if (diagnoseNoReturn(context, witness->getOriginalFunction(), invoker) ||
-        diagnoseUnsupportedControlFlow(
-            context, witness->getOriginalFunction(), invoker))
+    if (diagnoseNoReturn(context, orig, invoker) ||
+        diagnoseUnsupportedControlFlow(context, orig, invoker))
       return true;
 
     // Create empty VJP.

test/AutoDiff/SILGen/nil_coalescing.swift

@@ -0,0 +1,25 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s | %FileCheck %s
+
+import _Differentiation
+
+// CHECK: sil @test_nil_coalescing
+// CHECK: bb0(%{{.*}} : $*T, %[[ARG_OPT:.*]] : $*Optional<T>, %[[ARG_PB:.*]] :
+// CHECK:    $@noescape @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error any Error) for <T>):
+// CHECK: %[[ALLOC_OPT:.*]] = alloc_stack [lexical] $Optional<T>
+// CHECK: copy_addr %[[ARG_OPT]] to [init] %[[ALLOC_OPT]] : $*Optional<T>
+// CHECK: switch_enum_addr %[[ALLOC_OPT]] : $*Optional<T>, case #Optional.some!enumelt: {{.*}}, case #Optional.none!enumelt: {{.*}}
+// CHECK: try_apply %[[ARG_PB]](%{{.*}}) : $@noescape @callee_guaranteed @substituted <τ_0_0> () -> (@out τ_0_0, @error any Error) for <T>, normal {{.*}}, error {{.*}}
+//
+@_silgen_name("test_nil_coalescing")
+@derivative(of: ??)
+@usableFromInline
+func nilCoalescing<T: Differentiable>(optional: T?, defaultValue: @autoclosure () throws -> T)
+ rethrows -> (value: T, pullback: (T.TangentVector) -> Optional<T>.TangentVector)
+{
+ let hasValue = optional != nil
+ let value = try optional ?? defaultValue()
+ func pullback(_ v: T.TangentVector) -> Optional<T>.TangentVector {
+   return hasValue ? .init(v) : .zero
+ }
+ return (value, pullback)
+}