Utils: Inhibit load/store folding through phis for llvm.protected.field.ptr.

Protected pointer field loads/stores should be paired with the intrinsic
to avoid unnecessary address escapes.

Pull Request: llvm#151649

Author: pcc

llvm/include/llvm/Transforms/Utils/Local.h

@@ -182,6 +182,10 @@ LLVM_ABI bool EliminateDuplicatePHINodes(BasicBlock *BB);
 LLVM_ABI bool EliminateDuplicatePHINodes(BasicBlock *BB,
                                          SmallPtrSetImpl<PHINode *> &ToRemove);
 
+/// Returns whether it is allowed and beneficial for optimizations to transform
+/// phi(load(ptr)) into load(phi(ptr)) or a similar transformation for stores.
+bool shouldFoldLoadStoreWithPointerOperandThroughPhi(const Value *Ptr);
+
 /// This function is used to do simplification of a CFG.  For example, it
 /// adjusts branches to branches to eliminate the extra hop, it eliminates
 /// unreachable basic blocks, and does other peephole optimization of the CFG.

llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp

@@ -697,8 +697,7 @@ static bool isSafeAndProfitableToSinkLoad(LoadInst *L) {
 Instruction *InstCombinerImpl::foldPHIArgLoadIntoPHI(PHINode &PN) {
   LoadInst *FirstLI = cast<LoadInst>(PN.getIncomingValue(0));
 
-  // Can't forward swifterror through a phi.
-  if (FirstLI->getOperand(0)->isSwiftError())
+  if (!shouldFoldLoadStoreWithPointerOperandThroughPhi(FirstLI->getOperand(0)))
     return nullptr;
 
   // FIXME: This is overconservative; this transform is allowed in some cases
@@ -737,8 +736,7 @@ Instruction *InstCombinerImpl::foldPHIArgLoadIntoPHI(PHINode &PN) {
         LI->getPointerAddressSpace() != LoadAddrSpace)
       return nullptr;
 
-    // Can't forward swifterror through a phi.
-    if (LI->getOperand(0)->isSwiftError())
+    if (!shouldFoldLoadStoreWithPointerOperandThroughPhi(LI->getOperand(0)))
       return nullptr;
 
     // We can't sink the load if the loaded value could be modified between

llvm/lib/Transforms/Utils/Local.cpp

@@ -3846,10 +3846,7 @@ bool llvm::canReplaceOperandWithVariable(const Instruction *I, unsigned OpIdx) {
   if (Op->getType()->isMetadataTy())
     return false;
 
-  // swifterror pointers can only be used by a load, store, or as a swifterror
-  // argument; swifterror pointers are not allowed to be used in select or phi
-  // instructions.
-  if (Op->isSwiftError())
+  if (!shouldFoldLoadStoreWithPointerOperandThroughPhi(Op))
     return false;
 
   // Cannot replace alloca argument with phi/select.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

@@ -2134,6 +2134,22 @@ static bool replacingOperandWithVariableIsCheap(const Instruction *I,
   return !isa<IntrinsicInst>(I);
 }
 
+bool llvm::shouldFoldLoadStoreWithPointerOperandThroughPhi(const Value *Ptr) {
+  // swifterror pointers can only be used by a load or store; sinking a load
+  // or store would require introducing a select for the pointer operand,
+  // which isn't allowed for swifterror pointers.
+  if (Ptr->isSwiftError())
+    return false;
+
+  // Protected pointer field loads/stores should be paired with the intrinsic
+  // to avoid unnecessary address escapes.
+  if (auto *II = dyn_cast<IntrinsicInst>(Ptr))
+    if (II->getIntrinsicID() == Intrinsic::protected_field_ptr)
+      return false;
+
+  return true;
+}
+
 // All instructions in Insts belong to different blocks that all unconditionally
 // branch to a common successor. Analyze each instruction and return true if it
 // would be possible to sink them into their successor, creating one common

llvm/test/Transforms/Util/phi-protected-field-ptr.ll

@@ -0,0 +1,33 @@
+; RUN: opt -O2 -S < %s | FileCheck %s
+
+; Test that no optimization run at -O2 moves the loads into the exit block,
+; as this causes unnecessary address escapes with pointer field protection.
+
+target triple = "aarch64-unknown-linux-gnu"
+
+define ptr @phi_prot_ptr(i1 %sel, ptr %p1, ptr %p2) {
+  br i1 %sel, label %t, label %f
+
+; CHECK: t:
+t:
+  ; CHECK-NEXT: call
+  %protp1 = call ptr @llvm.protected.field.ptr(ptr %p1, i64 1, i1 true)
+  ; CHECK-NEXT: load
+  %load1 = load ptr, ptr %protp1
+  br label %exit
+
+; CHECK: f:
+f:
+  ; CHECK-NEXT: call
+  %protp2 = call ptr @llvm.protected.field.ptr(ptr %p2, i64 2, i1 true)
+  ; CHECK-NEXT: load
+  %load2 = load ptr, ptr %protp2
+  br label %exit
+
+; CHECK: exit:
+exit:
+  ; CHECK-NEXT: phi
+  %retval = phi ptr [ %load1, %t ], [ %load2, %f ]
+  ; CHECK-NEXT: ret
+  ret ptr %retval
+}