[InstCombine] Combine ptrauth intrin. callee into same-key bundle. #94707

ahmedbougacha · 2024-06-07T00:29:35Z

Try to optimize a call to the result of a ptrauth intrinsic, potentially into the ptrauth call bundle:

  call(ptrauth.resign(p)), ["ptrauth"()] ->  call p, ["ptrauth"()]
  call(ptrauth.sign(p)),   ["ptrauth"()] ->  call p

as long as the key/discriminator are the same in sign and auth-bundle, and we don't change the key in the bundle (to a potentially-invalid key.)

Generating a plain call to a raw unauthenticated pointer is generally undesirable, but if we ended up seeing a naked ptrauth.sign in the first place, we already have suspicious code. Unauthenticated calls are also easier to spot than naked signs, so let the indirect call shine.

Note that there is an arguably unsafe extension to this, where we don't bother checking that the key in bundle and intrinsic are the same (and also allow folding away an auth into a bundle.)

This can end up generating calls with a bundle that has an invalid key (which an informed frontend wouldn't have otherwise done), which can be problematic. The C that generates that is straightforward but arguably unreasonable. That wouldn't be an issue if we were to bite the bullet and make these fully AArch64-specific, allowing key knowledge to be embedded here.

I have a branch that does the "unsafe" combines in bce1917. I'm comfortable doing that in my toolchain but I imagine others may disagree here; open to people's thoughts.

Try to optimize a call to the result of a ptrauth intrinsic, potentially into the ptrauth call bundle: - call(ptrauth.resign(p)), ["ptrauth"()] -> call p, ["ptrauth"()] - call(ptrauth.sign(p)), ["ptrauth"()] -> call p as long as the key/discriminator are the same in sign and auth-bundle, and we don't change the key in the bundle (to a potentially-invalid key.) Generating a plain call to a raw unauthenticated pointer is generally undesirable, but if we ended up seeing a naked ptrauth.sign in the first place, we already have suspicious code. Unauthenticated calls are also easier to spot than naked signs, so let the indirect call shine.

llvmbot · 2024-06-12T00:05:16Z

@llvm/pr-subscribers-backend-aarch64

@llvm/pr-subscribers-llvm-transforms

Author: Ahmed Bougacha (ahmedbougacha)

Changes

Try to optimize a call to the result of a ptrauth intrinsic, potentially into the ptrauth call bundle:

  call(ptrauth.resign(p)), ["ptrauth"()] -&gt;  call p, ["ptrauth"()]
  call(ptrauth.sign(p)),   ["ptrauth"()] -&gt;  call p

as long as the key/discriminator are the same in sign and auth-bundle, and we don't change the key in the bundle (to a potentially-invalid key.)

Generating a plain call to a raw unauthenticated pointer is generally undesirable, but if we ended up seeing a naked ptrauth.sign in the first place, we already have suspicious code. Unauthenticated calls are also easier to spot than naked signs, so let the indirect call shine.

Note that there is an arguably unsafe extension to this, where we don't bother checking that the key in bundle and intrinsic are the same (and also allow folding away an auth into a bundle.)

This can end up generating calls with a bundle that has an invalid key (which an informed frontend wouldn't have otherwise done), which can be problematic. The C that generates that is straightforward but arguably unreasonable. That wouldn't be an issue if we were to bite the bullet and make these fully AArch64-specific, allowing key knowledge to be embedded here.

I have a branch that does the "unsafe" combines in bce1917. I'm comfortable doing that in my toolchain but I imagine others may disagree here; open to people's thoughts.

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

3 Files Affected:

(modified) llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp (+76)
(modified) llvm/lib/Transforms/InstCombine/InstCombineInternal.h (+8)
(added) llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll (+76)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
index 436cdbff75669..8c32f411320ab 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp
@@ -3665,6 +3665,78 @@ static IntrinsicInst *findInitTrampoline(Value *Callee) {
   return nullptr;
 }
 
+Instruction *InstCombinerImpl::foldPtrAuthIntrinsicCallee(CallBase &Call) {
+  Value *Callee = Call.getCalledOperand();
+  auto *IPC = dyn_cast<IntToPtrInst>(Callee);
+  if (!IPC || !IPC->isNoopCast(DL))
+    return nullptr;
+
+  IntrinsicInst *II = dyn_cast<IntrinsicInst>(IPC->getOperand(0));
+  if (!II)
+    return nullptr;
+
+  // Isolate the ptrauth bundle from the others.
+  std::optional<OperandBundleUse> PtrAuthBundleOrNone;
+  SmallVector<OperandBundleDef, 2> NewBundles;
+  for (unsigned BI = 0, BE = Call.getNumOperandBundles(); BI != BE; ++BI) {
+    OperandBundleUse Bundle = Call.getOperandBundleAt(BI);
+    if (Bundle.getTagID() == LLVMContext::OB_ptrauth)
+      PtrAuthBundleOrNone = Bundle;
+    else
+      NewBundles.emplace_back(Bundle);
+  }
+
+  Value *NewCallee = nullptr;
+  switch (II->getIntrinsicID()) {
+  default:
+    return nullptr;
+
+  // call(ptrauth.resign(p)), ["ptrauth"()] ->  call p, ["ptrauth"()]
+  // assuming the call bundle and the sign operands match.
+  case Intrinsic::ptrauth_resign: {
+    if (!PtrAuthBundleOrNone ||
+        II->getOperand(3) != PtrAuthBundleOrNone->Inputs[0] ||
+        II->getOperand(4) != PtrAuthBundleOrNone->Inputs[1])
+      return nullptr;
+
+    // Don't change the key used in the call; we don't know what's valid.
+    if (II->getOperand(1) != PtrAuthBundleOrNone->Inputs[0])
+      return nullptr;
+
+    Value *NewBundleOps[] = {II->getOperand(1), II->getOperand(2)};
+    NewBundles.emplace_back("ptrauth", NewBundleOps);
+    NewCallee = II->getOperand(0);
+    break;
+  }
+
+  // call(ptrauth.sign(p)), ["ptrauth"()] ->  call p
+  // assuming the call bundle and the sign operands match.
+  // Non-ptrauth indirect calls are undesirable, but so is ptrauth.sign.
+  case Intrinsic::ptrauth_sign: {
+    if (!PtrAuthBundleOrNone ||
+        II->getOperand(1) != PtrAuthBundleOrNone->Inputs[0] ||
+        II->getOperand(2) != PtrAuthBundleOrNone->Inputs[1])
+      return nullptr;
+    NewCallee = II->getOperand(0);
+    break;
+  }
+  }
+
+  if (!NewCallee)
+    return nullptr;
+
+  NewCallee = Builder.CreateBitOrPointerCast(NewCallee, Callee->getType());
+  CallBase *NewCall = nullptr;
+  if (auto *CI = dyn_cast<CallInst>(&Call)) {
+    NewCall = CallInst::Create(CI, NewBundles);
+  } else {
+    auto *IKI = cast<InvokeInst>(&Call);
+    NewCall = InvokeInst::Create(IKI, NewBundles);
+  }
+  NewCall->setCalledOperand(NewCallee);
+  return NewCall;
+}
+
 bool InstCombinerImpl::annotateAnyAllocSite(CallBase &Call,
                                             const TargetLibraryInfo *TLI) {
   // Note: We only handle cases which can't be driven from generic attributes
@@ -3812,6 +3884,10 @@ Instruction *InstCombinerImpl::visitCallBase(CallBase &Call) {
   if (IntrinsicInst *II = findInitTrampoline(Callee))
     return transformCallThroughTrampoline(Call, *II);
 
+  // Combine calls involving pointer authentication intrinsics.
+  if (Instruction *NewCall = foldPtrAuthIntrinsicCallee(Call))
+    return NewCall;
+
   if (isa<InlineAsm>(Callee) && !Call.doesNotThrow()) {
     InlineAsm *IA = cast<InlineAsm>(Callee);
     if (!IA->canThrow()) {
diff --git a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
index 984f02bcccad7..f290f446b424e 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
+++ b/llvm/lib/Transforms/InstCombine/InstCombineInternal.h
@@ -282,6 +282,14 @@ class LLVM_LIBRARY_VISIBILITY InstCombinerImpl final
   Instruction *transformCallThroughTrampoline(CallBase &Call,
                                               IntrinsicInst &Tramp);
 
+  /// Try to optimize a call to the result of a ptrauth intrinsic, potentially
+  /// into the ptrauth call bundle:
+  /// - call(ptrauth.resign(p)), ["ptrauth"()] ->  call p, ["ptrauth"()]
+  /// - call(ptrauth.sign(p)),   ["ptrauth"()] ->  call p
+  /// as long as the key/discriminator are the same in sign and auth-bundle,
+  /// and we don't change the key in the bundle (to a potentially-invalid key.)
+  Instruction *foldPtrAuthIntrinsicCallee(CallBase &Call);
+
   // Return (a, b) if (LHS, RHS) is known to be (a, b) or (b, a).
   // Otherwise, return std::nullopt
   // Currently it matches:
diff --git a/llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll b/llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll
new file mode 100644
index 0000000000000..900fd43f01602
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll
@@ -0,0 +1,76 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i32 @test_ptrauth_call_resign(ptr %p) {
+; CHECK-LABEL: @test_ptrauth_call_resign(
+; CHECK-NEXT:    [[V3:%.*]] = call i32 [[P:%.*]]() [ "ptrauth"(i32 1, i64 1234) ]
+; CHECK-NEXT:    ret i32 [[V3]]
+;
+  %v0 = ptrtoint ptr %p to i64
+  %v1 = call i64 @llvm.ptrauth.resign(i64 %v0, i32 1, i64 1234, i32 1, i64 5678)
+  %v2 = inttoptr i64 %v1 to i32()*
+  %v3 = call i32 %v2() [ "ptrauth"(i32 1, i64 5678) ]
+  ret i32 %v3
+}
+
+define i32 @test_ptrauth_call_resign_blend(ptr %pp) {
+; CHECK-LABEL: @test_ptrauth_call_resign_blend(
+; CHECK-NEXT:    [[V01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
+; CHECK-NEXT:    [[V6:%.*]] = call i32 [[V01]]() [ "ptrauth"(i32 1, i64 1234) ]
+; CHECK-NEXT:    ret i32 [[V6]]
+;
+  %v0 = load ptr, ptr %pp, align 8
+  %v1 = ptrtoint ptr %pp to i64
+  %v2 = ptrtoint ptr %v0 to i64
+  %v3 = call i64 @llvm.ptrauth.blend(i64 %v1, i64 5678)
+  %v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 1, i64 1234, i32 1, i64 %v3)
+  %v5 = inttoptr i64 %v4 to i32()*
+  %v6 = call i32 %v5() [ "ptrauth"(i32 1, i64 %v3) ]
+  ret i32 %v6
+}
+
+define i32 @test_ptrauth_call_resign_blend_2(ptr %pp) {
+; CHECK-LABEL: @test_ptrauth_call_resign_blend_2(
+; CHECK-NEXT:    [[V01:%.*]] = load ptr, ptr [[PP:%.*]], align 8
+; CHECK-NEXT:    [[V1:%.*]] = ptrtoint ptr [[PP]] to i64
+; CHECK-NEXT:    [[V3:%.*]] = call i64 @llvm.ptrauth.blend(i64 [[V1]], i64 5678)
+; CHECK-NEXT:    [[V6:%.*]] = call i32 [[V01]]() [ "ptrauth"(i32 0, i64 [[V3]]) ]
+; CHECK-NEXT:    ret i32 [[V6]]
+;
+  %v0 = load ptr, ptr %pp, align 8
+  %v1 = ptrtoint ptr %pp to i64
+  %v2 = ptrtoint ptr %v0 to i64
+  %v3 = call i64 @llvm.ptrauth.blend(i64 %v1, i64 5678)
+  %v4 = call i64 @llvm.ptrauth.resign(i64 %v2, i32 0, i64 %v3, i32 0, i64 1234)
+  %v5 = inttoptr i64 %v4 to i32()*
+  %v6 = call i32 %v5() [ "ptrauth"(i32 0, i64 1234) ]
+  ret i32 %v6
+}
+
+define i32 @test_ptrauth_call_sign(ptr %p) {
+; CHECK-LABEL: @test_ptrauth_call_sign(
+; CHECK-NEXT:    [[V3:%.*]] = call i32 [[P:%.*]]()
+; CHECK-NEXT:    ret i32 [[V3]]
+;
+  %v0 = ptrtoint ptr %p to i64
+  %v1 = call i64 @llvm.ptrauth.sign(i64 %v0, i32 2, i64 5678)
+  %v2 = inttoptr i64 %v1 to i32()*
+  %v3 = call i32 %v2() [ "ptrauth"(i32 2, i64 5678) ]
+  ret i32 %v3
+}
+
+define i32 @test_ptrauth_call_sign_otherbundle(ptr %p) {
+; CHECK-LABEL: @test_ptrauth_call_sign_otherbundle(
+; CHECK-NEXT:    [[V3:%.*]] = call i32 [[P:%.*]]() [ "somebundle"(ptr null), "otherbundle"(i64 0) ]
+; CHECK-NEXT:    ret i32 [[V3]]
+;
+  %v0 = ptrtoint ptr %p to i64
+  %v1 = call i64 @llvm.ptrauth.sign(i64 %v0, i32 2, i64 5678)
+  %v2 = inttoptr i64 %v1 to i32()*
+  %v3 = call i32 %v2() [ "somebundle"(ptr null), "ptrauth"(i32 2, i64 5678), "otherbundle"(i64 0) ]
+  ret i32 %v3
+}
+
+declare i64 @llvm.ptrauth.sign(i64, i32, i64)
+declare i64 @llvm.ptrauth.resign(i64, i32, i64, i32, i64)
+declare i64 @llvm.ptrauth.blend(i64, i64)

nikic

Looks fine from InstCombine perspective, but can't say anything about the actual logic here.

llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll

kovdan01

Changes are mostly OK as for me, but I think that several test cases are missing and also the code can be re-organized a little so we fail earlier for the majority of cases when such combining is not possible.

I'm OK with both adding tests as a part of this PR and submitting a follow-up patch with missing tests if adding them now is too time-consuming or there are other issues preventing that.

llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

kovdan01 · 2024-06-17T13:43:36Z

llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp

+
+  // call(ptrauth.sign(p)), ["ptrauth"()] ->  call p
+  // assuming the call bundle and the sign operands match.
+  // Non-ptrauth indirect calls are undesirable, but so is ptrauth.sign.


Yes, both of these are bad options - not sure though what is worse, but I'm happy with having unauthenticated indirect calls here if everyone else is happy with this as well.

It's actually interesting, can we have such an IR in "wild nature" when: 1) generating IR from C/C++ code and not manually writing it; 2) not using ptrauth intrinsics in code explicitly? If yes, could you please provide an example of such code?

I'd argue unauthenticated calls are better because they can be scanned for trivially, flagged for audit, rejected, etc.; silent raw signs are more insidious.

It's true we're not supposed to have these, but if nothing else, compilers do learn new IRGen every once in a while (maybe less often for c/c++ than other frontends), which can end up with this.

…in-callee-safe

…in-callee-safe Conflicts: llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp llvm/lib/Transforms/InstCombine/InstCombineInternal.h

ahmedbougacha force-pushed the users/ahmedbougacha/ptrauth-instcombine-intrin-callee-safe branch from 903f149 to e623f48 Compare June 12, 2024 00:03

ahmedbougacha marked this pull request as ready for review June 12, 2024 00:04

ahmedbougacha requested a review from nikic as a code owner June 12, 2024 00:04

ahmedbougacha requested review from asl and kovdan01 June 12, 2024 00:04

llvmbot added the llvm:transforms label Jun 12, 2024

ahmedbougacha added the backend:AArch64 label Jun 12, 2024

nikic reviewed Jun 12, 2024

View reviewed changes

llvm/test/Transforms/InstCombine/ptrauth-intrinsics-call.ll Outdated Show resolved Hide resolved

Replace typed pointers in casts with ptr.

5ea8d74

kovdan01 requested changes Jun 17, 2024

View reviewed changes

asl added this to the LLVM 21.x Release milestone Jun 2, 2025

asl assigned ahmedbougacha Jun 2, 2025

ahmedbougacha added 3 commits July 15, 2025 09:18

Merge branch 'main' into users/ahmedbougacha/ptrauth-instcombine-intr…

874ac2a

…in-callee-safe

Simplify a bit and address review feedback.

fa689d5

Add more tests.

a27020b

llvmbot added the llvm:instcombine Covers the InstCombine, InstSimplify and AggressiveInstCombine passes label Jul 15, 2025

Merge branch 'main' into users/ahmedbougacha/ptrauth-instcombine-intr…

66b9427

…in-callee-safe Conflicts: llvm/lib/Transforms/InstCombine/InstCombineCalls.cpp llvm/lib/Transforms/InstCombine/InstCombineInternal.h

ahmedbougacha merged commit 77bcab8 into llvm:main Jul 15, 2025
7 of 8 checks passed

github-project-automation bot moved this from In Progress to Done in Pointer Authentication Tasks Jul 15, 2025

ahmedbougacha deleted the users/ahmedbougacha/ptrauth-instcombine-intrin-callee-safe branch July 15, 2025 21:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[InstCombine] Combine ptrauth intrin. callee into same-key bundle. #94707

[InstCombine] Combine ptrauth intrin. callee into same-key bundle. #94707

Uh oh!

ahmedbougacha commented Jun 7, 2024 •

edited

Loading

Uh oh!

llvmbot commented Jun 12, 2024 •

edited

Loading

Uh oh!

nikic left a comment

Uh oh!

Uh oh!

kovdan01 left a comment •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kovdan01 Jun 17, 2024

Uh oh!

ahmedbougacha Jul 15, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

[InstCombine] Combine ptrauth intrin. callee into same-key bundle. #94707

[InstCombine] Combine ptrauth intrin. callee into same-key bundle. #94707

Uh oh!

Conversation

ahmedbougacha commented Jun 7, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

llvmbot commented Jun 12, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

nikic left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

kovdan01 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

kovdan01 Jun 17, 2024

Choose a reason for hiding this comment

Uh oh!

ahmedbougacha Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

ahmedbougacha commented Jun 7, 2024 •

edited

Loading

llvmbot commented Jun 12, 2024 •

edited

Loading

kovdan01 left a comment •

edited

Loading