[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-02 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko created 
https://github.com/llvm/llvm-project/pull/97472

In preparation for implementing hardening of BLRA* instructions, restrict thunk 
function generation to only the thunks being actually called from any function. 
As described in the existing comments, emitting all possible thunks for BLRAA 
and BLRAB instructions would mean adding about 1800 functions in total, most of 
which are likely not to be called.

This commit merges AArch64SLSHardening class into SLSBLRThunkInserter, so 
thunks can be created as needed while rewriting a machine function. The usages 
of TII, TRI and ST fields of AArch64SLSHardening class are replaced with 
requesting them in-place, as ThunkInserter assumes multiple "entry points" in 
contrast to the only runOnMachineFunction method of AArch64SLSHardening.

The runOnMachineFunction method essentially replaces pre-existing insertThunks 
implementation as there is no more need to insert all possible thunks 
unconditionally. Instead, thunks are created on first use from inside of 
insertThunks method.

>From a246cfe705b326c520d6b36882a17bd90b622e5d Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko 
Date: Fri, 28 Jun 2024 21:50:24 +0300
Subject: [PATCH] [AArch64] Only create called thunks when hardening against
 SLS

In preparation for implementing hardening of BLRA* instructions,
restrict thunk function generation to only the thunks being actually
called from any function. As described in the existing comments,
emitting all possible thunks for BLRAA and BLRAB instructions would
mean adding about 1800 functions in total, most of which are likely
not to be called.

This commit merges AArch64SLSHardening class into SLSBLRThunkInserter,
so thunks can be created as needed while rewriting a machine function.
The usages of TII, TRI and ST fields of AArch64SLSHardening class are
replaced with requesting them in-place, as ThunkInserter assumes
multiple "entry points" in contrast to the only runOnMachineFunction
method of AArch64SLSHardening.

The runOnMachineFunction method essentially replaces pre-existing
insertThunks implementation as there is no more need to insert all
possible thunks unconditionally. Instead, thunks are created on first
use from inside of insertThunks method.
---
 llvm/lib/Target/AArch64/AArch64.h |   1 -
 .../Target/AArch64/AArch64SLSHardening.cpp| 188 +++---
 .../Target/AArch64/AArch64TargetMachine.cpp   |   1 -
 llvm/test/CodeGen/AArch64/O0-pipeline.ll  |   1 -
 llvm/test/CodeGen/AArch64/O3-pipeline.ll  |   1 -
 .../AArch64/arm64-opt-remarks-lazy-bfi.ll |   8 -
 .../speculation-hardening-sls-blr-bti.mir |  20 --
 .../AArch64/speculation-hardening-sls-blr.mir |  20 +-
 8 files changed, 75 insertions(+), 165 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64.h 
b/llvm/lib/Target/AArch64/AArch64.h
index 6f2aeb83a451a..66ad701d83958 100644
--- a/llvm/lib/Target/AArch64/AArch64.h
+++ b/llvm/lib/Target/AArch64/AArch64.h
@@ -40,7 +40,6 @@ FunctionPass *createAArch64ISelDag(AArch64TargetMachine &TM,
 FunctionPass *createAArch64StorePairSuppressPass();
 FunctionPass *createAArch64ExpandPseudoPass();
 FunctionPass *createAArch64SLSHardeningPass();
-FunctionPass *createAArch64IndirectThunks();
 FunctionPass *createAArch64SpeculationHardeningPass();
 FunctionPass *createAArch64LoadStoreOptimizationPass();
 ModulePass *createAArch64LowerHomogeneousPrologEpilogPass();
diff --git a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp 
b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
index d35386eaeab12..b4ebd7d5377c2 100644
--- a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
+++ b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
@@ -13,20 +13,16 @@
 
 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
-#include "Utils/AArch64BaseInfo.h"
 #include "llvm/CodeGen/IndirectThunks.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
-#include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOperand.h"
-#include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/Pass.h"
-#include "llvm/Support/CodeGen.h"
-#include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Target/TargetMachine.h"
 #include 
 
@@ -36,38 +32,43 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
+static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+
 namespace {
 
-class AArch64SLSHardening : public MachineFunctionPass {
-public:
-  const TargetInstrInfo *TII;
-  const TargetRegisterInfo *TRI;
-  const AArch64Subtarget *ST;
+// Set of inserted thunks: bitmask with bits corresponding to
+// indexes in SLSBLRThunks array.
+typedef uint32_t ThunksSet;
 
-  static char ID;
-
-  AArch64SLSHardening() : MachineFunctionPass(ID) {
-

[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko updated 
https://github.com/llvm/llvm-project/pull/97472

>From a246cfe705b326c520d6b36882a17bd90b622e5d Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko 
Date: Fri, 28 Jun 2024 21:50:24 +0300
Subject: [PATCH 1/2] [AArch64] Only create called thunks when hardening
 against SLS

In preparation for implementing hardening of BLRA* instructions,
restrict thunk function generation to only the thunks being actually
called from any function. As described in the existing comments,
emitting all possible thunks for BLRAA and BLRAB instructions would
mean adding about 1800 functions in total, most of which are likely
not to be called.

This commit merges AArch64SLSHardening class into SLSBLRThunkInserter,
so thunks can be created as needed while rewriting a machine function.
The usages of TII, TRI and ST fields of AArch64SLSHardening class are
replaced with requesting them in-place, as ThunkInserter assumes
multiple "entry points" in contrast to the only runOnMachineFunction
method of AArch64SLSHardening.

The runOnMachineFunction method essentially replaces pre-existing
insertThunks implementation as there is no more need to insert all
possible thunks unconditionally. Instead, thunks are created on first
use from inside of insertThunks method.
---
 llvm/lib/Target/AArch64/AArch64.h |   1 -
 .../Target/AArch64/AArch64SLSHardening.cpp| 188 +++---
 .../Target/AArch64/AArch64TargetMachine.cpp   |   1 -
 llvm/test/CodeGen/AArch64/O0-pipeline.ll  |   1 -
 llvm/test/CodeGen/AArch64/O3-pipeline.ll  |   1 -
 .../AArch64/arm64-opt-remarks-lazy-bfi.ll |   8 -
 .../speculation-hardening-sls-blr-bti.mir |  20 --
 .../AArch64/speculation-hardening-sls-blr.mir |  20 +-
 8 files changed, 75 insertions(+), 165 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AArch64.h 
b/llvm/lib/Target/AArch64/AArch64.h
index 6f2aeb83a451a..66ad701d83958 100644
--- a/llvm/lib/Target/AArch64/AArch64.h
+++ b/llvm/lib/Target/AArch64/AArch64.h
@@ -40,7 +40,6 @@ FunctionPass *createAArch64ISelDag(AArch64TargetMachine &TM,
 FunctionPass *createAArch64StorePairSuppressPass();
 FunctionPass *createAArch64ExpandPseudoPass();
 FunctionPass *createAArch64SLSHardeningPass();
-FunctionPass *createAArch64IndirectThunks();
 FunctionPass *createAArch64SpeculationHardeningPass();
 FunctionPass *createAArch64LoadStoreOptimizationPass();
 ModulePass *createAArch64LowerHomogeneousPrologEpilogPass();
diff --git a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp 
b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
index d35386eaeab12..b4ebd7d5377c2 100644
--- a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
+++ b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
@@ -13,20 +13,16 @@
 
 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
-#include "Utils/AArch64BaseInfo.h"
 #include "llvm/CodeGen/IndirectThunks.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
-#include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstr.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineOperand.h"
-#include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/RegisterScavenging.h"
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/Pass.h"
-#include "llvm/Support/CodeGen.h"
-#include "llvm/Support/Debug.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Target/TargetMachine.h"
 #include 
 
@@ -36,38 +32,43 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
+static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+
 namespace {
 
-class AArch64SLSHardening : public MachineFunctionPass {
-public:
-  const TargetInstrInfo *TII;
-  const TargetRegisterInfo *TRI;
-  const AArch64Subtarget *ST;
+// Set of inserted thunks: bitmask with bits corresponding to
+// indexes in SLSBLRThunks array.
+typedef uint32_t ThunksSet;
 
-  static char ID;
-
-  AArch64SLSHardening() : MachineFunctionPass(ID) {
-initializeAArch64SLSHardeningPass(*PassRegistry::getPassRegistry());
+struct SLSBLRThunkInserter : ThunkInserter {
+public:
+  const char *getThunkPrefix() { return SLSBLRNamePrefix; }
+  bool mayUseThunk(const MachineFunction &MF) {
+// FIXME: ComdatThunks is only accumulated until the first thunk is 
created.
+ComdatThunks &= !MF.getSubtarget().hardenSlsNoComdat();
+// We are inserting barriers aside from thunk calls, so
+// check hardenSlsRetBr() as well.
+return MF.getSubtarget().hardenSlsBlr() ||
+   MF.getSubtarget().hardenSlsRetBr();
   }
+  ThunksSet insertThunks(MachineModuleInfo &MMI, MachineFunction &MF,
+ ThunksSet ExistingThunks);
+  void populateThunk(MachineFunction &MF);
 
-  bool runOnMachineFunction(MachineFunction &Fn) override;
+private:
+  bool ComdatThunks = true;
 
-  StringRef getPassName() const override { return AARCH64_SLS_HARDENING_NAME; }
+  bool hardenReturnsAndBRs(MachineModuleInfo &

[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits


@@ -36,38 +32,43 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
+static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+
 namespace {
 
-class AArch64SLSHardening : public MachineFunctionPass {
-public:
-  const TargetInstrInfo *TII;
-  const TargetRegisterInfo *TRI;
-  const AArch64Subtarget *ST;
+// Set of inserted thunks: bitmask with bits corresponding to
+// indexes in SLSBLRThunks array.
+typedef uint32_t ThunksSet;

atrosinenko wrote:

Yes, `ThunkSet` will be class.

https://github.com/llvm/llvm-project/pull/97472
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits


@@ -46,13 +40,5 @@ body: |
 
 
 ...

-name:__llvm_slsblr_thunk_x8
-tracksRegLiveness: true
-body: |
-  bb.0.entry:
-liveins: $x8
 
-BR $x8

atrosinenko wrote:

The above lines are **input** to `llc`, so I removed them to not conflict with 
the thunks that are inserted by the pass. Unfortunately, MIR syntax is quite 
verbose, so I added a few basic CHECK lines testing MIR output and duplicated 
the RUN lines to generate and check assembly output as well.

Now there is no separation between function-rewriting and thunk-inserting 
passes, and I assumed it is safe to expect that initially there are no thunks 
in the module. I guess the original version of thunk insertion logic would 
assert on pre-existing thunks as well (due to unknown thunk names with `.1` 
suffix).

https://github.com/llvm/llvm-project/pull/97472
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko edited 
https://github.com/llvm/llvm-project/pull/97472
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko commented:

@kbeyls thank you for the comments!

https://github.com/llvm/llvm-project/pull/97472
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits


@@ -36,38 +32,43 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
+static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+
 namespace {
 
-class AArch64SLSHardening : public MachineFunctionPass {
-public:
-  const TargetInstrInfo *TII;
-  const TargetRegisterInfo *TRI;
-  const AArch64Subtarget *ST;
+// Set of inserted thunks: bitmask with bits corresponding to
+// indexes in SLSBLRThunks array.
+typedef uint32_t ThunksSet;
 
-  static char ID;
-
-  AArch64SLSHardening() : MachineFunctionPass(ID) {
-initializeAArch64SLSHardeningPass(*PassRegistry::getPassRegistry());
+struct SLSBLRThunkInserter : ThunkInserter {

atrosinenko wrote:

Renamed it to `SLSHardeningInserter`, so the `AArch64SLSHardening` name is kept 
for the pass itself.

https://github.com/llvm/llvm-project/pull/97472
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko edited 
https://github.com/llvm/llvm-project/pull/97472
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [llvm] [AArch64][PAC] Support BLRA* instructions in SLS Hardening pass (PR #97605)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko created 
https://github.com/llvm/llvm-project/pull/97605

Make SLS Hardening pass handle BLRA* instructions the same way it handles BLR. 
The thunk names have the form

__llvm_slsblr_thunk_xNfor BLR thunks
__llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
__llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks

Now there are about 1800 possible thunk names, so do not rely on linear thunk 
function's name lookup and parse the name instead.

>From b389284b8e92f5bf09cea38f3f9a53974a84dc29 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko 
Date: Mon, 1 Jul 2024 20:13:54 +0300
Subject: [PATCH] [AArch64][PAC] Support BLRA* instructions in SLS Hardening
 pass

Make SLS Hardening pass handle BLRA* instructions the same way it
handles BLR. The thunk names have the form

__llvm_slsblr_thunk_xNfor BLR thunks
__llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
__llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks

Now there are about 1800 possible thunk names, so do not rely on linear
thunk function's name lookup and parse the name instead.
---
 .../Target/AArch64/AArch64SLSHardening.cpp| 326 --
 .../speculation-hardening-sls-blra.mir| 210 +++
 2 files changed, 432 insertions(+), 104 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/speculation-hardening-sls-blra.mir

diff --git a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp 
b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
index feb166f30127a..d93fe2a875845 100644
--- a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
+++ b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
@@ -13,6 +13,7 @@
 
 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
+#include "llvm/ADT/StringSwitch.h"
 #include "llvm/CodeGen/IndirectThunks.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -23,6 +24,7 @@
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Target/TargetMachine.h"
 #include 
 
@@ -32,17 +34,103 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
-static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+// Common name prefix of all thunks generated by this pass.
+//
+// The generic form is
+// __llvm_slsblr_thunk_xNfor BLR thunks
+// __llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
+// __llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks
+static constexpr StringRef CommonNamePrefix = "__llvm_slsblr_thunk_";
 
 namespace {
 
-// Set of inserted thunks: bitmask with bits corresponding to
-// indexes in SLSBLRThunks array.
-typedef uint32_t ThunksSet;
+struct ThunkKind {
+  enum ThunkKindId {
+ThunkBR,
+ThunkBRAA,
+ThunkBRAB,
+ThunkBRAAZ,
+ThunkBRABZ,
+  };
+
+  ThunkKindId Id;
+  StringRef NameInfix;
+  bool HasXmOperand;
+  bool NeedsPAuth;
+
+  // Opcode to perform indirect jump from inside the thunk.
+  unsigned BROpcode;
+
+  static const ThunkKind BR;
+  static const ThunkKind BRAA;
+  static const ThunkKind BRAB;
+  static const ThunkKind BRAAZ;
+  static const ThunkKind BRABZ;
+};
+
+// Set of inserted thunks.
+class ThunksSet {
+public:
+  static constexpr unsigned NumXRegisters = 32;
+
+  // Given Xn register, returns n.
+  static unsigned indexOfXReg(Register Xn);
+  // Given n, returns Xn register.
+  static Register xRegByIndex(unsigned N);
+
+  ThunksSet &operator|=(const ThunksSet &Other) {
+BLRThunks |= Other.BLRThunks;
+BLRAAZThunks |= Other.BLRAAZThunks;
+BLRABZThunks |= Other.BLRABZThunks;
+for (unsigned I = 0; I < NumXRegisters; ++I)
+  BLRAAThunks[I] |= Other.BLRAAThunks[I];
+for (unsigned I = 0; I < NumXRegisters; ++I)
+  BLRABThunks[I] |= Other.BLRABThunks[I];
+
+return *this;
+  }
+
+  bool get(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+uint32_t XnBit = 1u << indexOfXReg(Xn);
+return getBitmask(Kind, Xm) & XnBit;
+  }
+
+  void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+uint32_t XnBit = 1u << indexOfXReg(Xn);
+getBitmask(Kind, Xm) |= XnBit;
+  }
+
+private:
+  // Bitmasks representing operands used, with n-th bit corresponding to Xn
+  // register operand. If the instruction has a second operand (Xm), an array
+  // of bitmasks is used, indexed by m.
+  // Indexes corresponding to the forbidden x16, x17 and x30 registers are
+  // always unset, for simplicity there are no holes.
+  uint32_t BLRThunks = 0;
+  uint32_t BLRAAZThunks = 0;
+  uint32_t BLRABZThunks = 0;
+  uint32_t BLRAAThunks[NumXRegisters] = {};
+  uint32_t BLRABThunks[NumXRegisters] = {};
+
+  uint32_t &getBitmask(ThunkKind::ThunkKindId Kind, Register Xm) {
+switch (Kind) {
+case ThunkKind::ThunkBR:
+  return BLRThunks;
+case ThunkKind::ThunkBRAAZ:
+  return BLRAAZThunks;
+ca

[llvm-branch-commits] [llvm] [AArch64] Only create called thunks when hardening against SLS (PR #97472)

2024-07-03 Thread Anatoly Trosinenko via llvm-branch-commits


@@ -36,38 +32,43 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
+static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+
 namespace {
 
-class AArch64SLSHardening : public MachineFunctionPass {
-public:
-  const TargetInstrInfo *TII;
-  const TargetRegisterInfo *TRI;
-  const AArch64Subtarget *ST;
+// Set of inserted thunks: bitmask with bits corresponding to
+// indexes in SLSBLRThunks array.
+typedef uint32_t ThunksSet;

atrosinenko wrote:

Here is the PR: #97605

https://github.com/llvm/llvm-project/pull/97472
___
llvm-branch-commits mailing list
llvm-branch-commits@lists.llvm.org
https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits


[llvm-branch-commits] [llvm] [AArch64][PAC] Support BLRA* instructions in SLS Hardening pass (PR #97605)

2024-07-04 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko updated 
https://github.com/llvm/llvm-project/pull/97605

>From f49c32c8465e9e68d7345fa82ae1294cc2faf0e7 Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko 
Date: Mon, 1 Jul 2024 20:13:54 +0300
Subject: [PATCH] [AArch64][PAC] Support BLRA* instructions in SLS Hardening
 pass

Make SLS Hardening pass handle BLRA* instructions the same way it
handles BLR. The thunk names have the form

__llvm_slsblr_thunk_xNfor BLR thunks
__llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
__llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks

Now there are about 1800 possible thunk names, so do not rely on linear
thunk function's name lookup and parse the name instead.
---
 .../Target/AArch64/AArch64SLSHardening.cpp| 326 --
 .../speculation-hardening-sls-blra.mir| 210 +++
 2 files changed, 432 insertions(+), 104 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/speculation-hardening-sls-blra.mir

diff --git a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp 
b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
index 24f023a3d70e7..4ae2e48af337f 100644
--- a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
+++ b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
@@ -13,6 +13,7 @@
 
 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
+#include "llvm/ADT/StringSwitch.h"
 #include "llvm/CodeGen/IndirectThunks.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -23,6 +24,7 @@
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Target/TargetMachine.h"
 #include 
 
@@ -32,17 +34,103 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
-static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+// Common name prefix of all thunks generated by this pass.
+//
+// The generic form is
+// __llvm_slsblr_thunk_xNfor BLR thunks
+// __llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
+// __llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks
+static constexpr StringRef CommonNamePrefix = "__llvm_slsblr_thunk_";
 
 namespace {
 
-// Set of inserted thunks: bitmask with bits corresponding to
-// indexes in SLSBLRThunks array.
-typedef uint32_t ThunksSet;
+struct ThunkKind {
+  enum ThunkKindId {
+ThunkBR,
+ThunkBRAA,
+ThunkBRAB,
+ThunkBRAAZ,
+ThunkBRABZ,
+  };
+
+  ThunkKindId Id;
+  StringRef NameInfix;
+  bool HasXmOperand;
+  bool NeedsPAuth;
+
+  // Opcode to perform indirect jump from inside the thunk.
+  unsigned BROpcode;
+
+  static const ThunkKind BR;
+  static const ThunkKind BRAA;
+  static const ThunkKind BRAB;
+  static const ThunkKind BRAAZ;
+  static const ThunkKind BRABZ;
+};
+
+// Set of inserted thunks.
+class ThunksSet {
+public:
+  static constexpr unsigned NumXRegisters = 32;
+
+  // Given Xn register, returns n.
+  static unsigned indexOfXReg(Register Xn);
+  // Given n, returns Xn register.
+  static Register xRegByIndex(unsigned N);
+
+  ThunksSet &operator|=(const ThunksSet &Other) {
+BLRThunks |= Other.BLRThunks;
+BLRAAZThunks |= Other.BLRAAZThunks;
+BLRABZThunks |= Other.BLRABZThunks;
+for (unsigned I = 0; I < NumXRegisters; ++I)
+  BLRAAThunks[I] |= Other.BLRAAThunks[I];
+for (unsigned I = 0; I < NumXRegisters; ++I)
+  BLRABThunks[I] |= Other.BLRABThunks[I];
+
+return *this;
+  }
+
+  bool get(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+uint32_t XnBit = 1u << indexOfXReg(Xn);
+return getBitmask(Kind, Xm) & XnBit;
+  }
+
+  void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+uint32_t XnBit = 1u << indexOfXReg(Xn);
+getBitmask(Kind, Xm) |= XnBit;
+  }
+
+private:
+  // Bitmasks representing operands used, with n-th bit corresponding to Xn
+  // register operand. If the instruction has a second operand (Xm), an array
+  // of bitmasks is used, indexed by m.
+  // Indexes corresponding to the forbidden x16, x17 and x30 registers are
+  // always unset, for simplicity there are no holes.
+  uint32_t BLRThunks = 0;
+  uint32_t BLRAAZThunks = 0;
+  uint32_t BLRABZThunks = 0;
+  uint32_t BLRAAThunks[NumXRegisters] = {};
+  uint32_t BLRABThunks[NumXRegisters] = {};
+
+  uint32_t &getBitmask(ThunkKind::ThunkKindId Kind, Register Xm) {
+switch (Kind) {
+case ThunkKind::ThunkBR:
+  return BLRThunks;
+case ThunkKind::ThunkBRAAZ:
+  return BLRAAZThunks;
+case ThunkKind::ThunkBRABZ:
+  return BLRABZThunks;
+case ThunkKind::ThunkBRAA:
+  return BLRAAThunks[indexOfXReg(Xm)];
+case ThunkKind::ThunkBRAB:
+  return BLRABThunks[indexOfXReg(Xm)];
+}
+  }
+};
 
 struct SLSHardeningInserter : ThunkInserter {
 public:
-  const char *getThunkPrefix() { return SLSBLRNamePrefix; }
+  const char *getThunkPrefix() { return CommonNamePrefix.data(); }
   bool mayUseThunk

[llvm-branch-commits] [llvm] [AArch64][PAC] Support BLRA* instructions in SLS Hardening pass (PR #97605)

2024-07-04 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko updated 
https://github.com/llvm/llvm-project/pull/97605

>From 84e7eb3c36b99ac7f673d9a7ad0c88c469f7f45d Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko 
Date: Mon, 1 Jul 2024 20:13:54 +0300
Subject: [PATCH] [AArch64][PAC] Support BLRA* instructions in SLS Hardening
 pass

Make SLS Hardening pass handle BLRA* instructions the same way it
handles BLR. The thunk names have the form

__llvm_slsblr_thunk_xNfor BLR thunks
__llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
__llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks

Now there are about 1800 possible thunk names, so do not rely on linear
thunk function's name lookup and parse the name instead.
---
 .../Target/AArch64/AArch64SLSHardening.cpp| 326 --
 .../speculation-hardening-sls-blra.mir| 210 +++
 2 files changed, 432 insertions(+), 104 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/speculation-hardening-sls-blra.mir

diff --git a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp 
b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
index 00ba31b3e500d..4ae2e48af337f 100644
--- a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
+++ b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
@@ -13,6 +13,7 @@
 
 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
+#include "llvm/ADT/StringSwitch.h"
 #include "llvm/CodeGen/IndirectThunks.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -23,6 +24,7 @@
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Target/TargetMachine.h"
 #include 
 
@@ -32,17 +34,103 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
-static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+// Common name prefix of all thunks generated by this pass.
+//
+// The generic form is
+// __llvm_slsblr_thunk_xNfor BLR thunks
+// __llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
+// __llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks
+static constexpr StringRef CommonNamePrefix = "__llvm_slsblr_thunk_";
 
 namespace {
 
-// Set of inserted thunks: bitmask with bits corresponding to
-// indexes in SLSBLRThunks array.
-typedef uint32_t ThunksSet;
+struct ThunkKind {
+  enum ThunkKindId {
+ThunkBR,
+ThunkBRAA,
+ThunkBRAB,
+ThunkBRAAZ,
+ThunkBRABZ,
+  };
+
+  ThunkKindId Id;
+  StringRef NameInfix;
+  bool HasXmOperand;
+  bool NeedsPAuth;
+
+  // Opcode to perform indirect jump from inside the thunk.
+  unsigned BROpcode;
+
+  static const ThunkKind BR;
+  static const ThunkKind BRAA;
+  static const ThunkKind BRAB;
+  static const ThunkKind BRAAZ;
+  static const ThunkKind BRABZ;
+};
+
+// Set of inserted thunks.
+class ThunksSet {
+public:
+  static constexpr unsigned NumXRegisters = 32;
+
+  // Given Xn register, returns n.
+  static unsigned indexOfXReg(Register Xn);
+  // Given n, returns Xn register.
+  static Register xRegByIndex(unsigned N);
+
+  ThunksSet &operator|=(const ThunksSet &Other) {
+BLRThunks |= Other.BLRThunks;
+BLRAAZThunks |= Other.BLRAAZThunks;
+BLRABZThunks |= Other.BLRABZThunks;
+for (unsigned I = 0; I < NumXRegisters; ++I)
+  BLRAAThunks[I] |= Other.BLRAAThunks[I];
+for (unsigned I = 0; I < NumXRegisters; ++I)
+  BLRABThunks[I] |= Other.BLRABThunks[I];
+
+return *this;
+  }
+
+  bool get(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+uint32_t XnBit = 1u << indexOfXReg(Xn);
+return getBitmask(Kind, Xm) & XnBit;
+  }
+
+  void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+uint32_t XnBit = 1u << indexOfXReg(Xn);
+getBitmask(Kind, Xm) |= XnBit;
+  }
+
+private:
+  // Bitmasks representing operands used, with n-th bit corresponding to Xn
+  // register operand. If the instruction has a second operand (Xm), an array
+  // of bitmasks is used, indexed by m.
+  // Indexes corresponding to the forbidden x16, x17 and x30 registers are
+  // always unset, for simplicity there are no holes.
+  uint32_t BLRThunks = 0;
+  uint32_t BLRAAZThunks = 0;
+  uint32_t BLRABZThunks = 0;
+  uint32_t BLRAAThunks[NumXRegisters] = {};
+  uint32_t BLRABThunks[NumXRegisters] = {};
+
+  uint32_t &getBitmask(ThunkKind::ThunkKindId Kind, Register Xm) {
+switch (Kind) {
+case ThunkKind::ThunkBR:
+  return BLRThunks;
+case ThunkKind::ThunkBRAAZ:
+  return BLRAAZThunks;
+case ThunkKind::ThunkBRABZ:
+  return BLRABZThunks;
+case ThunkKind::ThunkBRAA:
+  return BLRAAThunks[indexOfXReg(Xm)];
+case ThunkKind::ThunkBRAB:
+  return BLRABThunks[indexOfXReg(Xm)];
+}
+  }
+};
 
 struct SLSHardeningInserter : ThunkInserter {
 public:
-  const char *getThunkPrefix() { return SLSBLRNamePrefix; }
+  const char *getThunkPrefix() { return CommonNamePrefix.data(); }
   bool mayUseThunk

[llvm-branch-commits] [llvm] [AArch64][PAC] Support BLRA* instructions in SLS Hardening pass (PR #97605)

2024-07-04 Thread Anatoly Trosinenko via llvm-branch-commits

https://github.com/atrosinenko updated 
https://github.com/llvm/llvm-project/pull/97605

>From 84e7eb3c36b99ac7f673d9a7ad0c88c469f7f45d Mon Sep 17 00:00:00 2001
From: Anatoly Trosinenko 
Date: Mon, 1 Jul 2024 20:13:54 +0300
Subject: [PATCH 1/2] [AArch64][PAC] Support BLRA* instructions in SLS
 Hardening pass

Make SLS Hardening pass handle BLRA* instructions the same way it
handles BLR. The thunk names have the form

__llvm_slsblr_thunk_xNfor BLR thunks
__llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
__llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks

Now there are about 1800 possible thunk names, so do not rely on linear
thunk function's name lookup and parse the name instead.
---
 .../Target/AArch64/AArch64SLSHardening.cpp| 326 --
 .../speculation-hardening-sls-blra.mir| 210 +++
 2 files changed, 432 insertions(+), 104 deletions(-)
 create mode 100644 llvm/test/CodeGen/AArch64/speculation-hardening-sls-blra.mir

diff --git a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp 
b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
index 00ba31b3e500d..4ae2e48af337f 100644
--- a/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
+++ b/llvm/lib/Target/AArch64/AArch64SLSHardening.cpp
@@ -13,6 +13,7 @@
 
 #include "AArch64InstrInfo.h"
 #include "AArch64Subtarget.h"
+#include "llvm/ADT/StringSwitch.h"
 #include "llvm/CodeGen/IndirectThunks.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFunction.h"
@@ -23,6 +24,7 @@
 #include "llvm/IR/DebugLoc.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/FormatVariadic.h"
 #include "llvm/Target/TargetMachine.h"
 #include 
 
@@ -32,17 +34,103 @@ using namespace llvm;
 
 #define AARCH64_SLS_HARDENING_NAME "AArch64 sls hardening pass"
 
-static const char SLSBLRNamePrefix[] = "__llvm_slsblr_thunk_";
+// Common name prefix of all thunks generated by this pass.
+//
+// The generic form is
+// __llvm_slsblr_thunk_xNfor BLR thunks
+// __llvm_slsblr_thunk_(aaz|abz)_xN  for BLRAAZ and BLRABZ thunks
+// __llvm_slsblr_thunk_(aa|ab)_xN_xM for BLRAA and BLRAB thunks
+static constexpr StringRef CommonNamePrefix = "__llvm_slsblr_thunk_";
 
 namespace {
 
-// Set of inserted thunks: bitmask with bits corresponding to
-// indexes in SLSBLRThunks array.
-typedef uint32_t ThunksSet;
+struct ThunkKind {
+  enum ThunkKindId {
+ThunkBR,
+ThunkBRAA,
+ThunkBRAB,
+ThunkBRAAZ,
+ThunkBRABZ,
+  };
+
+  ThunkKindId Id;
+  StringRef NameInfix;
+  bool HasXmOperand;
+  bool NeedsPAuth;
+
+  // Opcode to perform indirect jump from inside the thunk.
+  unsigned BROpcode;
+
+  static const ThunkKind BR;
+  static const ThunkKind BRAA;
+  static const ThunkKind BRAB;
+  static const ThunkKind BRAAZ;
+  static const ThunkKind BRABZ;
+};
+
+// Set of inserted thunks.
+class ThunksSet {
+public:
+  static constexpr unsigned NumXRegisters = 32;
+
+  // Given Xn register, returns n.
+  static unsigned indexOfXReg(Register Xn);
+  // Given n, returns Xn register.
+  static Register xRegByIndex(unsigned N);
+
+  ThunksSet &operator|=(const ThunksSet &Other) {
+BLRThunks |= Other.BLRThunks;
+BLRAAZThunks |= Other.BLRAAZThunks;
+BLRABZThunks |= Other.BLRABZThunks;
+for (unsigned I = 0; I < NumXRegisters; ++I)
+  BLRAAThunks[I] |= Other.BLRAAThunks[I];
+for (unsigned I = 0; I < NumXRegisters; ++I)
+  BLRABThunks[I] |= Other.BLRABThunks[I];
+
+return *this;
+  }
+
+  bool get(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+uint32_t XnBit = 1u << indexOfXReg(Xn);
+return getBitmask(Kind, Xm) & XnBit;
+  }
+
+  void set(ThunkKind::ThunkKindId Kind, Register Xn, Register Xm) {
+uint32_t XnBit = 1u << indexOfXReg(Xn);
+getBitmask(Kind, Xm) |= XnBit;
+  }
+
+private:
+  // Bitmasks representing operands used, with n-th bit corresponding to Xn
+  // register operand. If the instruction has a second operand (Xm), an array
+  // of bitmasks is used, indexed by m.
+  // Indexes corresponding to the forbidden x16, x17 and x30 registers are
+  // always unset, for simplicity there are no holes.
+  uint32_t BLRThunks = 0;
+  uint32_t BLRAAZThunks = 0;
+  uint32_t BLRABZThunks = 0;
+  uint32_t BLRAAThunks[NumXRegisters] = {};
+  uint32_t BLRABThunks[NumXRegisters] = {};
+
+  uint32_t &getBitmask(ThunkKind::ThunkKindId Kind, Register Xm) {
+switch (Kind) {
+case ThunkKind::ThunkBR:
+  return BLRThunks;
+case ThunkKind::ThunkBRAAZ:
+  return BLRAAZThunks;
+case ThunkKind::ThunkBRABZ:
+  return BLRABZThunks;
+case ThunkKind::ThunkBRAA:
+  return BLRAAThunks[indexOfXReg(Xm)];
+case ThunkKind::ThunkBRAB:
+  return BLRABThunks[indexOfXReg(Xm)];
+}
+  }
+};
 
 struct SLSHardeningInserter : ThunkInserter {
 public:
-  const char *getThunkPrefix() { return SLSBLRNamePrefix; }
+  const char *getThunkPrefix() { return CommonNamePrefix.data(); }
   bool mayUseT