[AutoDiff] Handle @inout_aliasable when converting associated functions.

When reapplying a `partial_apply` for a JVP/VJP, if the argument has convention `@inout_aliasable`, it is a noescape mutable capture and the underlying value should not be retained. This fixes a memory leak when a noescape closure being differentiated captures a mutable self.

Author: rxwei

lib/SILOptimizer/Mandatory/Differentiation.cpp

@@ -1951,19 +1951,35 @@ reapplyFunctionConversion(SILValue newFunc, SILValue oldFunc,
         builder.createDeallocStack(loc, alloc);
     };
     for (auto arg : pai->getArguments()) {
-      // Retain the argument since it's to be owned by the newly created
+      // Retain the argument if it's to be owned by the newly created
       // closure.
+      // Objects are to be retained.
       if (arg->getType().isObject()) {
         builder.createRetainValue(loc, arg, builder.getDefaultAtomicity());
         newArgs.push_back(arg);
-      } else if (arg->getType().isLoadable(builder.getFunction())) {
-        builder.createRetainValueAddr(loc, arg, builder.getDefaultAtomicity());
-        newArgs.push_back(arg);
-      } else {
-        auto *argCopy = builder.createAllocStack(loc, arg->getType());
-        copiedIndirectParams.push_back(argCopy);
-        builder.createCopyAddr(loc, arg, argCopy, IsNotTake, IsInitialization);
-        newArgs.push_back(argCopy);
+      }
+      // Addresses depend on argument conventions.
+      else {
+        auto conv = pai->getCalleeFunction()->getConventions();
+        auto argConv =
+            conv.getSILArgumentConvention(conv.getNumSILArguments() - 1);
+        // If the argument is an aliasable inout reference, do not retain the
+        // argument since it's a `@noescape` capture.
+        if (argConv == SILArgumentConvention::Indirect_InoutAliasable) {
+          newArgs.push_back(arg);
+        }
+        // Otherwise, retain/copy the underlying value.
+        else if (arg->getType().isLoadable(builder.getFunction())) {
+          builder.createRetainValueAddr(loc, arg,
+                                        builder.getDefaultAtomicity());
+          newArgs.push_back(arg);
+        } else {
+          auto *argCopy = builder.createAllocStack(loc, arg->getType());
+          copiedIndirectParams.push_back(argCopy);
+          builder.createCopyAddr(loc, arg, argCopy, IsNotTake,
+                                 IsInitialization);
+          newArgs.push_back(argCopy);
+        }
       }
     }
     auto innerNewFunc = reapplyFunctionConversion(

test/AutoDiff/closures.swift

@@ -13,6 +13,27 @@ func diffableClosureInStruct(s: Foo) {
 // CHECK:   retain_value [[CLOSURE]] : $@differentiable @callee_guaranteed (Float) -> Float
 // CHECK:   autodiff_function_extract [original] [[CLOSURE]] : $@differentiable @callee_guaranteed (Float) -> Float
 
+struct InoutAliasableCapture {
+  var x: Float = .zero
+  mutating func foo() {
+    func capturesMutableSelf(t: Float) -> Float {
+      self.x = .zero
+      return t
+    }
+    _ = gradient(at: .zero, in: capturesMutableSelf)
+  }
+}
+
+// CHECK-LABEL: @{{.*}}InoutAliasableCapture{{.*}}foo{{.*}} : $@convention(method) (@inout InoutAliasableCapture) -> () {
+// CHECK: bb0([[SELF:%.*]] : $*InoutAliasableCapture):
+// CHECK:   [[JVP:%.*]] = function_ref @{{.*}}capturesMutableSelf{{.*}}__jvp_src_0_wrt_0 : $@convention(thin) (Float, @inout_aliasable InoutAliasableCapture) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK-NOT:  retain_value_addr [[SELF]]
+// CHECK-NOT:  copy_addr [[SELF]]
+// CHECK:   [[JVP_CAPTURED:%.*]] = partial_apply [callee_guaranteed] [[JVP]]([[SELF]]) : $@convention(thin) (Float, @inout_aliasable InoutAliasableCapture) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK:   [[VJP:%.*]] = function_ref @{{.*}}capturesMutableSelf{{.*}}__vjp_src_0_wrt_0 : $@convention(thin) (Float, @inout_aliasable InoutAliasableCapture) -> (Float, @owned @callee_guaranteed (Float) -> Float)
+// CHECK-NOT:  retain_value_addr [[SELF]]
+// CHECK-NOT:  copy_addr [[SELF]]
+// CHECK:   [[VJP_CAPTURED:%.*]] = partial_apply [callee_guaranteed] [[VJP]]([[SELF]]) : $@convention(thin) (Float, @inout_aliasable InoutAliasableCapture) -> (Float, @owned @callee_guaranteed (Float) -> Float)
 
 public func closureCaptureMutable() {
   var val: Float = 10

test/AutoDiff/leakchecking.swift

@@ -207,6 +207,21 @@ LeakCheckingTests.testWithLeakChecking("ClosureCaptureLeakChecking") {
       return model.applied(to: x)
     }
   }
+
+  do {
+    struct Foo {
+      let x: Tracked<Float> = .zero
+      var y: Tracked<Float> = .zero
+      mutating func differentiateSomethingThatCapturesSelf() {
+        _ = x.gradient { x in
+          self.y += .zero
+          return .zero
+        }
+      }
+    }
+    var foo = Foo()
+    foo.differentiateSomethingThatCapturesSelf()
+  }
 }
 
 LeakCheckingTests.testWithLeakChecking("ControlFlowWithTrivialUnconditionalMath") {