SILOptimizer: support for Dictionary literals generated in the data section.

eeckstein · eeckstein · commit a4025445880e · 2019-04-23T09:02:53.000-07:00
Actually: generate the array of (key, value) tuples in the data section, which is then passed to Dictionary.init(dictionaryLiteral:)
We already do this for simple arrays, e.g. arrays with trivial element types.
The only change needed for dictionary literals is to support tuple types in the ObjectOutliner.

The effect of this optimization is a significant reduction in code size for dictionary literals - and an increase in data size.
But in most cases there is a considerable net win for code+data size in total.
diff --git a/lib/SILOptimizer/Transforms/ObjectOutliner.cpp b/lib/SILOptimizer/Transforms/ObjectOutliner.cpp
@@ -27,23 +27,26 @@ class ObjectOutliner {
   NominalTypeDecl *ArrayDecl = nullptr;
   int GlobIdx = 0;
 
+  // Instructions to be deleted.
+  llvm::SmallVector<SILInstruction *, 4> ToRemove;
+
   bool isCOWType(SILType type) {
     return type.getNominalOrBoundGenericNominal() == ArrayDecl;
   }
 
-  bool isValidUseOfObject(SILInstruction *Val, bool isCOWObject,
+  bool isValidUseOfObject(SILInstruction *Val,
+                          bool isCOWObject,
                           ApplyInst **FindStringCall = nullptr);
 
   bool getObjectInitVals(SILValue Val,
                          llvm::DenseMap<VarDecl *, StoreInst *> &MemberStores,
                          llvm::SmallVectorImpl<StoreInst *> &TailStores,
+                         unsigned NumTailTupleElements,
                          ApplyInst **FindStringCall);
-  bool handleTailAddr(int TailIdx, SILInstruction *I,
+  bool handleTailAddr(int TailIdx, SILInstruction *I, unsigned NumTailTupleElements,
                       llvm::SmallVectorImpl<StoreInst *> &TailStores);
 
-  bool
-  optimizeObjectAllocation(AllocRefInst *ARI,
-                           llvm::SmallVector<SILInstruction *, 4> &ToRemove);
+  bool optimizeObjectAllocation(AllocRefInst *ARI);
   void replaceFindStringCall(ApplyInst *FindStringCall);
 
 public:
@@ -60,23 +63,27 @@ bool ObjectOutliner::run(SILFunction *F) {
   for (auto &BB : *F) {
     auto Iter = BB.begin();
 
-    // We can't remove instructions willy-nilly as we iterate because
-    // that might cause a pointer to the next instruction to become
-    // garbage, causing iterator invalidations (and crashes).
-    // Instead, we collect in a list the instructions we want to remove
-    // and erase the BB they belong to at the end of the loop, once we're
-    // sure it's safe to do so.
-    llvm::SmallVector<SILInstruction *, 4> ToRemove;
-
     while (Iter != BB.end()) {
       SILInstruction *I = &*Iter;
       Iter++;
       if (auto *ARI = dyn_cast<AllocRefInst>(I)) {
-        hasChanged |= optimizeObjectAllocation(ARI, ToRemove);
+        unsigned GarbageSize = ToRemove.size();
+
+        // Try to replace the alloc_ref with a static object.
+        if (optimizeObjectAllocation(ARI)) {
+          hasChanged = true;
+        } else {
+          // No transformation was made. Restore the original state of the garbage list.
+          assert(GarbageSize <= ToRemove.size());
+          ToRemove.resize(GarbageSize);
+        }
       }
     }
+    // Delaying the deallocation of instructions avoids problems with iterator invalidation in the
+    // instruction loop above.
     for (auto *I : ToRemove)
       I->eraseFromParent();
+    ToRemove.clear();
   }
   return hasChanged;
 }
@@ -177,6 +184,13 @@ bool ObjectOutliner::isValidUseOfObject(SILInstruction *I, bool isCOWObject,
     BuiltinValueKind K = BI->getBuiltinInfo().ID;
     if (K == BuiltinValueKind::ICMP_EQ || K == BuiltinValueKind::ICMP_NE)
       return true;
+    if (K == BuiltinValueKind::DestroyArray) {
+      // We must not try to delete the tail allocated values. Although this would be a no-op
+      // (because we only handle trivial types), it would be semantically wrong to apply this
+      // builtin on the outlined object.
+      ToRemove.push_back(BI);
+      return true;
+    }
     return false;
   }
 
@@ -194,14 +208,28 @@ bool ObjectOutliner::isValidUseOfObject(SILInstruction *I, bool isCOWObject,
 
 /// Handle the address of a tail element.
 bool ObjectOutliner::handleTailAddr(int TailIdx, SILInstruction *TailAddr,
-                              llvm::SmallVectorImpl<StoreInst *> &TailStores) {
-  if (TailIdx >= 0 && TailIdx < (int)TailStores.size()) {
-    if (auto *SI = dyn_cast<StoreInst>(TailAddr)) {
-      if (!isValidInitVal(SI->getSrc()) || TailStores[TailIdx])
-        return false;
-      TailStores[TailIdx] = SI;
+                                    unsigned NumTailTupleElements,
+                                    llvm::SmallVectorImpl<StoreInst *> &TailStores) {
+  if (NumTailTupleElements > 0) {
+    if (auto *TEA = dyn_cast<TupleElementAddrInst>(TailAddr)) {
+      unsigned TupleIdx = TEA->getFieldNo();
+      assert(TupleIdx < NumTailTupleElements);
+      for (Operand *Use : TEA->getUses()) {
+        if (!handleTailAddr(TailIdx * NumTailTupleElements + TupleIdx, Use->getUser(), 0,
+                            TailStores))
+          return false;
+      }
       return true;
     }
+  } else {
+    if (TailIdx >= 0 && TailIdx < (int)TailStores.size()) {
+      if (auto *SI = dyn_cast<StoreInst>(TailAddr)) {
+        if (!isValidInitVal(SI->getSrc()) || TailStores[TailIdx])
+          return false;
+        TailStores[TailIdx] = SI;
+        return true;
+      }
+    }
   }
   return isValidUseOfObject(TailAddr, /*isCOWObject*/false);
 }
@@ -210,12 +238,13 @@ bool ObjectOutliner::handleTailAddr(int TailIdx, SILInstruction *TailAddr,
 bool ObjectOutliner::getObjectInitVals(SILValue Val,
                         llvm::DenseMap<VarDecl *, StoreInst *> &MemberStores,
                         llvm::SmallVectorImpl<StoreInst *> &TailStores,
+                        unsigned NumTailTupleElements,
                         ApplyInst **FindStringCall) {
   for (Operand *Use : Val->getUses()) {
     SILInstruction *User = Use->getUser();
     if (auto *UC = dyn_cast<UpcastInst>(User)) {
       // Upcast is transparent.
-      if (!getObjectInitVals(UC, MemberStores, TailStores, FindStringCall))
+      if (!getObjectInitVals(UC, MemberStores, TailStores, NumTailTupleElements, FindStringCall))
         return false;
     } else if (auto *REA = dyn_cast<RefElementAddrInst>(User)) {
       // The address of a stored property.
@@ -242,11 +271,11 @@ bool ObjectOutliner::getObjectInitVals(SILValue Val,
             TailIdx = Index->getValue().getZExtValue();
 
           for (Operand *IAUse : IA->getUses()) {
-            if (!handleTailAddr(TailIdx, IAUse->getUser(), TailStores))
+            if (!handleTailAddr(TailIdx, IAUse->getUser(), NumTailTupleElements, TailStores))
               return false;
           }
         // Without an index_addr it's the first tail element.
-        } else if (!handleTailAddr(/*TailIdx*/0, TailUser, TailStores)) {
+        } else if (!handleTailAddr(/*TailIdx*/0, TailUser, NumTailTupleElements, TailStores)) {
           return false;
         }
       }
@@ -280,9 +309,7 @@ class GlobalVariableMangler : public Mangle::ASTMangler {
 ///     func getarray() -> [Int] {
 ///       return [1, 2, 3]
 ///     }
-bool ObjectOutliner::optimizeObjectAllocation(
-    AllocRefInst *ARI, llvm::SmallVector<SILInstruction *, 4> &ToRemove) {
-
+bool ObjectOutliner::optimizeObjectAllocation(AllocRefInst *ARI) {
   if (ARI->isObjC())
     return false;
 
@@ -316,13 +343,30 @@ bool ObjectOutliner::optimizeObjectAllocation(
   llvm::SmallVector<VarDecl *, 16> Fields;
   getFields(Cl, Fields);
 
-  // Get the initialization stores of the object's properties and tail
-  // allocated elements. Also check if there are any "bad" uses of the object.
   llvm::DenseMap<VarDecl *, StoreInst *> MemberStores;
+
+  // A store for each element of the tail allocated array. In case of a tuple, there is a store
+  // for each tuple element. For example, a 3 element array of 2-element tuples
+  //     [ (i0, i1), (i2, i3), (i4, i5) ]
+  // results in following store instructions, collected in TailStores:
+  //     [ store i0, store i1, store i2, store i3, store i4, store i5 ]
   llvm::SmallVector<StoreInst *, 16> TailStores;
-  TailStores.resize(NumTailElems);
+
+  unsigned NumStores = NumTailElems;
+  unsigned NumTailTupleElems = 0;
+  if (auto Tuple = TailType.getAs<TupleType>()) {
+    NumTailTupleElems = Tuple->getNumElements();
+    if (NumTailTupleElems == 0)
+      return false;
+    NumStores *= NumTailTupleElems;
+  }
+
+  TailStores.resize(NumStores);
   ApplyInst *FindStringCall = nullptr;
-  if (!getObjectInitVals(ARI, MemberStores, TailStores, &FindStringCall))
+
+  // Get the initialization stores of the object's properties and tail
+  // allocated elements. Also check if there are any "bad" uses of the object.
+  if (!getObjectInitVals(ARI, MemberStores, TailStores, NumTailTupleElems, &FindStringCall))
     return false;
 
   // Is there a store for all the class properties?
@@ -372,13 +416,32 @@ bool ObjectOutliner::optimizeObjectAllocation(
                            cast<SingleValueInstruction>(MemberStore->getSrc())));
     ToRemove.push_back(MemberStore);
   }
-  // Create the initializers for the tail elements.
   unsigned NumBaseElements = ObjectArgs.size();
-  for (StoreInst *TailStore : TailStores) {
-    ObjectArgs.push_back(Cloner.clone(
-                           cast<SingleValueInstruction>(TailStore->getSrc())));
-    ToRemove.push_back(TailStore);
+
+  // Create the initializers for the tail elements.
+  if (NumTailTupleElems == 0) {
+    // The non-tuple element case.
+    for (StoreInst *TailStore : TailStores) {
+      ObjectArgs.push_back(Cloner.clone(
+                             cast<SingleValueInstruction>(TailStore->getSrc())));
+      ToRemove.push_back(TailStore);
+    }
+  } else {
+    // The elements are tuples: combine NumTailTupleElems elements from TailStores to a single tuple
+    // instruction.
+    for (unsigned EIdx = 0; EIdx < NumTailElems; EIdx++) {
+      SmallVector<SILValue, 8> TupleElems;
+      for (unsigned TIdx = 0; TIdx < NumTailTupleElems; TIdx++) {
+        StoreInst *TailStore = TailStores[EIdx * NumTailTupleElems + TIdx];
+        SILValue V = Cloner.clone(cast<SingleValueInstruction>(TailStore->getSrc()));
+        TupleElems.push_back(V);
+        ToRemove.push_back(TailStore);
+      }
+      auto *TI = Cloner.getBuilder().createTuple(ARI->getLoc(), TailType, TupleElems);
+      ObjectArgs.push_back(TI);
+    }
   }
+
   // Create the initializer for the object itself.
   SILBuilder StaticInitBuilder(Glob);
   StaticInitBuilder.createObject(ArtificialUnreachableLocation(),
diff --git a/test/SILOptimizer/static_arrays.swift b/test/SILOptimizer/static_arrays.swift
@@ -50,6 +50,22 @@
 // CHECK:         object {{.*}} ({{[^,]*}}, [tail_elems] {{[^,]*}}, {{[^,]*}})
 // CHECK-NEXT:  }
 
+// CHECK-LABEL: outlined variable #0 of returnDictionary()
+// CHECK-NEXT:  sil_global private @{{.*}}returnDictionary{{.*}} = {
+// CHECK-DAG:     integer_literal $Builtin.Int{{[0-9]+}}, 5
+// CHECK-DAG:     integer_literal $Builtin.Int{{[0-9]+}}, 4
+// CHECK-DAG:     integer_literal $Builtin.Int{{[0-9]+}}, 2
+// CHECK-DAG:     integer_literal $Builtin.Int{{[0-9]+}}, 1
+// CHECK-DAG:     integer_literal $Builtin.Int{{[0-9]+}}, 6
+// CHECK-DAG:     integer_literal $Builtin.Int{{[0-9]+}}, 3
+// CHECK:         object {{.*}} ({{[^,]*}}, [tail_elems]
+// CHECK-NEXT:  }
+
+// CHECK-LABEL: outlined variable #0 of returnStringDictionary()
+// CHECK-NEXT:  sil_global private @{{.*}}returnStringDictionary{{.*}} = {
+// CHECK:         object {{.*}} ({{[^,]*}}, [tail_elems]
+// CHECK-NEXT:  }
+
 // CHECK-LABEL: sil_global private @{{.*}}main{{.*}} = {
 // CHECK-DAG:     integer_literal $Builtin.Int{{[0-9]+}}, 100
 // CHECK-DAG:     integer_literal $Builtin.Int{{[0-9]+}}, 101
@@ -121,6 +137,22 @@ func arrayWithEmptyElements() -> [Empty] {
   return [Empty()]
 }
 
+// CHECK-LABEL: sil {{.*}}returnDictionary{{.*}} : $@convention(thin) () -> @owned Dictionary<Int, Int> {
+// CHECK:   global_value @{{.*}}returnDictionary{{.*}}
+// CHECK:   return
+@inline(never)
+public func returnDictionary() -> [Int:Int] {
+  return [1:2, 3:4, 5:6]
+}
+
+// CHECK-LABEL: sil {{.*}}returnStringDictionary{{.*}} : $@convention(thin) () -> @owned Dictionary<String, String> {
+// CHECK:   global_value @{{.*}}returnStringDictionary{{.*}}
+// CHECK:   return
+@inline(never)
+public func returnStringDictionary() -> [String:String] {
+  return ["1":"2", "3":"4", "5":"6"]
+}
+
 // CHECK-OUTPUT:      [100, 101, 102]
 print(globalVariable)
 // CHECK-OUTPUT-NEXT: 11
@@ -136,9 +168,13 @@ storeArray()
 // CHECK-OUTPUT-NEXT: [227, 228]
 print(gg!)
 
+let dict = returnDictionary()
+// CHECK-OUTPUT-NEXT: dict 3: 2, 4, 6
+print("dict \(dict.count): \(dict[1]!), \(dict[3]!), \(dict[5]!)")
 
-
-
+let sdict = returnStringDictionary()
+// CHECK-OUTPUT-NEXT: sdict 3: 2, 4, 6
+print("sdict \(sdict.count): \(sdict["1"]!), \(sdict["3"]!), \(sdict["5"]!)")
 
 
 public class SwiftClass {}
