EscapeAnalysis: make the use-point analysis more precise

In canEscapeToUsePoint only check the content node if it's a reference (see comment why this is needed).
For all other node types, especially addresses, handle defer edges by propagating use-point infomation backward in the graph.
This makes escape analysis more precise with address types, e.g. don't consider an inout address to escape to an apply if just the loaded value is passed to an apply argument.

Author: eeckstein

include/swift/SILOptimizer/Analysis/EscapeAnalysis.h

@@ -411,6 +411,16 @@ class EscapeAnalysis : public BottomUpIPAnalysis {
         return false;
       return true;
     }
+
+    /// Returns true if the node has any predecessors with link to this node
+    /// with a defer-edge.
+    bool hasDeferPredecessors() const {
+      for (Predecessor Pred : Preds) {
+        if (Pred.getInt() == EdgeType::Defer)
+          return true;
+      }
+      return false;
+    }
 
     friend class CGNodeMap;
     friend class ConnectionGraph;
@@ -631,19 +641,7 @@ class EscapeAnalysis : public BottomUpIPAnalysis {
     }
 
     /// Adds an argument/instruction in which the node's value is used.
-    int addUsePoint(CGNode *Node, SILNode *User) {
-      if (Node->getEscapeState() >= EscapeState::Global)
-        return -1;
-
-      User = User->getRepresentativeSILNodeInObject();
-      int Idx = (int)UsePoints.size();
-      assert(UsePoints.count(User) == 0 && "value is already a use-point");
-      UsePoints[User] = Idx;
-      UsePointTable.push_back(User);
-      assert(UsePoints.size() == UsePointTable.size());
-      Node->setUsePointBit(Idx);
-      return Idx;
-    }
+    int addUsePoint(CGNode *Node, SILNode *User);
 
     /// Specifies that the node's value escapes to global or unidentified
     /// memory.

lib/SILOptimizer/Analysis/EscapeAnalysis.cpp

@@ -376,6 +376,27 @@ updatePointsTo(CGNode *InitialNode, CGNode *pointsTo) {
   clearWorkListFlags(WorkList);
 }
 
+int EscapeAnalysis::ConnectionGraph::addUsePoint(CGNode *Node, SILNode *User) {
+  // For escaping nodes we have to make worst case assumptions anyway. So no
+  // need to store use point information.
+  if (Node->getEscapeState() >= EscapeState::Global &&
+      // ... except it has defer predecessor edges, which are (potentially) not
+      // escaping: use-points are propagated in reverse direction along defer
+      // edges!
+      !Node->hasDeferPredecessors()) {
+    return -1;
+  }
+
+  User = User->getRepresentativeSILNodeInObject();
+  int Idx = (int)UsePoints.size();
+  assert(UsePoints.count(User) == 0 && "value is already a use-point");
+  UsePoints[User] = Idx;
+  UsePointTable.push_back(User);
+  assert(UsePoints.size() == UsePointTable.size());
+  Node->setUsePointBit(Idx);
+  return Idx;
+}
+
 void EscapeAnalysis::ConnectionGraph::propagateEscapeStates() {
   bool Changed = false;
   do {
@@ -449,6 +470,10 @@ void EscapeAnalysis::ConnectionGraph::computeUsePoints() {
       for (CGNode *Def : Node->defersTo) {
         Changed |= Def->mergeUsePoints(Node);
       }
+      for (Predecessor Pred : Node->Preds) {
+        if (Pred.getInt() == EdgeType::Defer)
+          Changed |= Pred.getPointer()->mergeUsePoints(Node);
+      }
     }
   } while (Changed);
 }
@@ -1858,23 +1883,28 @@ bool EscapeAnalysis::canEscapeToUsePoint(SILValue V, SILNode *UsePoint,
   if (ConGraph->isUsePoint(UsePoint, Node))
     return true;
 
-  assert(isPointer(V) && "should not have a node for a non-pointer");
-
-  // Check if the object "content" can escape to the called function.
-  // This will catch cases where V is a reference and a pointer to a stored
-  // property escapes.
-  // It's also important in case of a pointer assignment, e.g.
-  //    V = V1
-  //    apply(V1)
-  // In this case the apply is only a use-point for V1 and V1's content node.
-  // As V1's content node is the same as V's content node, we also make the
-  // check for the content node.
-  CGNode *ContentNode = ConGraph->getContentNode(Node);
-  if (ContentNode->escapesInsideFunction(false))
-    return true;
+  if (V->getType().hasReferenceSemantics()) {
+    // In case V is the result of getUnderlyingObject, we also have to check
+    // the content node. Example:
+    //
+    //   %a = ref_element %r
+    //   apply %f(%a)
+    //
+    // The reference %r does not actually escape to the function call. But in
+    // case this API is called like
+    //
+    //   canEscapeToUsePoint(getUnderlyingObject(applyArgument))
+    //
+    // it would yield false, although the projected object address is passed to
+    // the function.
+
+    CGNode *ContentNode = ConGraph->getContentNode(Node);
+    if (ContentNode->escapesInsideFunction(false))
+      return true;
 
-  if (ConGraph->isUsePoint(UsePoint, ContentNode))
-    return true;
+    if (ConGraph->isUsePoint(UsePoint, ContentNode))
+      return true;
+  }
 
   return false;
 }

