Merge pull request #33232 from eeckstein/string-optimization

StringOptimization: handle static let variables for String constant folding.

Author: eeckstein

include/swift/SIL/SILGlobalVariable.h

@@ -260,8 +260,7 @@ SILFunction *getCalleeOfOnceCall(BuiltinInst *BI);
 /// Given an addressor, AddrF, find the call to the global initializer if
 /// present, otherwise return null. If an initializer is returned, then
 /// `CallToOnce` is initialized to the corresponding builtin "once" call.
-SILFunction *findInitializer(SILModule *Module, SILFunction *AddrF,
-                             BuiltinInst *&CallToOnce);
+SILFunction *findInitializer(SILFunction *AddrF, BuiltinInst *&CallToOnce);
 
 /// Helper for getVariableOfGlobalInit(), so GlobalOpts can deeply inspect and
 /// rewrite the initialization pattern.

include/swift/SILOptimizer/PassManager/Passes.def

@@ -299,6 +299,8 @@ PASS(SILCombine, "sil-combine",
      "Combine SIL Instructions via Peephole Optimization")
 PASS(SILDebugInfoGenerator, "sil-debuginfo-gen",
      "Generate Debug Information with Source Locations into Textual SIL")
+PASS(EarlySROA, "early-sroa",
+     "Scalar Replacement of Aggregate Stack Objects on high-level SIL")
 PASS(SROA, "sroa",
      "Scalar Replacement of Aggregate Stack Objects")
 PASS(SROABBArgs, "sroa-bb-args",

lib/SIL/IR/SILGlobalVariable.cpp

@@ -220,9 +220,9 @@ SILGlobalVariable *swift::getVariableOfGlobalInit(SILFunction *AddrF) {
   // and the globalinit_func is called by "once" from a single location,
   // continue; otherwise bail.
   BuiltinInst *CallToOnce;
-  auto *InitF = findInitializer(&AddrF->getModule(), AddrF, CallToOnce);
+  auto *InitF = findInitializer(AddrF, CallToOnce);
 
-  if (!InitF || !InitF->getName().startswith("globalinit_"))
+  if (!InitF)
     return nullptr;
 
   // If the globalinit_func is trivial, continue; otherwise bail.
@@ -247,8 +247,8 @@ SILFunction *swift::getCalleeOfOnceCall(BuiltinInst *BI) {
 }
 
 // Find the globalinit_func by analyzing the body of the addressor.
-SILFunction *swift::findInitializer(SILModule *Module, SILFunction *AddrF,
-                             BuiltinInst *&CallToOnce) {
+SILFunction *swift::findInitializer(SILFunction *AddrF,
+                                    BuiltinInst *&CallToOnce) {
   // We only handle a single SILBasicBlock for now.
   if (AddrF->size() != 1)
     return nullptr;
@@ -272,7 +272,10 @@ SILFunction *swift::findInitializer(SILModule *Module, SILFunction *AddrF,
   }
   if (!CallToOnce)
     return nullptr;
-  return getCalleeOfOnceCall(CallToOnce);
+  SILFunction *callee = getCalleeOfOnceCall(CallToOnce);
+  if (!callee->getName().startswith("globalinit_"))
+    return nullptr;
+  return callee;
 }
 
 SILGlobalVariable *

lib/SILOptimizer/IPO/GlobalOpt.cpp

@@ -431,9 +431,8 @@ bool SILGlobalOpt::optimizeInitializer(SILFunction *AddrF,
   // If the addressor contains a single "once" call, it calls globalinit_func,
   // and the globalinit_func is called by "once" from a single location,
   // continue; otherwise bail.
-  auto *InitF = findInitializer(Module, AddrF, CallToOnce);
-  if (!InitF || !InitF->getName().startswith("globalinit_") ||
-      InitializerCount[InitF] > 1)
+  auto *InitF = findInitializer(AddrF, CallToOnce);
+  if (!InitF || InitializerCount[InitF] > 1)
     return false;
 
   // If the globalinit_func is trivial, continue; otherwise bail.

lib/SILOptimizer/PassManager/PassPipeline.cpp

@@ -222,7 +222,7 @@ void addHighLevelLoopOptPasses(SILPassPipelinePlan &P) {
   // Perform classic SSA optimizations for cleanup.
   P.addLowerAggregateInstrs();
   P.addSILCombine();
-  P.addSROA();
+  P.addEarlySROA();
   P.addMem2Reg();
   P.addDCE();
   P.addSILCombine();
@@ -290,7 +290,11 @@ void addFunctionPasses(SILPassPipelinePlan &P,
   P.addLowerAggregateInstrs();
 
   // Split up operations on stack-allocated aggregates (struct, tuple).
-  P.addSROA();
+  if (OpLevel == OptimizationLevelKind::HighLevel) {
+    P.addEarlySROA();
+  } else {
+    P.addSROA();
+  }
 
   // Promote stack allocations to values.
   P.addMem2Reg();
@@ -489,6 +493,8 @@ static void addHighLevelFunctionPipeline(SILPassPipelinePlan &P) {
   addFunctionPasses(P, OptimizationLevelKind::HighLevel);
 
   addHighLevelLoopOptPasses(P);
+  
+  P.addStringOptimization();
 }
 
 // After "high-level" function passes have processed the entire call tree, run
@@ -525,8 +531,6 @@ static void addSerializePipeline(SILPassPipelinePlan &P) {
 static void addMidLevelFunctionPipeline(SILPassPipelinePlan &P) {
   P.startPipeline("MidLevel,Function", true /*isFunctionPassPipeline*/);
 
-  P.addStringOptimization();
-
   addFunctionPasses(P, OptimizationLevelKind::MidLevel);
 
   // Specialize partially applied functions with dead arguments as a preparation

lib/SILOptimizer/Transforms/SILMem2Reg.cpp

@@ -375,15 +375,15 @@ static void replaceLoad(LoadInst *LI, SILValue val, AllocStackInst *ASI) {
 
 static void replaceDestroy(DestroyAddrInst *DAI, SILValue NewValue) {
   SILFunction *F = DAI->getFunction();
+  auto Ty = DAI->getOperand()->getType();
 
-  assert(DAI->getOperand()->getType().isLoadable(*F) &&
+  assert(Ty.isLoadable(*F) &&
          "Unexpected promotion of address-only type!");
 
-  assert(NewValue && "Expected a value to release!");
+  assert(NewValue || (Ty.is<TupleType>() && Ty.getAs<TupleType>()->getNumElements() == 0));
 
   SILBuilderWithScope Builder(DAI);
 
-  auto Ty = DAI->getOperand()->getType();
   auto &TL = F->getTypeLowering(Ty);
 
   bool expand = shouldExpand(DAI->getModule(),

lib/SILOptimizer/Transforms/SILSROA.cpp

@@ -301,18 +301,33 @@ void SROAMemoryUseAnalyzer::chopUpAlloca(std::vector<AllocStackInst *> &Worklist
   eraseFromParentWithDebugInsts(AI);
 }
 
-static bool runSROAOnFunction(SILFunction &Fn) {
+/// Returns true, if values of \ty should be ignored, because \p ty is known
+/// by a high-level SIL optimization. Values of that type must not be split
+/// so that those high-level optimizations can analyze the code.
+static bool isSemanticType(ASTContext &ctxt, SILType ty) {
+  NominalTypeDecl *stringDecl = ctxt.getStringDecl();
+  if (ty.getStructOrBoundGenericStruct() == stringDecl)
+    return true;
+  return false;
+}
+
+static bool runSROAOnFunction(SILFunction &Fn, bool splitSemanticTypes) {
   std::vector<AllocStackInst *> Worklist;
   bool Changed = false;
+  ASTContext &ctxt = Fn.getModule().getASTContext();
 
   // For each basic block BB in Fn...
   for (auto &BB : Fn)
     // For each instruction in BB...
     for (auto &I : BB)
       // If the instruction is an alloc stack inst, add it to the worklist.
-      if (auto *AI = dyn_cast<AllocStackInst>(&I))
+      if (auto *AI = dyn_cast<AllocStackInst>(&I)) {
+        if (!splitSemanticTypes && isSemanticType(ctxt, AI->getElementType()))
+          continue;
+
         if (shouldExpand(Fn.getModule(), AI->getElementType()))
           Worklist.push_back(AI);
+      }
 
   while (!Worklist.empty()) {
     AllocStackInst *AI = Worklist.back();
@@ -332,6 +347,11 @@ static bool runSROAOnFunction(SILFunction &Fn) {
 namespace {
 class SILSROA : public SILFunctionTransform {
 
+  bool splitSemanticTypes;
+  
+public:
+  SILSROA(bool splitSemanticTypes) : splitSemanticTypes(splitSemanticTypes) { }
+
   /// The entry point to the transformation.
   void run() override {
     SILFunction *F = getFunction();
@@ -343,7 +363,7 @@ class SILSROA : public SILFunctionTransform {
     LLVM_DEBUG(llvm::dbgs() << "***** SROA on function: " << F->getName()
                             << " *****\n");
 
-    if (runSROAOnFunction(*F))
+    if (runSROAOnFunction(*F, splitSemanticTypes))
       invalidateAnalysis(SILAnalysis::InvalidationKind::Instructions);
   }
 
@@ -352,5 +372,9 @@ class SILSROA : public SILFunctionTransform {
 
 
 SILTransform *swift::createSROA() {
-  return new SILSROA();
+  return new SILSROA(/*splitSemanticTypes*/ true);
+}
+
+SILTransform *swift::createEarlySROA() {
+  return new SILSROA(/*splitSemanticTypes*/ false);
 }

lib/SILOptimizer/Transforms/StringOptimization.cpp

@@ -19,6 +19,7 @@
 #include "swift/SILOptimizer/Analysis/ValueTracking.h"
 #include "swift/SIL/SILFunction.h"
 #include "swift/SIL/SILBasicBlock.h"
+#include "swift/SIL/SILGlobalVariable.h"
 #include "swift/SIL/SILBuilder.h"
 #include "swift/AST/SemanticAttrs.h"
 #include "swift/AST/ParameterList.h"
@@ -58,13 +59,15 @@ class StringOptimization {
     StringRef str;
 
     /// Negative means: not constant
-    int numCodeUnits = -1;
-    
-    /// Not 0 for the empty-string initializer which reserves a capacity.
     int reservedCapacity = 0;
 
-    bool isConstant() const { return numCodeUnits >= 0; }
+    StringInfo(StringRef str, int reservedCapacity = 0) :
+      str(str), reservedCapacity(reservedCapacity) { }
+    
+    bool isConstant() const { return reservedCapacity >= 0; }
     bool isEmpty() const { return isConstant() && str.empty(); }
+    
+    static StringInfo unknown() { return StringInfo(StringRef(), -1); }
   };
 
   /// The stdlib's String type.
@@ -96,6 +99,8 @@ class StringOptimization {
   static void invalidateModifiedObjects(SILInstruction *inst,
                             llvm::DenseMap<SILValue, SILValue> &storedStrings);
   static StringInfo getStringInfo(SILValue value);
+  static StringInfo getStringFromStaticLet(SILValue value);
+
   static Optional<int> getIntConstant(SILValue value);
   static void replaceAppendWith(ApplyInst *appendCall, SILValue newValue,
                                 bool copyNewValue);
@@ -368,32 +373,31 @@ void StringOptimization::invalidateModifiedObjects(SILInstruction *inst,
 
 /// Returns information about value if it's a constant string.
 StringOptimization::StringInfo StringOptimization::getStringInfo(SILValue value) {
-  // Start with a non-constant result.
-  StringInfo result;
+  if (!value)
+    return StringInfo::unknown();
 
-  auto *apply = dyn_cast_or_null<ApplyInst>(value);
-  if (!apply)
-    return result;
+  auto *apply = dyn_cast<ApplyInst>(value);
+  if (!apply) {
+    return getStringFromStaticLet(value);
+  }
 
   SILFunction *callee = apply->getReferencedFunctionOrNull();
   if (!callee)
-    return result;
+    return StringInfo::unknown();
 
   if (callee->hasSemanticsAttr(semantics::STRING_INIT_EMPTY)) {
     // An empty string initializer.
-    result.numCodeUnits = 0;
-    return result;
+    return StringInfo("");
   }
 
   if (callee->hasSemanticsAttr(semantics::STRING_INIT_EMPTY_WITH_CAPACITY)) {
     // An empty string initializer with initial capacity.
-    result.numCodeUnits = 0;
-    result.reservedCapacity = std::numeric_limits<int>::max();
+    int reservedCapacity = std::numeric_limits<int>::max();
     if (apply->getNumArguments() > 0) {
       if (Optional<int> capacity = getIntConstant(apply->getArgument(0)))
-        result.reservedCapacity = capacity.getValue();
+        reservedCapacity = capacity.getValue();
     }
-    return result;
+    return StringInfo("", reservedCapacity);
   }
 
   if (callee->hasSemanticsAttr(semantics::STRING_MAKE_UTF8)) {
@@ -404,13 +408,79 @@ StringOptimization::StringInfo StringOptimization::getStringInfo(SILValue value)
     auto *intLiteral = dyn_cast<IntegerLiteralInst>(lengthVal);
     if (intLiteral && stringLiteral &&
         // For simplicity, we only support UTF8 string literals.
-        stringLiteral->getEncoding() == StringLiteralInst::Encoding::UTF8) {
-      result.str = stringLiteral->getValue();
-      result.numCodeUnits = intLiteral->getValue().getSExtValue();
-      return result;
+        stringLiteral->getEncoding() == StringLiteralInst::Encoding::UTF8 &&
+        // This passed number of code units should always match the size of the
+        // string in the string literal. Just to be on the safe side, check it.
+        intLiteral->getValue() == stringLiteral->getValue().size()) {
+      return StringInfo(stringLiteral->getValue());
     }
   }
-  return result;
+  return StringInfo::unknown();
+}
+
+/// Return the string if \p value is a load from a global static let, which is
+/// initialized with a String constant.
+StringOptimization::StringInfo
+StringOptimization::getStringFromStaticLet(SILValue value) {
+  // Match the pattern
+  //   %ptr_to_global = apply %addressor()
+  //   %global_addr = pointer_to_address %ptr_to_global
+  //   %value = load %global_addr
+  auto *load = dyn_cast<LoadInst>(value);
+  if (!load)
+    return StringInfo::unknown();
+ 
+  auto *pta = dyn_cast<PointerToAddressInst>(load->getOperand());
+  if (!pta)
+    return StringInfo::unknown();
+    
+  auto *addressorCall = dyn_cast<ApplyInst>(pta->getOperand());
+  if (!addressorCall)
+    return StringInfo::unknown();
+    
+  SILFunction *addressorFunc = addressorCall->getReferencedFunctionOrNull();
+  if (!addressorFunc)
+    return StringInfo::unknown();
+    
+  // The addressor function has a builtin.once call to the initializer.
+  BuiltinInst *onceCall = nullptr;
+  SILFunction *initializer = findInitializer(addressorFunc, onceCall);
+  if (!initializer)
+    return StringInfo::unknown();
+  
+  if (initializer->size() != 1)
+    return StringInfo::unknown();
+
+  // Match the pattern
+  //   %addr = global_addr @staticStringLet
+  //   ...
+  //   %str = apply %stringInitializer(...)
+  //   store %str to %addr
+  GlobalAddrInst *gAddr = nullptr;
+  for (SILInstruction &inst : initializer->front()) {
+    if (auto *ga = dyn_cast<GlobalAddrInst>(&inst)) {
+      if (gAddr)
+        return StringInfo::unknown();
+      gAddr = ga;
+    }
+  }
+  if (!gAddr || !gAddr->getReferencedGlobal()->isLet())
+    return StringInfo::unknown();
+  
+  Operand *gUse = gAddr->getSingleUse();
+  auto *store = dyn_cast<StoreInst>(gUse->getUser());
+  if (!store || store->getDest() != gAddr)
+    return StringInfo::unknown();
+    
+  SILValue initVal = store->getSrc();
+  
+  // This check is probably not needed, but let's be on the safe side:
+  // it prevents an infinite recursion if the initializer of the global is
+  // itself a load of another global, and so on.
+  if (isa<LoadInst>(initVal))
+    return StringInfo::unknown();
+
+  return getStringInfo(initVal);
 }
 
 /// Returns the constant integer value if \a value is an Int or Bool struct with

test/SILOptimizer/string_optimization.swift

@@ -12,6 +12,8 @@ import Foundation
 
 struct Outer {
   struct Inner { }
+
+  static let staticString = "static"
 }
 
 // More types are tested in test/stdlib/TypeName.swift and
@@ -36,6 +38,15 @@ public func testFoldCompleteInterpolation() -> String {
   return "-\([Int].self) \(s) \("cool")+"
 }
 
+// CHECK-LABEL: sil [noinline] @$s4test0A13FoldStaticLetSSyF
+// CHECK-NOT: apply
+// CHECK-NOT: bb1
+// CHECK: } // end sil function '$s4test0A13FoldStaticLetSSyF'
+@inline(never)
+public func testFoldStaticLet() -> String {
+  return "-\(Outer.staticString)+"
+}
+
 // CHECK-LABEL: sil [noinline] @$s4test0A19UnqualifiedTypeNameSSyF 
 // CHECK-NOT: apply
 // CHECK-NOT: bb1
@@ -92,6 +103,9 @@ printEmbeeded(testTypeNameInterpolation())
 // CHECK-OUTPUT: <-Array<Int> is cool+>
 printEmbeeded(testFoldCompleteInterpolation())
 
+// CHECK-OUTPUT: <-static+>
+printEmbeeded(testFoldStaticLet())
+
 // CHECK-OUTPUT: <Inner>
 printEmbeeded(testUnqualifiedTypeName())