SILOptimizer: fix partial_apply optimization. (#31552)

`partial_apply` can be rewritten to `thin_to_thick_function` only if the
specialized callee is `@convention(thin)`.

This condition is newly exercised by the differentiation transform:
`{JVP,VJP}Emitter::visitApplyInst` generates argument-less `partial_apply`
with `@convention(method)` callees.

Resolves SR-12732.

Author: dan-zheng

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -460,15 +460,20 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
   // arguments are dead?
   std::pair<SILFunction *, SILFunction *> GenericSpecialized;
   SILOptFunctionBuilder FuncBuilder(*this);
-  if (auto *NewFunc = getSpecializedWithDeadParams(FuncBuilder,
-          PAI, SubstF, PAI->getNumArguments(), GenericSpecialized)) {
-    rewritePartialApply(PAI, NewFunc);
-    if (GenericSpecialized.first) {
-      // Notify the pass manager about the new function.
-      addFunctionToPassManagerWorklist(GenericSpecialized.first,
-                                       GenericSpecialized.second);
+  if (auto *NewFunc = getSpecializedWithDeadParams(FuncBuilder, PAI, SubstF,
+                                                   PAI->getNumArguments(),
+                                                   GenericSpecialized)) {
+    // `partial_apply` can be rewritten to `thin_to_thick_function` only if the
+    // specialized callee is `@convention(thin)`.
+    if (NewFunc->getRepresentation() == SILFunctionTypeRepresentation::Thin) {
+      rewritePartialApply(PAI, NewFunc);
+      if (GenericSpecialized.first) {
+        // Notify the pass manager about the new function.
+        addFunctionToPassManagerWorklist(GenericSpecialized.first,
+                                         GenericSpecialized.second);
+      }
+      return true;
     }
-    return true;
   }
 
   // Second possibility: Are all partially applied arguments constant?

test/AutoDiff/compiler_crashers_fixed/sr12732-optimize-partial-apply-convention-thin-only.swift

@@ -0,0 +1,22 @@
+// RUN: %target-build-swift -Osize %s
+// REQUIRES: asserts
+
+// SR-12732: Fix `partial_apply` optimization.
+
+// Do not rewrite `partial_apply` to `thin_to_thick_function` if the specialized
+// callee is not `@convention(thin)`.
+
+import DifferentiationUnittest
+
+func callback(_ x: inout Tracked<Float>.TangentVector) {}
+
+@differentiable
+func caller(_ x: Tracked<Float>) -> Tracked<Float> {
+  return x.withDerivative(callback)
+}
+
+// SIL verification failed: operand of thin_to_thick_function must be thin: opFTy->getRepresentation() == SILFunctionType::Representation::Thin
+// Verifying instruction:
+//      // function_ref specialized Differentiable._vjpWithDerivative(_:)
+//   %10 = function_ref @$s16_Differentiation14DifferentiablePAAE18_vjpWithDerivativeyx5value_13TangentVectorQzAGc8pullbacktyAGzcF0A8Unittest7TrackedVySfG_Tg5 : $@convention(method) (@guaranteed @callee_guaranteed @substituted <τ_0_0> (@inout τ_0_0) -> () for <Tracked<Float>>, @in_guaranteed Tracked<Float>) -> (@out Tracked<Float>, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> @out τ_0_1 for <Tracked<Float>, Tracked<Float>>) // user: %11
+// ->   %11 = thin_to_thick_function %10 : $@convention(method) (@guaranteed @callee_guaranteed @substituted <τ_0_0> (@inout τ_0_0) -> () for <Tracked<Float>>, @in_guaranteed Tracked<Float>) -> (@out Tracked<Float>, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> @out τ_0_1 for <Tracked<Float>, Tracked<Float>>) to $@callee_guaranteed (@guaranteed @callee_guaranteed @substituted <τ_0_0> (@inout τ_0_0) -> () for <Tracked<Float>>, @in_guaranteed Tracked<Float>) -> (@out Tracked<Float>, @owned @callee_guaranteed @substituted <τ_0_0, τ_0_1> (@in_guaranteed τ_0_0) -> @out τ_0_1 for <Tracked<Float>, Tracked<Float>>) // user: %12

test/AutoDiff/stdlib/derivative_customization.swift

@@ -1,8 +1,6 @@
 // RUN: %target-run-simple-swift
 // REQUIRES: executable_test
 
-// REQUIRES: SR12732
-
 import DifferentiationUnittest
 import StdlibUnittest
 

test/AutoDiff/validation-test/custom_derivatives.swift

@@ -1,8 +1,6 @@
 // RUN: %target-run-simple-swift
 // REQUIRES: executable_test
 
-// REQUIRES: SR12732
-
 import StdlibUnittest
 #if os(macOS) || os(iOS) || os(watchOS) || os(tvOS)
 import Darwin.C

test/AutoDiff/validation-test/derivative_registration.swift

@@ -1,8 +1,6 @@
 // RUN: %target-run-simple-swift
 // REQUIRES: executable_test
 
-// REQUIRES: SR12732
-
 import StdlibUnittest
 import DifferentiationUnittest