Support differentiating closures with contexts (#20186)

Author: rxwei

lib/SILOptimizer/Mandatory/TFDifferentiation.cpp

@@ -226,6 +226,8 @@ static FunctionRefInst *findReferenceToVisibleFunction(SILValue value) {
     return findReferenceToVisibleFunction(thinToThick->getOperand());
   if (auto *convertFn = dyn_cast<ConvertFunctionInst>(inst))
     return findReferenceToVisibleFunction(convertFn->getOperand());
+  if (auto *partialApply = dyn_cast<PartialApplyInst>(inst))
+    return findReferenceToVisibleFunction(partialApply->getCallee());
   return nullptr;
 }
 
@@ -1568,7 +1570,7 @@ reapplyFunctionConversion(SILValue newFunc, SILValue oldFunc,
         newFunc, oldFunc, pai->getCallee(), builder, loc, substituteOperand);
     return builder.createPartialApply(
         loc, innerNewFunc, pai->getSubstitutionMap(), newArgs,
-        pai->getOrigCalleeType()->getCalleeConvention());
+        ParameterConvention::Direct_Guaranteed);
   }
   llvm_unreachable("Unhandled function convertion instruction");
 }
@@ -4004,7 +4006,7 @@ void AdjointEmitter::accumulateMaterializedAdjointsIndirect(
     // Ensure the witness method is linked.
     getModule().lookUpFunctionInWitnessTable(confRef, declRef);
     auto subMap =
-      SubstitutionMap::getProtocolSubstitutions(proto, adjointASTTy, confRef);
+        SubstitutionMap::getProtocolSubstitutions(proto, adjointASTTy, confRef);
     // %1 = metatype $T.Type
     auto metatypeType =
         CanMetatypeType::get(adjointASTTy, MetatypeRepresentation::Thick);

test/AutoDiff/autodiff_e2e_basic.swift

@@ -11,7 +11,7 @@ func adjointId(_ x: Float, originalValue: Float, seed: Float) -> Float {
 
 _ = #gradient(id)(2)
 
-// CHECK: @{{.*}}id{{.*}}__grad_src_0_wrt_0
+// CHECK-LABEL: @{{.*}}id{{.*}}__grad_src_0_wrt_0
 // CHECK-LABEL: @{{.*}}id{{.*}}__grad_src_0_wrt_0_s_p
 
 #if os(macOS) || os(iOS) || os(tvOS) || os(watchOS)
@@ -41,7 +41,7 @@ let x = #gradient(sigmoid)(3)
 let (value: y, gradient: z) = #valueAndGradient(sigmoid)(4)
 print(x * z)
 
-// CHECK: @{{.*}}sigmoid{{.*}}__grad_src_0_wrt_0
+// CHECK-LABEL: @{{.*}}sigmoid{{.*}}__grad_src_0_wrt_0
 // CHECK: @{{.*}}sigmoid{{.*}}__grad_src_0_wrt_0_s_p
 // CHECK: @{{.*}}sigmoid{{.*}}__grad_src_0_wrt_0_p
 
@@ -51,7 +51,6 @@ public func publicFunc(_ x: Float) -> Float {
 }
 _ = #gradient(publicFunc)
 
-// CHECK: sil non_abi @{{.*}}publicFunc{{.*}}__grad_src_0_wrt_0
-// CHECK: sil non_abi @{{.*}}publicFunc{{.*}}__primal_src_0_wrt_0
-// CHECK: sil non_abi @{{.*}}publicFunc{{.*}}__adjoint_src_0_wrt_0
-
+// CHECK-LABEL: @{{.*}}publicFunc{{.*}}__grad_src_0_wrt_0
+// CHECK: @{{.*}}publicFunc{{.*}}__primal_src_0_wrt_0
+// CHECK: @{{.*}}publicFunc{{.*}}__adjoint_src_0_wrt_0

test/AutoDiff/closures.swift

@@ -0,0 +1,15 @@
+// RUN: %target-swift-frontend -emit-sil %s | %FileCheck %s
+
+public func closureCapture() {
+  let val: Float = 10
+  let clo: (Float) -> Float = { x in
+    val * x
+  }
+  _ = #gradient(clo)
+}
+
+// CHECK-LABEL: @{{.*}}closureCapture{{.*}}
+// CHECK: [[ORIG_FN:%.*]] = function_ref @{{.*}}closureCapture{{.*}}
+// CHECK: [[PARTIAL_APPLIED:%.*]] = partial_apply [callee_guaranteed] [[ORIG_FN]]
+// CHECK: [[GRAD_FN:%.*]] = function_ref @{{.*}}closureCapture{{.*}}___grad
+// CHECK: [[PARTIAL_APPLIED_GRAD:%.*]] = partial_apply [callee_guaranteed] [[GRAD_FN]]

test/AutoDiff/simple_math.swift

@@ -56,4 +56,13 @@ SimpleMathTests.test("ResultSelection") {
   expectEqual((0, 1), #gradient(foo, result: .1)(3, 3))
 }
 
+
+SimpleMathTests.test("CaptureGlobal") {
+  let z: Float = 10
+  func foo(_ x: Float) -> Float {
+    return z * x
+  }
+  expectEqual(10, #gradient(foo)(0))
+}
+
 runAllTests()