[AddressLowering] Allow use-projection into phi.

Resolves numerous regressions when enabling AddressLowering early in the
pipeline.

Previously, if a value incoming into a phi had storage which itself was
a use-projection out of some other storage, PhiStorageOptimizer bailed
out.  The result was unnecessary "moves" (i.e. `copy_addr [take] [init]`
instructions).

Here, this bailout is removed.  In order to do this, it is necessary to
find (1) all values whose storage recursively project out of an incoming
value (such a value may have storage which is either a use _or_ a def
projection) and (2) the block which dominates the defs of all these
values.

Together, these values are used to compute liveness to determine
interference.  Previously, the live region was that between the uses of
an incoming value and its defining block.  Now, it is that between the
uses of any of the values found in (1) and the dominating block found in
(2).

Author: nate-chandler

lib/SILOptimizer/Mandatory/AddressLowering.cpp

@@ -1446,7 +1446,7 @@ void OpaqueStorageAllocation::allocatePhi(PhiValue phi) {
   // The phi operand projections are computed first to give them priority. Then
   // we determine if the phi itself can share storage with one of its users.
   CoalescedPhi coalescedPhi;
-  coalescedPhi.coalesce(phi, pass.valueStorageMap);
+  coalescedPhi.coalesce(phi, pass.valueStorageMap, pass.domInfo);
 
   SmallVector<SILValue, 4> coalescedValues;
   coalescedValues.reserve(coalescedPhi.getCoalescedOperands().size());

lib/SILOptimizer/Mandatory/PhiStorageOptimizer.cpp

@@ -87,6 +87,8 @@
 
 #include "PhiStorageOptimizer.h"
 #include "swift/SIL/BasicBlockDatastructures.h"
+#include "swift/SIL/Dominance.h"
+#include "swift/SIL/NodeDatastructures.h"
 #include "swift/SIL/SILBasicBlock.h"
 #include "swift/SIL/SILInstruction.h"
 
@@ -101,35 +103,46 @@ namespace swift {
 class PhiStorageOptimizer {
   PhiValue phi;
   const ValueStorageMap &valueStorageMap;
+  DominanceInfo *domInfo;
 
   CoalescedPhi &coalescedPhi;
 
   BasicBlockSet occupiedBlocks;
 
 public:
   PhiStorageOptimizer(PhiValue phi, const ValueStorageMap &valueStorageMap,
-                      CoalescedPhi &coalescedPhi)
-      : phi(phi), valueStorageMap(valueStorageMap), coalescedPhi(coalescedPhi),
-        occupiedBlocks(getFunction()) {}
+                      CoalescedPhi &coalescedPhi, DominanceInfo *domInfo)
+      : phi(phi), valueStorageMap(valueStorageMap), domInfo(domInfo),
+        coalescedPhi(coalescedPhi), occupiedBlocks(getFunction()) {}
 
   SILFunction *getFunction() const { return phi.phiBlock->getParent(); }
 
   void optimize();
 
 protected:
-  bool hasUseProjection(SILInstruction *defInst);
-  bool canCoalesceValue(SILValue incomingVal);
+  using ProjectedValues = StackList<SILValue>;
+  bool findUseProjections(SILInstruction *defInst,
+                          SmallVectorImpl<Operand *> &operands);
+  bool findDefProjections(SILValue value,
+                          SmallVectorImpl<SILValue> *projecteds);
+  SILBasicBlock *canCoalesceValue(SILValue incomingVal,
+                                  ProjectedValues *projectedValues);
   void tryCoalesceOperand(SILBasicBlock *incomingPred);
-  bool recordUseLiveness(SILValue incomingVal, BasicBlockSetVector &liveBlocks);
+  bool recordUseLiveness(ProjectedValues &incomingVals,
+                         SILBasicBlock *dominatingBlock,
+                         BasicBlockSetVector &liveBlocks);
+  SILBasicBlock *getDominatingBlock(SILValue incomingVal,
+                                    ProjectedValues *projectedValues);
 };
 
 } // namespace swift
 
 void CoalescedPhi::coalesce(PhiValue phi,
-                            const ValueStorageMap &valueStorageMap) {
+                            const ValueStorageMap &valueStorageMap,
+                            DominanceInfo *domInfo) {
   assert(empty() && "attempt to recoalesce the same phi");
 
-  PhiStorageOptimizer(phi, valueStorageMap, *this).optimize();
+  PhiStorageOptimizer(phi, valueStorageMap, *this, domInfo).optimize();
 }
 
 /// Optimize phi storage by coalescing phi operands.