lib/SILOptimizer/Transforms/StackPromotion.cpp

@@ -127,8 +127,9 @@ bool StackPromotion::tryPromoteAlloc(AllocRefInst *ARI, EscapeAnalysis *EA,
     if (SILInstruction *I = dyn_cast<SILInstruction>(UsePoint)) {
       UsePoints.push_back(I);
     } else {
-      // Also block arguments can be use points.
-      SILBasicBlock *UseBB = cast<SILPhiArgument>(UsePoint)->getParent();
+      // Also block arguments can be use points (even function arguments, as
+      // use-points are propagated backwards along defer edges).
+      SILBasicBlock *UseBB = cast<SILArgument>(UsePoint)->getParent();
       // For simplicity we just add the first instruction of the block as use
       // point.
       UsePoints.push_back(&UseBB->front());

test/SILOptimizer/escape_analysis.sil

@@ -119,7 +119,7 @@ bb0(%0 : $*Y, %1 : $X):
 // CHECK-LABEL: CG of deferringEdge
 // CHECK-NEXT:    Arg %0 Esc: A, Succ: (%3.1)
 // CHECK-NEXT:    Arg %1 Esc: A, Succ:
-// CHECK-NEXT:    Val %3 Esc: %3, Succ: (%3.1), %0
+// CHECK-NEXT:    Val %3 Esc: %0,%3, Succ: (%3.1), %0
 // CHECK-NEXT:    Con %3.1 Esc: A, Succ: (%3.2)
 // CHECK-NEXT:    Con %3.2 Esc: A, Succ: %1
 // CHECK-NEXT:    Ret Esc: R, Succ: %0
@@ -153,8 +153,8 @@ bb0:
 // CHECK-NEXT:    Val %1 Esc: A, Succ: (%1.1)
 // CHECK-NEXT:    Con %1.1 Esc: A, Succ: (%11.1)
 // CHECK-NEXT:    Val %4 Esc: A, Succ: (%1.1)
-// CHECK-NEXT:    Val %7 Esc: %11, Succ: (%1.1)
-// CHECK-NEXT:    Val %11 Esc: %11, Succ: (%1.1), %7, %11.1
+// CHECK-NEXT:    Val %7 Esc: %0,%11, Succ: (%1.1)
+// CHECK-NEXT:    Val %11 Esc: %0,%11, Succ: (%1.1), %7, %11.1
 // CHECK-NEXT:    Con %11.1 Esc: A, Succ: (%1.1), %0, %1, %4
 // CHECK-NEXT:    Ret Esc: R, Succ: %11.1
 // CHECK-NEXT:  End
@@ -212,7 +212,7 @@ bb0(%0 : $LinkedNode):
 
 // CHECK-LABEL: CG of loadNext
 // CHECK-NEXT:    Arg %0 Esc: A, Succ: (%2.1)
-// CHECK-NEXT:    Val %2 Esc: %2, Succ: (%2.1), %0, %2.2
+// CHECK-NEXT:    Val %2 Esc: %0,%2, Succ: (%2.1), %0, %2.2
 // CHECK-NEXT:    Con %2.1 Esc: A, Succ: (%2.2)
 // CHECK-NEXT:    Con %2.2 Esc: A, Succ: (%2.1)
 // CHECK-NEXT:    Ret Esc: R, Succ: %2.2
@@ -406,12 +406,12 @@ sil @call_copy_addr_content : $@convention(thin) () -> () {
 // CHECK-LABEL: CG of test_partial_apply
 // CHECK-NEXT:    Arg %1 Esc: G, Succ:
 // CHECK-NEXT:    Arg %2 Esc: A, Succ: (%6.3)
-// CHECK-NEXT:    Val %3 Esc: %14,%15,%17, Succ: (%6.1)
-// CHECK-NEXT:    Val %6 Esc: %14,%15,%16, Succ: (%6.1)
+// CHECK-NEXT:    Val %3 Esc: %14,%15,%16,%17, Succ: (%6.1)
+// CHECK-NEXT:    Val %6 Esc: %14,%15,%16,%17, Succ: (%6.1)
 // CHECK-NEXT:    Con %6.1 Esc: %14,%15,%16,%17, Succ: (%6.2)
-// CHECK-NEXT:    Con %6.2 Esc: %14,%15,%16,%17, Succ: %2
+// CHECK-NEXT:    Con %6.2 Esc: %2,%14,%15,%16,%17, Succ: %2
 // CHECK-NEXT:    Con %6.3 Esc: G, Succ:
-// CHECK-NEXT:    Val %12 Esc: %14,%15, Succ: %3, %6
+// CHECK-NEXT:    Val %12 Esc: %14,%15,%16,%17, Succ: %3, %6
 // CHECK-NEXT:  End
 sil @test_partial_apply : $@convention(thin) (Int64, @owned X, @owned Y) -> Int64 {
 bb0(%0 : $Int64, %1 : $X, %2 : $Y):
@@ -443,7 +443,7 @@ bb0(%0 : $Int64, %1 : $X, %2 : $Y):
 // CHECK-NEXT:    Con %2.1 Esc: A, Succ: (%2.2)
 // CHECK-NEXT:    Con %2.2 Esc: A, Succ: (%2.3)
 // CHECK-NEXT:    Con %2.3 Esc: G, Succ:
-// CHECK-NEXT:    Val %7 Esc: %8, Succ: %2
+// CHECK-NEXT:    Val %7 Esc: %2,%8, Succ: %2
 // CHECK-NEXT:  End
 sil @closure1 : $@convention(thin) (@owned X, @owned <τ_0_0> { var τ_0_0 } <Int64>, @owned <τ_0_0> { var τ_0_0 } <Y>) -> Int64 {
 bb0(%0 : $X, %1 : $<τ_0_0> { var τ_0_0 } <Int64>, %2 : $<τ_0_0> { var τ_0_0 } <Y>):
@@ -761,7 +761,7 @@ sil @unknown_throwing_func : $@convention(thin) (@owned X) -> (@owned X, @error
 // CHECK-NEXT:    Arg %0 Esc: A, Succ: (%1.3)
 // CHECK-NEXT:    Val %1 Esc: %6, Succ: (%1.1)
 // CHECK-NEXT:    Con %1.1 Esc: %6, Succ: (%1.2)
-// CHECK-NEXT:    Con %1.2 Esc: %6, Succ: %0
+// CHECK-NEXT:    Con %1.2 Esc: %0,%6, Succ: %0
 // CHECK-NEXT:    Con %1.3 Esc: G, Succ:
 // CHECK-NEXT:    Val %5 Esc: %6, Succ: %1
 // CHECK-NEXT:  End
@@ -969,7 +969,7 @@ bb0(%0 : $Builtin.Int64, %1 : $X, %2 : $X, %3 : $X):
 // CHECK-NEXT:    Arg %0 Esc: A, Succ:
 // CHECK-NEXT:    Val %1 Esc: , Succ: (%1.1)
 // CHECK-NEXT:    Con %1.1 Esc: , Succ: (%1.2)
-// CHECK-NEXT:    Con %1.2 Esc: , Succ: %0
+// CHECK-NEXT:    Con %1.2 Esc: %0, Succ: %0
 // CHECK-NEXT:  End
 sil @test_existential_addr : $@convention(thin) (@owned Pointer) -> () {
 bb0(%0 : $Pointer):
@@ -1181,7 +1181,7 @@ bb0(%0 : $*Array<X>):
 // CHECK-LABEL: CG of arraysemantics_get_element_address
 // CHECK-NEXT:    Arg %0 Esc: A, Succ: (%0.1)
 // CHECK-NEXT:    Con %0.1 Esc: A, Succ:
-// CHECK-NEXT:    Val %4 Esc: , Succ: %0.1
+// CHECK-NEXT:    Val %4 Esc: %0,%4, Succ: %0.1
 // CHECK-NEXT:  End
 sil @arraysemantics_get_element_address : $@convention(thin) (Array<X>) -> () {
 bb0(%0 : $Array<X>):