Merge pull request #4457 from eeckstein/partial-apply-opt

Add an optimization to eliminate a partial_apply if all applied arguments are dead in the applied function.

Author: eeckstein

include/swift/SILOptimizer/Analysis/CallerAnalysis.h

@@ -19,19 +19,11 @@
 #include "swift/SIL/SILModule.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/TinyPtrVector.h"
 
 namespace swift {
 
-/// NOTE: this can be extended to contain the callsites of the function.
-struct CallerAnalysisFunctionInfo {
-  /// A list of all the functions this function calls.
-  llvm::SetVector<SILFunction *> Callees;
-  /// A list of all the callers this function has.
-  llvm::SetVector<SILFunction *> Callers;
-};
-
 /// CallerAnalysis relies on keeping the Caller/Callee relation up-to-date
 /// lazily. i.e. when a function is invalidated, instead of recomputing the
 /// function it calls right away, its kept in a recompute list and
@@ -44,13 +36,57 @@ struct CallerAnalysisFunctionInfo {
 /// to run every function through the a sequence of function passes which might
 /// invalidate the functions and make the computed list incomplete. So
 /// O(n) * O(n) = O(n^2).
+///
+/// In addition of caller information this analysis also provides information
+/// about partial applies of a function.
 class CallerAnalysis : public SILAnalysis {
 
+public:
+
+  /// NOTE: this can be extended to contain the callsites of the function.
+  class FunctionInfo {
+    friend class CallerAnalysis;
+
+    /// A list of all the functions this function calls or partially applies.
+    llvm::SetVector<SILFunction *> Callees;
+    /// A list of all the callers this function has.
+    llvm::SmallSet<SILFunction *, 4> Callers;
+
+    /// The number of partial applied arguments of this function.
+    /// Specifically, it stores the minimum number of partial applied arguments
+    /// of each function which contain one or multiple partial_applys of this
+    /// function.
+    /// This is a little bit off-topic because a partial_apply is not really
+    /// a "call" of this function.
+    llvm::DenseMap<SILFunction *, int> PartialAppliers;
+
+  public:
+    /// Returns true if this function has at least one caller.
+    bool hasCaller() const {
+      return !Callers.empty();
+    }
+
+    /// Returns non zero if this function is partially applied anywhere.
+    /// The return value is the minimum number of partially applied arguments.
+    /// Usually all partial applies of a function partially apply the same
+    /// number of arguments anyway.
+    int getMinPartialAppliedArgs() const {
+      int minArgs = 0;
+      for (auto Iter : PartialAppliers) {
+        int numArgs = Iter.second;
+        if (minArgs == 0 || numArgs < minArgs)
+          minArgs = numArgs;
+      }
+      return minArgs;
+    }
+  };
+  
+private:
   /// Current module we are analyzing.
   SILModule &Mod;
 
   /// A map between all the functions and their callsites in the module.
-  llvm::DenseMap<SILFunction *, CallerAnalysisFunctionInfo> CallInfo;
+  llvm::DenseMap<SILFunction *, FunctionInfo> FuncInfos;
 
   /// A list of functions that needs to be recomputed.
   llvm::SetVector<SILFunction *> RecomputeFunctionList;
@@ -74,7 +110,7 @@ class CallerAnalysis : public SILAnalysis {
   CallerAnalysis(SILModule *M) : SILAnalysis(AnalysisKind::Caller), Mod(*M) {
     // Make sure we compute everything first time called.
     for (auto &F : Mod) {
-      CallInfo.FindAndConstruct(&F);
+      FuncInfos.FindAndConstruct(&F);
       RecomputeFunctionList.insert(&F);
     }
   }
@@ -89,7 +125,7 @@ class CallerAnalysis : public SILAnalysis {
 
   virtual void invalidate(SILFunction *F, InvalidationKind K) {
     // Should we invalidate based on the invalidation kind.
-    bool shouldInvalidate = K & InvalidationKind::Calls;
+    bool shouldInvalidate = K & InvalidationKind::CallsAndInstructions;
     if (!shouldInvalidate)
       return;
 
@@ -111,20 +147,18 @@ class CallerAnalysis : public SILAnalysis {
     if (!shouldInvalidate)
       return;
 
-    CallInfo.clear();
+    FuncInfos.clear();
     RecomputeFunctionList.clear();
     for (auto &F : Mod) {
       RecomputeFunctionList.insert(&F);
     }
   }
 
-  /// Return true if the function has a caller inside current module.
-  bool hasCaller(SILFunction *F) {
+  const FunctionInfo &getCallerInfo(SILFunction *F) {
     // Recompute every function in the invalidated function list and empty the
     // list.
     processRecomputeFunctionList();
-    auto Iter = CallInfo.FindAndConstruct(F);
-    return !Iter.second.Callers.empty();
+    return FuncInfos[F];
   }
 };
 

include/swift/SILOptimizer/PassManager/Passes.def

@@ -81,6 +81,8 @@ PASS(RedundantOverflowCheckRemoval, "remove-redundant-overflow-checks",
      "Removes redundant overflow checks")
 PASS(DCE, "dce",
      "Eliminate dead code")
+PASS(DeadArgSignatureOpt, "dead-arg-signature-opt",
+     "Create function with removed dead arguments")
 PASS(DeadFunctionElimination, "sil-deadfuncelim",
      "Remove unused functions")
 PASS(DeadObjectElimination, "deadobject-elim",

lib/SILOptimizer/Analysis/CallerAnalysis.cpp

@@ -28,25 +28,42 @@ void CallerAnalysis::processFunctionCallSites(SILFunction *F) {
           continue;
 
         // Update the callee information for this function.
-        CallerAnalysisFunctionInfo &CallerInfo
-                               = CallInfo.FindAndConstruct(F).second;
+        FunctionInfo &CallerInfo = FuncInfos[F];
         CallerInfo.Callees.insert(CalleeFn);
 
         // Update the callsite information for the callee.
-        CallerAnalysisFunctionInfo &CalleeInfo
-                               = CallInfo.FindAndConstruct(CalleeFn).second;
+        FunctionInfo &CalleeInfo = FuncInfos[CalleeFn];
         CalleeInfo.Callers.insert(F);
-      }   
+        continue;
+      }
+      if (auto *PAI = dyn_cast<PartialApplyInst>(&II)) {
+        SILFunction *CalleeFn = PAI->getCalleeFunction();
+        if (!CalleeFn)
+          continue;
+
+        // Update the callee information for this function.
+        FunctionInfo &CallerInfo = FuncInfos[F];
+        CallerInfo.Callees.insert(CalleeFn);
+        
+        // Update the partial-apply information for the callee.
+        FunctionInfo &CalleeInfo = FuncInfos[CalleeFn];
+        int &minAppliedArgs = CalleeInfo.PartialAppliers[F];
+        int numArgs = (int)PAI->getNumArguments();
+        if (minAppliedArgs == 0 || numArgs < minAppliedArgs) {
+          minAppliedArgs = numArgs;
+        }
+        continue;
+      }
     }   
   }   
 }
 
 void CallerAnalysis::invalidateExistingCalleeRelation(SILFunction *F) {
-  CallerAnalysisFunctionInfo &CallerInfo = CallInfo.FindAndConstruct(F).second;
+  FunctionInfo &CallerInfo = FuncInfos[F];
   for (auto Callee : CallerInfo.Callees) {
-    CallerAnalysisFunctionInfo &CalleeInfo
-                                    = CallInfo.FindAndConstruct(Callee).second;
-    CalleeInfo.Callers.remove(F);
+    FunctionInfo &CalleeInfo = FuncInfos[Callee];
+    CalleeInfo.Callers.erase(F);
+    CalleeInfo.PartialAppliers.erase(F);
   }
 }
 

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -31,6 +31,8 @@ STATISTIC(NumCapturesPropagated, "Number of constant captures propagated");
 namespace {
 /// Propagate constants through closure captures by specializing the partially
 /// applied function.
+/// Also optimize away partial_apply instructions where all partially applied
+/// arguments are dead.
 class CapturePropagation : public SILFunctionTransform
 {
 public:
@@ -258,15 +260,11 @@ void CapturePropagation::rewritePartialApply(PartialApplyInst *OrigPAI,
                                              SILFunction *SpecialF) {
   SILBuilderWithScope Builder(OrigPAI);
   auto FuncRef = Builder.createFunctionRef(OrigPAI->getLoc(), SpecialF);
-  auto NewPAI = Builder.createPartialApply(OrigPAI->getLoc(),
-                                           FuncRef,
-                                           SpecialF->getLoweredType(),
-                                           ArrayRef<Substitution>(),
-                                           ArrayRef<SILValue>(),
-                                           OrigPAI->getType());
-  OrigPAI->replaceAllUsesWith(NewPAI);
+  auto *T2TF = Builder.createThinToThickFunction(OrigPAI->getLoc(),
+                                                 FuncRef, OrigPAI->getType());
+  OrigPAI->replaceAllUsesWith(T2TF);
   recursivelyDeleteTriviallyDeadInstructions(OrigPAI, true);
-  DEBUG(llvm::dbgs() << "  Rewrote caller:\n" << *NewPAI);
+  DEBUG(llvm::dbgs() << "  Rewrote caller:\n" << *T2TF);
 }
 
 /// For now, we conservative only specialize if doing so can eliminate dynamic
@@ -286,6 +284,92 @@ static bool isProfitable(SILFunction *Callee) {
   return false;
 }
 
+/// Returns true if block \p BB only contains a return or throw of the first
+/// block argument and side-effect-free instructions.
+static bool isArgReturnOrThrow(SILBasicBlock *BB) {
+  for (SILInstruction &I : *BB) {
+    if (isa<ReturnInst>(&I) || isa<ThrowInst>(&I)) {
+      SILValue RetVal = I.getOperand(0);
+      if (BB->getNumBBArg() == 1 && RetVal == BB->getBBArg(0))
+        return true;
+      return false;
+    }
+    if (I.mayHaveSideEffects() || isa<TermInst>(&I))
+      return false;
+  }
+  llvm_unreachable("should have seen a terminator instruction");
+}
+
+/// Checks if \p Orig is a thunk which calls another function but without
+/// passing the trailing \p numDeadParams dead parameters.
+static SILFunction *getSpecializedWithDeadParams(SILFunction *Orig,
+                                                 int numDeadParams) {
+  SILBasicBlock &EntryBB = *Orig->begin();
+  unsigned NumArgs = EntryBB.getNumBBArg();
+  SILModule &M = Orig->getModule();
+  
+  // Check if all dead parameters have trivial types. We don't support non-
+  // trivial types because it's very hard to find places where we can release
+  // those parameters (as a replacement for the removed partial_apply).
+  // TODO: maybe we can skip this restrication when we have semantic ARC.
+  for (unsigned Idx = NumArgs - numDeadParams; Idx < NumArgs; ++Idx) {
+    SILType ArgTy = EntryBB.getBBArg(Idx)->getType();
+    if (!ArgTy.isTrivial(M))
+      return nullptr;
+  }
+  SILFunction *Specialized = nullptr;
+  SILValue RetValue;
+  
+  // Check all instruction of the entry block.
+  for (SILInstruction &I : EntryBB) {
+    if (auto FAS = FullApplySite::isa(&I)) {
+      
+      // Check if this is the call of the specialized function.
+      // As the original function is not generic, also the specialized function
+      // must be not generic.
+      if (FAS.hasSubstitutions())
+        return nullptr;
+      // Is it the only call?
+      if (Specialized)
+        return nullptr;
+      
+      Specialized = FAS.getReferencedFunction();
+      if (!Specialized)
+        return nullptr;
+
+      // Check if parameters are passes 1-to-1
+      unsigned NumArgs = FAS.getNumArguments();
+      if (EntryBB.getNumBBArg() - numDeadParams != NumArgs)
+        return nullptr;
+
+      for (unsigned Idx = 0; Idx < NumArgs; ++Idx) {
+        if (FAS.getArgument(Idx) != (ValueBase *)EntryBB.getBBArg(Idx))
+          return nullptr;
+      }
+
+      if (TryApplyInst *TAI = dyn_cast<TryApplyInst>(&I)) {
+        // Check the normal and throw blocks of the try_apply.
+        if (isArgReturnOrThrow(TAI->getNormalBB()) &&
+            isArgReturnOrThrow(TAI->getErrorBB()))
+          return Specialized;
+        return nullptr;
+      }
+      assert(isa<ApplyInst>(&I) && "unknown FullApplySite instruction");
+      RetValue = &I;
+      continue;
+    }
+    if (auto *RI = dyn_cast<ReturnInst>(&I)) {
+      // Check if we return the result of the apply.
+      if (RI->getOperand() != RetValue)
+        return nullptr;
+      continue;
+    }
+    if (I.mayHaveSideEffects() || isa<TermInst>(&I))
+      return nullptr;
+  }
+  return Specialized;
+}
+
 bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
   // Check if the partial_apply has generic substitutions.
   // FIXME: We could handle generic thunks if it's worthwhile.
@@ -295,15 +379,26 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
   SILFunction *SubstF = PAI->getReferencedFunction();
   if (!SubstF)
     return false;
+  if (SubstF->isExternalDeclaration())
+    return false;
 
   assert(!SubstF->getLoweredFunctionType()->isPolymorphic() &&
          "cannot specialize generic partial apply");
 
+  // First possibility: Is it a partial_apply where all partially applied
+  // arguments are dead?
+  if (SILFunction *NewFunc = getSpecializedWithDeadParams(SubstF,
+                                                    PAI->getNumArguments())) {
+    rewritePartialApply(PAI, NewFunc);
+    return true;
+  }
+
+  // Second possibility: Are all partially applied arguments constant?
   for (auto Arg : PAI->getArguments()) {
     if (!isConstant(Arg))
       return false;
   }
-  if (SubstF->isExternalDeclaration() || !isProfitable(SubstF))
+  if (!isProfitable(SubstF))
     return false;
 
   DEBUG(llvm::dbgs() << "Specializing closure for constant arguments:\n"

lib/SILOptimizer/PassManager/Passes.cpp

@@ -307,6 +307,11 @@ void swift::runSILOptimizationPasses(SILModule &Module) {
   // Run an iteration of the mid-level SSA passes.
   PM.setStageName("MidLevel");
   AddSSAPasses(PM, OptimizationLevelKind::MidLevel);
+  
+  // Specialy partially applied functions with dead arguments as a preparation
+  // for CapturePropagation.
+  PM.addDeadArgSignatureOpt();
+
   PM.runOneIteration();
   PM.resetAndRemoveTransformations();