diff --git a/llvm/include/llvm/MC/MCDwarf.h b/llvm/include/llvm/MC/MCDwarf.h --- a/llvm/include/llvm/MC/MCDwarf.h +++ b/llvm/include/llvm/MC/MCDwarf.h @@ -688,6 +688,7 @@ 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 --- a/llvm/include/llvm/MC/MCStreamer.h +++ b/llvm/include/llvm/MC/MCStreamer.h @@ -926,6 +926,7 @@ 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 --- a/llvm/lib/MC/MCAsmStreamer.cpp +++ b/llvm/lib/MC/MCAsmStreamer.cpp @@ -321,6 +321,7 @@ 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; @@ -2027,6 +2028,12 @@ 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 --- a/llvm/lib/MC/MCDwarf.cpp +++ b/llvm/lib/MC/MCDwarf.cpp @@ -1604,6 +1604,8 @@ 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 --- a/llvm/lib/MC/MCParser/AsmParser.cpp +++ b/llvm/lib/MC/MCParser/AsmParser.cpp @@ -540,6 +540,7 @@ DK_PSEUDO_PROBE, DK_LTO_DISCARD, DK_LTO_SET_CONDITIONAL, + DK_CFI_MTE_TAGGED_FRAME, DK_END }; @@ -5560,6 +5561,7 @@ 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 --- a/llvm/lib/MC/MCStreamer.cpp +++ b/llvm/lib/MC/MCStreamer.cpp @@ -252,6 +252,13 @@ 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 --- a/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp +++ b/llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp @@ -1285,6 +1285,17 @@ 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 --- 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 @@ .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 --- a/llvm/lib/Target/AArch64/AArch64InstrInfo.td +++ b/llvm/lib/Target/AArch64/AArch64InstrInfo.td @@ -2434,6 +2434,10 @@ // 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.h b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h --- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h +++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h @@ -175,6 +175,8 @@ /// 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; @@ -403,6 +405,7 @@ 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/AArch64MachineFunctionInfo.cpp b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp --- a/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp +++ b/llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp @@ -89,6 +89,8 @@ 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/AsmParser/AArch64AsmParser.cpp b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp --- a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp +++ b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp @@ -190,6 +190,7 @@ bool parseDirectiveUnreq(SMLoc L); bool parseDirectiveCFINegateRAState(); bool parseDirectiveCFIBKeyFrame(); + bool parseDirectiveCFIMTETaggedFrame(); bool parseDirectiveVariantPCS(SMLoc L); @@ -6038,6 +6039,8 @@ 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") @@ -6523,6 +6526,15 @@ 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 --- /dev/null +++ b/llvm/test/CodeGen/AArch64/stack-tagging-cfi.ll @@ -0,0 +1,33 @@ +; 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 @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 +}