Merge pull request #34898 from gottesmm/pr-d3b6d903b097410b535d20940005a1fa67d3a1bf

[auto-diff] Fix a bunch of places in the *Cloners where we were not closing borrow scopes.

Author: gottesmm

lib/SIL/IR/SILBuilder.cpp

@@ -637,15 +637,18 @@ DebugValueAddrInst *SILBuilder::createDebugValueAddr(SILLocation Loc,
 
 void SILBuilder::emitScopedBorrowOperation(SILLocation loc, SILValue original,
                                            function_ref<void(SILValue)> &&fun) {
-  if (original->getType().isAddress()) {
-    original = createLoadBorrow(loc, original);
+  SILValue value = original;
+  if (value->getType().isAddress()) {
+    value = createLoadBorrow(loc, value);
   } else {
-    original = createBeginBorrow(loc, original);
+    value = emitBeginBorrowOperation(loc, value);
   }
 
-  fun(original);
+  fun(value);
 
-  createEndBorrow(loc, original);
+  // If we actually inserted a borrowing operation... insert the end_borrow.
+  if (value != original)
+    createEndBorrow(loc, value);
 }
 
 CheckedCastBranchInst *SILBuilder::createCheckedCastBranch(

lib/SILOptimizer/Differentiation/JVPCloner.cpp

@@ -519,11 +519,13 @@ class JVPCloner::Implementation final
           return;
         }
       }
-      auto borrowedDiffFunc = builder.emitBeginBorrowOperation(loc, origCallee);
-      jvpValue = builder.createDifferentiableFunctionExtract(
-          loc, NormalDifferentiableFunctionTypeComponent::JVP,
-          borrowedDiffFunc);
-      jvpValue = builder.emitCopyValueOperation(loc, jvpValue);
+      builder.emitScopedBorrowOperation(
+          loc, origCallee, [&](SILValue borrowedDiffFunc) {
+            jvpValue = builder.createDifferentiableFunctionExtract(
+                loc, NormalDifferentiableFunctionTypeComponent::JVP,
+                borrowedDiffFunc);
+            jvpValue = builder.emitCopyValueOperation(loc, jvpValue);
+          });
     }
 
     // If JVP has not yet been found, emit an `differentiable_function`
@@ -614,11 +616,13 @@ class JVPCloner::Implementation final
       // Record the `differentiable_function` instruction.
       context.getDifferentiableFunctionInstWorklist().push_back(diffFuncInst);
 
-      auto borrowedADFunc = builder.emitBeginBorrowOperation(loc, diffFuncInst);
-      auto extractedJVP = builder.createDifferentiableFunctionExtract(
-          loc, NormalDifferentiableFunctionTypeComponent::JVP, borrowedADFunc);
-      jvpValue = builder.emitCopyValueOperation(loc, extractedJVP);
-      builder.emitEndBorrowOperation(loc, borrowedADFunc);
+      builder.emitScopedBorrowOperation(
+          loc, diffFuncInst, [&](SILValue borrowedADFunc) {
+            auto extractedJVP = builder.createDifferentiableFunctionExtract(
+                loc, NormalDifferentiableFunctionTypeComponent::JVP,
+                borrowedADFunc);
+            jvpValue = builder.emitCopyValueOperation(loc, extractedJVP);
+          });
       builder.emitDestroyValueOperation(loc, diffFuncInst);
     }
 

lib/SILOptimizer/Differentiation/VJPCloner.cpp

@@ -495,11 +495,13 @@ class VJPCloner::Implementation final
             loc, origCallee, SILType::getPrimitiveObjectType(origFnUnsubstType),
             /*withoutActuallyEscaping*/ false);
       }
