[MC] Add 'G' to augmentation string for MTE instrumented functions

[fmayer]

Summary

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.

[fmayer]

naming

[fmayer]

respect nounwind

[fmayer]

comment

[fmayer]

comment

[fmayer]

test

[MaskRay]

The lld/ELF change should be in a separate patch.
It's the convention and I am usually picky on how tests should be made.

[MaskRay]

This was agreed on in https://lists.llvm.org/pipermail/llvm-dev/2020-May/141345.html.

This agreement is insufficient for the broad Linux community.

Linux Standard Base Core Specification, Generic Part defines augmentation string characters. Essentially this needs some GCC agreement, otherwise there could be potential future conflict.

[fmayer]

In D127007#3557500, @MaskRay wrote:

This was agreed on in https://lists.llvm.org/pipermail/llvm-dev/2020-May/141345.html.

This agreement is insufficient for the broad Linux community.

Linux Standard Base Core Specification, Generic Part defines augmentation string characters. Essentially this needs some GCC agreement, otherwise there could be potential future conflict.

Thanks! I was wondering about this. I will get broader agreement prior to moving forward with this.

[fmayer]

Reached out to GCC here: https://gcc.gnu.org/pipermail/gcc/2022-June/238813.html

[fmayer]

split off lld / dwarf parser change.

[MaskRay]

Thanks for starting the GCC thread! I think the implementation is reasonable.

[llvm] Add 'G' to augmentation string for MTE instrumented functions.

[MC] may be more specific. The convention is to avoid trailing dot in the subject line.

This was agreed on in https://lists.llvm.org/pipermail/llvm-dev/2020-May/141345.html.

Avoid trailing . in an URI. Many tools consider . part of the URI.

Try summarizing the main point of the thread and place it in the commit message (summary).

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

1294

There are too many blank lines.

Use

if (getFunctionCFISectionType(*MF) != CFISection::None)
  OutStreamer->emitCFIMTETaggedFrame();

llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

6532

if (parseEOL())
  return true;

I've updated other unexpected *** in ... directive to use the prevailing style expected newline.

llvm/test/CodeGen/AArch64/stack-tagging-cfi.ll

3 ↗

