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

[Zachary Yedidia via llvm-commits]

https://github.com/zyedidia created https://github.com/llvm/llvm-project/pull/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.

>From 26d91b1976c84acc288f070c6dfacb07766e25fe Mon Sep 17 00:00:00 2001
From: Zachary Yedidia <zyedidia at gmail.com>
Date: Sat, 3 May 2025 18:29:29 -0700
Subject: [PATCH 1/2] Emit .arch and .arch_extension during assembly

When assembling to text, this patch ensures that .arch and
.arch_extension directives are preserved in the output. This prevents
errors related to unavailable extensions in the assembly output.
---
 llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp    | 7 +++++--
 .../Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp    | 8 ++++++++
 .../Target/AArch64/MCTargetDesc/AArch64TargetStreamer.h   | 3 +++
 3 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
index 28b4cbb5efed8..1a46d68f486bd 100644
--- a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
+++ b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
@@ -7152,9 +7152,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)
@@ -7201,6 +7201,7 @@ bool AArch64AsmParser::parseDirectiveArch(SMLoc L) {
   }
   FeatureBitset Features = ComputeAvailableFeatures(STI.getFeatureBits());
   setAvailableFeatures(Features);
+  getTargetStreamer().emitDirectiveArch(Name);
   return false;
 }
 
@@ -7234,6 +7235,8 @@ bool AArch64AsmParser::parseDirectiveArchExtension(SMLoc L) {
     STI.ClearFeatureBitsTransitively(It->Features);
   FeatureBitset Features = ComputeAvailableFeatures(STI.getFeatureBits());
   setAvailableFeatures(Features);
+
+  getTargetStreamer().emitDirectiveArchExtension(Name);
   return false;
 }
 
diff --git a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp b/llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp
index b12a12436db81..411a9da5e4362 100644
--- a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64ELFStreamer.cpp
+++ b/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";
   }
diff --git a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64TargetStreamer.h b/llvm/lib/Target/AArch64/MCTargetDesc/AArch64TargetStreamer.h
index 43e099f919999..d0f857243fed1 100644
--- a/llvm/lib/Target/AArch64/MCTargetDesc/AArch64TargetStreamer.h
+++ b/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) {}

>From 747bb90e4188a675bdebd04f7d32bf56633700d6 Mon Sep 17 00:00:00 2001
From: Zachary Yedidia <zyedidia at gmail.com>
Date: Sat, 3 May 2025 18:50:03 -0700
Subject: [PATCH 2/2] Add unit test

---
 .../Target/AArch64/AsmParser/AArch64AsmParser.cpp   |  1 +
 llvm/test/MC/AArch64/arch_directive.s               | 13 +++++++++++++
 2 files changed, 14 insertions(+)
 create mode 100644 llvm/test/MC/AArch64/arch_directive.s

diff --git a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
index 1a46d68f486bd..87c7923101e7d 100644
--- a/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
+++ b/llvm/lib/Target/AArch64/AsmParser/AArch64AsmParser.cpp
@@ -7201,6 +7201,7 @@ bool AArch64AsmParser::parseDirectiveArch(SMLoc L) {
   }
   FeatureBitset Features = ComputeAvailableFeatures(STI.getFeatureBits());
   setAvailableFeatures(Features);
+
   getTargetStreamer().emitDirectiveArch(Name);
   return false;
 }
diff --git a/llvm/test/MC/AArch64/arch_directive.s b/llvm/test/MC/AArch64/arch_directive.s
new file mode 100644
index 0000000000000..dd4edf43924b9
--- /dev/null
+++ b/llvm/test/MC/AArch64/arch_directive.s
@@ -0,0 +1,13 @@
+# 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]
+.text
+.arch armv8-a+lse
+cas x0, x1, [x2]
+.arch armv8-a
+.arch_extension lse
+cas x0, x1, [x2]