From 0593ce5f0bbb7b9aebc5f15f60876b0bc7aad652 Mon Sep 17 00:00:00 2001 From: Florian Mayer Date: Thu, 2 Jun 2022 18:05:02 -0700 Subject: [PATCH] [MC] Add 'G' to augmentation string for MTE instrumented functions This was agreed on in https://lists.llvm.org/pipermail/llvm-dev/2020-May/141345.html The thread proposed two options * add a character to augmentation string and handle in libuwind * use a separate personality function. It was determined that this is the simpler and better option. This is part of ARM's Aarch64 ABI: https://github.com/ARM-software/abi-aa/blob/main/aadwarf64/aadwarf64.rst#id22 The next step after this is teaching libunwind to untag when this augmentation character is set. Reviewed By: MaskRay, eugenis Differential Revision: https://reviews.llvm.org/D127007 --- llvm/include/llvm/MC/MCDwarf.h | 1 + llvm/include/llvm/MC/MCStreamer.h | 1 + llvm/lib/MC/MCAsmStreamer.cpp | 7 ++++ llvm/lib/MC/MCDwarf.cpp | 2 + llvm/lib/MC/MCParser/AsmParser.cpp | 2 + llvm/lib/MC/MCStreamer.cpp | 7 ++++ llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp | 11 ++++++ llvm/lib/Target/AArch64/AArch64FrameLowering.cpp | 5 +++ llvm/lib/Target/AArch64/AArch64InstrInfo.td | 4 ++ .../Target/AArch64/AArch64MachineFunctionInfo.cpp | 2 + .../Target/AArch64/AArch64MachineFunctionInfo.h | 3 ++ .../Target/AArch64/AsmParser/AArch64AsmParser.cpp | 12 ++++++ llvm/test/CodeGen/AArch64/stack-tagging-cfi.ll | 43 ++++++++++++++++++++++ 13 files changed, 100 insertions(+) create mode 100644 llvm/test/CodeGen/AArch64/stack-tagging-cfi.ll diff --git a/llvm/include/llvm/MC/MCDwarf.h b/llvm/include/llvm/MC/MCDwarf.h index a4f08b4..d0eb06e 100644 --- a/llvm/include/llvm/MC/MCDwarf.h +++ b/llvm/include/llvm/MC/MCDwarf.h @@ -688,6 +688,7 @@ struct MCDwarfFrameInfo { bool IsSimple = false; unsigned RAReg = static_cast(INT_MAX); bool IsBKeyFrame = false; + bool IsMTETaggedFrame = false; }; class MCDwarfFrameEmitter { diff --git a/llvm/include/llvm/MC/MCStreamer.h b/llvm/include/llvm/MC/MCStreamer.h index 47594ab..68c4ba8 100644 --- a/llvm/include/llvm/MC/MCStreamer.h +++ b/llvm/include/llvm/MC/MCStreamer.h @@ -926,6 +926,7 @@ public: unsigned CUID = 0); virtual void emitCFIBKeyFrame(); + virtual void emitCFIMTETaggedFrame(); /// This implements the DWARF2 '.loc fileno lineno ...' assembler /// directive. diff --git a/llvm/lib/MC/MCAsmStreamer.cpp b/llvm/lib/MC/MCAsmStreamer.cpp index 3d0853f..e48ea22 100644 --- a/llvm/lib/MC/MCAsmStreamer.cpp +++ b/llvm/lib/MC/MCAsmStreamer.cpp @@ -319,6 +319,7 @@ public: void emitIdent(StringRef IdentString) override; void emitCFIBKeyFrame() override; + void emitCFIMTETaggedFrame() override; void emitCFISections(bool EH, bool Debug) override; void emitCFIDefCfa(int64_t Register, int64_t Offset) override; void emitCFIDefCfaOffset(int64_t Offset) override; @@ -2025,6 +2026,12 @@ void MCAsmStreamer::emitCFIBKeyFrame() { EmitEOL(); } +void MCAsmStreamer::emitCFIMTETaggedFrame() { + MCStreamer::emitCFIMTETaggedFrame(); + OS << "\t.cfi_mte_tagged_frame"; + EmitEOL(); +} + void MCAsmStreamer::emitWinCFIStartProc(const MCSymbol *Symbol, SMLoc Loc) { MCStreamer::emitWinCFIStartProc(Symbol, Loc); diff --git a/llvm/lib/MC/MCDwarf.cpp b/llvm/lib/MC/MCDwarf.cpp index 0f45f2f..4db409a 100644 --- a/llvm/lib/MC/MCDwarf.cpp +++ b/llvm/lib/MC/MCDwarf.cpp @@ -1604,6 +1604,8 @@ const MCSymbol &FrameEmitterImpl::EmitCIE(const MCDwarfFrameInfo &Frame) { Augmentation += "S"; if (Frame.IsBKeyFrame) Augmentation += "B"; + if (Frame.IsMTETaggedFrame) + Augmentation += "G"; Streamer.emitBytes(Augmentation); } Streamer.emitInt8(0); diff --git a/llvm/lib/MC/MCParser/AsmParser.cpp b/llvm/lib/MC/MCParser/AsmParser.cpp index 35d75a9..ccc8e80 100644 --- a/llvm/lib/MC/MCParser/AsmParser.cpp +++ b/llvm/lib/MC/MCParser/AsmParser.cpp @@ -540,6 +540,7 @@ private: DK_PSEUDO_PROBE, DK_LTO_DISCARD, DK_LTO_SET_CONDITIONAL, + DK_CFI_MTE_TAGGED_FRAME, DK_END }; @@ -5556,6 +5557,7 @@ void AsmParser::initializeDirectiveKindMap() { DirectiveKindMap[".cfi_register"] = DK_CFI_REGISTER; DirectiveKindMap[".cfi_window_save"] = DK_CFI_WINDOW_SAVE; DirectiveKindMap[".cfi_b_key_frame"] = DK_CFI_B_KEY_FRAME; + DirectiveKindMap[".cfi_mte_tagged_frame"] = DK_CFI_MTE_TAGGED_FRAME; DirectiveKindMap[".macros_on"] = DK_MACROS_ON; DirectiveKindMap[".macros_off"] = DK_MACROS_OFF; DirectiveKindMap[".macro"] = DK_MACRO; diff --git a/llvm/lib/MC/MCStreamer.cpp b/llvm/lib/MC/MCStreamer.cpp index 066a7cd..64b2d26 100644 --- a/llvm/lib/MC/MCStreamer.cpp +++ b/llvm/lib/MC/MCStreamer.cpp @@ -252,6 +252,13 @@ void MCStreamer::emitCFIBKeyFrame() { CurFrame->IsBKeyFrame = true; } +void MCStreamer::emitCFIMTETaggedFrame() { + MCDwarfFrameInfo *CurFrame = getCurrentDwarfFrameInfo(); + if (!CurFrame) + return; + CurFrame->IsMTETaggedFrame = true; +} + void MCStreamer::emitDwarfLocDirective(unsigned FileNo, unsigned Line, unsigned Column, unsigned Flags, unsigned Isa, unsigned Discriminator, diff --git a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp index d0363d0..ef4af4d 100644 --- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp +++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp @@ -1285,6 +1285,17 @@ void AArch64AsmPrinter::emitInstruction(const MachineInstr *MI) { return; } + case AArch64::EMITMTETAGGED: { + ExceptionHandling ExceptionHandlingType = MAI->getExceptionHandlingType(); + if (ExceptionHandlingType != ExceptionHandling::DwarfCFI && + ExceptionHandlingType != ExceptionHandling::ARM) + return; + + if (getFunctionCFISectionType(*MF) != CFISection::None) + OutStreamer->emitCFIMTETaggedFrame(); + return; + } + // Tail calls use pseudo instructions so they have the proper code-gen // attributes (isCall, isReturn, etc.). We lower them to the real // instruction here. diff --git a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp index 9f8e153..78babdf 100644 --- a/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp +++ b/llvm/lib/Target/AArch64/AArch64FrameLowering.cpp @@ -192,6 +192,7 @@ #include "AArch64Subtarget.h" #include "AArch64TargetMachine.h" #include "MCTargetDesc/AArch64AddressingModes.h" +#include "MCTargetDesc/AArch64MCTargetDesc.h" #include "llvm/ADT/ScopeExit.h" #include "llvm/ADT/SmallVector.h" #include "llvm/ADT/Statistic.h" @@ -1427,6 +1428,10 @@ void AArch64FrameLowering::emitPrologue(MachineFunction &MF, .setMIFlags(MachineInstr::FrameSetup); } } + if (EmitCFI && MFnI.isMTETagged()) { + BuildMI(MBB, MBBI, DL, TII->get(AArch64::EMITMTETAGGED)) + .setMIFlag(MachineInstr::FrameSetup); + } // We signal the presence of a Swift extended frame to external tools by // storing FP with 0b0001 in bits 63:60. In normal userland operation a simple diff --git a/llvm/lib/Target/AArch64/AArch64InstrInfo.td b/llvm/lib/Target/AArch64/AArch64InstrInfo.td index 39e05b3..00c0d74 100644 --- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td +++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td @@ -2438,6 +2438,10 @@ def TLSDESCCALL : Pseudo<(outs), (ins i64imm:$sym), []>, Sched<[]> { // augmentation string. def EMITBKEY : Pseudo<(outs), (ins), []>, Sched<[]> {} +// Pseudo instruction to tell the streamer to emit a 'G' character into the +// augmentation string. +def EMITMTETAGGED : Pseudo<(outs), (ins), []>, Sched<[]> {} + // FIXME: maybe the scratch register used shouldn't be fixed to X1? // FIXME: can "hasSideEffects be dropped? // This gets lowered to an instruction sequence which takes 16 bytes diff --git a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp index edca145..8de3cba 100644 --- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp +++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp @@ -89,6 +89,8 @@ AArch64FunctionInfo::AArch64FunctionInfo(MachineFunction &MF_) : MF(&MF_) { const Function &F = MF->getFunction(); std::tie(SignReturnAddress, SignReturnAddressAll) = GetSignReturnAddress(F); SignWithBKey = ShouldSignWithBKey(F); + // TODO: skip functions that have no instrumented allocas for optimization + IsMTETagged = F.hasFnAttribute(Attribute::SanitizeMemTag); if (!F.hasFnAttribute("branch-target-enforcement")) { if (const auto *BTE = mdconst::extract_or_null( diff --git a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h index 17d01ac..df047fb 100644 --- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h +++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h @@ -175,6 +175,8 @@ class AArch64FunctionInfo final : public MachineFunctionInfo { /// The stack slot where the Swift asynchronous context is stored. int SwiftAsyncContextFrameIdx = std::numeric_limits::max(); + bool IsMTETagged = false; + /// True if the function need unwind information. mutable Optional NeedsDwarfUnwindInfo; @@ -408,6 +410,7 @@ public: bool shouldSignReturnAddress(bool SpillsLR) const; bool shouldSignWithBKey() const { return SignWithBKey; } + bool isMTETagged() const { return IsMTETagged; } bool branchTargetEnforcement() const { return BranchTargetEnforcement; } diff --git a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp index 58247fc..a64e0a8 100644 --- a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp +++ b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp @@ -190,6 +190,7 @@ private: bool parseDirectiveUnreq(SMLoc L); bool parseDirectiveCFINegateRAState(); bool parseDirectiveCFIBKeyFrame(); + bool parseDirectiveCFIMTETaggedFrame(); bool parseDirectiveVariantPCS(SMLoc L); @@ -6038,6 +6039,8 @@ bool AArch64AsmParser::ParseDirective(AsmToken DirectiveID) { parseDirectiveCFINegateRAState(); else if (IDVal == ".cfi_b_key_frame") parseDirectiveCFIBKeyFrame(); + else if (IDVal == ".cfi_mte_tagged_frame") + parseDirectiveCFIMTETaggedFrame(); else if (IDVal == ".arch_extension") parseDirectiveArchExtension(Loc); else if (IDVal == ".variant_pcs") @@ -6518,6 +6521,15 @@ bool AArch64AsmParser::parseDirectiveCFIBKeyFrame() { return false; } +/// parseDirectiveCFIMTETaggedFrame +/// ::= .cfi_mte_tagged_frame +bool AArch64AsmParser::parseDirectiveCFIMTETaggedFrame() { + if (parseEOL()) + return true; + getStreamer().emitCFIMTETaggedFrame(); + return false; +} + /// parseDirectiveVariantPCS /// ::= .variant_pcs symbolname bool AArch64AsmParser::parseDirectiveVariantPCS(SMLoc L) { diff --git a/llvm/test/CodeGen/AArch64/stack-tagging-cfi.ll b/llvm/test/CodeGen/AArch64/stack-tagging-cfi.ll new file mode 100644 index 0000000..7e1ab1c --- /dev/null +++ b/llvm/test/CodeGen/AArch64/stack-tagging-cfi.ll @@ -0,0 +1,43 @@ +; RUN: llc < %s -mtriple=aarch64 -mattr=+mte | FileCheck %s + +declare void @use32(ptr) + +define void @WithUnwind() sanitize_memtag { +entry: +; CHECK-LABEL: WithUnwind: +; CHECK: .cfi_mte_tagged_frame +; CHECK: stg + %x = alloca i32, align 4 + call void @use32(i32* %x) + ret void +} + +define void @NoUnwind() sanitize_memtag nounwind { +entry: +; CHECK-LABEL: NoUnwind: +; CHECK-NOT: .cfi_mte_tagged_frame +; CHECK: stg + %x = alloca i32, align 4 + call void @use32(i32* %x) + ret void +} + +define void @NoUnwindUwTable() sanitize_memtag nounwind uwtable { +entry: +; CHECK-LABEL: NoUnwindUwTable: +; CHECK: .cfi_mte_tagged_frame +; CHECK: stg + %x = alloca i32, align 4 + call void @use32(i32* %x) + ret void +} + +define void @NoMemtag() { +entry: +; CHECK-LABEL: NoMemtag: +; CHECK-NOT: .cfi_mte_tagged_frame +; CHECK-NOT: stg + %x = alloca i32, align 4 + call void @use32(i32* %x) + ret void +} -- 2.7.4