[llvm] [LFI][AArch64] Add rewrites for control flow (PR #192602)
via llvm-commits
llvm-commits at lists.llvm.org
Thu Apr 16 23:59:30 PDT 2026
llvmbot wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-backend-aarch64
Author: Zachary Yedidia (zyedidia)
<details>
<summary>Changes</summary>
Adds LFI rewrites for control flow instructions (indirect branches and returns). Indirect branches must go through `x28`, which is always guaranteed to hold a sandbox address. Modifications to `x30` must guard `x30` afterwards, to uphold the invariant that `x30` always holds a sandbox address. As a result, bare return instructions can be used without any additional rewrites.
---
Full diff: https://github.com/llvm/llvm-project/pull/192602.diff
6 Files Affected:
- (modified) llvm/docs/LFI.rst (-4)
- (modified) llvm/lib/Target/AArch64/AArch64InstrInfo.cpp (+22-4)
- (modified) llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.cpp (+69)
- (modified) llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.h (+12)
- (added) llvm/test/MC/AArch64/LFI/branch.s (+21)
- (added) llvm/test/MC/AArch64/LFI/return.s (+25)
``````````diff
diff --git a/llvm/docs/LFI.rst b/llvm/docs/LFI.rst
index cf99cc5faae38..62a7ce8feabec 100644
--- a/llvm/docs/LFI.rst
+++ b/llvm/docs/LFI.rst
@@ -140,8 +140,6 @@ In the following assembly rewrites, some shorthand is used.
Control flow
~~~~~~~~~~~~
-**Note**: not yet implemented.
-
Indirect branches get rewritten to branch through register ``x28``, which must
always contain an address within the sandbox. An ``add`` is used to safely
update ``x28`` with the destination address. Since ``ret`` uses ``x30`` by
@@ -271,8 +269,6 @@ before moving it back into ``sp`` with a safe ``add``.
Link register modification
~~~~~~~~~~~~~~~~~~~~~~~~~~~
-**Note**: not yet implemented.
-
When the link register is modified, we write the modified value to a
temporary, before loading it back into ``x30`` with a safe ``add``.
diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
index 3ca851ef83d27..d9be4b2c0e6f0 100644
--- a/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
+++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.cpp
@@ -119,12 +119,30 @@ static std::optional<unsigned> getLFIInstSizeInBytes(const MachineInstr &MI) {
case AArch64::SVC:
// SVC expands to 4 instructions.
return 16;
+ case AArch64::BR:
+ case AArch64::BLR:
+ // Indirect branches/calls expand to 2 instructions (guard + br/blr).
+ return 8;
+ case AArch64::RET:
+ // RET through LR is not rewritten, but RET through another register
+ // expands to 2 instructions (guard + ret).
+ if (MI.getOperand(0).getReg() != AArch64::LR)
+ return 8;
+ return 4;
default:
- // Default case: instructions that don't cause expansion.
- // - TP accesses in LFI are a single load/store, so no expansion.
- // - All remaining instructions are not rewritten.
- return std::nullopt;
+ break;
}
+
+ // Instructions that explicitly modify LR expand to 2 instructions (original
+ // + guard).
+ for (const MachineOperand &MO : MI.explicit_operands())
+ if (MO.isReg() && MO.isDef() && MO.getReg() == AArch64::LR)
+ return 8;
+
+ // Default case: instructions that don't cause expansion.
+ // - TP accesses in LFI are a single load/store, so no expansion.
+ // - All remaining instructions are not rewritten.
+ return std::nullopt;
}
/// GetInstSize - Return the number of bytes of code the specified
diff --git a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.cpp b/llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.cpp
index 92d7d66992e51..0d84de66f08e5 100644
--- a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.cpp
+++ b/llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.cpp
@@ -111,6 +111,58 @@ void AArch64MCLFIRewriter::emitMov(MCRegister Dest, MCRegister Src,
emitInst(Inst, Out, STI);
}
+// {br,blr} xN
+// ->
+// add x28, x27, wN, uxtw
+// {br,blr} x28
+void AArch64MCLFIRewriter::rewriteIndirectBranch(const MCInst &Inst,
+ MCStreamer &Out,
+ const MCSubtargetInfo &STI) {
+ if (!Inst.getOperand(0).isReg())
+ return error(Inst, "unsupported instruction: expected target register");
+ MCRegister BranchReg = Inst.getOperand(0).getReg();
+
+ // Guard the branch target through X28.
+ emitAddMask(LFIAddrReg, BranchReg, Out, STI);
+ emitBranch(Inst.getOpcode(), LFIAddrReg, Out, STI);
+}
+
+void AArch64MCLFIRewriter::rewriteCall(const MCInst &Inst, MCStreamer &Out,
+ const MCSubtargetInfo &STI) {
+ if (Inst.getOperand(0).isReg())
+ rewriteIndirectBranch(Inst, Out, STI);
+ else
+ emitInst(Inst, Out, STI);
+}
+
+// ret xN (where xN != x30)
+// ->
+// add x28, x27, wN, uxtw
+// ret x28
+//
+// ret (x30) is safe since x30 is always within the sandbox.
+void AArch64MCLFIRewriter::rewriteReturn(const MCInst &Inst, MCStreamer &Out,
+ const MCSubtargetInfo &STI) {
+ if (Inst.getNumOperands() == 0 || !Inst.getOperand(0).isReg())
+ return error(Inst, "unsupported instruction: expected target register");
+ // RET through LR is safe since LR is always within sandbox.
+ if (Inst.getOperand(0).getReg() != AArch64::LR)
+ rewriteIndirectBranch(Inst, Out, STI);
+ else
+ emitInst(Inst, Out, STI);
+}
+
+// modify x30
+// ->
+// modify x30
+// add x30, x27, w30, uxtw
+void AArch64MCLFIRewriter::rewriteLRModification(const MCInst &Inst,
+ MCStreamer &Out,
+ const MCSubtargetInfo &STI) {
+ emitInst(Inst, Out, STI);
+ emitAddMask(AArch64::LR, AArch64::LR, Out, STI);
+}
+
// svc #0
// ->
// mov x26, x30
@@ -192,6 +244,23 @@ void AArch64MCLFIRewriter::doRewriteInst(const MCInst &Inst, MCStreamer &Out,
return;
}
+ // Control flow.
+ if (isReturn(Inst))
+ return rewriteReturn(Inst, Out, STI);
+
+ if (isIndirectBranch(Inst))
+ return rewriteIndirectBranch(Inst, Out, STI);
+
+ if (isCall(Inst))
+ return rewriteCall(Inst, Out, STI);
+
+ if (isBranch(Inst))
+ return emitInst(Inst, Out, STI);
+
+ // Link register modification.
+ if (mayModifyRegister(Inst, AArch64::LR))
+ return rewriteLRModification(Inst, Out, STI);
+
emitInst(Inst, Out, STI);
}
diff --git a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.h b/llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.h
index 049860c1ab6cc..915a6b86b2e17 100644
--- a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.h
+++ b/llvm/lib/Target/AArch64/MCTargetDesc/AArch64MCLFIRewriter.h
@@ -67,6 +67,18 @@ class AArch64MCLFIRewriter : public MCLFIRewriter {
void doRewriteInst(const MCInst &Inst, MCStreamer &Out,
const MCSubtargetInfo &STI);
+ // Control flow.
+ void rewriteIndirectBranch(const MCInst &Inst, MCStreamer &Out,
+ const MCSubtargetInfo &STI);
+ void rewriteCall(const MCInst &Inst, MCStreamer &Out,
+ const MCSubtargetInfo &STI);
+ void rewriteReturn(const MCInst &Inst, MCStreamer &Out,
+ const MCSubtargetInfo &STI);
+
+ // Link register modification.
+ void rewriteLRModification(const MCInst &Inst, MCStreamer &Out,
+ const MCSubtargetInfo &STI);
+
// System instructions.
void rewriteSyscall(const MCInst &Inst, MCStreamer &Out,
const MCSubtargetInfo &STI);
diff --git a/llvm/test/MC/AArch64/LFI/branch.s b/llvm/test/MC/AArch64/LFI/branch.s
new file mode 100644
index 0000000000000..684033f14d696
--- /dev/null
+++ b/llvm/test/MC/AArch64/LFI/branch.s
@@ -0,0 +1,21 @@
+// RUN: llvm-mc -triple aarch64_lfi %s | FileCheck %s
+
+foo:
+
+b foo
+// CHECK: b foo
+
+br x0
+// CHECK: add x28, x27, w0, uxtw
+// CHECK-NEXT: br x28
+
+blr x0
+// CHECK: add x28, x27, w0, uxtw
+// CHECK-NEXT: blr x28
+
+ret
+// CHECK: ret
+
+ret x0
+// CHECK: add x28, x27, w0, uxtw
+// CHECK-NEXT: ret x28
diff --git a/llvm/test/MC/AArch64/LFI/return.s b/llvm/test/MC/AArch64/LFI/return.s
new file mode 100644
index 0000000000000..7e7431df57aa0
--- /dev/null
+++ b/llvm/test/MC/AArch64/LFI/return.s
@@ -0,0 +1,25 @@
+// RUN: llvm-mc -triple aarch64_lfi %s | FileCheck %s
+
+mov x30, x0
+ret
+// CHECK: mov x30, x0
+// CHECK-NEXT: add x30, x27, w30, uxtw
+// CHECK-NEXT: ret
+
+ldr x30, [sp]
+ret
+// CHECK: ldr x30, [sp]
+// CHECK-NEXT: add x30, x27, w30, uxtw
+// CHECK-NEXT: ret
+
+ldp x29, x30, [sp]
+ret
+// CHECK: ldp x29, x30, [sp]
+// CHECK-NEXT: add x30, x27, w30, uxtw
+// CHECK-NEXT: ret
+
+ldp x30, x29, [sp]
+ret
+// CHECK: ldp x30, x29, [sp]
+// CHECK-NEXT: add x30, x27, w30, uxtw
+// CHECK-NEXT: ret
``````````
</details>
https://github.com/llvm/llvm-project/pull/192602
More information about the llvm-commits
mailing list