Merge pull request #18077 from eeckstein/closure-specializer

ClosureSpecializer: extend the benefit-analysis so that specialization is also done in case a closure is passed through multiple call-levels.

Author: eeckstein

include/swift/SIL/SILInstruction.h

@@ -7683,11 +7683,17 @@ class ApplySite {
     }
   }
 
+  /// Returns true if \p oper is an argument operand and not the callee
+  /// operand.
+  bool isArgumentOperand(const Operand &oper) {
+    return oper.getOperandNumber() >= getOperandIndexOfFirstArgument();
+  }
+
   // Translate the index of the argument to the full apply or partial_apply into
   // to the corresponding index into the arguments of the called function.
   unsigned getCalleeArgIndex(const Operand &oper) {
     assert(oper.getUser() == Inst);
-    assert(oper.getOperandNumber() >= getOperandIndexOfFirstArgument());
+    assert(isArgumentOperand(oper));
 
     unsigned appliedArgIdx =
         oper.getOperandNumber() - getOperandIndexOfFirstArgument();

lib/SILOptimizer/IPO/ClosureSpecializer.cpp

@@ -908,6 +908,37 @@ static void markReabstractionPartialApplyAsUsed(
   llvm_unreachable("Unexpect instruction");
 }
 
+/// Returns true if the \p closureArgIdx argument of \p callee is called in
+/// \p callee or any function called by callee.
+static bool isClosureAppliedIn(SILFunction *Callee, unsigned closureArgIdx,
+                               SmallPtrSetImpl<SILFunction *> &HandledFuncs) {
+  // Limit the number of recursive calls to not go into exponential behavior in
+  // corner cases.
+  const int RecursionBudget = 8;
+
+  SILValue Arg = Callee->getArgument(closureArgIdx);
+  for (Operand *ArgUse : Arg->getUses()) {
+    if (auto UserAI = FullApplySite::isa(ArgUse->getUser())) {
+      if (UserAI.getCallee() == Arg)
+        return true;
+
+      assert(UserAI.isArgumentOperand(*ArgUse) &&
+             "any other non-argument operands than the callee?");
+
+      SILFunction *ApplyCallee = UserAI.getReferencedFunction();
+      if (ApplyCallee && !ApplyCallee->isExternalDeclaration() &&
+          HandledFuncs.count(ApplyCallee) == 0 &&
+          HandledFuncs.size() < RecursionBudget) {
+        HandledFuncs.insert(ApplyCallee);
+        if (isClosureAppliedIn(UserAI.getReferencedFunction(),
+                               UserAI.getCalleeArgIndex(*ArgUse), HandledFuncs))
+          return true;
+      }
+    }
+  }
+  return false;
+}
+
 bool SILClosureSpecializerTransform::gatherCallSites(
     SILFunction *Caller,
     llvm::SmallVectorImpl<ClosureInfo*> &ClosureCandidates,
@@ -1018,43 +1049,23 @@ bool SILClosureSpecializerTransform::gatherCallSites(
         if (mayBindDynamicSelf(ApplyCallee))
           return CFGChanged;
 
-        // Ok, we know that we can perform the optimization but not whether or
-        // not the optimization is profitable. Find the index of the argument
-        // corresponding to our partial apply.
-        Optional<unsigned> ClosureIndex;
-        for (unsigned i = 0, e = AI.getNumArguments(); i != e; ++i) {
-          if (AI.getArgument(i) != Use->get())
-            continue;
-          ClosureIndex = i;
-          DEBUG(llvm::dbgs() << "    Found callsite with closure argument at "
-                << i << ": " << *AI.getInstruction());
-          break;
-        }
-
-        // If we did not find an index, there is nothing further to do,
-        // continue.
-        if (!ClosureIndex.hasValue())
+        // Check if the closure is passed as an argument (and not called).
+        if (!AI.isArgumentOperand(*Use))
           continue;
 
-        // Make sure that the Closure is invoked in the Apply's callee. We only
-        // want to perform closure specialization if we know that we will be
-        // able to change a partial_apply into an apply.
-        //
-        // TODO: Maybe just call the function directly instead of moving the
-        // partial apply?
-        SILValue Arg = ApplyCallee->getArgument(ClosureIndex.getValue());
-        if (std::none_of(Arg->use_begin(), Arg->use_end(),
-                         [&Arg](Operand *Op) -> bool {
-                           auto UserAI = FullApplySite::isa(Op->getUser());
-                           return UserAI && UserAI.getCallee() == Arg;
-                         })) {
+        unsigned ClosureIndex = AI.getCalleeArgIndex(*Use);
+
+        // Ok, we know that we can perform the optimization but not whether or
+        // not the optimization is profitable. Check if the closure is actually
+        // called in the callee (or in a function called by the callee).
+        SmallPtrSet<SILFunction *, 8> HandledFuncs;
+        if (!isClosureAppliedIn(ApplyCallee, ClosureIndex, HandledFuncs))
           continue;
-        }
 
         unsigned firstParamArgIdx =
             AI.getSubstCalleeConv().getSILArgIndexOfFirstParam();
-        assert(ClosureIndex.getValue() >= firstParamArgIdx);
-        auto ClosureParamIndex = ClosureIndex.getValue() - firstParamArgIdx;
+        assert(ClosureIndex >= firstParamArgIdx);
+        auto ClosureParamIndex = ClosureIndex - firstParamArgIdx;
 
         auto ParamInfo = AI.getSubstCalleeType()->getParameters();
         SILParameterInfo ClosureParamInfo = ParamInfo[ClosureParamIndex];
@@ -1097,7 +1108,7 @@ bool SILClosureSpecializerTransform::gatherCallSites(
         // Now we know that CSDesc is profitable to specialize. Add it to our
         // call site list.
         CInfo->CallSites.push_back(
-          CallSiteDescriptor(CInfo, AI, ClosureIndex.getValue(),
+          CallSiteDescriptor(CInfo, AI, ClosureIndex,
                              ClosureParamInfo, std::move(NonFailureExitBBs)));
       }
       if (CInfo) {

test/SILOptimizer/closure_specialize_simple.sil

@@ -26,6 +26,23 @@ bb2:
   return %2 : $Builtin.Int1
 }
 
+// CHECK-LABEL: sil shared @$S37simple_partial_apply_2nd_level_caller0a1_b1_C4_funBi1_Tf1c_n : $@convention(thin) (Builtin.Int1) -> Builtin.Int1 {
+// CHECK: bb0([[CAPTURED_ARG:%.*]] : $Builtin.Int1):
+// CHECK: [[SPECIALIZED_CALLEE:%.*]] = function_ref @$S27simple_partial_apply_caller0a1_b1_C4_funBi1_Tf1c_n  :
+// CHECK: [[RET:%.*]]= apply [[SPECIALIZED_CALLEE]]([[CAPTURED_ARG]])
+// CHECK: return [[RET]]
+sil @simple_partial_apply_2nd_level_caller : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1 {
+bb0(%0 : $@callee_owned (Builtin.Int1) -> Builtin.Int1):
+  br bb1
+
+bb1:
+  %1 = function_ref @simple_partial_apply_caller : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
+  %2 = apply %1(%0) : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
+  cond_br undef, bb1, bb2
+
+bb2:
+  return %2 : $Builtin.Int1
+}
 sil @simple_partial_apply_caller_decl : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
 
 sil @simple_multiple_partial_apply_caller : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1, @owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1 {
@@ -177,8 +194,10 @@ bb2:
 }
 
 // CHECK-LABEL: sil @loop_driver : $@convention(thin) (Builtin.Int1, Builtin.Int1) -> () {
-// CHECK: [[SPECIALIZED_FUN:%.*]] = function_ref @$S27simple_partial_apply_caller0a1_b1_C4_funBi1_Tf1c_n : $@convention(thin) (Builtin.Int1) -> Builtin.Int1
+// CHECK-DAG: [[SPECIALIZED_FUN:%.*]] = function_ref @$S27simple_partial_apply_caller0a1_b1_C4_funBi1_Tf1c_n : $@convention(thin) (Builtin.Int1) -> Builtin.Int1
+// CHECK-DAG: [[SPECIALIZED_FUN2:%.*]] = function_ref @$S37simple_partial_apply_2nd_level_caller0a1_b1_C4_funBi1_Tf1c_n : $@convention(thin) (Builtin.Int1) -> Builtin.Int1
 // CHECK: apply [[SPECIALIZED_FUN]]
+// CHECK: apply [[SPECIALIZED_FUN2]]
 
 // We can't call this one b/c it is just a declaration.
 // CHECK: [[UNSPECIALIZED_FUN_DECL:%.*]] = function_ref @simple_partial_apply_caller_decl : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
@@ -218,6 +237,9 @@ bb0(%0 : $Builtin.Int1, %1 : $Builtin.Int1):
   %4 = function_ref @simple_partial_apply_caller : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
   %5 = apply %4(%3) : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
 
+  %51 = function_ref @simple_partial_apply_2nd_level_caller : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
+  %52 = apply %51(%3) : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
+
   %6 = partial_apply %2(%0) : $@convention(thin) (Builtin.Int1, Builtin.Int1) -> Builtin.Int1
   %7 = function_ref @simple_partial_apply_caller_decl : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1
   %8 = apply %7(%6) : $@convention(thin) (@owned @callee_owned (Builtin.Int1) -> Builtin.Int1) -> Builtin.Int1