[pred-memopt] Rather than extracting values at the load site, extract at the store site.

Previously in PredMemOpts, we would insert any extracts at the load site, i.e.:

  store %foo to %0
  ...
  %1 = struct_element_addr %0, $Type, $Type.field
  %2 = load %1
  ...
  apply %use(%2)

would transform to:

  store %foo to %0
  ...
  %2 = struct_extract %foo
  apply %use(%2)

This framework will not work with Ownership enabled since the value stored is
considered consumed by the store. This change fixes the issue by moving such
non-destructive extracts to occur while %foo is considered live, i.e. before the
store:

  %2 = struct_extract %foo
  store %foo to %0
  ...
  apply %use(%2)

This means that we have to store insertion points for each store that provides
us with available values and insert the extracts at those points. This creates
some complications in the case where we have multiple stores since we need to
deal with phi nodes. Rather than dealing with it by hand, we just insert the
extracts at each point and then use the SSA updater to insert the relevant phi
nodes.

rdar://31521023

Author: gottesmm

include/swift/SIL/SILBuilder.h

@@ -2049,6 +2049,35 @@ class SILBuilderWithScope : public SILBuilder {
   }
 };
 
+class SavedInsertionPointRAII {
+  SILBuilder &Builder;
+  PointerUnion<SILInstruction *, SILBasicBlock *> SavedIP;
+
+public:
+  SavedInsertionPointRAII(SILBuilder &B, SILInstruction *NewIP)
+      : Builder(B), SavedIP(&*B.getInsertionPoint()) {
+    Builder.setInsertionPoint(NewIP);
+  }
+
+  SavedInsertionPointRAII(SILBuilder &B, SILBasicBlock *NewIP)
+      : Builder(B), SavedIP(B.getInsertionBB()) {
+    Builder.setInsertionPoint(NewIP);
+  }
+
+  SavedInsertionPointRAII(const SavedInsertionPointRAII &) = delete;
+  SavedInsertionPointRAII &operator=(const SavedInsertionPointRAII &) = delete;
+  SavedInsertionPointRAII(SavedInsertionPointRAII &&) = delete;
+  SavedInsertionPointRAII &operator=(SavedInsertionPointRAII &&) = delete;
+
+  ~SavedInsertionPointRAII() {
+    if (SavedIP.is<SILInstruction *>()) {
+      Builder.setInsertionPoint(SavedIP.get<SILInstruction *>());
+    } else {
+      Builder.setInsertionPoint(SavedIP.get<SILBasicBlock *>());
+    }
+  }
+};
+
 } // end swift namespace
 
 #endif

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -12,11 +12,12 @@
 
 #define DEBUG_TYPE "predictable-memopt"
 
-#include "swift/SILOptimizer/PassManager/Passes.h"
 #include "DIMemoryUseCollector.h"
 #include "swift/SIL/SILBuilder.h"
-#include "swift/SILOptimizer/Utils/Local.h"
+#include "swift/SILOptimizer/PassManager/Passes.h"
 #include "swift/SILOptimizer/PassManager/Transforms.h"
