[llvm] [SystemZ] Implement .machine (push|pop) directives (PR #137302)

[Dominik Steenken via llvm-commits]

https://github.com/dominik-steenken updated https://github.com/llvm/llvm-project/pull/137302

>From 35949c374eb2f2babf59c20e0c9394e4bb142e3e Mon Sep 17 00:00:00 2001
From: Dominik Steenken <dost at de.ibm.com>
Date: Thu, 24 Apr 2025 11:38:22 +0200
Subject: [PATCH] [SystemZ] Implement .machine (push|pop) directives

The `.machine push` and `.machine pop` directives were missing from
the SystemZ Backend Asm Parser. This commit adds them, and expands the
corresponding test to test proper operation.
---
 .../SystemZ/AsmParser/SystemZAsmParser.cpp    | 55 +++++++++++++++----
 .../MC/SystemZ/machine-directive-invalid.s    | 10 +++-
 llvm/test/MC/SystemZ/machine-directive.s      |  6 +-
 3 files changed, 56 insertions(+), 15 deletions(-)

diff --git a/llvm/lib/Target/SystemZ/AsmParser/SystemZAsmParser.cpp b/llvm/lib/Target/SystemZ/AsmParser/SystemZAsmParser.cpp
index 6d9a7a73f72db..bf8be064cbb4e 100644
--- a/llvm/lib/Target/SystemZ/AsmParser/SystemZAsmParser.cpp
+++ b/llvm/lib/Target/SystemZ/AsmParser/SystemZAsmParser.cpp
@@ -33,6 +33,7 @@
 #include "llvm/Support/Casting.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/SMLoc.h"
+#include "llvm/TargetParser/SubtargetFeature.h"
 #include <algorithm>
 #include <cassert>
 #include <cstddef>
@@ -410,6 +411,18 @@ class SystemZAsmParser : public MCTargetAsmParser {
 
 private:
   MCAsmParser &Parser;
+
+  // A vector to contain the stack of machine specs created by `.machine push`.
+  SmallVector<StringRef> MachineStack;
+  // Specifies the current CPU. It is initialized to "default", which represents
+  // the feature bit set in use before any `.machine` directive is issued.
+  StringRef CurrentCPU = "default";
+  // The feature bits associated with the "default" cpu. Derived from the
+  // SubtargetInfo handed to the constructor of the AsmParser. This is
+  // used to recreate the feature bits if the "default" cpu is saved by a
+  // `.machine push` and then restored by a `.machine pop`.
+  const FeatureBitset &DefaultBits;
+
   enum RegisterGroup {
     RegGR,
     RegFP,
@@ -494,16 +507,16 @@ class SystemZAsmParser : public MCTargetAsmParser {
 
 public:
   SystemZAsmParser(const MCSubtargetInfo &sti, MCAsmParser &parser,
-                   const MCInstrInfo &MII,
-                   const MCTargetOptions &Options)
-    : MCTargetAsmParser(Options, sti, MII), Parser(parser) {
+                   const MCInstrInfo &MII, const MCTargetOptions &Options)
+      : MCTargetAsmParser(Options, sti, MII), Parser(parser),
+        DefaultBits(sti.getFeatureBits()) {
     MCAsmParserExtension::Initialize(Parser);
 
     // Alias the .word directive to .short.
     parser.addAliasForDirective(".word", ".short");
 
     // Initialize the set of available features.
-    setAvailableFeatures(ComputeAvailableFeatures(getSTI().getFeatureBits()));
+    setAvailableFeatures(ComputeAvailableFeatures(DefaultBits));
   }
 
   // Override MCTargetAsmParser.
@@ -1382,17 +1395,39 @@ bool SystemZAsmParser::parseDirectiveMachine(SMLoc L) {
       Parser.getTok().isNot(AsmToken::String))
     return TokError("unexpected token in '.machine' directive");
 
-  StringRef CPU = Parser.getTok().getIdentifier();
+  StringRef Id = Parser.getTok().getIdentifier();
+
+  // Parse push and pop directives first
+  StringRef CPU = "";
+  if (Id == "push") {
+    if (CurrentCPU.empty())
+      llvm_unreachable(".machine push without a cpu to save");
+    MachineStack.push_back(CurrentCPU);
+  } else if (Id == "pop") {
+    if (MachineStack.empty())
+      return TokError("pop without corresponding push in '.machine' directive");
+    CPU = MachineStack.back();
+    MachineStack.pop_back();
+  } else {
+    CPU = Id;
+  }
+
   Parser.Lex();
   if (parseEOL())
     return true;
 
-  MCSubtargetInfo &STI = copySTI();
-  STI.setDefaultFeatures(CPU, /*TuneCPU*/ CPU, "");
-  setAvailableFeatures(ComputeAvailableFeatures(STI.getFeatureBits()));
-
-  getTargetStreamer().emitMachine(CPU);
+  // check if we need to change cpu
+  if (CPU == "default") {
+    setAvailableFeatures(DefaultBits);
+    CurrentCPU = "default";
+  } else if (!CPU.empty()) {
+    MCSubtargetInfo &STI = copySTI();
+    STI.setDefaultFeatures(CPU, /*TuneCPU*/ CPU, "");
+    setAvailableFeatures(ComputeAvailableFeatures(STI.getFeatureBits()));
+    CurrentCPU = CPU;
 
+    getTargetStreamer().emitMachine(CPU);
+  }
   return false;
 }
 
diff --git a/llvm/test/MC/SystemZ/machine-directive-invalid.s b/llvm/test/MC/SystemZ/machine-directive-invalid.s
index 8b3147ecc62d4..bd4889912a5df 100644
--- a/llvm/test/MC/SystemZ/machine-directive-invalid.s
+++ b/llvm/test/MC/SystemZ/machine-directive-invalid.s
@@ -1,10 +1,14 @@
+// NOTE: Assertions have been autogenerated by utils/update_mc_test_checks.py UTC_ARGS: --version 5
 # RUN: not llvm-mc -triple=s390x %s 2>&1 | FileCheck %s
 
-# CHECK: [[#@LINE+1]]:9: error: unexpected token in '.machine' directive
 .machine
+// CHECK: :[[@LINE-1]]:9: error: unexpected token in '.machine' directive
+
+.machine pop
+// CHECK: :[[@LINE-1]]:10: error: pop without corresponding push in '.machine' directive
 
-# CHECK: [[#@LINE+1]]:10: error: unexpected token in '.machine' directive
 .machine 42
+// CHECK: :[[@LINE-1]]:10: error: unexpected token in '.machine' directive
 
-# CHECK: [[#@LINE+1]]:13: error: expected newline
 .machine z13+
+// CHECK: :[[@LINE-1]]:13: error: expected newline
diff --git a/llvm/test/MC/SystemZ/machine-directive.s b/llvm/test/MC/SystemZ/machine-directive.s
index aa71c6aea5572..8f4cb3f49678b 100644
--- a/llvm/test/MC/SystemZ/machine-directive.s
+++ b/llvm/test/MC/SystemZ/machine-directive.s
@@ -11,10 +11,12 @@
 # CHECK: .machine z13
 # CHECK: vgbm	%v0, 3
 
+.machine push
 .machine z13
+.machine push
 vgbm    %v0, 0
 .machine zEC12
 vgbm    %v0, 1
-.machine z13
+.machine pop
 vgbm    %v0, 3
-
+.machine pop