[llvm-branch-commits] [llvm] [CGP][PAC] Flip PHI and blends when all immediate modifiers are the same (PR #150226)

[Oliver Hunt via llvm-branch-commits]

================
@@ -8309,6 +8316,74 @@ bool CodeGenPrepare::optimizeExtractElementInst(Instruction *Inst) {
   return false;
 }
 
+// Given the instruction Inst, rewrite its discriminator operand Disc if it is
+// a PHI node with all incoming values being @llvm.ptrauth.blend(addr, imm)
+// with the same immediate modifier.
+bool CodeGenPrepare::optimizePtrauthInst(Instruction *Inst, Value *Disc) {
+  // GVN PRE, SimplifyCFG and possibly other passes may hoist or sink the call
+  // to @llvm.ptrauth.blend intrinsic, introducing multiple duplicate call
+  // instructions hidden behind a PHI node.
+  //
+  // To enforce the specific immediate modifier being blended into the
+  // discriminator even if the other, address component was reloaded from the
+  // stack, the AArch64 backend defines a number of pseudo instructions
+  // representing auth, sign and other operations. These pseudo instructions
+  // have separate register and immediate operands absorbing the arguments of
+  // blend intrinsic in case of a common `op(ptr, key_id, blend(addr, imm))`
+  // code pattern.
+  //
+  // To help the instruction selector, this function detects the discriminators
+  // represented by a PHI node with all incoming values being `blend`s with
+  // the same integer operand - each such discriminator is rewritten as a single
+  // blend, whose address operand is a PHI node.
+
+  PHINode *P = dyn_cast<PHINode>(Disc);
+  if (!P)
+    return false;
+
+  // Checks if V is blend(something, imm), returns imm if it is.
+  auto GetImmModifier = [](Value *V) -> std::optional<uint64_t> {
+    auto *II = dyn_cast<IntrinsicInst>(V);
+    if (!II || II->getIntrinsicID() != Intrinsic::ptrauth_blend)
+      return std::nullopt;
+    auto *ImmModifier = dyn_cast<ConstantInt>(II->getArgOperand(1));
+    if (!ImmModifier)
----------------
ojhunt wrote:

I was about to ask "when is there no second argument" because apparently reading `dyn_cast` was too hard for me :D

but is this limitation strictly required? is the concern spilling a variable extra discriminator along the lines of?
(pseudo ir)
```
...
%a = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 %v1)
...
%b=  call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 %v1)
...
%r = phi(%a,%b)
```
getting turned into
```
%v1 = ...
; spill %v1
...
; load %v1
%x = phi(%a, %b)
%r = call i64 @llvm.ptrauth.blend(i64 %x, i64 %v1)
```



https://github.com/llvm/llvm-project/pull/150226