SILCombine: handle convert_escape_to_noescape in the apply-of-convert-function optimization

lib/SILOptimizer/SILCombiner/SILCombinerApplyVisitors.cpp

@@ -135,8 +135,9 @@ SILCombiner::optimizeApplyOfConvertFunctionInst(FullApplySite AI,
   if (auto *TTI = dyn_cast<ThinToThickFunctionInst>(funcOper))
     funcOper = TTI->getOperand();
 
-  auto *FRI = dyn_cast<FunctionRefInst>(funcOper);
-  if (!FRI)
+  if (!isa<FunctionRefInst>(funcOper) &&
+      // Optimizing partial_apply will then enable the partial_apply -> apply peephole.
+      !isa<PartialApplyInst>(funcOper))
     return nullptr;
 
   // Grab our relevant callee types...
@@ -151,8 +152,8 @@ SILCombiner::optimizeApplyOfConvertFunctionInst(FullApplySite AI,
   // relevant types from the ConvertFunction function type and AI.
   Builder.setCurrentDebugScope(AI.getDebugScope());
   OperandValueArrayRef Ops = AI.getArguments();
-  SILFunctionConventions substConventions(SubstCalleeTy, FRI->getModule());
-  SILFunctionConventions convertConventions(ConvertCalleeTy, FRI->getModule());
+  SILFunctionConventions substConventions(SubstCalleeTy, CFI->getModule());
+  SILFunctionConventions convertConventions(ConvertCalleeTy, CFI->getModule());
   auto context = AI.getFunction()->getTypeExpansionContext();
   auto oldOpRetTypes = substConventions.getIndirectSILResultTypes(context);
   auto newOpRetTypes = convertConventions.getIndirectSILResultTypes(context);
@@ -229,7 +230,7 @@ SILCombiner::optimizeApplyOfConvertFunctionInst(FullApplySite AI,
       Builder.createBranch(AI.getLoc(), TAI->getNormalBB(), branchArgs);
     }
 
-    return Builder.createTryApply(AI.getLoc(), FRI, SubstitutionMap(), Args,
+    return Builder.createTryApply(AI.getLoc(), funcOper, SubstitutionMap(), Args,
                                   normalBB, TAI->getErrorBB(),
                                   TAI->getApplyOptions());
   }
@@ -239,9 +240,9 @@ SILCombiner::optimizeApplyOfConvertFunctionInst(FullApplySite AI,
   // otherwise, we would be creating malformed SIL).
   ApplyOptions Options = AI.getApplyOptions();
   Options -= ApplyFlags::DoesNotThrow;
-  if (FRI->getFunctionType()->hasErrorResult())
+  if (funcOper->getType().castTo<SILFunctionType>()->hasErrorResult())
     Options |= ApplyFlags::DoesNotThrow;
-  ApplyInst *NAI = Builder.createApply(AI.getLoc(), FRI, SubstitutionMap(),
+  ApplyInst *NAI = Builder.createApply(AI.getLoc(), funcOper, SubstitutionMap(),
                                        Args, Options);
   SILInstruction *result = NAI;
 
@@ -1455,7 +1456,11 @@ SILInstruction *SILCombiner::visitApplyInst(ApplyInst *AI) {
   if (isa<PartialApplyInst>(AI->getCallee()))
     return nullptr;
 
-  if (auto *CFI = dyn_cast<ConvertFunctionInst>(AI->getCallee()))
+  SILValue callee = AI->getCallee();
+  if (auto *cee = dyn_cast<ConvertEscapeToNoEscapeInst>(callee)) {
+    callee = cee->getOperand();
+  }
+  if (auto *CFI = dyn_cast<ConvertFunctionInst>(callee))
     return optimizeApplyOfConvertFunctionInst(AI, CFI);
 
   if (tryOptimizeKeypath(AI))

lib/SILOptimizer/SILCombiner/SILCombinerMiscVisitors.cpp

@@ -2367,6 +2367,11 @@ SILCombiner::visitDifferentiableFunctionExtractInst(DifferentiableFunctionExtrac
   // match the type of the original `differentiable_function_extract`,
   // create a `convert_function`.
   if (newValue->getType() != DFEI->getType()) {
+    CanSILFunctionType opTI = newValue->getType().castTo<SILFunctionType>();
+    CanSILFunctionType resTI = DFEI->getType().castTo<SILFunctionType>();
+    if (!opTI->isABICompatibleWith(resTI, *DFEI->getFunction()).isCompatible())
+      return nullptr;
+
     std::tie(newValue, std::ignore) =
       castValueToABICompatibleType(&Builder, DFEI->getLoc(),
                                    newValue,

test/AutoDiff/e2e_optimizations.swift

@@ -87,14 +87,10 @@ func test_gradient_float_loop() {
 }
 
 // Check whether `apply`s are inlined.
-// Currently, the VJP is inlined but the pullback is not.
 // CHECK-LABEL: sil hidden @test_gradient_float_loop : $@convention(thin) () -> ()
-// CHECK: [[PB_FN_REF:%.*]] = function_ref @{{.*}}24test_gradient_float_loopyyFS2fcfU_TJrSpSr : $@convention(thin) (Float) -> (Float, @owned @callee_guaranteed (Float) -> Float)
-// CHECK: [[GRADIENT_RESULT:%.*]] = apply [[PB_FN_REF]]
-// CHECK: [[EXTRACT:%.*]] = tuple_extract [[GRADIENT_RESULT]]
-// CHECK: [[GRADIENT_RESULT2:%.*]] = apply [[EXTRACT]]
+// CHECK:  = function_ref @${{.*24test_gradient_float_loopyyFS2fcfU_TJrSpSr|sSf16_DifferentiationE12_vjpMultiply3lhs3rhsSf5value_Sf_SftSfc8pullbacktSf_SftFZSf_SftSfcfU_}}
 // CHECK: [[FN_REF:%.*]] = function_ref @$s9blackHoleSf_Tg5 : $@convention(thin) (Float) -> Float
-// CHECK-NEXT: apply [[FN_REF:%.*]]([[GRADIENT_RESULT2]])
+// CHECK-NEXT: apply [[FN_REF:%.*]]
 // CHECK-NOT: apply
 // CHECK-LABEL: } // end sil function 'test_gradient_float_loop'
 func array_loop(_ array: [Float]) -> Float {

test/SILOptimizer/dictionary_lookup_with_default.swift

@@ -0,0 +1,19 @@
+// RUN: %target-swift-frontend  %s -O -module-name=test -emit-sil | %FileCheck %s
+
+// REQUIRES: swift_stdlib_no_asserts,optimized_stdlib
+
+public struct S {
+    var x: Int
+
+    @inline(never)
+    mutating func doSomething() { }
+}
+
+// Check that all partial_applys can be optimized away so that no closure context needs to be allocated.
+
+// CHECK-LABEL: sil @$s4test6testit_1xySDySiAA1SVGz_SitF :
+// CHECK-NOT:     partial_apply
+// CHECK:       } // end sil function '$s4test6testit_1xySDySiAA1SVGz_SitF'
+public func testit(_ data: inout [Int: S], x: Int) {
+  data[x, default: S(x: x)].doSomething()
+}

test/SILOptimizer/sil_combine.sil

@@ -1442,6 +1442,28 @@ entry(%a : $AnotherClass):
   return %r : $MyNSObj
 }
 
+sil @createMyInt : $@convention(thin) (Builtin.Int32) -> @out MyInt
+
+// CHECK-LABEL: sil @convert_function_of_closure :
+// CHECK:         [[F:%.*]] = function_ref @createMyInt 
+// CHECK:         [[S:%.*]] = alloc_stack $MyInt
+// CHECK:         apply [[F]]([[S]], %0)
+// CHECK-NOT:     strong_release
+// CHECK:       } // end sil function 'convert_function_of_closure'
+sil @convert_function_of_closure : $@convention(thin) (Builtin.Int32) -> () {
+bb0(%0 : $Builtin.Int32):
+  %1 = function_ref @createMyInt : $@convention(thin) (Builtin.Int32) -> @out MyInt
+  %2 = partial_apply [callee_guaranteed] %1(%0) : $@convention(thin) (Builtin.Int32) -> @out MyInt
+  %3 = convert_function %2 : $@callee_guaranteed () -> @out MyInt to $@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <MyInt>
+  %4 = convert_escape_to_noescape %3 : $@callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <MyInt> to $@noescape @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <MyInt>
+  %5 = alloc_stack $MyInt
+  %6 = apply %4(%5) : $@noescape @callee_guaranteed @substituted <τ_0_0> () -> @out τ_0_0 for <MyInt>
+  dealloc_stack %5 : $*MyInt
+  strong_release %2 : $@callee_guaranteed () -> @out MyInt
+  %9 = tuple ()
+  return %9 : $()
+}
+
 // CHECK-LABEL: sil @upcast_formation : $@convention(thin) (@inout E, E, @inout B) -> B {
 // CHECK: bb0
 // CHECK-NEXT: upcast