[AddressLowering] Sink projections to uses.

Visit the tree of projections in post-order, sinking each to the lowest
position in the least common ancestor of all uses.

Author: nate-chandler

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -137,6 +137,7 @@
 #include "PhiStorageOptimizer.h"
 #include "swift/AST/Decl.h"
 #include "swift/Basic/BlotSetVector.h"
+#include "swift/Basic/FrozenMultiMap.h"
 #include "swift/Basic/Range.h"
 #include "swift/SIL/BasicBlockUtils.h"
 #include "swift/SIL/DebugUtils.h"
@@ -522,6 +523,9 @@ struct AddressLoweringState {
   // Handle moves from a phi's operand storage to the phi storage.
   std::unique_ptr<PhiRewriter> phiRewriter;
 
+  // Projections created for uses, recorded in order to be sunk.
+  SmallVector<SILValue, 16> useProjections;
+
   AddressLoweringState(SILFunction *function, DominanceInfo *domInfo,
                        DeadEndBlocks *deBlocks)
       : function(function), loweredFnConv(getLoweredFnConv(function)),
@@ -1147,6 +1151,8 @@ class OpaqueStorageAllocation {
   /// jointly postdominate.
   void finalizeOpaqueStorage();
 
+  void sinkProjections();
+
 protected:
   void allocateValue(SILValue value);
   bool findProjectionIntoUseImpl(SILValue value,
@@ -1501,6 +1507,54 @@ AllocStackInst *OpaqueStorageAllocation::createStackAllocation(SILValue value) {
   return alloc;
 }
 
+namespace {
+enum class SinkResult {
+  NoUsers,
+  Unmoved,
+  Moved,
+};
+SinkResult sinkToUses(SingleValueInstruction *svi, DominanceInfo *domInfo) {
+  // Fast paths for 0 and 1 users.
+
+  if (svi->use_begin() == svi->use_end()) {
+    return SinkResult::NoUsers;
+  }
+
+  if (auto *use = svi->getSingleUse()) {
+    auto *user = use->getUser();
+    if (user == svi->getNextInstruction())
+      return SinkResult::Unmoved;
+    svi->moveBefore(user);
+    return SinkResult::Moved;
+  }
+
+  // Compute the lca and sink the instruction to before its first user or its
+  // end if there are none.
+
+  SILBasicBlock *lca = domInfo->getLeastCommonAncestorOfUses(svi);
+
+  // The lca may contain a user.  Look for the user to insert before it.
+
+  InstructionSet userSet(svi->getFunction());
+  for (auto user : svi->getUsers()) {
+    userSet.insert(user);
+  }
+
+  for (auto &instruction : *lca) {
+    if (userSet.contains(&instruction)) {
+      if (&instruction == svi->getNextInstruction())
+        return SinkResult::Unmoved;
+      svi->moveBefore(&instruction);
+      return SinkResult::Moved;
+    }
+  }
+
+  // No user was found in the lca, move to before the end.
+  svi->moveBefore(&lca->back());
+  return SinkResult::Moved;
+}
+} // end anonymous namespace
+
 void OpaqueStorageAllocation::finalizeOpaqueStorage() {
   SmallVector<SILBasicBlock *, 4> boundary;
   for (auto maybeAlloc : allocs) {
@@ -1531,6 +1585,79 @@ void OpaqueStorageAllocation::finalizeOpaqueStorage() {
   }
 }
 
+void OpaqueStorageAllocation::sinkProjections() {
+  // First, sink use projections to their uses.
+  //
+  // Done in reverse order because outer projections are materialized first and
+  // so appear in `useProjections` before inner projections, and inner
+  // projections must be sunk first.
+  for (auto projection : llvm::reverse(pass.useProjections)) {
+    assert(projection);
+    auto *svi = dyn_cast<SingleValueInstruction>(projection);
+    assert(svi);
+    auto sank = sinkToUses(svi, pass.domInfo);
+    if (sank == SinkResult::NoUsers) {
+      pass.deleter.forceDelete(svi);
+    }
+  }
+
+  // Second, sink all storage from the valueStorageMap.
+  //
+  // Calculate the (multi-rooted) tree of projections represented by their
+  // indices in `valueStorageMap` (a).  Then visit the projections in
+  // post-order (b), in order to visit all inner projections before visiting
+  // outer projections.
+
+  // (a) Assemble the tree.
+  SmallFrozenMultiMap<uint32_t, uint32_t, 32> projections;
+  SmallVector<uint32_t, 16> roots;
+  uint32_t index = 0;
+  for (auto &valueAndStorage : pass.valueStorageMap) {
+    auto &storage = valueAndStorage.storage;
+    if (storage.isProjection()) {
+      projections.insert(storage.projectedStorageID, index);
+    } else {
+      roots.push_back(index);
+    }
+    ++index;
+  }
+  projections.setFrozen();
+
+  auto sinkProjectionAtIndex = [&](uint32_t index) {
+    auto pair = pass.valueStorageMap[index];
+    auto addr = pair.storage.getMaterializedAddress();
+    if (!pair.storage.isProjection())
+      return;
+    auto *inst = dyn_cast<SingleValueInstruction>(addr);
+    if (!inst)
+      return;
+    auto sank = sinkToUses(inst, pass.domInfo);
+    if (sank == SinkResult::NoUsers) {
+      pass.deleter.forceDelete(inst);
+    }
+  };
+
+  // (b) Visit the projections in post-order.
+  for (auto root : roots) {
+    SmallVector<std::pair<uint32_t, size_t>, 32> stack;
+    stack.push_back({root, 0});
+    while (!stack.empty()) {
+      while (stack.back().second <
+             projections.find(stack.back().first)
+                 .transform([](auto &i) { return i.size(); })
+                 .value_or(0)) {
+        auto child =
+            projections.find(stack.back().first).value()[stack.back().second];
+        stack.back().second++;
+        stack.push_back({child, 0});
+      }
+      auto index = stack.back().first;
+      sinkProjectionAtIndex(index);
+      stack.pop_back();
+    }
+  }
+}
+
 //===----------------------------------------------------------------------===//
 //                           AddressMaterialization
 //
@@ -1590,6 +1717,8 @@ class AddressMaterialization {
                                    SILValue elementValue, unsigned fieldIdx);
 
   SILValue materializeProjectionIntoUse(Operand *operand, bool intoPhiOperand);
+  SILValue materializeProjectionIntoUseImpl(Operand *operand,
+                                            bool intoPhiOperand);
 
   SILValue materializeComposingUser(SingleValueInstruction *user,
                                     bool intoPhiOperand) {
@@ -1771,12 +1900,20 @@ SILValue AddressMaterialization::materializeTupleExtract(
       elementValue->getType().getAddressType());
 }
 
+SILValue
+AddressMaterialization::materializeProjectionIntoUse(Operand *operand,
+                                                     bool intoPhiOperand) {
+  auto projection = materializeProjectionIntoUseImpl(operand, intoPhiOperand);
+  pass.useProjections.push_back(projection);
+  return projection;
+}
+
 /// Recursively materialize the address of a subobject that is a member of the
 /// operand's user. The operand's user must be an aggregate struct, tuple, enum,
 /// init_existential_value.
 SILValue
-AddressMaterialization::materializeProjectionIntoUse(Operand *operand,
-                                                     bool intoPhiOperand) {
+AddressMaterialization::materializeProjectionIntoUseImpl(Operand *operand,
+                                                         bool intoPhiOperand) {
   SILInstruction *user = operand->getUser();
   switch (user->getKind()) {
   default:
@@ -4183,6 +4320,8 @@ void AddressLowering::runOnFunction(SILFunction *function) {
 
   allocator.finalizeOpaqueStorage();
 
+  allocator.sinkProjections();
+
   deleteRewrittenInstructions(pass);
 
   StackNesting::fixNesting(function);

lib/SILOptimizer/Mandatory/AddressLowering.h

@@ -270,6 +270,12 @@ class ValueStorageMap {
     return hashIter->second;
   }
 
+  ValueStoragePair &operator[](uint32_t index) { return valueVector[index]; }
+
+  ValueStoragePair const &operator[](uint32_t index) const {
+    return valueVector[index];
+  }
+
   ValueStorage &getStorage(SILValue value) {
     return valueVector[getOrdinal(value)].storage;
   }