(On Diff #434211)

i32* => ptr

7 ↗

(On Diff #434211)

[fmayer]

In D127007#3559248, @MaskRay wrote:

Thanks for starting the GCC thread! I think the implementation is reasonable.

Turns out this is already part of ARM's ABI:
https://github.com/ARM-software/abi-aa/blob/main/aadwarf64/aadwarf64.rst#id22

[fmayer]

address comments

[MaskRay]

Looks great! Hope that @eugenis can check the test interop with msan.

[fmayer]

format

[eugenis]

Should we have a test for the augmentation string, through llvm-readelf?

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

1289

As we've seen recently, sometimes even nounwind frames get unwound. Is it possible to relax some of these conditions to emit the "G" in as many cases as possible? This is the situation where MTE is more demanding to the unwind correctness, so I'd rather err on the safe side.

[fmayer]

In D127007#3561986, @eugenis wrote:

Should we have a test for the augmentation string, through llvm-readelf?

We have that test the follow-up CL that does the change to lld (as requested to be split by MaskRay).

[fmayer]

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

1289

I'm not sure. This is exactly the same logic as for the B-Key, so either we change both or neither. B-Key is just as important to correctness as this is.

What do you mean "nounwind frames get unwound"? AFAICT they do get skipped over, so the 'G' wouldn't help.

[eugenis]

In D127007#3562002, @fmayer wrote:

In D127007#3561986, @eugenis wrote:

Should we have a test for the augmentation string, through llvm-readelf?

We have that test the follow-up CL that does the change to lld (as requested to be split by MaskRay).

Hmm ok. This tests does not really need to involve lld at all, but that's fine.

[fmayer]

In D127007#3562017, @eugenis wrote:

In D127007#3562002, @fmayer wrote:

In D127007#3561986, @eugenis wrote:

Should we have a test for the augmentation string, through llvm-readelf?

We have that test the follow-up CL that does the change to lld (as requested to be split by MaskRay).

Hmm ok. This tests does not really need to involve lld at all, but that's fine.

It does, because we need LLD so we can get an ELF file (which we then can readelf or, as the test does objdump).

[eugenis]

You can llvm-objdump --dwarf=frames a .o file.

[fmayer]

In D127007#3562022, @fmayer wrote:

In D127007#3562017, @eugenis wrote:

In D127007#3562002, @fmayer wrote:

In D127007#3561986, @eugenis wrote:

Should we have a test for the augmentation string, through llvm-readelf?

We have that test the follow-up CL that does the change to lld (as requested to be split by MaskRay).

Hmm ok. This tests does not really need to involve lld at all, but that's fine.

It does, because we need LLD so we can get an ELF file (which we then can readelf or, as the test does objdump).

Good point. The change to the DWARF parser for that is also in the follow up change though, but I changed the test to both check the .o and the .so files.

[fmayer]

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

1289

I think the only thing we can do is force unwind tables for all frames that have stack instrumentation, but that would be quite a big hammer.

I'm not sure that is the right thing to do, where the real solution would be for people to fix their builds to not mix exception and noexception code in this way. Given that this mode is (at least for now) intended for testing, to me that seems like the better thing for people to do, as all sorts of stuff can go wrong when throwing exceptions through no-exception code (PAC as mentioned above, anything that needs RAII, &c).

[MaskRay]

It may be better to move the DebugInfo change from the lld patch into this patch.

[fmayer]

In D127007#3562088, @MaskRay wrote:

It may be better to move the DebugInfo change from the lld patch into this patch.

I split it into D127171. Can also merge into here, no strong opinions

[fmayer]

add another test

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

1289

OK I think we have resolved this confusion. Compiling with -fno-exceptions does still generate IR with uwtables but also with`nounwind`. In this case the uwtable does still get generated, but not the cleanup actions. Added a test to that effect.

[eugenis]

LGTM

lld/ELF/EhFrame.cpp

for (char c : aug) {
     if (c == 'R')
       return readByte();
     if (c == 'z')
       skipLeb128();
     else if (c == 'L')
       readByte();
     else if (c == 'P')
       skipAugP();
-    else if (c != 'B' && c != 'S')
+    else if (c != 'B' && c != 'S' && c != 'G')
       failOn(aug.data(), "unknown .eh_frame augmentation string: " + aug);
   }
   return DW_EH_PE_absptr;
 }
 
 bool EhReader::hasLSDA() {
   StringRef aug = getAugmentation();
   for (char c : aug) {
     if (c == 'L')
       return true;
     if (c == 'z')
       skipLeb128();
     else if (c == 'P')
       skipAugP();
     else if (c == 'R')
       readByte();
-    else if (c != 'B' && c != 'S')
+    else if (c != 'B' && c != 'S' && c != 'G')
       failOn(aug.data(), "unknown .eh_frame augmentation string: " + aug);
   }
   return false;
 }

llvm/include/llvm/MC/MCDwarf.h

struct MCDwarfFrameInfo {
   unsigned CurrentCfaRegister = 0;
   unsigned PersonalityEncoding = 0;
   unsigned LsdaEncoding = 0;
   uint32_t CompactUnwindEncoding = 0;
   bool IsSignalFrame = false;
   bool IsSimple = false;
   unsigned RAReg = static_cast<unsigned>(INT_MAX);
   bool IsBKeyFrame = false;
+  bool IsMTETaggedFrame = false;
 };
 
 class MCDwarfFrameEmitter {
 public:
   //
   // This emits the frame info section.
   //
   static void Emit(MCObjectStreamer &streamer, MCAsmBackend *MAB, bool isEH);
   static void EmitAdvanceLoc(MCObjectStreamer &Streamer, uint64_t AddrDelta);
   static void EncodeAdvanceLoc(MCContext &Context, uint64_t AddrDelta,
                                raw_ostream &OS);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_MC_MCDWARF_H

llvm/include/llvm/MC/MCStreamer.h

public:
 
   /// Specify the "root" file of the compilation, using the ".file 0" extension.
   virtual void emitDwarfFile0Directive(StringRef Directory, StringRef Filename,
                                        Optional<MD5::MD5Result> Checksum,
                                        Optional<StringRef> Source,
                                        unsigned CUID = 0);
 
   virtual void emitCFIBKeyFrame();
+  virtual void emitCFIMTETaggedFrame();
 
   /// This implements the DWARF2 '.loc fileno lineno ...' assembler
   /// directive.
   virtual void emitDwarfLocDirective(unsigned FileNo, unsigned Line,
                                      unsigned Column, unsigned Flags,
                                      unsigned Isa, unsigned Discriminator,
                                      StringRef FileName);
 

llvm/lib/DebugInfo/DWARF/DWARFDebugFrame.cpp

if (Id == getCIEId(IsDWARF64, IsEH)) {
             AugmentationLength = Data.getULEB128(&Offset);
             StartAugmentationOffset = Offset;
             EndAugmentationOffset = Offset + *AugmentationLength;
             break;
           case 'B':
             // B-Key is used for signing functions associated with this
             // augmentation string
             break;
+          case 'G':
+            break;
           }
         }
 
         if (AugmentationLength.hasValue()) {
           if (Offset != EndAugmentationOffset)
             return createStringError(errc::invalid_argument,
                                      "parsing augmentation data at 0x%" PRIx64
                                      " failed",

llvm/lib/MC/MCAsmStreamer.cpp

public:
 
   void emitCVStringTableDirective() override;
   void emitCVFileChecksumsDirective() override;
   void emitCVFileChecksumOffsetDirective(unsigned FileNo) override;
   void emitCVFPOData(const MCSymbol *ProcSym, SMLoc L) override;
 
   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;
   void emitCFIDefCfaRegister(int64_t Register) override;
   void emitCFILLVMDefAspaceCfa(int64_t Register, int64_t Offset,
                                int64_t AddressSpace) override;
   void emitCFIOffset(int64_t Register, int64_t Offset) override;
   void emitCFIPersonality(const MCSymbol *Sym, unsigned Encoding) override;

 }
 
 void MCAsmStreamer::emitCFIBKeyFrame() {
   MCStreamer::emitCFIBKeyFrame();
   OS << "\t.cfi_b_key_frame";
   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);
 
   OS << ".seh_proc ";
   Symbol->print(OS, MAI);
   EmitEOL();
 }
 

llvm/lib/MC/MCDwarf.cpp

if (Frame.Personality)
       Augmentation += "P";
     if (Frame.Lsda)
       Augmentation += "L";
     Augmentation += "R";
     if (Frame.IsSignalFrame)
       Augmentation += "S";
     if (Frame.IsBKeyFrame)
       Augmentation += "B";
+    if (Frame.IsMTETaggedFrame)
+      Augmentation += "G";
     Streamer.emitBytes(Augmentation);
   }
   Streamer.emitInt8(0);
 
   if (CIEVersion >= 4) {
     // Address Size
     Streamer.emitInt8(context.getAsmInfo()->getCodePointerSize());
 

llvm/lib/MC/MCParser/AsmParser.cpp

enum DirectiveKind {
     DK_ERROR,
     DK_WARNING,
     DK_PRINT,
     DK_ADDRSIG,
     DK_ADDRSIG_SYM,
     DK_PSEUDO_PROBE,
     DK_LTO_DISCARD,
     DK_LTO_SET_CONDITIONAL,
+    DK_CFI_MTE_TAGGED_FRAME,
     DK_END
   };
 
   /// Maps directive name --> DirectiveKind enum, for
   /// directives parsed by this class.
   StringMap<DirectiveKind> DirectiveKindMap;
 
   // Codeview def_range type parsing.
void AsmParser::initializeDirectiveKindMap() {
   DirectiveKindMap[".cfi_restore"] = DK_CFI_RESTORE;
   DirectiveKindMap[".cfi_escape"] = DK_CFI_ESCAPE;
   DirectiveKindMap[".cfi_return_column"] = DK_CFI_RETURN_COLUMN;
   DirectiveKindMap[".cfi_signal_frame"] = DK_CFI_SIGNAL_FRAME;
   DirectiveKindMap[".cfi_undefined"] = DK_CFI_UNDEFINED;
   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;
   DirectiveKindMap[".exitm"] = DK_EXITM;
   DirectiveKindMap[".endm"] = DK_ENDM;
   DirectiveKindMap[".endmacro"] = DK_ENDMACRO;
   DirectiveKindMap[".purgem"] = DK_PURGEM;
   DirectiveKindMap[".err"] = DK_ERR;

llvm/lib/MC/MCStreamer.cpp

 
 void MCStreamer::emitCFIBKeyFrame() {
   MCDwarfFrameInfo *CurFrame = getCurrentDwarfFrameInfo();
   if (!CurFrame)
     return;
   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,
                                        StringRef FileName) {
   getContext().setCurrentDwarfLoc(FileNo, Line, Column, Flags, Isa,
                                   Discriminator);
 }
 

llvm/lib/Target/AArch64/AArch64AsmPrinter.cpp

case AArch64::EMITBKEY: {
 
       if (getFunctionCFISectionType(*MF) == CFISection::None)
         return;
 
       OutStreamer->emitCFIBKeyFrame();
       return;
   }
 
+  case AArch64::EMITMTETAGGED: {
+      ExceptionHandling ExceptionHandlingType = MAI->getExceptionHandlingType();
+      if (ExceptionHandlingType != ExceptionHandling::DwarfCFI &&
+          ExceptionHandlingType != ExceptionHandling::ARM)
+        return;
+
+      if (getFunctionCFISectionType(*MF) == CFISection::None)
+        return;
+
+      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.
   case AArch64::TCRETURNri:
   case AArch64::TCRETURNriBTI:
   case AArch64::TCRETURNriALL: {
     MCInst TmpInst;
     TmpInst.setOpcode(AArch64::BR);

llvm/lib/Target/AArch64/AArch64FrameLowering.cpp

 
 #include "AArch64FrameLowering.h"
 #include "AArch64InstrInfo.h"
 #include "AArch64MachineFunctionInfo.h"
 #include "AArch64RegisterInfo.h"
 #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"
 #include "llvm/CodeGen/LivePhysRegs.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineInstr.h"
if (MFnI.shouldSignReturnAddress()) {
     if (EmitCFI) {
       unsigned CFIIndex =
           MF.addFrameInst(MCCFIInstruction::createNegateRAState(nullptr));
       BuildMI(MBB, MBBI, DL, TII->get(TargetOpcode::CFI_INSTRUCTION))
           .addCFIIndex(CFIIndex)
           .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
   // ORR is sufficient, it is assumed a Swift kernel would initialize the TBI
   // bits so that is still true.
   if (HasFP && AFI->hasSwiftAsyncContext()) {
     switch (MF.getTarget().Options.SwiftAsyncFramePointer) {
     case SwiftAsyncFramePointerMode::DeploymentBased:

llvm/lib/Target/AArch64/AArch64InstrInfo.td

 def TLSDESCCALL : Pseudo<(outs), (ins i64imm:$sym), []>, Sched<[]> {
   let AsmString = ".tlsdesccall $sym";
 }
 
 // Pseudo instruction to tell the streamer to emit a 'B' character into the
 // 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
 let isCall = 1, Defs = [LR, X0, X1], hasSideEffects = 1, Size = 16,
     isCodeGenOnly = 1 in
 def TLSDESC_CALLSEQ
     : Pseudo<(outs), (ins i64imm:$sym),
              [(AArch64tlsdesc_callseq tglobaltlsaddr:$sym)]>,

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.h

class AArch64FunctionInfo final : public MachineFunctionInfo {
   /// Whether this function has an extended frame record [Ctx, FP, LR]. If so,
   /// bit 60 of the in-memory FP will be 1 to enable other tools to detect the
   /// extended record.
   bool HasSwiftAsyncContext = false;
 
   /// The stack slot where the Swift asynchronous context is stored.
   int SwiftAsyncContextFrameIdx = std::numeric_limits<int>::max();
 
+  bool IsMTETagged = false;
+
   /// True if the function need unwind information.
   mutable Optional<bool> NeedsDwarfUnwindInfo;
 
   /// True if the function need asynchronous unwind information.
   mutable Optional<bool> NeedsAsyncDwarfUnwindInfo;
 
 public:
   explicit AArch64FunctionInfo(MachineFunction &MF);
#endif
   void setCalleeSaveBaseToFrameRecordOffset(int Offset) {
     CalleeSaveBaseToFrameRecordOffset = Offset;
   }
 
   bool shouldSignReturnAddress() const;
   bool shouldSignReturnAddress(bool SpillsLR) const;
 
   bool shouldSignWithBKey() const { return SignWithBKey; }
+  bool isMTETagged() const { return IsMTETagged; }
 
   bool branchTargetEnforcement() const { return BranchTargetEnforcement; }
 
   void setHasSwiftAsyncContext(bool HasContext) {
     HasSwiftAsyncContext = HasContext;
   }
   bool hasSwiftAsyncContext() const { return HasSwiftAsyncContext; }
 

llvm/lib/Target/AArch64/AArch64MachineFunctionInfo.cpp

AArch64FunctionInfo::AArch64FunctionInfo(MachineFunction &MF) : MF(MF) {
   // If we already know that the function doesn't have a redzone, set
   // HasRedZone here.
   if (MF.getFunction().hasFnAttribute(Attribute::NoRedZone))
     HasRedZone = false;
 
   const Function &F = MF.getFunction();
   std::tie(SignReturnAddress, SignReturnAddressAll) = GetSignReturnAddress(F);
   SignWithBKey = ShouldSignWithBKey(F);
+  // TODO: should we 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<ConstantInt>(
             F.getParent()->getModuleFlag("branch-target-enforcement")))
       BranchTargetEnforcement = BTE->getZExtValue();
     return;
   }
 

llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp

private:
 
   bool parseDirectiveLOH(StringRef LOH, SMLoc L);
   bool parseDirectiveLtorg(SMLoc L);
 
   bool parseDirectiveReq(StringRef Name, SMLoc L);
   bool parseDirectiveUnreq(SMLoc L);
   bool parseDirectiveCFINegateRAState();
   bool parseDirectiveCFIBKeyFrame();
+  bool parseDirectiveCFIMTETaggedFrame();
 
   bool parseDirectiveVariantPCS(SMLoc L);
 
   bool parseDirectiveSEHAllocStack(SMLoc L);
   bool parseDirectiveSEHPrologEnd(SMLoc L);
   bool parseDirectiveSEHSaveR19R20X(SMLoc L);
   bool parseDirectiveSEHSaveFPLR(SMLoc L);
   bool parseDirectiveSEHSaveFPLRX(SMLoc L);
bool AArch64AsmParser::ParseDirective(AsmToken DirectiveID) {
   else if (IDVal == ".unreq")
     parseDirectiveUnreq(Loc);
   else if (IDVal == ".inst")
     parseDirectiveInst(Loc);
   else if (IDVal == ".cfi_negate_ra_state")
     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")
     parseDirectiveVariantPCS(Loc);
   else if (IsMachO) {
     if (IDVal == MCLOHDirectiveName())
       parseDirectiveLOH(IDVal, Loc);
     else

 bool AArch64AsmParser::parseDirectiveCFIBKeyFrame() {
   if (parseToken(AsmToken::EndOfStatement,
                  "unexpected token in '.cfi_b_key_frame'"))
     return true;
   getStreamer().emitCFIBKeyFrame();
   return false;
 }
 
+/// parseDirectiveCFIMTETaggedFrame
+/// ::= .cfi_mte_tagged_frame
+bool AArch64AsmParser::parseDirectiveCFIMTETaggedFrame() {
+  if (parseToken(AsmToken::EndOfStatement,
+                 "unexpected token in '.cfi_mte_tagged_frame'"))
+    return true;
+  getStreamer().emitCFIMTETaggedFrame();
+  return false;
+}
+
 /// parseDirectiveVariantPCS
 /// ::= .variant_pcs symbolname
 bool AArch64AsmParser::parseDirectiveVariantPCS(SMLoc L) {
   StringRef Name;
   if (getParser().parseIdentifier(Name))
     return TokError("expected symbol name");
   if (parseEOL())
     return true;

llvm/test/CodeGen/AArch64/stack-tagging-asm.ll

+; RUN: llc < %s -mtriple=aarch64 -mattr=+mte > %t.s
+; RUN: cat %t.s | FileCheck --check-prefix=ASM %s
+; RUN: llc -filetype=obj < %s -mtriple=aarch64 -mattr=+mte > %t.o
+; RUN: clang --target=aarch64-linux+mte -nostdlib -fuse-ld=lld -shared %t.o -o %t.so
+; RUN: llvm-readelf -u %t.so | FileCheck --check-prefix=READELF %s
+
+declare void @use32(i32*)
+
+define void @WithUnwind() sanitize_memtag {
+entry:
+; ASM-LABEL: WithUnwind
+; ASM: .cfi_mte_tagged_frame
+; ASM: stg
+; READELF: augmentation: zRG
+  %x = alloca i32, align 4
+  call void @use32(i32* %x)
+  ret void
+}
+
+define void @NoUnwind() sanitize_memtag nounwind {
+entry:
+; ASM-LABEL: NoUnwind
+; ASM-NOT: .cfi_mte_tagged_frame
+; ASM: stg
+  %x = alloca i32, align 4
+  call void @use32(i32* %x)
+  ret void
+}