[AutoDiff] Fix cloned curry thunk verification error. (#28662)

Differentiation has special support for canonicalizing `differentiable_function`
instructions of curry thunk applications: curry thunks are cloned to a version
returning a `differentiable_function` of the original return value.

In cloned curry thunks, the `differentiable_function` instruction is now
created right after the return value's defining instruction, not at the end of
the function. This avoids `alloc_stack`/`dealloc_stack` ordering issues.

Resolves TF-1039.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -1392,12 +1392,14 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction(
 
       auto thunkTy = thunk->getLoweredFunctionType();
       auto thunkResult = thunkTy->getSingleResult();
-      if (auto resultFnTy = thunkResult.getInterfaceType()->getAs<SILFunctionType>()) {
-        // Construct new curry thunk type with `@differentiable` result.
-        auto diffableResultFnTy = resultFnTy->getWithExtInfo(
-            resultFnTy->getExtInfo()
-                .withDifferentiabilityKind(DifferentiabilityKind::Normal));
-        auto newThunkResult = thunkResult.getWithInterfaceType(diffableResultFnTy);
+      if (auto resultFnTy =
+              thunkResult.getInterfaceType()->getAs<SILFunctionType>()) {
+        // Construct new curry thunk type with `@differentiable` function
+        // result.
+        auto diffResultFnTy = resultFnTy->getWithExtInfo(
+            resultFnTy->getExtInfo().withDifferentiabilityKind(
+                DifferentiabilityKind::Normal));
+        auto newThunkResult = thunkResult.getWithInterfaceType(diffResultFnTy);
         auto thunkType = SILFunctionType::get(
             thunkTy->getSubstGenericSignature(), thunkTy->getExtInfo(),
             thunkTy->getCoroutineKind(), thunkTy->getCalleeConvention(),
@@ -1425,12 +1427,18 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction(
           cloner.run();
           auto *retInst =
               cast<ReturnInst>(newThunk->findReturnBB()->getTerminator());
-          SILBuilder thunkBuilder(retInst);
-          auto *dfi = context.createDifferentiableFunction(thunkBuilder, loc,
-                                                   parameterIndices,
-                                                   retInst->getOperand());
+          auto returnValue = retInst->getOperand();
+          // Create `differentiable_function` instruction directly after the
+          // defining instruction (e.g. `partial_apply`) of the returned value.
+          // Note: `differentiable_function` is not created at the end of the
+          // new thunk to avoid `alloc_stack`/`dealloc_stack` ordering issues.
+          SILBuilder dfiBuilder(
+              std::next(returnValue->getDefiningInstruction()->getIterator()));
+          auto *dfi = context.createDifferentiableFunction(
+              dfiBuilder, loc, parameterIndices, returnValue);
           context.setResultIndex(dfi, resultIndex);
-          thunkBuilder.createReturn(loc, dfi);
+          dfiBuilder.setInsertionPoint(newThunk->findReturnBB());
+          dfiBuilder.createReturn(loc, dfi);
           retInst->eraseFromParent();
 
           context.recordGeneratedFunction(newThunk);
@@ -1450,12 +1458,8 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction(
         auto *newApply = builder.createApply(
             ai->getLoc(), newThunkRef, ai->getSubstitutionMap(), newArgs,
             ai->isNonThrowing());
-        for (auto arg : newArgsToDestroy) {
-          if (arg->getType().isObject())
-            builder.emitDestroyValueOperation(loc, arg);
-          else
-            builder.emitDestroyAddr(loc, arg);
-        }
+        for (auto arg : newArgsToDestroy)
+          builder.emitDestroyOperation(loc, arg);
         for (auto *alloc : newBuffersToDealloc)
           builder.createDeallocStack(loc, alloc);
         return newApply;

test/AutoDiff/compiler_crashers_fixed/tf1039-cloned-curry-thunk-verification.swift

@@ -0,0 +1,53 @@
+// RUN: %target-swift-emit-sil %s -verify
+// REQUIRES: asserts
+
+// TF-1039: Cloned curry thunks generated by differentiation should create a
+// `differentiable_function` instruction before any `dealloc_stack` instructions
+// to prevent `alloc_stack`/`dealloc_stack` ordering issues.
+
+protocol P {
+  @differentiable
+  func foo(_ x: Float) -> Float
+}
+extension P {
+  @derivative(of: foo)
+  func vjpFoo(_ x: Float) -> (value: Float, pullback: (Float) -> Float) {
+    return (x, { $0 })
+  }
+}
+struct S: P {
+  @differentiable
+  func foo(_ x: Float) -> Float { x }
+}
+func foo<T: P>(_ x: T) {
+  // Curry thunk emitted here for `x.foo`.
+  _ = gradient(at: 1, in: x.foo)
+}
+
+// SIL verification failed: stack dealloc does not match most recent stack alloc: op == state.Stack.back()
+// Verifying instruction:
+//      %2 = alloc_stack $τ_0_0                     // users: %7, %5, %9, %8, %3
+// ->   dealloc_stack %2 : $*τ_0_0                  // id: %9
+// In function:
+// // AD__$s5curry1PP3fooyS2fFTc__differentiable_curry_thunk_src_0_wrt_0
+// sil shared [thunk] @AD__$s5curry1PP3fooyS2fFTc__differentiable_curry_thunk_src_0_wrt_0 : $@convention(thin) <τ_0_0 where τ_0_0 : P> (@in_guaranteed τ_0_0) -> @owned @differentiable @callee_guaranteed (Float) -> Float {
+// // %0                                             // user: %3
+// bb0(%0 : $*τ_0_0):
+//   %1 = witness_method $τ_0_0, #P.foo!1 : <Self where Self : P> (Self) -> (Float) -> Float : $@convention(witness_method: P) <τ_0_0 where τ_0_0 : P> (Float, @in_guaranteed τ_0_0) -> Float // user: %8
+//   %2 = alloc_stack $τ_0_0                        // users: %7, %5, %9, %8, %3
+//   copy_addr %0 to [initialization] %2 : $*τ_0_0  // id: %3
+//   %4 = alloc_stack $τ_0_0                        // users: %15, %11, %5
+//   copy_addr %2 to [initialization] %4 : $*τ_0_0  // id: %5
+//   %6 = alloc_stack $τ_0_0                        // users: %14, %13, %7
+//   copy_addr %2 to [initialization] %6 : $*τ_0_0  // id: %7
+//   %8 = partial_apply [callee_guaranteed] %1<τ_0_0>(%2) : $@convention(witness_method: P) <τ_0_0 where τ_0_0 : P> (Float, @in_guaranteed τ_0_0) -> Float // user: %16
+//   dealloc_stack %2 : $*τ_0_0                     // id: %9
+//   %10 = witness_method $τ_0_0, #P.foo!1.jvp.SU : <Self where Self : P> (Self) -> (Float) -> Float : $@convention(witness_method: P) <τ_0_0 where τ_0_0 : P> (Float, @in_guaranteed τ_0_0) -> (Float, @owned @callee_guaranteed (Float) -> Float) // user: %11
+//   %11 = partial_apply [callee_guaranteed] %10<τ_0_0>(%4) : $@convention(witness_method: P) <τ_0_0 where τ_0_0 : P> (Float, @in_guaranteed τ_0_0) -> (Float, @owned @callee_guaranteed (Float) -> Float) // user: %16
+//   %12 = witness_method $τ_0_0, #P.foo!1.vjp.SU : <Self where Self : P> (Self) -> (Float) -> Float : $@convention(witness_method: P) <τ_0_0 where τ_0_0 : P> (Float, @in_guaranteed τ_0_0) -> (Float, @owned @callee_guaranteed (Float) -> Float) // user: %13
+//   %13 = partial_apply [callee_guaranteed] %12<τ_0_0>(%6) : $@convention(witness_method: P) <τ_0_0 where τ_0_0 : P> (Float, @in_guaranteed τ_0_0) -> (Float, @owned @callee_guaranteed (Float) -> Float) // user: %16
+//   dealloc_stack %6 : $*τ_0_0                     // id: %14
+//   dealloc_stack %4 : $*τ_0_0                     // id: %15
+//   %16 = differentiable_function [parameters 0] %8 : $@callee_guaranteed (Float) -> Float with_derivative {%11 : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float), %13 : $@callee_guaranteed (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)} // user: %17
+//   return %16 : $@differentiable @callee_guaranteed (Float) -> Float // id: %17
+// } // end sil function 'AD__$s5curry1PP3fooyS2fFTc__differentiable_curry_thunk_src_0_wrt_0'