[AArch64] Emit .arch and .arch_extension during assembly to textual asm (#138433)

This patch ensures that when assembling to text (`-filetype asm`),
`.arch` and `.arch_extension` directives are preserved in the output.
This prevents errors related to unavailable extensions in the assembly
output (see #117221 for an example).

Fixes #117221.

Author: zyedidia

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

@@ -7189,9 +7189,9 @@ static SMLoc incrementLoc(SMLoc L, int Offset) {
 bool AArch64AsmParser::parseDirectiveArch(SMLoc L) {
   SMLoc CurLoc = getLoc();
 
+  StringRef Name = getParser().parseStringToEndOfStatement().trim();
   StringRef Arch, ExtensionString;
-  std::tie(Arch, ExtensionString) =
-      getParser().parseStringToEndOfStatement().trim().split('+');
+  std::tie(Arch, ExtensionString) = Name.split('+');
 
   const AArch64::ArchInfo *ArchInfo = AArch64::parseArch(Arch);
   if (!ArchInfo)
@@ -7238,6 +7238,8 @@ bool AArch64AsmParser::parseDirectiveArch(SMLoc L) {
   }
   FeatureBitset Features = ComputeAvailableFeatures(STI.getFeatureBits());
   setAvailableFeatures(Features);
+
+  getTargetStreamer().emitDirectiveArch(Name);
   return false;
 }
 
@@ -7246,12 +7248,13 @@ bool AArch64AsmParser::parseDirectiveArch(SMLoc L) {
 bool AArch64AsmParser::parseDirectiveArchExtension(SMLoc L) {
   SMLoc ExtLoc = getLoc();
 
-  StringRef Name = getParser().parseStringToEndOfStatement().trim();
+  StringRef FullName = getParser().parseStringToEndOfStatement().trim();
 
   if (parseEOL())
     return true;
 
   bool EnableFeature = true;
+  StringRef Name = FullName;
   if (Name.starts_with_insensitive("no")) {
     EnableFeature = false;
     Name = Name.substr(2);
@@ -7271,6 +7274,8 @@ bool AArch64AsmParser::parseDirectiveArchExtension(SMLoc L) {
     STI.ClearFeatureBitsTransitively(It->Features);
   FeatureBitset Features = ComputeAvailableFeatures(STI.getFeatureBits());
   setAvailableFeatures(Features);
+
+  getTargetStreamer().emitDirectiveArchExtension(FullName);
   return false;
 }
 

llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp

@@ -55,6 +55,14 @@ class AArch64TargetAsmStreamer : public AArch64TargetStreamer {
     OS << "\t.variant_pcs\t" << Symbol->getName() << "\n";
   }
 
+  void emitDirectiveArch(StringRef Name) override {
+    OS << "\t.arch\t" << Name << "\n";
+  }
+
+  void emitDirectiveArchExtension(StringRef Name) override {
+    OS << "\t.arch_extension\t" << Name << "\n";
+  }
+
   void emitARM64WinCFIAllocStack(unsigned Size) override {
     OS << "\t.seh_stackalloc\t" << Size << "\n";
   }

llvm/lib/Target/AArch64/MCTargetDesc/AArch64TargetStreamer.h

@@ -55,6 +55,9 @@ class AArch64TargetStreamer : public MCTargetStreamer {
   /// Callback used to implement the .variant_pcs directive.
   virtual void emitDirectiveVariantPCS(MCSymbol *Symbol) {};
 
+  virtual void emitDirectiveArch(StringRef Name) {};
+  virtual void emitDirectiveArchExtension(StringRef Name) {};
+
   virtual void emitARM64WinCFIAllocStack(unsigned Size) {}
   virtual void emitARM64WinCFISaveR19R20X(int Offset) {}
   virtual void emitARM64WinCFISaveFPLR(int Offset) {}

llvm/test/MC/AArch64/arch_directive.s

@@ -0,0 +1,15 @@
+# RUN: llvm-mc -assemble -triple=aarch64- %s | FileCheck %s
+# CHECK: .text
+# CHECK-NEXT: .arch armv8-a+lse
+# CHECK-NEXT: cas x0, x1, [x2]
+# CHECK-NEXT: .arch armv8-a
+# CHECK-NEXT: .arch_extension lse
+# CHECK-NEXT: cas x0, x1, [x2]
+# CHECK-NEXT: .arch_extension nolse
+.text
+.arch armv8-a+lse
+cas x0, x1, [x2]
+.arch armv8-a
+.arch_extension lse
+cas x0, x1, [x2]
+.arch_extension nolse