diff --git a/llvm/lib/CodeGen/CodeGenPrepare.cpp b/llvm/lib/CodeGen/CodeGenPrepare.cpp index c21058ca51344..336470801e610 100644 --- a/llvm/lib/CodeGen/CodeGenPrepare.cpp +++ b/llvm/lib/CodeGen/CodeGenPrepare.cpp @@ -444,6 +444,7 @@ class CodeGenPrepare { bool optimizeSwitchPhiConstants(SwitchInst *SI); bool optimizeSwitchInst(SwitchInst *SI); bool optimizeExtractElementInst(Instruction *Inst); + bool optimizePtrauthInst(Instruction *Inst, Value *Disc); bool dupRetToEnableTailCallOpts(BasicBlock *BB, ModifyDT &ModifiedDT); bool fixupDbgVariableRecord(DbgVariableRecord &I); bool fixupDbgVariableRecordsOnInst(Instruction &I); @@ -2765,6 +2766,12 @@ bool CodeGenPrepare::optimizeCallInst(CallInst *CI, ModifyDT &ModifiedDT) { return optimizeGatherScatterInst(II, II->getArgOperand(0)); case Intrinsic::masked_scatter: return optimizeGatherScatterInst(II, II->getArgOperand(1)); + case Intrinsic::ptrauth_auth: + case Intrinsic::ptrauth_sign: + return optimizePtrauthInst(II, II->getArgOperand(2)); + case Intrinsic::ptrauth_resign: + return optimizePtrauthInst(II, II->getArgOperand(2)) || + optimizePtrauthInst(II, II->getArgOperand(4)); } SmallVector PtrOps; @@ -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(Disc); + if (!P) + return false; + + // Checks if V is blend(something, imm), returns imm if it is. + auto GetImmModifier = [](Value *V) -> std::optional { + auto *II = dyn_cast(V); + if (!II || II->getIntrinsicID() != Intrinsic::ptrauth_blend) + return std::nullopt; + auto *ImmModifier = dyn_cast(II->getArgOperand(1)); + if (!ImmModifier) + return std::nullopt; + return ImmModifier->getZExtValue(); + }; + + // Check that all incoming values of P are blends with the same immediate + // modifier. + std::optional ImmModifier = GetImmModifier(*P->op_begin()); + if (!ImmModifier || + !llvm::all_of(P->operands(), [&ImmModifier, GetImmModifier](Value *V) { + return ImmModifier == GetImmModifier(V); + })) + return false; + + // At this point, P is a PHI node, with all the incoming values being a blend + // operations with the same immediate modifier, but possibly different address + // modifiers. Replace Disc argument of Inst with a single ptrauth_blend + // whose address modifier if provided by a PHI node. + + IRBuilder<> Builder(Inst->getContext()); + Type *Int64Ty = Builder.getInt64Ty(); + + Builder.SetInsertPoint(P); + PHINode *AddrDiscPhi = Builder.CreatePHI(Int64Ty, P->getNumIncomingValues()); + for (auto [Val, BB] : llvm::zip_equal(P->operands(), P->blocks())) { + Value *AddrDisc = cast(Val)->getArgOperand(0); + AddrDiscPhi->addIncoming(AddrDisc, BB); + } + + Builder.SetInsertPoint(Inst); + Value *ImmModifValue = ConstantInt::get(Int64Ty, *ImmModifier); + Value *Blend = Builder.CreateIntrinsic(Intrinsic::ptrauth_blend, + {AddrDiscPhi, ImmModifValue}); + Inst->replaceUsesOfWith(Disc, Blend); + return true; +} + /// For the instruction sequence of store below, F and I values /// are bundled together as an i64 value before being stored into memory. /// Sometimes it is more efficient to generate separate stores for F and I, diff --git a/llvm/test/Transforms/CodeGenPrepare/AArch64/ptrauth.ll b/llvm/test/Transforms/CodeGenPrepare/AArch64/ptrauth.ll new file mode 100644 index 0000000000000..7f06137c891de --- /dev/null +++ b/llvm/test/Transforms/CodeGenPrepare/AArch64/ptrauth.ll @@ -0,0 +1,142 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: -p --version 5 +; RUN: opt -mtriple=aarch64-linux-pauthtest -S -passes='require,function(codegenprepare)' < %s | FileCheck %s +; RUN: opt -mtriple=arm64e-apple-darwin -S -passes='require,function(codegenprepare)' < %s | FileCheck %s + +define i64 @choose_disc(i1 %cond, i64 %addr.1, i64 %addr.2, i64 %ptr) { +; CHECK-LABEL: define i64 @choose_disc(i1 %cond, i64 %addr.1, i64 %addr.2, i64 %ptr) { +; CHECK-NEXT: entry: +; CHECK-NEXT: br i1 %cond, label %bb.blend.1, label %bb.blend.2 +; CHECK: bb.blend.1: +; CHECK-NEXT: %disc.1 = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 1234) +; CHECK-NEXT: br label %bb.user +; CHECK: bb.blend.2: +; CHECK-NEXT: %disc.2 = call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 1234) +; CHECK-NEXT: br label %bb.user +; CHECK: bb.user: +; CHECK-NEXT: %0 = phi i64 [ %addr.1, %bb.blend.1 ], [ %addr.2, %bb.blend.2 ] +; CHECK-NEXT: %disc = phi i64 [ %disc.1, %bb.blend.1 ], [ %disc.2, %bb.blend.2 ] +; CHECK-NEXT: %1 = call i64 @llvm.ptrauth.blend(i64 %0, i64 1234) +; CHECK-NEXT: %result = call i64 @llvm.ptrauth.sign(i64 %ptr, i32 0, i64 %1) +; CHECK-NEXT: ret i64 %result +; +entry: + br i1 %cond, label %bb.blend.1, label %bb.blend.2 +bb.blend.1: + %disc.1 = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 1234) + br label %bb.user +bb.blend.2: + %disc.2 = call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 1234) + br label %bb.user +bb.user: + %disc = phi i64 [ %disc.1, %bb.blend.1 ], [ %disc.2, %bb.blend.2 ] + %result = call i64 @llvm.ptrauth.sign(i64 %ptr, i32 0, i64 %disc) + ret i64 %result +} + +define i64 @choose_disc_different_imm(i1 %cond, i64 %addr.1, i64 %addr.2, i64 %ptr) { +; CHECK-LABEL: define i64 @choose_disc_different_imm(i1 %cond, i64 %addr.1, i64 %addr.2, i64 %ptr) { +; CHECK-NEXT: entry: +; CHECK-NEXT: br i1 %cond, label %bb.blend.1, label %bb.blend.2 +; CHECK: bb.blend.1: +; CHECK-NEXT: %disc.1 = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 1234) +; CHECK-NEXT: br label %bb.user +; CHECK: bb.blend.2: +; CHECK-NEXT: %disc.2 = call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 42) +; CHECK-NEXT: br label %bb.user +; CHECK: bb.user: +; CHECK-NEXT: %disc = phi i64 [ %disc.1, %bb.blend.1 ], [ %disc.2, %bb.blend.2 ] +; CHECK-NEXT: %result = call i64 @llvm.ptrauth.sign(i64 %ptr, i32 0, i64 %disc) +; CHECK-NEXT: ret i64 %result +; +entry: + br i1 %cond, label %bb.blend.1, label %bb.blend.2 +bb.blend.1: + %disc.1 = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 1234) + br label %bb.user +bb.blend.2: + %disc.2 = call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 42) + br label %bb.user +bb.user: + %disc = phi i64 [ %disc.1, %bb.blend.1 ], [ %disc.2, %bb.blend.2 ] + %result = call i64 @llvm.ptrauth.sign(i64 %ptr, i32 0, i64 %disc) + ret i64 %result +} + +define i64 @choose_disc_multi_edge(i32 %arg, i64 %addr.1, i64 %addr.2, i64 %ptr) { +; CHECK-LABEL: define i64 @choose_disc_multi_edge(i32 %arg, i64 %addr.1, i64 %addr.2, i64 %ptr) { +; CHECK-NEXT: entry: +; CHECK-NEXT: %cond = icmp eq i32 %arg, 0 +; CHECK-NEXT: br i1 %cond, label %bb.blend.1, label %bb.blend.2 +; CHECK: bb.blend.1: +; CHECK-NEXT: %disc.1 = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 1234) +; CHECK-NEXT: switch i32 %arg, label %bb.blend.2 [ +; CHECK-NEXT: i32 0, label %bb.user +; CHECK-NEXT: i32 3, label %bb.user +; CHECK-NEXT: ] +; CHECK: bb.blend.2: +; CHECK-NEXT: %disc.2 = call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 1234) +; CHECK-NEXT: br label %bb.user +; CHECK: bb.user: +; CHECK-NEXT: %0 = phi i64 [ %addr.1, %bb.blend.1 ], [ %addr.1, %bb.blend.1 ], [ %addr.2, %bb.blend.2 ] +; CHECK-NEXT: %disc = phi i64 [ %disc.1, %bb.blend.1 ], [ %disc.1, %bb.blend.1 ], [ %disc.2, %bb.blend.2 ] +; CHECK-NEXT: %1 = call i64 @llvm.ptrauth.blend(i64 %0, i64 1234) +; CHECK-NEXT: %result = call i64 @llvm.ptrauth.sign(i64 %ptr, i32 0, i64 %1) +; CHECK-NEXT: ret i64 %result +; +entry: + %cond = icmp eq i32 %arg, 0 + br i1 %cond, label %bb.blend.1, label %bb.blend.2 +bb.blend.1: + %disc.1 = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 1234) + switch i32 %arg, label %bb.blend.2 [ + i32 0, label %bb.user + i32 2, label %bb.blend.2 + i32 3, label %bb.user + ] +bb.blend.2: + %disc.2 = call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 1234) + br label %bb.user +bb.user: + %disc = phi i64 [ %disc.1, %bb.blend.1 ], [ %disc.1, %bb.blend.1 ], [ %disc.2, %bb.blend.2 ] + %result = call i64 @llvm.ptrauth.sign(i64 %ptr, i32 0, i64 %disc) + ret i64 %result +} + +define i64 @choose_disc_phi_placement(i1 %cond, i64 %addr.1, i64 %addr.2, i64 %ptr) { +; CHECK-LABEL: define i64 @choose_disc_phi_placement(i1 %cond, i64 %addr.1, i64 %addr.2, i64 %ptr) { +; CHECK-NEXT: entry: +; CHECK-NEXT: br i1 %cond, label %bb.blend.1, label %bb.blend.2 +; CHECK: bb.blend.1: +; CHECK-NEXT: %disc.1 = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 1234) +; CHECK-NEXT: br label %bb.phi +; CHECK: bb.blend.2: +; CHECK-NEXT: %disc.2 = call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 1234) +; CHECK-NEXT: br label %bb.phi +; CHECK: bb.phi: +; CHECK-NEXT: %0 = phi i64 [ %addr.1, %bb.blend.1 ], [ %addr.2, %bb.blend.2 ] +; CHECK-NEXT: %disc = phi i64 [ %disc.1, %bb.blend.1 ], [ %disc.2, %bb.blend.2 ] +; CHECK-NEXT: br i1 %cond, label %bb.user, label %bb.ret0 +; CHECK: bb.user: +; CHECK-NEXT: %1 = call i64 @llvm.ptrauth.blend(i64 %0, i64 1234) +; CHECK-NEXT: %result = call i64 @llvm.ptrauth.sign(i64 %ptr, i32 0, i64 %1) +; CHECK-NEXT: ret i64 %result +; CHECK: bb.ret0: +; CHECK-NEXT: ret i64 0 +; +entry: + br i1 %cond, label %bb.blend.1, label %bb.blend.2 +bb.blend.1: + %disc.1 = call i64 @llvm.ptrauth.blend(i64 %addr.1, i64 1234) + br label %bb.phi +bb.blend.2: + %disc.2 = call i64 @llvm.ptrauth.blend(i64 %addr.2, i64 1234) + br label %bb.phi +bb.phi: + %disc = phi i64 [ %disc.1, %bb.blend.1 ], [ %disc.2, %bb.blend.2 ] + br i1 %cond, label %bb.user, label %bb.ret0 +bb.user: + %result = call i64 @llvm.ptrauth.sign(i64 %ptr, i32 0, i64 %disc) + ret i64 %result +bb.ret0: + ret i64 0 +}