[AutoDiff] Destroy unused pullback direct results. (#28207)

In `PullbackEmitter::visitApplyInst`, destroy unused pullback direct results.
This is needed for VJPs extracted from `@differentiable` function callees,
where the `@differentiable` function's differentiation parameter indices are
a superset of the active `apply` parameter indices.

Resolves TF-953.

Author: dan-zheng

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -6923,8 +6923,8 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
         ? *applyInfo.originalPullbackType
         : pullbackType;
     for (auto indRes : actualPullbackType->getIndirectFormalResults()) {
-      auto *alloc =
-          builder.createAllocStack(loc, remapType(indRes.getSILStorageInterfaceType()));
+      auto *alloc = builder.createAllocStack(
+          loc, remapType(indRes.getSILStorageInterfaceType()));
       pullbackIndirectResults.push_back(alloc);
       args.push_back(alloc);
     }
@@ -6975,13 +6975,22 @@ class PullbackEmitter final : public SILInstructionVisitor<PullbackEmitter> {
           addToAdjointBuffer(bb, origArg, tmpBuf, loc);
           builder.emitDestroyAddrAndFold(loc, tmpBuf);
           builder.createDeallocStack(loc, tmpBuf);
-        }
-        else {
+        } else {
           recordTemporary(tan);
           addAdjointValue(bb, origArg, makeConcreteAdjointValue(tan), loc);
         }
       }
     }
+    // Destroy unused pullback direct results. Needed for pullback results from
+    // VJPs extracted from `@differentiable` function callees, where the
+    // `@differentiable` function's differentiation parameter indices are a
+    // superset of the active `apply` parameter indices.
+    while (allResultsIt != allResults.end()) {
+      auto unusedPullbackDirectResult = *allResultsIt++;
+      if (unusedPullbackDirectResult->getType().isAddress())
+        continue;
+      builder.emitDestroyValueOperation(loc, unusedPullbackDirectResult);
+    }
     // Destroy and deallocate pullback indirect results.
     for (auto *alloc : llvm::reverse(pullbackIndirectResults)) {
       builder.emitDestroyAddrAndFold(loc, alloc);

test/AutoDiff/compiler_crashers/tf928-pullback-ownership-memory-leak.swift renamed to test/AutoDiff/compiler_crashers_fixed/tf928-pullback-ownership-memory-leak.swift

@@ -1,5 +1,4 @@
-// RUN: not --crash %target-swift-emit-sil %s
-// REQUIRES: asserts
+// RUN: %target-swift-emit-sil %s
 
 // TF-928: Ownership verification error in pullback function generated by the
 // differentiation transform.
@@ -27,6 +26,7 @@ func TF_928(
   _ x: Tracked<Float>
 ) {
   _ = pullback(at: x) { x in lossFunction(x, Tracked<Float>()) }
+  _ = pullback(at: x) { x in lossFunction(Tracked<Float>(), x) }
 }
 
 // Function: 'AD__$s4main6TF_928yyAA7TrackedVySfGAE_AEtXF_AEtFA2EcfU___pullback_src_0_wrt_0'