@@ -168,7 +181,8 @@ void CoalescedPhi::coalesce(PhiValue phi,
 void PhiStorageOptimizer::optimize() {
   // The single incoming value case always projects storage.
   if (auto *predecessor = phi.phiBlock->getSinglePredecessorBlock()) {
-    if (canCoalesceValue(phi.getValue()->getIncomingPhiValue(predecessor))) {
+    if (canCoalesceValue(phi.getValue()->getIncomingPhiValue(predecessor),
+                         /*projectedValues=*/nullptr)) {
       // Storage will always be allocated for the phi.  The optimization
       // attempts to let incoming values reuse the phi's storage.  This isn't
       // always possible, even in the single incoming value case.  For example,
@@ -183,21 +197,53 @@ void PhiStorageOptimizer::optimize() {
   }
 }
 
-// Return true if any of \p defInst's operands are composing use projections
-// into \p defInst's storage.
-bool PhiStorageOptimizer::hasUseProjection(SILInstruction *defInst) {
+// Return \p defInst's operand which is a composing use projection into \p
+// defInst's storage, or nullptr if none exists.
+bool PhiStorageOptimizer::findUseProjections(
+    SILInstruction *defInst, SmallVectorImpl<Operand *> &operands) {
+  auto *svi = dyn_cast<SingleValueInstruction>(defInst);
+  if (!svi) {
+    return false;
+  }
+  auto ordinal = valueStorageMap.getOrdinal(svi);
+  bool found = false;
   for (Operand &oper : defInst->getAllOperands()) {
-    if (valueStorageMap.isComposingUseProjection(&oper))
-      return true;
+    if (valueStorageMap.isComposingUseProjection(&oper) &&
+        valueStorageMap.getStorage(oper.get()).projectedStorageID == ordinal) {
+      found = true;
+      operands.push_back(&oper);
+    }
+  }
+  return found;
+}
+
+bool PhiStorageOptimizer::findDefProjections(
+    SILValue value, SmallVectorImpl<SILValue> *projecteds) {
+  bool found = false;
+  for (auto user : value->getUsers()) {
+    for (auto result : user->getResults()) {
+      auto *storage = valueStorageMap.getStorageOrNull(result);
+      if (!storage)
+        continue;
+      if (storage->isDefProjection) {
+        assert(storage->projectedStorageID ==
+               valueStorageMap.getOrdinal(value));
+        projecteds->push_back(result);
+        found = true;
+      }
+    }
   }
-  return false;
+  return found;
 }
 
-// Return true in \p incomingVal can be coalesced with this phi ignoring
-// possible interference. Simply determine whether storage reuse is possible.
+// Determine whether storage reuse is possible with \p incomingVal.  Ignore
+// possible interference.  If reuse is possible, return the block which
+// dominates all defs whose storage is projected out of incomingVal.
 //
 // Precondition: \p incomingVal is an operand of this phi.
-bool PhiStorageOptimizer::canCoalesceValue(SILValue incomingVal) {
+SILBasicBlock *
+PhiStorageOptimizer::canCoalesceValue(SILValue incomingVal,
+                                      ProjectedValues *projectedValues) {
   // A Phi must not project from storage that was initialized on a path that
   // reaches the phi because other uses of the storage may interfere with the
   // phi. A phi may, however, be a composing use projection.
@@ -215,31 +261,20 @@ bool PhiStorageOptimizer::canCoalesceValue(SILValue incomingVal) {
   // could be handled, but would require by recursively following uses across
   // projections when computing liveness.
   if (!incomingStorage.isAllocated() || incomingStorage.isProjection())
-    return false;
+    return nullptr;
 
-  auto *defInst = incomingVal->getDefiningInstruction();
-  if (!defInst) {
-    // Indirect function arguments were replaced by loads.
-    assert(!isa<SILFunctionArgument>(incomingVal));
+  if (PhiValue(incomingVal)) {
     // Do not coalesce a phi with other phis. This would require liveness
     // analysis of the whole phi web before coalescing phi operands.
-    return false;
+    return nullptr;
   }
 
   // Don't coalesce an incoming value unless it's storage is from a stack
   // allocation, which can be replaced with another alloc_stack.
   if (!isa<AllocStackInst>(incomingStorage.storageAddress))
-    return false;
-
-  // Make sure that the incomingVal is not coalesced with any of its operands.
-  //
-  // Handling incomingValues whose operands project into them would require by
-  // recursively finding the set of value definitions and their dominating defBB
-  // instead of simply incomingVal->getParentBlock().
-  if (hasUseProjection(defInst))
-    return false;
+    return nullptr;
 
-  return true;
+  return getDominatingBlock(incomingVal, projectedValues);
 }
 
 // Process a single incoming phi operand. Compute the value's liveness while
@@ -248,11 +283,13 @@ void PhiStorageOptimizer::tryCoalesceOperand(SILBasicBlock *incomingPred) {
   Operand *incomingOper = phi.getOperand(incomingPred);
   SILValue incomingVal = incomingOper->get();
 
-  if (!canCoalesceValue(incomingVal))
+  ProjectedValues projectedValues(incomingPred->getFunction());
+  auto *dominatingBlock = canCoalesceValue(incomingVal, &projectedValues);
+  if (!dominatingBlock)
     return;
 
   BasicBlockSetVector liveBlocks(getFunction());
-  if (!recordUseLiveness(incomingVal, liveBlocks))
+  if (!recordUseLiveness(projectedValues, dominatingBlock, liveBlocks))
     return;
 
   for (auto *block : liveBlocks) {
@@ -262,39 +299,100 @@ void PhiStorageOptimizer::tryCoalesceOperand(SILBasicBlock *incomingPred) {
   coalescedPhi.coalescedOperands.push_back(incomingOper);
 }
 
-// Record liveness generated by uses of \p incomingVal.
+/// The block which dominates all the defs whose storage is projected out of the
+/// storage for \p incomingVal.
+///
+/// Populates \p projectedValues with the values whose storage is recursively
+/// projected out of the storage for incomingVal.
+SILBasicBlock *
+PhiStorageOptimizer::getDominatingBlock(SILValue incomingVal,
+                                        ProjectedValues *projectedValues) {
+  assert(domInfo);
+
+  // Recursively find the set of value definitions and the dominating LCA.
+  ValueWorklist values(incomingVal->getFunction());
+
+  values.push(incomingVal);
+
+  SILBasicBlock *lca = incomingVal->getParentBlock();
+  auto updateLCA = [&](SILBasicBlock *other) {
+    lca = domInfo->findNearestCommonDominator(lca, other);
+  };
+
+  while (auto value = values.pop()) {
+    if (projectedValues)
+      projectedValues->push_back(value);
+    auto *defInst = value->getDefiningInstruction();
+    if (!defInst) {
+      assert(!PhiValue(value));
+      updateLCA(value->getParentBlock());
+      continue;
+    }
+    SmallVector<Operand *, 2> operands;
+    if (findUseProjections(defInst, operands)) {
+      // Any operand whose storage is a use-projection out of `value`'s storage
+      // dominates `value` (i.e. `defInst`), so skip updating the LCA here.
+      for (auto *operand : operands) {
+        values.pushIfNotVisited(operand->get());
+      }
+      continue;
+    }
+    SmallVector<SILValue, 2> projecteds;
+    if (findDefProjections(value, &projecteds)) {
+      // Walking up the projection chain from this point, every subsequent
+      // projection must be a def projection [projection_chain_structure].
+      // Every such projection is dominated by `value` (i.e. defInst).
+      //
+      // If the walk were only updating the LCA, it could stop here.
+      //
+      // It is also collecting values whose storage is projected out of the
+      // phi, however, so the walk must continue.
+      updateLCA(defInst->getParent());
+      for (auto projected : projecteds) {
+        values.pushIfNotVisited(projected);
+      }
+      continue;
+    }
+    updateLCA(defInst->getParent());
+  }
+  return lca;
+}
+
+// Record liveness generated by uses of \p projectedVals.
 //
 // Return true if no interference was detected along the way.
-bool PhiStorageOptimizer::recordUseLiveness(SILValue incomingVal,
+bool PhiStorageOptimizer::recordUseLiveness(ProjectedValues &projectedVals,
+                                            SILBasicBlock *dominatingBlock,
                                             BasicBlockSetVector &liveBlocks) {
   assert(liveBlocks.empty());
 
-  // Stop liveness traversal at defBB.
-  SILBasicBlock *defBB = incomingVal->getParentBlock();
-  for (auto *use : incomingVal->getUses()) {
-    StackList<SILBasicBlock *> liveBBWorklist(getFunction());
+  for (auto projectedVal : projectedVals) {
+    for (auto *use : projectedVal->getUses()) {
+      StackList<SILBasicBlock *> liveBBWorklist(getFunction());
+
+      // If \p liveBB is already occupied by another value, return
+      // false. Otherwise, mark \p liveBB live and push it onto liveBBWorklist.
+      auto visitLiveBlock = [&](SILBasicBlock *liveBB) {
+        assert(liveBB != phi.phiBlock && "phi operands are consumed");
 
-    // If \p liveBB is already occupied by another value, return
-    // false. Otherwise, mark \p liveBB live and push it onto liveBBWorklist.
-    auto visitLiveBlock = [&](SILBasicBlock *liveBB) {
-      assert(liveBB != phi.phiBlock && "phi operands are consumed");
+        if (occupiedBlocks.contains(liveBB))
+          return false;
 
-      if (occupiedBlocks.contains(liveBB))
+        // Stop liveness traversal at dominatingBlock.
+        if (liveBlocks.insert(liveBB) && liveBB != dominatingBlock) {
+          liveBBWorklist.push_back(liveBB);
+        }
+        return true;
+      };
+      if (!visitLiveBlock(use->getUser()->getParent()))
         return false;
 
-      if (liveBlocks.insert(liveBB) && liveBB != defBB) {
-        liveBBWorklist.push_back(liveBB);
-      }
-      return true;
-    };
-    if (!visitLiveBlock(use->getUser()->getParent()))
-      return false;
-
-    while (!liveBBWorklist.empty()) {
-      auto *succBB = liveBBWorklist.pop_back_val();
-      for (auto *predBB : succBB->getPredecessorBlocks()) {
-        if (!visitLiveBlock(predBB))
-          return false;
+      while (!liveBBWorklist.empty()) {
+        auto *succBB = liveBBWorklist.pop_back_val();
+        for (auto *predBB : succBB->getPredecessorBlocks()) {
+          if (!visitLiveBlock(predBB))
+            return false;
+        }
       }
     }
   }

lib/SILOptimizer/Mandatory/PhiStorageOptimizer.h

@@ -24,6 +24,8 @@
 
 namespace swift {
 
+class DominanceInfo;
+
 class CoalescedPhi {
   friend class PhiStorageOptimizer;
 
@@ -37,7 +39,8 @@ class CoalescedPhi {
   CoalescedPhi(CoalescedPhi &&) = default;
   CoalescedPhi &operator=(CoalescedPhi &&) = default;
 
-  void coalesce(PhiValue phi, const ValueStorageMap &valueStorageMap);
+  void coalesce(PhiValue phi, const ValueStorageMap &valueStorageMap,
+                DominanceInfo *domInfo);
 
   bool empty() const { return coalescedOperands.empty(); }
 

test/SILOptimizer/address_lowering.sil

@@ -1431,10 +1431,8 @@ nopers(%error : @owned $any Error):
 
 // CHECK-LABEL: sil [ossa] @f262_testTryApplyIntoPhi : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[ADDR:%[^,]+]] :
-// CHECK:         [[STACK:%[^,]+]] = alloc_stack
-// CHECK:         try_apply undef<R>([[STACK]]) {{.*}}, normal [[NORMAL:bb[0-9]+]]
+// CHECK:         try_apply undef<R>([[ADDR]]) {{.*}}, normal [[NORMAL:bb[0-9]+]]
 // CHECK:       [[NORMAL]]({{%[^,]+}} : $()):
-// CHECK:         copy_addr [take] [[STACK]] to [init] [[ADDR]]
 // CHECK:         br [[EXIT:bb[0-9]+]]                                          
 // CHECK:       [[EXIT]]:                                              
 // CHECK:         return {{%[^,]+}} : $()                                 
@@ -2268,15 +2266,13 @@ bb3:
 // Test the result being stored into an @out enum.
 // CHECK-LABEL: sil [ossa] @test_checked_cast_br4 : {{.*}} {
 // CHECK:       {{bb[0-9]+}}([[OUT_ADDR:%[^,]+]] : $*Optional<T>, [[INSTANCE:%[^,]+]] : @owned $TestGeneric<Element>):
-// CHECK:         [[TEMP:%[^,]+]] = alloc_stack $Optional<T>
 // CHECK:         [[CAST_SOURCE_ADDR:%[^,]+]] = alloc_stack $TestGeneric<Element>
 // CHECK:         store [[INSTANCE]] to [init] [[CAST_SOURCE_ADDR]]
-// CHECK:         [[CAST_DEST_ADDR:%[^,]+]] = init_enum_data_addr [[TEMP]] : $*Optional<T>, #Optional.some!enumelt
+// CHECK:         [[CAST_DEST_ADDR:%[^,]+]] = init_enum_data_addr [[OUT_ADDR]] : $*Optional<T>, #Optional.some!enumelt
 // CHECK:         checked_cast_addr_br take_on_success TestGeneric<Element> in [[CAST_SOURCE_ADDR]] : $*TestGeneric<Element> to T in [[CAST_DEST_ADDR]] : $*T, [[SUCCESS:bb[0-9]+]]
 // CHECK:       [[SUCCESS]]:
 // CHECK:         dealloc_stack [[CAST_SOURCE_ADDR]]
-// CHECK:         inject_enum_addr [[TEMP]] : $*Optional<T>, #Optional.some!enumelt
-// CHECK:         copy_addr [take] [[TEMP]] to [init] [[OUT_ADDR]]
+// CHECK:         inject_enum_addr [[OUT_ADDR]] : $*Optional<T>, #Optional.some!enumelt
 // CHECK-LABEL: } // end sil function 'test_checked_cast_br4'
 sil [ossa] @test_checked_cast_br4 : $@convention(method) <Element><T> (@owned TestGeneric<Element>) -> @out Optional<T> {
 bb0(%3 : @owned $TestGeneric<Element>):