Merge pull request #84704 from eeckstein/closure-specialization

ClosureSpecialization: support for OSSA and a big overhaul

include/swift/AST/SILOptions.h

@@ -303,10 +303,6 @@ public:
   /// Are we building in embedded Swift + -no-allocations?
   bool NoAllocations = false;
 
-  /// Should we use the experimental Swift based closure-specialization
-  /// optimization pass instead of the existing C++ one.
-  bool EnableExperimentalSwiftBasedClosureSpecialization = false;
-
   /// The name of the file to which the backend should save optimization
   /// records.
   std::string OptRecordFile;

include/swift/Demangling/Demangle.h

@@ -117,6 +117,7 @@ enum class FunctionSigSpecializationParamKind : unsigned {
   InOutToOut = 8,
   ConstantPropKeyPath = 9,
   ConstantPropStruct = 10,
+  ClosurePropPreviousArg = 11,
 
   // Option Set Flags use bits 6-31. This gives us 26 bits to use for option
   // flags.

include/swift/Option/FrontendOptions.td

@@ -395,10 +395,6 @@ def public_autolink_library :
 // HIDDEN FLAGS
 let Flags = [FrontendOption, NoDriverOption, HelpHidden] in {
 
-def enable_experimental_swift_based_closure_specialization :
-  Flag<["-"], "experimental-swift-based-closure-specialization">,
-  HelpText<"Use the experimental Swift based closure-specialization optimization pass instead of the existing C++ one">;
-
 def checked_async_objc_bridging : Joined<["-"], "checked-async-objc-bridging=">,
   HelpText<"Control whether checked continuations are used when bridging "
            "async calls from Swift to ObjC: 'on', 'off' ">;

include/swift/SIL/SILBridging.h

@@ -243,8 +243,6 @@ BridgedParameterInfoArray SILFunctionType_getParameters(BridgedCanType);
 
 BRIDGED_INLINE bool SILFunctionType_hasSelfParam(BridgedCanType);
 
-BRIDGED_INLINE bool SILFunctionType_isTrivialNoescape(BridgedCanType);
-
 SWIFT_IMPORT_UNSAFE BRIDGED_INLINE
 BridgedYieldInfoArray SILFunctionType_getYields(BridgedCanType);
 
@@ -819,6 +817,8 @@ struct BridgedInstruction {
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedGenericSpecializationInformation ApplyInst_getSpecializationInfo() const;
   BRIDGED_INLINE bool TryApplyInst_getNonAsync() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedGenericSpecializationInformation TryApplyInst_getSpecializationInfo() const;
+  BRIDGED_INLINE bool BeginApplyInst_getNonThrowing() const;
+  BRIDGED_INLINE bool BeginApplyInst_getNonAsync() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclRef ClassMethodInst_getMember() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedDeclRef WitnessMethodInst_getMember() const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedCanType WitnessMethodInst_getLookupType() const;
@@ -1232,6 +1232,8 @@ struct BridgedBuilder{
                                                                                BridgedType type) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUncheckedAddrCast(BridgedValue op,
                                                                                 BridgedType type) const;
+  SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUncheckedValueCast(BridgedValue op,
+                                                                                 BridgedType type) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createUpcast(BridgedValue op, BridgedType type) const;
   SWIFT_IMPORT_UNSAFE BRIDGED_INLINE BridgedInstruction createCheckedCastAddrBranch(
       BridgedValue source, BridgedCanType sourceFormalType,
@@ -1548,6 +1550,7 @@ struct BridgedCloner {
   void recordClonedInstruction(BridgedInstruction origInst, BridgedInstruction clonedInst) const;
   void recordFoldedValue(BridgedValue orig, BridgedValue mapped) const;
   BridgedInstruction clone(BridgedInstruction inst) const;
+  void setInsertionBlockIfNotSet(BridgedBasicBlock block) const;
 };
 
 struct BridgedTypeSubstCloner {

include/swift/SIL/SILBridgingImpl.h

@@ -260,10 +260,6 @@ bool SILFunctionType_hasSelfParam(BridgedCanType funcTy) {
   return funcTy.unbridged()->castTo<swift::SILFunctionType>()->hasSelfParam();
 }
 
-bool SILFunctionType_isTrivialNoescape(BridgedCanType funcTy) {
-  return funcTy.unbridged()->castTo<swift::SILFunctionType>()->isTrivialNoEscape();
-}
-
 BridgedYieldInfoArray SILFunctionType_getYields(BridgedCanType funcTy) {
   return {funcTy.unbridged()->castTo<swift::SILFunctionType>()->getYields()};
 }
@@ -1364,6 +1360,14 @@ BridgedGenericSpecializationInformation BridgedInstruction::TryApplyInst_getSpec
   return {getAs<swift::TryApplyInst>()->getSpecializationInfo()};
 }
 
+bool BridgedInstruction::BeginApplyInst_getNonThrowing() const {
+  return getAs<swift::BeginApplyInst>()->isNonThrowing();
+}
+
+bool BridgedInstruction::BeginApplyInst_getNonAsync() const {
+  return getAs<swift::BeginApplyInst>()->isNonAsync();
+}
+
 BridgedDeclRef BridgedInstruction::ClassMethodInst_getMember() const {
   return getAs<swift::ClassMethodInst>()->getMember();
 }
@@ -2316,6 +2320,11 @@ BridgedInstruction BridgedBuilder::createUncheckedAddrCast(BridgedValue op, Brid
                                               type.unbridged())};
 }
 
+BridgedInstruction BridgedBuilder::createUncheckedValueCast(BridgedValue op, BridgedType type) const {
+  return {unbridged().createUncheckedValueCast(regularLoc(), op.getSILValue(),
+                                              type.unbridged())};
+}
+
 BridgedInstruction BridgedBuilder::createUpcast(BridgedValue op, BridgedType type) const {
   return {unbridged().createUpcast(regularLoc(), op.getSILValue(),
                                    type.unbridged())};

include/swift/SIL/SILCloner.h

@@ -217,6 +217,8 @@ public:
                          ArrayRef<SILValue> entryArgs,
                          bool replaceOriginalFunctionInPlace = false);
 
+  void cloneFunctionBody(SILFunction *F);
+
   /// Clone all blocks in this function and all instructions in those
   /// blocks.
   ///
@@ -835,9 +837,31 @@ void SILCloner<ImplClass>::cloneFunctionBody(SILFunction *F,
   commonFixUp(F);
 }
 
+template <typename ImplClass>
+void SILCloner<ImplClass>::cloneFunctionBody(SILFunction *F) {
+  assert(!Builder.getFunction().empty() && "Expect the entry block to already be created");
+
+  assert(F != &Builder.getFunction() && "Must clone into a new function.");
+  assert(BBMap.empty() && "This API does not allow clients to map blocks.");
+
+  SILBasicBlock *clonedEntryBB = Builder.getFunction().getEntryBlock();
+  BBMap.insert(std::make_pair(&*F->begin(), clonedEntryBB));
+
+  Builder.setInsertionPoint(clonedEntryBB);
+
+  // This will layout all newly cloned blocks immediate after clonedEntryBB.
+  visitBlocksDepthFirst(&*F->begin());
+
+  commonFixUp(F);
+}
+
 template <typename ImplClass>
 void SILCloner<ImplClass>::cloneFunction(SILFunction *origF) {
   SILFunction *newF = &Builder.getFunction();
+  if (!newF->empty()) {
+    cloneFunctionBody(origF);
+    return;
+  }
 
   auto *newEntryBB = newF->createBasicBlock();
 
@@ -861,23 +885,12 @@ template <typename ImplClass>
 void SILCloner<ImplClass>::cloneFunctionBody(
     SILFunction *F, SILBasicBlock *clonedEntryBB, ArrayRef<SILValue> entryArgs,
     llvm::function_ref<SILValue(SILValue)> entryArgIndexToOldArgIndex) {
-  assert(F != clonedEntryBB->getParent() && "Must clone into a new function.");
-  assert(BBMap.empty() && "This API does not allow clients to map blocks.");
   assert(ValueMap.empty() && "Stale ValueMap.");
-
   assert(entryArgs.size() == F->getArguments().size());
   for (unsigned i = 0, e = entryArgs.size(); i != e; ++i) {
     ValueMap[entryArgIndexToOldArgIndex(entryArgs[i])] = entryArgs[i];
   }
-
-  BBMap.insert(std::make_pair(&*F->begin(), clonedEntryBB));
-
-  Builder.setInsertionPoint(clonedEntryBB);
-
-  // This will layout all newly cloned blocks immediate after clonedEntryBB.
-  visitBlocksDepthFirst(&*F->begin());
-
-  commonFixUp(F);
+  cloneFunctionBody(F);
 }
 
 template<typename ImplClass>

include/swift/SILOptimizer/OptimizerBridging.h

@@ -187,9 +187,15 @@ struct BridgedPassContext {
   bool specializeAppliesInFunction(BridgedFunction function, bool isMandatory) const;
   BridgedOwnedString mangleOutlinedVariable(BridgedFunction function) const;
   BridgedOwnedString mangleAsyncRemoved(BridgedFunction function) const;
+
+  struct ClosureArgMangling {
+    SwiftInt argIdx;
+    OptionalBridgedInstruction inst;
+    SwiftInt otherArgIdx;
+  };
+
   BridgedOwnedString mangleWithDeadArgs(BridgedArrayRef bridgedDeadArgIndices, BridgedFunction function) const;
-  BridgedOwnedString mangleWithClosureArgs(BridgedArrayRef closureArgIndices,
-                                                               BridgedFunction applySiteCallee) const;
+  BridgedOwnedString mangleWithClosureArgs(BridgedArrayRef closureArgManglings, BridgedFunction applySiteCallee) const;
   BridgedOwnedString mangleWithConstCaptureArgs(BridgedArrayRef bridgedConstArgs,
                                                 BridgedFunction applySiteCallee) const;
   BridgedOwnedString mangleWithBoxToStackPromotedArgs(BridgedArrayRef bridgedPromotedArgIndices,

include/swift/SILOptimizer/PassManager/Passes.def

@@ -153,9 +153,8 @@ PASS(TempLValueElimination, "temp-lvalue-elimination",
 PASS(LoopInvariantCodeMotion, "loop-invariant-code-motion",
      "New Loop Invariant Code Motion")
 
-// NOTE - ExperimentalSwiftBasedClosureSpecialization and AutodiffClosureSpecialization are a WIP
-PASS(ExperimentalSwiftBasedClosureSpecialization, "experimental-swift-based-closure-specialization",
-     "General closure-specialization pass written in Swift")
+PASS(ClosureSpecialization, "closure-specialization",
+     "Specialize functions with closure arguments")
 PASS(AutodiffClosureSpecialization, "autodiff-closure-specialization",
      "Autodiff specific closure-specialization pass")
 

include/swift/SILOptimizer/Utils/SpecializationMangler.h

@@ -68,6 +68,7 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
     BoxToValue = 3,
     BoxToStack = 4,
     InOutToOut = 5,
+    ClosurePropPreviousArg = 6, // the same closure as a previous `ClosureProp` argument
 
     First_Option = 0,
     Last_Option = 31,
@@ -82,8 +83,20 @@ class FunctionSignatureSpecializationMangler : public SpecializationMangler {
     LastOptionSetEntry = 32768,
   };
 
-  using ArgInfo = std::pair<ArgumentModifierIntBase,
-                            NullablePtr<SILInstruction>>;
+  struct ArgInfo {
+    ArgumentModifier kind;
+    union {
+      SILInstruction *inst;
+      unsigned otherArgIdx; // only for `ClosurePropPreviousArg`
+    };
+
+    ArgInfo(ArgumentModifier kind) : kind(kind), inst(nullptr) {}
+
+    void append(ArgumentModifier modifier) {
+      kind = ArgumentModifier(ArgumentModifierIntBase(kind) | ArgumentModifierIntBase(modifier));
+    }
+  };
+
   // Information for each SIL argument in the original function before
   // specialization. This includes SIL indirect result argument required for
   // the original function type at the current stage of compilation.
@@ -99,9 +112,8 @@ public:
   void setArgumentConstantProp(unsigned OrigArgIdx, SILInstruction *constInst);
   void appendStringAsIdentifier(StringRef str);
 
-  void setArgumentClosureProp(unsigned OrigArgIdx, PartialApplyInst *PAI);
-  void setArgumentClosureProp(unsigned OrigArgIdx,
-                              ThinToThickFunctionInst *TTTFI);
+  void setArgumentClosureProp(unsigned OrigArgIdx, SILInstruction *closure);
+  void setArgumentClosurePropPreviousArg(unsigned OrigArgIdx, unsigned otherArgIdx);
   void setArgumentDead(unsigned OrigArgIdx);
   void setArgumentOwnedToGuaranteed(unsigned OrigArgIdx);
   void setArgumentGuaranteedToOwned(unsigned OrigArgIdx);
@@ -120,8 +132,7 @@ public:
 private:
   void mangleConstantProp(SILInstruction *constInst);
   void mangleClosureProp(SILInstruction *Inst);
-  void mangleArgument(ArgumentModifierIntBase ArgMod,
-                      NullablePtr<SILInstruction> Inst);
+  void mangleArgument(ArgInfo argInfo);
   void mangleReturnValue(ReturnValueModifierIntBase RetMod);
 };