CapturePropagation: specialize closures which capture a constant keypath

This optimizes keypath-closures, like
```
   a.map { \.x }
```
It results in a significant performance improvement for such code patterns.

rdar://87968067

Author: eeckstein

include/swift/SIL/SILInstruction.h

@@ -78,6 +78,7 @@ class StringLiteralExpr;
 class ValueDecl;
 class VarDecl;
 class FunctionRefBaseInst;
+class SILPrintContext;
 
 template <typename ImplClass> class SILClonerWithScopes;
 
@@ -3489,7 +3490,9 @@ class KeyPathPatternComponent {
 
   void incrementRefCounts() const;
   void decrementRefCounts() const;
-  
+
+  void print(SILPrintContext &ctxt) const;
+
   void Profile(llvm::FoldingSetNodeID &ID);
 };
 

include/swift/SILOptimizer/Utils/SpecializationMangler.h

@@ -95,7 +95,9 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
   FunctionSignatureSpecializationMangler(SpecializationPass Pass,
                                          IsSerialized_t Serialized,
                                          SILFunction *F);
-  void setArgumentConstantProp(unsigned OrigArgIdx, LiteralInst *LI);
+  void setArgumentConstantProp(unsigned OrigArgIdx, SILInstruction *constInst);
+  void appendStringAsIdentifier(StringRef str);
+
   void setArgumentClosureProp(unsigned OrigArgIdx, PartialApplyInst *PAI);
   void setArgumentClosureProp(unsigned OrigArgIdx,
                               ThinToThickFunctionInst *TTTFI);
@@ -112,7 +114,7 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
   std::string mangle();
 
 private:
-  void mangleConstantProp(LiteralInst *LI);
+  void mangleConstantProp(SILInstruction *constInst);
   void mangleClosureProp(SILInstruction *Inst);
   void mangleArgument(ArgumentModifierIntBase ArgMod,
                       NullablePtr<SILInstruction> Inst);

lib/SIL/IR/SILPrinter.cpp

@@ -3855,6 +3855,11 @@ void SILSpecializeAttr::print(llvm::raw_ostream &OS) const {
   }
 }
 
+void KeyPathPatternComponent::print(SILPrintContext &ctxt) const {
+  SILPrinter printer(ctxt);
+  printer.printKeyPathPatternComponent(*this);
+}
+
 //===----------------------------------------------------------------------===//
 // SILPrintContext members
 //===----------------------------------------------------------------------===//

lib/SILOptimizer/IPO/CapturePropagation.cpp

@@ -50,25 +50,38 @@ class CapturePropagation : public SILFunctionTransform
 };
 } // end anonymous namespace
 
