[AutoDiff] Fix 'autodiff_function_extract' operand ownership kind and memory leaks. (#27199)

The `autodiff_function_extract` instruction behaves like `tuple_extract`, where it extracts some element from an aggregate. Its operand should have the same ownership kind as that of `tuple_extract`. That is, it should be defined as `CONSTANT_OR_TRIVIAL_OWNERSHIP_INST(Guaranteed, AutoDiffFunctionInst)` in ValueOwnershipKindClassifier.

However, this is currently defined wrongly as `FORWARDING_OWNERSHIP_INST(AutoDiffFunctionExtract)`, which caused a bug in the differentiation transform to be uncaught: VJPEmitter and JVPEmitter in the differentiation transform is performing `autodiff_function_extract` on an `@owned` `@differentiable` function, which caused associated functions that are not extracted to be not released:

```
%f = autodiff_function %original
%f_vjp = autodiff_function_extract [vjp] %f
... // %f is not released, and not caught by ownership verification!
```

After we fix the operand ownership kind for `autodiff_function_extract`, all these cases are now caught by ownership verification. The reproducer in [TF-795](https://bugs.swift.org/browse/TF-795) and most differentiation tests are failing to compile because ownership verification caught the bug in AD-generated code. The existing AD test suite is serving as good test cases for this ownership error.

To fix this, VJPEmitter and JVPEmitter are now changed to emit borrows of `@differentiable` functions and copies of associated functions and property destroying the `@differentiable` function:

```
%f = autodiff_function %original
%f_borrowed = begin_borrow %f
%f_vjp_extracted = autodiff_function_extract [vjp] %f_borrowed
%f_vjp = copy_value %f_vjp_extracted
end_borrow %f_borrowed
destroy_value %f
```

Fixes [TF-795](https://bugs.swift.org/browse/TF-795).

Author: rxwei

lib/SIL/ValueOwnership.cpp

@@ -163,6 +163,8 @@ CONSTANT_OWNERSHIP_INST(Unowned, ValueToBridgeObject)
   }
 CONSTANT_OR_TRIVIAL_OWNERSHIP_INST(Guaranteed, StructExtract)
 CONSTANT_OR_TRIVIAL_OWNERSHIP_INST(Guaranteed, TupleExtract)
+// SWIFT_ENABLE_TENSORFLOW
+CONSTANT_OR_TRIVIAL_OWNERSHIP_INST(Guaranteed, AutoDiffFunctionExtract)
 // OpenExistentialValue opens the boxed value inside an existential
 // CoW box. The semantics of an existential CoW box implies that we
 // can only consume the projected value inside the box if the box is
@@ -254,7 +256,6 @@ FORWARDING_OWNERSHIP_INST(SelectEnum)
 FORWARDING_OWNERSHIP_INST(Enum)
 // SWIFT_ENABLE_TENSORFLOW
 FORWARDING_OWNERSHIP_INST(AutoDiffFunction)
-FORWARDING_OWNERSHIP_INST(AutoDiffFunctionExtract)
 #undef FORWARDING_OWNERSHIP_INST
 
 ValueOwnershipKind

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -3749,9 +3749,14 @@ class VJPEmitter final
       context.getResultIndices()[autoDiffFuncInst] =
           activeResultIndices.front();
 
-      vjpValue = getBuilder().createAutoDiffFunctionExtract(
+      auto borrowedADFunc =
+          builder.emitBeginBorrowOperation(loc, autoDiffFuncInst);
+      auto extractedVJP = getBuilder().createAutoDiffFunctionExtract(
           loc, AutoDiffFunctionExtractInst::Extractee::VJP,
-          /*differentiationOrder*/ 1, autoDiffFuncInst);
+          /*differentiationOrder*/ 1, borrowedADFunc);
+      vjpValue = builder.emitCopyValueOperation(loc, extractedVJP);
+      builder.emitEndBorrowOperation(loc, borrowedADFunc);
+      builder.emitDestroyValueOperation(loc, autoDiffFuncInst);
     }
 
     // Record desired/actual VJP indices.
@@ -4850,7 +4855,6 @@ class JVPEmitter final
     // on the remapped original function operand and `autodiff_function_extract`
     // the JVP. The actual JVP functions will be populated in the
     // `autodiff_function` during the transform main loop.
-    SILValue differentiableFunc;
     if (!jvpValue) {
       // FIXME: Handle indirect differentiation invokers. This may require some
       // redesign: currently, each original function + attribute pair is mapped
@@ -4911,16 +4915,20 @@ class JVPEmitter final
       auto *autoDiffFuncInst =
           context.createAutoDiffFunction(builder, loc, indices.parameters,
                                          /*differentiationOrder*/ 1, original);
-      differentiableFunc = autoDiffFuncInst;
 
       // Record the `autodiff_function` instruction.
       context.getAutoDiffFunctionInsts().push_back(autoDiffFuncInst);
       context.getResultIndices()[autoDiffFuncInst] =
           activeResultIndices.front();
 
-      jvpValue = builder.createAutoDiffFunctionExtract(
+      auto borrowedADFunc =
+          builder.emitBeginBorrowOperation(loc, autoDiffFuncInst);
+      auto extractedJVP = builder.createAutoDiffFunctionExtract(
           loc, AutoDiffFunctionExtractInst::Extractee::JVP,
-          /*differentiationOrder*/ 1, autoDiffFuncInst);
+          /*differentiationOrder*/ 1, borrowedADFunc);
+      jvpValue = builder.emitCopyValueOperation(loc, extractedJVP);
+      builder.emitEndBorrowOperation(loc, borrowedADFunc);
+      builder.emitDestroyValueOperation(loc, autoDiffFuncInst);
     }
 
     // Call the JVP using the original parameters.