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

[pcc]

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

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Peter Collingbourne (pcc)

Changes

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

---

Full diff: https://github.com/llvm/llvm-project/pull/151649.diff

5 Files Affected:

• (modified) llvm/include/llvm/Transforms/Utils/Local.h (+4)
• (modified) llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp (+2-4)
• (modified) llvm/lib/Transforms/Utils/Local.cpp (+1-4)
• (modified) llvm/lib/Transforms/Utils/SimplifyCFG.cpp (+16)
• (added) llvm/test/Transforms/Util/phi-protected-field-ptr.ll (+33)

diff --git a/llvm/include/llvm/Transforms/Utils/Local.h b/llvm/include/llvm/Transforms/Utils/Local.h
index 3f5f4278a2766..557f610e10851 100644
--- a/llvm/include/llvm/Transforms/Utils/Local.h
+++ b/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.
diff --git a/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp b/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp
index 6477141ab095f..f156883d9dcee 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombinePHI.cpp
+++ b/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
diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index babd7f6b3a058..e2bf56b774493 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/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.
diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 674de573f7919..6b7a05b1bd57c 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/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
diff --git a/llvm/test/Transforms/Util/phi-protected-field-ptr.ll b/llvm/test/Transforms/Util/phi-protected-field-ptr.ll
new file mode 100644
index 0000000000000..2c667125edda3
--- /dev/null
+++ b/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
+}

[nikic]

This is going to apply not to just loads and stores.

[pcc]

Okay, gave it a better name.

[nikic]

Needs to either drop the triple (more likely) or add REQUIRES.

[pcc]

Done

[nikic]

I think my general preference would be to not add a new API for this and just call canReplaceOperandWithVariable() here.

[pcc]

Sounds good, done.

[nikic]

LGTM