-static LiteralInst *getConstant(SILValue V) {
+static SILInstruction *getConstant(SILValue V) {
   if (auto I = dyn_cast<ThinToThickFunctionInst>(V))
     return getConstant(I->getOperand());
   if (auto I = dyn_cast<ConvertFunctionInst>(V))
     return getConstant(I->getOperand());
-  return dyn_cast<LiteralInst>(V);
-}
 
-static bool isOptimizableConstant(SILValue V) {
-  // We do not optimize string literals of length > 32 since we would need to
-  // encode them into the symbol name for uniqueness.
-  if (auto *SLI = dyn_cast<StringLiteralInst>(V))
-    return SLI->getValue().size() <= 32;
-  return true;
-}
+  if (auto *SLI = dyn_cast<StringLiteralInst>(V)) {
+    // We do not optimize string literals of length > 32 since we would need to
+    // encode them into the symbol name for uniqueness.
+    if (SLI->getValue().size() > 32)
+      return nullptr;
+    return SLI;
+  }
+
+  if (auto *lit = dyn_cast<LiteralInst>(V))
+    return lit;
 
-static bool isConstant(SILValue V) {
-  V = getConstant(V);
-  return V && isOptimizableConstant(V);
+  if (auto *kp = dyn_cast<KeyPathInst>(V)) {
+    // We could support operands, if they are constants, to enable propagation
+    // of subscript keypaths. This would require to add the operands in the
+    // mangling scheme.
+    // But currently it's not worth it because we do not optimize subscript
+    // keypaths in SILCombine.
+    if (kp->getNumOperands() != 0)
+      return nullptr;
+    if (!kp->hasPattern())
+      return nullptr;
+    if (kp->getSubstitutions().hasAnySubstitutableParams())
+      return nullptr;
+    return kp;
+  }
+  return nullptr;
 }
 
 static std::string getClonedName(PartialApplyInst *PAI, IsSerialized_t Serialized,
@@ -194,11 +207,15 @@ void CapturePropagationCloner::cloneClosure(
   // Replace the rest of the old arguments with constants.
   getBuilder().setInsertionPoint(ClonedEntryBB);
   IsCloningConstant = true;
+  llvm::SmallVector<KeyPathInst *, 8> toDestroy;
   for (SILValue PartialApplyArg : PartialApplyArgs) {
-    assert(isConstant(PartialApplyArg) &&
+    assert(getConstant(PartialApplyArg) &&
            "expected a constant arg to partial apply");
 
     cloneConstValue(PartialApplyArg);
+    if (auto *kp = dyn_cast<KeyPathInst>(getMappedValue(PartialApplyArg))) {
+      toDestroy.push_back(kp);
+    }
 
     // The PartialApplyArg from the caller is now mapped to its cloned
     // instruction.  Also map the original argument to the cloned instruction.
@@ -213,6 +230,22 @@ void CapturePropagationCloner::cloneClosure(
   // Visit original BBs in depth-first preorder, starting with the
   // entry block, cloning all instructions and terminators.
   cloneFunctionBody(OrigF, ClonedEntryBB, entryArgs);
+  
+  // Destroy all the inserted keypaths at the function exits.
+  for (KeyPathInst *kpToDestroy : toDestroy) {
+    SILLocation loc = RegularLocation::getAutoGeneratedLocation();
+    for (SILBasicBlock &clonedBB : CloneF) {
+      TermInst *term = clonedBB.getTerminator();
+      if (term->isFunctionExiting()) {
+        SILBuilder builder(term);
+        if (CloneF.hasOwnership()) {
+          builder.createDestroyValue(loc, kpToDestroy);
+        } else {
+          builder.createStrongRelease(loc, kpToDestroy, builder.getDefaultAtomicity());
+        }
+      }
+    }
+  }
 }
 
 CanSILFunctionType getPartialApplyInterfaceResultType(PartialApplyInst *PAI) {
@@ -305,6 +338,20 @@ void CapturePropagation::rewritePartialApply(PartialApplyInst *OrigPAI,
   LLVM_DEBUG(llvm::dbgs() << "  Rewrote caller:\n" << *T2TF);
 }
 
+static bool isKeyPathFunction(FullApplySite FAS, SILValue keyPath) {
+  SILFunction *callee = FAS.getReferencedFunctionOrNull();
+  if (!callee)
+    return false;
+  if (callee->getName() == "swift_setAtWritableKeyPath" ||
+      callee->getName() == "swift_setAtReferenceWritableKeyPath") {
+    return FAS.getArgument(1) == keyPath;
+  }
+  if (callee->getName() == "swift_getAtKeyPath") {
+    return FAS.getArgument(2) == keyPath;
+  }
+  return false;
+}
+
 /// For now, we conservative only specialize if doing so can eliminate dynamic
 /// dispatch.
 ///
@@ -316,6 +363,8 @@ static bool isProfitable(SILFunction *Callee) {
       if (FullApplySite FAS = FullApplySite::isa(Operand->getUser())) {
         if (FAS.getCallee() == Operand->get())
           return true;
+        if (isKeyPathFunction(FAS, Arg))
+          return true;
       }
     }
   }
@@ -477,9 +526,35 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
   }
 
   // Second possibility: Are all partially applied arguments constant?
-  for (auto Arg : PAI->getArguments()) {
-    if (!isConstant(Arg))
+  llvm::SmallVector<SILInstruction *, 8> toDelete;
+  for (const Operand &argOp : PAI->getArgumentOperands()) {
+    SILInstruction *constInst = getConstant(argOp.get());
+    if (!constInst)
       return false;
+    if (auto *kp = dyn_cast<KeyPathInst>(constInst)) {
+      auto argConv = ApplySite(PAI).getArgumentConvention(argOp).Value;
+      // Only handle the common case of a guaranteed keypath arguments. That
+      // refers to the callee function.
+      if (argConv != SILArgumentConvention::Direct_Guaranteed)
+        return false;
+      
+      // For escaping closures:
+      // To keep things simple, we don't do a liferange analysis to insert
+      // compensating destroys of the keypath.
+      // Instead we require that the PAI is the only use of the keypath (= the
+      // common case). This allows us to just delete the now unused keypath
+      // instruction.
+      //
+      // For non-escaping closures:
+      // The keypath is not consumed by the PAI. We don't need todelete the
+      // keypath instruction in this pass, but let dead-object-elimination clean
+      // it up later.
+      if (!PAI->isOnStack()) {
+        if (getSingleNonDebugUser(kp) != PAI)
+          return false;
+        toDelete.push_back(kp);
+      }
+    }
   }
   if (!isProfitable(SubstF))
     return false;
@@ -491,6 +566,8 @@ bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
   SILFunction *NewF = specializeConstClosure(PAI, SubstF);
   rewritePartialApply(PAI, NewF);
 
+  recursivelyDeleteTriviallyDeadInstructions(toDelete, /*force*/ true);
+
   addFunctionToPassManagerWorklist(NewF, SubstF);
   return true;
 }

lib/SILOptimizer/Utils/SpecializationMangler.cpp

@@ -14,8 +14,10 @@
 #include "swift/AST/GenericEnvironment.h"
 #include "swift/AST/GenericSignature.h"
 #include "swift/AST/SubstitutionMap.h"
+#include "swift/Basic/MD5Stream.h"
 #include "swift/Demangling/ManglingMacros.h"
 #include "swift/SIL/SILGlobalVariable.h"
+#include "llvm/ADT/StringExtras.h"
 
 using namespace swift;
 using namespace Mangle;
@@ -68,10 +70,10 @@ void FunctionSignatureSpecializationMangler::setArgumentClosureProp(
 }
 
 void FunctionSignatureSpecializationMangler::setArgumentConstantProp(
-    unsigned OrigArgIdx, LiteralInst *LI) {
+    unsigned OrigArgIdx, SILInstruction *constInst) {
   auto &Info = OrigArgs[OrigArgIdx];
   Info.first = ArgumentModifierIntBase(ArgumentModifier::ConstantProp);
-  Info.second = LI;
+  Info.second = constInst;
 }
 
 void FunctionSignatureSpecializationMangler::setArgumentOwnedToGuaranteed(
@@ -122,41 +124,41 @@ setReturnValueOwnedToUnowned() {
 }
 
 void
-FunctionSignatureSpecializationMangler::mangleConstantProp(LiteralInst *LI) {
+FunctionSignatureSpecializationMangler::mangleConstantProp(SILInstruction *constInst) {
   // Append the prefix for constant propagation 'p'.
   ArgOpBuffer << 'p';
 
   // Then append the unique identifier of our literal.
-  switch (LI->getKind()) {
+  switch (constInst->getKind()) {
   default:
     llvm_unreachable("unknown literal");
   case SILInstructionKind::PreviousDynamicFunctionRefInst:
   case SILInstructionKind::DynamicFunctionRefInst:
   case SILInstructionKind::FunctionRefInst: {
     SILFunction *F =
-        cast<FunctionRefBaseInst>(LI)->getInitiallyReferencedFunction();
+        cast<FunctionRefBaseInst>(constInst)->getInitiallyReferencedFunction();
     ArgOpBuffer << 'f';
     appendIdentifier(F->getName());
     break;
   }
   case SILInstructionKind::GlobalAddrInst: {
-    SILGlobalVariable *G = cast<GlobalAddrInst>(LI)->getReferencedGlobal();
+    SILGlobalVariable *G = cast<GlobalAddrInst>(constInst)->getReferencedGlobal();
     ArgOpBuffer << 'g';
     appendIdentifier(G->getName());
     break;
   }
   case SILInstructionKind::IntegerLiteralInst: {
-    APInt apint = cast<IntegerLiteralInst>(LI)->getValue();
+    APInt apint = cast<IntegerLiteralInst>(constInst)->getValue();
     ArgOpBuffer << 'i' << apint;
     break;
   }
   case SILInstructionKind::FloatLiteralInst: {
-    APInt apint = cast<FloatLiteralInst>(LI)->getBits();
+    APInt apint = cast<FloatLiteralInst>(constInst)->getBits();
     ArgOpBuffer << 'd' << apint;
     break;
   }
   case SILInstructionKind::StringLiteralInst: {
-    StringLiteralInst *SLI = cast<StringLiteralInst>(LI);
+    StringLiteralInst *SLI = cast<StringLiteralInst>(constInst);
     StringRef V = SLI->getValue();
     assert(V.size() <= 32 && "Cannot encode string of length > 32");
     std::string VBuffer;
@@ -175,7 +177,43 @@ FunctionSignatureSpecializationMangler::mangleConstantProp(LiteralInst *LI) {
     }
     break;
   }
+  case SILInstructionKind::KeyPathInst: {
+    // Mangle a keypath instruction by creating a MD5 hash of the printed
+    // instruction. Everything else would be too complicated.
+  
+    auto *kp = cast<KeyPathInst>(constInst);
+    KeyPathPattern *pattern = kp->getPattern();
+    
+    MD5Stream md5Stream;
+    SILPrintContext printCtxt(md5Stream);
+    for (auto &component : pattern->getComponents()) {
+      component.print(printCtxt);
+    }
+    llvm::MD5::MD5Result md5Hash;
+    md5Stream.final(md5Hash);
+    SmallString<32> resultStr;
+    llvm::MD5::stringifyResult(md5Hash, resultStr);
+    appendStringAsIdentifier(resultStr);
+
+    // Also, mangle the involved types.
+    appendType(pattern->getRootType(), nullptr);
+    appendType(pattern->getValueType(), nullptr);
+
+    ArgOpBuffer << 'k';
+    break;
+  }
+  }
+}
+
+void
+FunctionSignatureSpecializationMangler::appendStringAsIdentifier(StringRef str) {
+  std::string buffer;
+  if (!str.empty() && (isDigit(str[0]) || str[0] == '_')) {
+    buffer = "_";
+    buffer.append(str.data(), str.size());
+    str = buffer;
   }
+  appendIdentifier(str);
 }
 
 void
@@ -206,7 +244,7 @@ FunctionSignatureSpecializationMangler::mangleClosureProp(SILInstruction *Inst)
 void FunctionSignatureSpecializationMangler::mangleArgument(
     ArgumentModifierIntBase ArgMod, NullablePtr<SILInstruction> Inst) {
   if (ArgMod == ArgumentModifierIntBase(ArgumentModifier::ConstantProp)) {
-    mangleConstantProp(cast<LiteralInst>(Inst.get()));
+    mangleConstantProp(Inst.get());
     return;
   }