[AutoDiff] fix ownership instructions (SR-13973)

Author: Marc Rasi

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -324,7 +324,8 @@ static void copyParameterArgumentsForApply(
     // Objects are to be retained.
     if (arg->getType().isObject()) {
       auto newArg = arg;
-      if (newArg.getOwnershipKind() != OwnershipKind::None)
+      if (!copyBuilder.hasOwnership() ||
+          newArg.getOwnershipKind() != OwnershipKind::None)
         newArg = copyBuilder.emitCopyValueOperation(loc, arg);
       collectNewArg(newArg);
       continue;
@@ -500,7 +501,8 @@ emitDerivativeFunctionReference(
           builder.emitBeginBorrowOperation(original.getLoc(), original);
       SILValue derivativeFn = builder.createDifferentiableFunctionExtract(
           borrowedDiffFunc.getLoc(), kind, borrowedDiffFunc);
-      if (derivativeFn.getOwnershipKind() != OwnershipKind::None)
+      if (!builder.hasOwnership() ||
+          derivativeFn.getOwnershipKind() != OwnershipKind::None)
         derivativeFn =
             builder.emitCopyValueOperation(original.getLoc(), derivativeFn);
       builder.emitEndBorrowOperation(original.getLoc(), borrowedDiffFunc);
@@ -867,7 +869,8 @@ static void emitFatalError(ADContext &context, SILFunction *f,
   auto loc = f->getLocation();
   // Destroy all owned arguments to pass ownership verification.
   for (auto *arg : entry->getArguments())
-    if (arg->getOwnershipKind() == OwnershipKind::Owned)
+    if (!builder.hasOwnership() ||
+        arg->getOwnershipKind() == OwnershipKind::Owned)
       builder.emitDestroyOperation(loc, arg);
   // Fatal error with a nice message.
   auto neverResultInfo =
@@ -1213,7 +1216,8 @@ SILValue DifferentiationTransformer::promoteToDifferentiableFunction(
     builder.createDeallocStack(loc, buf);
 
   // If our original copy does not have none ownership, copy it.
-  if (origFnOperand.getOwnershipKind() != OwnershipKind::None)
+  if (!builder.hasOwnership() ||
+      origFnOperand.getOwnershipKind() != OwnershipKind::None)
     origFnOperand = builder.emitCopyValueOperation(loc, origFnOperand);
   auto *newDiffFn = context.createDifferentiableFunction(
       builder, loc, parameterIndices, resultIndices, origFnOperand,
@@ -1229,7 +1233,8 @@ SILValue DifferentiationTransformer::promoteToLinearFunction(
   // with an undef transpose function operand. Eventually, a legitimate
   // transpose function operand should be created and used.
   auto origFnOperand = lfi->getOriginalFunction();
-  if (origFnOperand.getOwnershipKind() != OwnershipKind::None)
+  if (!builder.hasOwnership() ||
+      origFnOperand.getOwnershipKind() != OwnershipKind::None)
     origFnOperand = builder.emitCopyValueOperation(loc, origFnOperand);
   auto *parameterIndices = lfi->getParameterIndices();
   auto originalType = origFnOperand->getType().castTo<SILFunctionType>();

test/AutoDiff/validation-test/address_sanitizer.swift

@@ -0,0 +1,24 @@
+// RUN: %empty-directory(%t)
+// RUN: %target-build-swift %s -sanitize=address -o %t/sr13973
+// RUN: %target-run %t/sr13973
+
+// REQUIRES: executable_test
+// REQUIRES: asan_runtime
+
+import _Differentiation
+
+struct SR13973 {
+  let x: Float = 0
+
+  @differentiable
+  func errorVector(_ t: Float) -> Float {
+    return t
+  }
+}
+
+func sr13973() {
+  let s = SR13973()
+  _ = valueWithPullback(at: 0, in: s.errorVector)
+}
+
+sr13973()