[AutoDiff upstream] handle differentiable_function in DiagnoseInvalidEscapingCaptures (#30909)

Author: marcrasi

lib/SILOptimizer/Mandatory/DiagnoseInvalidEscapingCaptures.cpp

@@ -78,6 +78,12 @@ static bool checkNoEscapePartialApplyUse(Operand *oper, FollowUse followUses) {
     return false;
   }
 
+  // Look through `differentiable_function`.
+  if (auto *DFI = dyn_cast<DifferentiableFunctionInst>(user)) {
+    followUses(DFI);
+    return false;
+  }
+
   // @noescape block storage can be passed as an Optional (Nullable).
   if (auto *EI = dyn_cast<EnumInst>(user)) {
     followUses(EI);

test/AutoDiff/SILOptimizer/differentiation.swift

@@ -0,0 +1,14 @@
+// RUN: %target-swift-frontend -emit-sil -verify %s
+
+// Tests that differentiation features interact correctly with non-differentiation SILOptimizer
+// passes.
+
+import _Differentiation
+
+// - MARK: DiagnoseInvalidEscapingCaptures
+
+func nonEscapingUseOfDifferentiableFunction(_ f: @differentiable (Float) -> Float) {}
+func testDiagnoseInvalidEscapingCaptures(_ f: @differentiable (Float) -> Float) {
+  // Should not be diagnosed as invalid escaping capture.
+  nonEscapingUseOfDifferentiableFunction { f($0) }
+}