-      auto borrowedDiffFunc = builder.emitBeginBorrowOperation(loc, origCallee);
-      vjpValue = builder.createDifferentiableFunctionExtract(
-          loc, NormalDifferentiableFunctionTypeComponent::VJP,
-          borrowedDiffFunc);
-      vjpValue = builder.emitCopyValueOperation(loc, vjpValue);
+      builder.emitScopedBorrowOperation(
+          loc, origCallee, [&](SILValue borrowedDiffFunc) {
+            vjpValue = builder.createDifferentiableFunctionExtract(
+                loc, NormalDifferentiableFunctionTypeComponent::VJP,
+                borrowedDiffFunc);
+            vjpValue = builder.emitCopyValueOperation(loc, vjpValue);
+          });
       auto vjpFnType = vjpValue->getType().castTo<SILFunctionType>();
       auto vjpFnUnsubstType = vjpFnType->getUnsubstitutedType(getModule());
       if (vjpFnType != vjpFnUnsubstType) {
@@ -601,11 +603,14 @@ class VJPCloner::Implementation final
       // Record the `differentiable_function` instruction.
       context.getDifferentiableFunctionInstWorklist().push_back(diffFuncInst);
 
-      auto borrowedADFunc = builder.emitBeginBorrowOperation(loc, diffFuncInst);
-      auto extractedVJP = getBuilder().createDifferentiableFunctionExtract(
-          loc, NormalDifferentiableFunctionTypeComponent::VJP, borrowedADFunc);
-      vjpValue = builder.emitCopyValueOperation(loc, extractedVJP);
-      builder.emitEndBorrowOperation(loc, borrowedADFunc);
+      builder.emitScopedBorrowOperation(
+          loc, diffFuncInst, [&](SILValue borrowedADFunc) {
+            auto extractedVJP =
+                getBuilder().createDifferentiableFunctionExtract(
+                    loc, NormalDifferentiableFunctionTypeComponent::VJP,
+                    borrowedADFunc);
+            vjpValue = builder.emitCopyValueOperation(loc, extractedVJP);
+          });
       builder.emitDestroyValueOperation(loc, diffFuncInst);
     }
 

test/AutoDiff/SILOptimizer/derivative_sil.swift

@@ -33,7 +33,6 @@ func foo(_ x: Float) -> Float {
 // CHECK-SIL:   [[ADD_VJP_REF:%.*]] = differentiability_witness_function [vjp] [parameters 0 1] [results 0] @add
 // CHECK-SIL:   [[ADD_DIFF_FN:%.*]] = differentiable_function [parameters 0 1] [results 0] [[ADD_ORIG_REF]] : $@convention(method) (Float, Float, @thin Float.Type) -> Float with_derivative {[[ADD_JVP_REF]] : $@convention(method) (Float, Float, @thin Float.Type) -> (Float, @owned @callee_guaranteed (Float, Float) -> Float), [[ADD_VJP_REF]] : $@convention(method) (Float, Float, @thin Float.Type) -> (Float, @owned @callee_guaranteed (Float) -> (Float, Float))}
 // CHECK-SIL:   [[ADD_JVP_FN:%.*]] = differentiable_function_extract [jvp] [[ADD_DIFF_FN]]
-// CHECK-SIL:   end_borrow [[ADD_DIFF_FN]]
 // CHECK-SIL:   [[ADD_RESULT:%.*]] = apply [[ADD_JVP_FN]]([[X]], [[X]], {{.*}})
 // CHECK-SIL:   ([[ORIG_RES:%.*]], [[ADD_DF:%.*]]) = destructure_tuple [[ADD_RESULT]]
 // CHECK-SIL:   [[DF_STRUCT:%.*]] = struct $_AD__foo_bb0__DF__src_0_wrt_0 ([[ADD_DF]] : $@callee_guaranteed (Float, Float) -> Float)
@@ -58,7 +57,6 @@ func foo(_ x: Float) -> Float {
 // CHECK-SIL:   [[ADD_VJP_REF:%.*]] = differentiability_witness_function [vjp] [parameters 0 1] [results 0] @add
 // CHECK-SIL:   [[ADD_DIFF_FN:%.*]] = differentiable_function [parameters 0 1] [results 0] [[ADD_ORIG_REF]] : $@convention(method) (Float, Float, @thin Float.Type) -> Float with_derivative {[[ADD_JVP_REF]] : $@convention(method) (Float, Float, @thin Float.Type) -> (Float, @owned @callee_guaranteed (Float, Float) -> Float), [[ADD_VJP_REF]] : $@convention(method) (Float, Float, @thin Float.Type) -> (Float, @owned @callee_guaranteed (Float) -> (Float, Float))}
 // CHECK-SIL:   [[ADD_VJP_FN:%.*]] = differentiable_function_extract [vjp] [[ADD_DIFF_FN]]
-// CHECK-SIL:   end_borrow [[ADD_DIFF_FN]]
 // CHECK-SIL:   [[ADD_RESULT:%.*]] = apply [[ADD_VJP_FN]]([[X]], [[X]], {{.*}})
 // CHECK-SIL:   ([[ORIG_RES:%.*]], [[ADD_PB:%.*]]) = destructure_tuple [[ADD_RESULT]]
 // CHECK-SIL:   [[PB_STRUCT:%.*]] = struct $_AD__foo_bb0__PB__src_0_wrt_0 ([[ADD_PB]] : $@callee_guaranteed (Float) -> (Float, Float))