[pmo] Teach PMO that in ossa, load [take] invalidates an available value.

Example would be a case where we dynamically control if the load [take] occurs
and the load [take] happens with conditional control flow.

This was exposed by the throwing_initializer and failable_initializer
interpreter tests in the test of PolarBear(n:before:during).

Author: gottesmm

lib/SILOptimizer/Mandatory/PredictableMemOpt.cpp

@@ -147,7 +147,8 @@ static unsigned computeSubelement(SILValue Pointer,
       continue;
     }
 
-    
+    // This fails when we visit unchecked_take_enum_data_addr. We should just
+    // add support for enums.
     assert(isa<InitExistentialAddrInst>(Pointer) &&
            "Unknown access path instruction");
     // Cannot promote loads and stores from within an existential projection.
@@ -233,7 +234,7 @@ struct AvailableValue {
     InsertionPoints.set_union(Other.InsertionPoints);
   }
 
-  void addInsertionPoint(StoreInst *I) & { InsertionPoints.insert(I); }
+  void addInsertionPoint(StoreInst *si) & { InsertionPoints.insert(si); }
 
   AvailableValue emitStructExtract(SILBuilder &B, SILLocation Loc, VarDecl *D,
                                    unsigned SubElementNumber) const {
@@ -845,8 +846,18 @@ AvailableValueAggregator::addMissingDestroysForCopiedValues(
 namespace {
 
 /// Given a piece of memory, the memory's uses, and destroys perform a single
-/// round of optimistic dataflow switching to intersection when a back edge is
-/// encountered.
+/// round of semi-optimistic backwards dataflow for each use. The result is the
+/// set of available values that reach the specific use of the field in the
+/// allocated object.
+///
+/// The general form of the algorithm is that in our constructor, we analyze our
+/// uses and determine available values. Then users call computeAvailableValues
+/// which looks backwards up the control flow graph for available values that we
+/// can use.
+///
+/// NOTE: The reason why we say that the algorithm is semi-optimistic is that we
+/// assume that all incoming elements into a loopheader will be the same. If we
+/// find a conflict, we record it and fail.
 class AvailableValueDataflowContext {
   /// The base memory we are performing dataflow upon.
   AllocationInst *TheMemory;
@@ -862,9 +873,21 @@ class AvailableValueDataflowContext {
   /// The set of blocks with local definitions.
   ///
   /// We use this to determine if we should visit a block or look at a block's
-  /// predecessors during dataflow.
+  /// predecessors during dataflow for an available value.
   llvm::SmallPtrSet<SILBasicBlock *, 32> HasLocalDefinition;
 
+  /// The set of blocks that have definitions which specifically "kill" the
+  /// given value. If a block is in this set, there must be an instruction in
+  /// LoadTakeUse whose parent is the block. This is just used to speed up
+  /// computation.
+  ///
+  /// NOTE: These are not considered escapes.
+  llvm::SmallPtrSet<SILBasicBlock *, 32> HasLocalKill;
+
+  /// This is a set of load takes that we are tracking. HasLocalKill is the set
+  /// of parent blocks of these instructions.
+  llvm::SmallPtrSet<SILInstruction *, 8> LoadTakeUses;
+
   /// This is a map of uses that are not loads (i.e., they are Stores,
   /// InOutUses, and Escapes), to their entry in Uses.
   llvm::SmallDenseMap<SILInstruction *, unsigned, 16> NonLoadUses;
@@ -925,10 +948,36 @@ AvailableValueDataflowContext::AvailableValueDataflowContext(
     auto &Use = Uses[ui];
     assert(Use.Inst && "No instruction identified?");
 
-    // Keep track of all the uses that aren't loads.
-    if (Use.Kind == PMOUseKind::Load)
-      continue;
+    // If we have a load...
+    if (Use.Kind == PMOUseKind::Load) {
+      // Skip load borrow use and open_existential_addr.
+      if (isa<LoadBorrowInst>(Use.Inst) || isa<OpenExistentialAddrInst>(Use.Inst))
+        continue;
+
+      // That is not a load take, continue. Otherwise, stash the load [take].
+      if (auto *LI = dyn_cast<LoadInst>(Use.Inst)) {
+        if (LI->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
+          LoadTakeUses.insert(LI);
+          HasLocalKill.insert(LI->getParent());
+        }
+        continue;
+      }
+
+      // If we have a copy_addr as our load, it means we are processing a source
+      // of the value. If the copy_addr is taking from the source, we need to
+      // treat it like a load take use.
+      if (auto *CAI = dyn_cast<CopyAddrInst>(Use.Inst)) {
+        if (CAI->isTakeOfSrc() == IsTake) {
+          LoadTakeUses.insert(CAI);
+          HasLocalKill.insert(CAI->getParent());
+        }
+        continue;
+      }
+
+      llvm_unreachable("Unhandled SILInstructionKind for PMOUseKind::Load?!");
+    }
 
+    // Keep track of all the uses that aren't loads.
     NonLoadUses[Use.Inst] = ui;
     HasLocalDefinition.insert(Use.Inst->getParent());
 
@@ -945,50 +994,149 @@ AvailableValueDataflowContext::AvailableValueDataflowContext(
   HasLocalDefinition.insert(TheMemory->getParent());
 }
 
-void AvailableValueDataflowContext::updateAvailableValues(
-    SILInstruction *Inst, SmallBitVector &RequiredElts,
-    SmallVectorImpl<AvailableValue> &Result,
-    SmallBitVector &ConflictingValues) {
-  // Handle store.
-  if (auto *SI = dyn_cast<StoreInst>(Inst)) {
-    unsigned StartSubElt = computeSubelement(SI->getDest(), TheMemory);
-    assert(StartSubElt != ~0U && "Store within enum projection not handled");
-    SILType ValTy = SI->getSrc()->getType();
+// This function takes in the current (potentially uninitialized) available
+// values for theMemory and for the subset of AvailableValues corresponding to
+// \p address either:
+//
+// 1. If uninitialized, optionally initialize the available value with a new
+//    SILValue. It is optional since in certain cases, (for instance when
+//    invalidating one just wants to skip empty available values).
+//
+// 2. Given an initialized value, either add the given instruction as an
+//    insertion point or state that we have a conflict.
+static inline void updateAvailableValuesHelper(
+    SingleValueInstruction *theMemory, SILInstruction *inst, SILValue address,
+    SmallBitVector &requiredElts, SmallVectorImpl<AvailableValue> &result,
+    SmallBitVector &conflictingValues,
+    function_ref<Optional<AvailableValue>(unsigned)> defaultFunc,
+    function_ref<bool(AvailableValue &, unsigned)> isSafeFunc) {
+  auto &mod = theMemory->getModule();
+  unsigned startSubElt = computeSubelement(address, theMemory);
+
+  // TODO: Is this needed now?
+  assert(startSubElt != ~0U && "Store within enum projection not handled");
+  for (unsigned i :
+       range(getNumSubElements(address->getType().getObjectType(), mod))) {
+    // If this element is not required, don't fill it in.
+    if (!requiredElts[startSubElt + i])
+      continue;
 
-    for (unsigned i : range(getNumSubElements(ValTy, getModule()))) {
-      // If this element is not required, don't fill it in.
-      if (!RequiredElts[StartSubElt+i]) continue;
+    // At this point we know that we will either mark the value as conflicting
+    // or give it a value.
+    requiredElts[startSubElt + i] = false;
+
+    // First see if we have an entry at all.
+    auto &entry = result[startSubElt + i];
+
+    // If we don't...
+    if (!entry) {
+      // and we are told to initialize it, do so.
+      if (auto defaultValue = defaultFunc(i)) {
+        entry = std::move(defaultValue.getValue());
+      } else {
+        // Otherwise, mark this as a conflicting value. There is some available
+        // value here, we just do not know what it is at this point. This
+        // ensures that if we visit a kill where we do not have an entry yet, we
+        // properly invalidate our state.
+        conflictingValues[startSubElt + i] = true;
+      }
+      continue;
+    }
 
-      // This element is now provided.
-      RequiredElts[StartSubElt + i] = false;
+    // Check if our caller thinks that the value currently in entry is
+    // compatible with \p inst. If not, mark the values as conflicting and
+    // continue.
+    if (!isSafeFunc(entry, i)) {
+      conflictingValues[startSubElt + i] = true;
+      continue;
+    }
 
-      // If there is no result computed for this subelement, record it.  If
-      // 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, SI};
-        continue;
-      }
+    // Otherwise, we found another insertion point for our available
+    // value. Today this will always be a Store.
+    entry.addInsertionPoint(cast<StoreInst>(inst));
+  }
+}
 
-      // 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;
-        continue;
-      }
+void AvailableValueDataflowContext::updateAvailableValues(
+    SILInstruction *Inst, SmallBitVector &RequiredElts,
+    SmallVectorImpl<AvailableValue> &Result,
+    SmallBitVector &ConflictingValues) {
 
-      Entry.addInsertionPoint(SI);
+  // If we are visiting a load [take], it invalidates the underlying available
+  // values.
+  //
+  // NOTE: Since we are always looking back from the instruction to promote,
+  // when we attempt to promote the load [take] itself, we will never hit this
+  // code since.
+  if (auto *LI = dyn_cast<LoadInst>(Inst)) {
+    // First see if this is a load inst that we are tracking.
+    if (LoadTakeUses.count(LI)) {
+      updateAvailableValuesHelper(TheMemory, LI, LI->getOperand(), RequiredElts,
+                                  Result, ConflictingValues,
+                                  /*default*/
+                                  [](unsigned) -> Optional<AvailableValue> {
+                                    // We never initialize values. We only
+                                    // want to invalidate.
+                                    return None;
+                                  },
+                                  /*isSafe*/
+                                  [](AvailableValue &, unsigned) -> bool {
+                                    // Always assume values conflict.
+                                    return false;
+                                  });
+      return;
     }
+  }
 
+  // Handle store.
+  if (auto *SI = dyn_cast<StoreInst>(Inst)) {
+    updateAvailableValuesHelper(
+        TheMemory, SI, SI->getDest(), RequiredElts, Result, ConflictingValues,
+        /*default*/
+        [&](unsigned ResultOffset) -> Optional<AvailableValue> {
+          Optional<AvailableValue> Result;
+          Result.emplace(SI->getSrc(), ResultOffset, SI);
+          return Result;
+        },
+        /*isSafe*/
+        [&](AvailableValue &Entry, unsigned ResultOffset) -> bool {
+          // TODO: This is /really/, /really/, conservative. This basically
+          // means that if we do not have an identical store, we will not
+          // promote.
+          return Entry.getValue() == SI->getSrc() &&
+                 Entry.getSubElementNumber() == ResultOffset;
+        });
     return;
   }
-  
-  // If we get here with a copy_addr, it must be storing into the element. Check
-  // to see if any loaded subelements are being used, and if so, explode the
-  // copy_addr to its individual pieces.
+
+  // If we got here from an apply, we must either be initializing the element
+  // via an @out parameter or we are trying to model an invalidating load of the
+  // value (e.x.: indirect_in, indirect_inout).
+
+  // If we get here with a copy_addr, we must either be storing into the element
+  // or tracking some sort of take of the src. First check if we are taking (in
+  // which case, we just track invalidation of src) and continue. Otherwise we
+  // must be storing into the copy_addr so see which loaded subelements are
+  // being used, and if so, explode the copy_addr to its individual pieces.
   if (auto *CAI = dyn_cast<CopyAddrInst>(Inst)) {
+    // If we have a load take use, we must be tracking a store of CAI.
+    if (LoadTakeUses.count(CAI)) {
+      updateAvailableValuesHelper(TheMemory, CAI, CAI->getSrc(), RequiredElts,
+                                  Result, ConflictingValues,
+                                  /*default*/
+                                  [](unsigned) -> Optional<AvailableValue> {
+                                    // We never give values default initialized
+                                    // values. We only want to invalidate.
+                                    return None;
+                                  },
+                                  /*isSafe*/
+                                  [](AvailableValue &, unsigned) -> bool {
+                                    // Always assume values conflict.
+                                    return false;
+                                  });
+      return;
+    }
+
     unsigned StartSubElt = computeSubelement(CAI->getDest(), TheMemory);
     assert(StartSubElt != ~0U && "Store within enum projection not handled");
     SILType ValTy = CAI->getDest()->getType();
@@ -1082,20 +1230,21 @@ void AvailableValueDataflowContext::computeAvailableValuesFrom(
         &VisitedBlocks,
     SmallBitVector &ConflictingValues) {
   assert(!RequiredElts.none() && "Scanning with a goal of finding nothing?");
-  
+
   // If there is a potential modification in the current block, scan the block
-  // to see if the store or escape is before or after the load.  If it is
-  // before, check to see if it produces the value we are looking for.
-  if (HasLocalDefinition.count(BB)) {
+  // to see if the store, escape, or load [take] is before or after the load. If
+  // it is before, check to see if it produces the value we are looking for.
+  bool shouldCheckBlock =
+      HasLocalDefinition.count(BB) || HasLocalKill.count(BB);
+  if (shouldCheckBlock) {
     for (SILBasicBlock::iterator BBI = StartingFrom; BBI != BB->begin();) {
       SILInstruction *TheInst = &*std::prev(BBI);
-
       // If this instruction is unrelated to the element, ignore it.
-      if (!NonLoadUses.count(TheInst)) {
+      if (!NonLoadUses.count(TheInst) && !LoadTakeUses.count(TheInst)) {
         --BBI;
         continue;
       }
-      
+
       // Given an interesting instruction, incorporate it into the set of
       // results, and filter down the list of demanded subelements that we still
       // need.
@@ -1112,7 +1261,6 @@ void AvailableValueDataflowContext::computeAvailableValuesFrom(
     }
   }
 
-  
   // Otherwise, we need to scan up the CFG looking for available values.
   for (auto PI = BB->pred_begin(), E = BB->pred_end(); PI != E; ++PI) {
     SILBasicBlock *PredBB = *PI;
@@ -1174,6 +1322,9 @@ void AvailableValueDataflowContext::explodeCopyAddr(CopyAddrInst *CAI) {
 
   // Update our internal state for this being gone.
   NonLoadUses.erase(CAI);
+  LoadTakeUses.erase(CAI);
+  // NOTE: We do not need to update HasLocalKill since the copy_addr
+  // and the loads/stores will have the same parent block.
 
   // Remove the copy_addr from Uses.  A single copy_addr can appear multiple
   // times if the source and dest are to elements within a single aggregate, but
@@ -1237,6 +1388,12 @@ void AvailableValueDataflowContext::explodeCopyAddr(CopyAddrInst *CAI) {
       // "assign" operation on the destination of the copyaddr.
       if (LoadUse.isValid() &&
           getAccessPathRoot(NewInst->getOperand(0)) == TheMemory) {
+        if (auto *LI = dyn_cast<LoadInst>(NewInst)) {
+          if (LI->getOwnershipQualifier() == LoadOwnershipQualifier::Take) {
+            LoadTakeUses.insert(LI);
+            HasLocalKill.insert(LI->getParent());
+          }
+        }
         LoadUse.Inst = NewInst;
         Uses.push_back(LoadUse);
       }