+#include "swift/SILOptimizer/Utils/Local.h"
+#include "swift/SILOptimizer/Utils/SILSSAUpdater.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Support/Compiler.h"
@@ -91,14 +92,14 @@ static SILValue getAccessPathRoot(SILValue Pointer) {
 /// If this pointer is to within an existential projection, it returns ~0U.
 static unsigned computeSubelement(SILValue Pointer,
                                   SingleValueInstruction *RootInst) {
-  unsigned SubEltNumber = 0;
+  unsigned SubElementNumber = 0;
   SILModule &M = RootInst->getModule();
 
   while (1) {
     // If we got to the root, we're done.
     if (RootInst == Pointer)
-      return SubEltNumber;
-    
+      return SubElementNumber;
+
     if (auto *PBI = dyn_cast<ProjectBoxInst>(Pointer)) {
       Pointer = PBI->getOperand();
       continue;
@@ -114,7 +115,7 @@ static unsigned computeSubelement(SILValue Pointer,
 
       // Keep track of what subelement is being referenced.
       for (unsigned i = 0, e = TEAI->getFieldNo(); i != e; ++i) {
-        SubEltNumber += getNumSubElements(TT.getTupleElementType(i), M);
+        SubElementNumber += getNumSubElements(TT.getTupleElementType(i), M);
       }
       Pointer = TEAI->getOperand();
       continue;
@@ -127,7 +128,7 @@ static unsigned computeSubelement(SILValue Pointer,
       StructDecl *SD = SEAI->getStructDecl();
       for (auto *D : SD->getStoredProperties()) {
         if (D == SEAI->getField()) break;
-        SubEltNumber += getNumSubElements(ST.getFieldType(D, M), M);
+        SubElementNumber += getNumSubElements(ST.getFieldType(D, M), M);
       }
 
       Pointer = SEAI->getOperand();
@@ -149,14 +150,26 @@ static unsigned computeSubelement(SILValue Pointer,
 namespace {
 
 struct AvailableValue {
+  /// If this gets too expensive in terms of copying, we can use an arena and a
+  /// FrozenPtrSet like we do in ARC.
+  using SetVector = llvm::SmallSetVector<SILInstruction *, 1>;
+
   SILValue Value;
   unsigned SubElementNumber;
+  SetVector InsertionPoints;
 
 public:
   AvailableValue() = default;
 
-  AvailableValue(SILValue Value, unsigned SubElementNumber)
-      : Value(Value), SubElementNumber(SubElementNumber) {}
+  /// Main initializer for available values.
+  ///
+  /// *NOTE* We assume that all available values start with a singular insertion
+  /// point and insertion points are added by merging.
+  AvailableValue(SILValue Value, unsigned SubElementNumber,
+                 SILInstruction *InsertPoint)
+      : Value(Value), SubElementNumber(SubElementNumber), InsertionPoints() {
+    InsertionPoints.insert(InsertPoint);
+  }
 
   /// Deleted copy constructor. This is a move only type.
   AvailableValue(const AvailableValue &) = delete;
@@ -166,15 +179,17 @@ struct AvailableValue {
 
   /// Move constructor.
   AvailableValue(AvailableValue &&Other)
-      : Value(nullptr), SubElementNumber(~0) {
+      : Value(nullptr), SubElementNumber(~0), InsertionPoints() {
     std::swap(Value, Other.Value);
     std::swap(SubElementNumber, Other.SubElementNumber);
+    std::swap(InsertionPoints, Other.InsertionPoints);
   }
 
   /// Move operator.
   AvailableValue &operator=(AvailableValue &&Other) {
     std::swap(Value, Other.Value);
     std::swap(SubElementNumber, Other.SubElementNumber);
+    std::swap(InsertionPoints, Other.InsertionPoints);
     return *this;
   }
 
@@ -191,24 +206,70 @@ struct AvailableValue {
   SILValue getValue() const { return Value; }
   SILType getType() const { return Value->getType(); }
   unsigned getSubElementNumber() const { return SubElementNumber; }
+  ArrayRef<SILInstruction *> getInsertionPoints() const {
+    return InsertionPoints.getArrayRef();
+  }
+
+  void mergeInsertionPoints(const AvailableValue &Other) & {
+    assert(Value == Other.Value && SubElementNumber == Other.SubElementNumber);
+    InsertionPoints.set_union(Other.InsertionPoints);
+  }
+
+  void addInsertionPoint(SILInstruction *I) & { InsertionPoints.insert(I); }
 
   /// TODO: This needs a better name.
   AvailableValue emitStructExtract(SILBuilder &B, SILLocation Loc, VarDecl *D,
-                                   unsigned SubEltNumber) const {
+                                   unsigned SubElementNumber) const {
     SILValue NewValue = B.emitStructExtract(Loc, Value, D);
-    return {NewValue, SubEltNumber};
+    return {NewValue, SubElementNumber, InsertionPoints};
   }
 
   /// TODO: This needs a better name.
   AvailableValue emitTupleExtract(SILBuilder &B, SILLocation Loc,
-                                  unsigned EltNo, unsigned SubEltNumber) const {
+                                  unsigned EltNo,
+                                  unsigned SubElementNumber) const {
     SILValue NewValue = B.emitTupleExtract(Loc, Value, EltNo);
-    return {NewValue, SubEltNumber};
+    return {NewValue, SubElementNumber, InsertionPoints};
   }
+
+  void dump() const __attribute__((used));
+  void print(llvm::raw_ostream &os) const;
+
+private:
+  /// Private constructor for use by emitStructExtract and emitTupleExtract.
+  AvailableValue(SILValue Value, unsigned SubElementNumber,
+                 const SetVector &InsertPoints)
+      : Value(Value), SubElementNumber(SubElementNumber),
+        InsertionPoints(InsertPoints) {}
 };
 
 } // end anonymous namespace
 
+void AvailableValue::dump() const { print(llvm::dbgs()); }
+
+void AvailableValue::print(llvm::raw_ostream &os) const {
+  os << "Available Value Dump. Value: ";
+  if (getValue()) {
+    os << getValue();
+  } else {
+    os << "NoValue;\n";
+  }
+  os << "SubElementNumber: " << getSubElementNumber() << "\n";
+  os << "Insertion Points:\n";
+  for (auto *I : getInsertionPoints()) {
+    os << *I;
+  }
+}
+
+namespace llvm {
+
+llvm::raw_ostream &operator<<(llvm::raw_ostream &os, const AvailableValue &V) {
+  V.print(os);
+  return os;
+}
+
+} // end llvm namespace
+
 //===----------------------------------------------------------------------===//
 //                           Subelement Extraction
 //===----------------------------------------------------------------------===//
@@ -301,7 +362,7 @@ class AvailableValueAggregator {
   SILValue aggregateValues(SILType LoadTy, SILValue Address, unsigned FirstElt);
 
   void print(llvm::raw_ostream &os) const;
-  void dump() const;
+  void dump() const __attribute__((used));
 
 private:
   SILValue aggregateFullyAvailableValue(SILType LoadTy, unsigned FirstElt);
@@ -321,8 +382,7 @@ void AvailableValueAggregator::print(llvm::raw_ostream &os) const {
   os << "Available Value List, N = " << AvailableValueList.size()
      << ". Elts:\n";
   for (auto &V : AvailableValueList) {
-    os << "Value: " << V.getValue()
-       << "SubElementNumber: " << V.getSubElementNumber() << "\n";
+    os << V;
   }
 }
 
@@ -440,14 +500,30 @@ SILValue AvailableValueAggregator::handlePrimitiveValue(SILType LoadTy,
     return B.createLoad(Loc, Address, LoadOwnershipQualifier::Unqualified);
   }
 
-  // If we have an available value, we know that we know that the available
-  // value is already being consumed by the store. This means that we must
-  // insert a copy of EltVal after we extract it if we do not have a trivial
-  // value. We use SILBuilder::emit*Operation to handle both trivial/non-trivial
-  // cases without needing to introduce control flow here.
+  // If we have 1 insertion point, just extract the value and return.
+  //
+  // This saves us from having to spend compile time in the SSA updater in this
+  // case.
+  ArrayRef<SILInstruction *> InsertPts = Val.getInsertionPoints();
+  if (InsertPts.size() == 1) {
+    SavedInsertionPointRAII SavedInsertPt(B, InsertPts[0]);
+    SILValue EltVal = nonDestructivelyExtractSubElement(Val, B, Loc);
+    assert(EltVal->getType() == LoadTy && "Subelement types mismatch");
+    return EltVal;
+  }
+
+  // If we have an available value, then we want to extract the subelement from
+  // the borrowed aggregate before each insertion point.
+  SILSSAUpdater Updater;
+  Updater.Initialize(LoadTy);
+  for (auto *I : Val.getInsertionPoints()) {
+    SavedInsertionPointRAII SavedInsertPt(B, I);
+    SILValue EltVal = nonDestructivelyExtractSubElement(Val, B, Loc);
+    Updater.AddAvailableValue(I->getParent(), EltVal);
+  }
 
-  // Then extract the subelement from the borrowed aggregate.
-  SILValue EltVal = nonDestructivelyExtractSubElement(Val, B, Loc);
+  // Finally, grab the value from the SSA updater.
+  SILValue EltVal = Updater.GetValueInMiddleOfBlock(B.getInsertionBB());
   assert(EltVal->getType() == LoadTy && "Subelement types mismatch");
   return EltVal;
 }
@@ -587,6 +663,7 @@ void AllocOptimize::updateAvailableValues(
     SILInstruction *Inst, llvm::SmallBitVector &RequiredElts,
     SmallVectorImpl<AvailableValue> &Result,
     llvm::SmallBitVector &ConflictingValues) {
+
   // Handle store and assign.
   if (auto *SI = dyn_cast<StoreInst>(Inst)) {
     unsigned StartSubElt = computeSubelement(SI->getDest(), TheMemory);
@@ -601,12 +678,19 @@ void AllocOptimize::updateAvailableValues(
       // there already is a result, check it for conflict.  If there is no
       // conflict, then we're ok.
       auto &Entry = Result[StartSubElt+i];
-      if (!Entry)
-        Entry = {SI->getSrc(), i};
-      else if (Entry.getValue() != SI->getSrc() ||
-               Entry.getSubElementNumber() != i)
-        ConflictingValues[StartSubElt+i] = true;
-      
+      if (!Entry) {
+        Entry = {SI->getSrc(), i, Inst};
+      } else {
+        // TODO: This is /really/, /really/, conservative. This basically means
+        // that if we do not have an identical store, we will not promote.
+        if (Entry.getValue() != SI->getSrc() ||
+            Entry.getSubElementNumber() != i) {
+          ConflictingValues[StartSubElt + i] = true;
+        } else {
+          Entry.addInsertionPoint(Inst);
+        }
+      }
+
       // This element is now provided.
       RequiredElts[StartSubElt+i] = false;
     }
@@ -682,8 +766,8 @@ void AllocOptimize::computeAvailableValues(
   if (!ConflictingValues.none())
     for (unsigned i = 0, e = Result.size(); i != e; ++i)
       if (ConflictingValues[i])
-        Result[i] = { SILValue(), 0U };
-  
+        Result[i] = {};
+
   return;
 }