[llvm] Reland [AMDGPU] Serialize WWM_REG vreg flag (#110229) (PR #112492)

[Akshat Oke via llvm-commits]

https://github.com/optimisan updated https://github.com/llvm/llvm-project/pull/112492

>From 1ed2b069364303403c4760dfe61d08c0bb7b17ae Mon Sep 17 00:00:00 2001
From: Akshat Oke <76596238+optimisan at users.noreply.github.com>
Date: Mon, 14 Oct 2024 14:37:21 +0530
Subject: [PATCH 1/2] Reland [AMDGPU] Serialize WWM_REG vreg flag (#110229)

---
 llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp    | 11 +++++++++++
 llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp         | 10 ++++++++++
 llvm/lib/Target/AMDGPU/SIRegisterInfo.h           |  8 ++++++++
 .../MIR/AMDGPU/machine-function-info-no-ir.mir    | 15 +++++++++++++++
 4 files changed, 44 insertions(+)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index e4cc522194f2a9..affb8a2654c1ab 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -1718,6 +1718,17 @@ bool GCNTargetMachine::parseMachineFunctionInfo(
     MFI->reserveWWMRegister(ParsedReg);
   }
 
+  for (const auto &[_, Info] : PFS.VRegInfosNamed) {
+    for (uint8_t Flag : Info->Flags) {
+      MFI->setFlag(Info->VReg, Flag);
+    }
+  }
+  for (const auto &[_, Info] : PFS.VRegInfos) {
+    for (uint8_t Flag : Info->Flags) {
+      MFI->setFlag(Info->VReg, Flag);
+    }
+  }
+
   auto parseAndCheckArgument = [&](const std::optional<yaml::SIArgument> &A,
                                    const TargetRegisterClass &RC,
                                    ArgDescriptor &Arg, unsigned UserSGPRs,
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
index 8de16974a3e79b..1929df12124f94 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -3860,3 +3860,13 @@ SIRegisterInfo::getNumUsedPhysRegs(const MachineRegisterInfo &MRI,
       return getHWRegIndex(Reg) + 1;
   return 0;
 }
+
+SmallVector<StringLiteral>
+SIRegisterInfo::getVRegFlagsOfReg(Register Reg,
+                                  const MachineFunction &MF) const {
+  SmallVector<StringLiteral> RegFlags;
+  const SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
+  if (FuncInfo->checkFlag(Reg, AMDGPU::VirtRegFlag::WWM_REG))
+    RegFlags.push_back("WWM_REG");
+  return RegFlags;
+}
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
index e12a41371c7fd7..8e481e3ac23043 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
@@ -462,6 +462,14 @@ class SIRegisterInfo final : public AMDGPUGenRegisterInfo {
   // Does not go inside function calls.
   unsigned getNumUsedPhysRegs(const MachineRegisterInfo &MRI,
                               const TargetRegisterClass &RC) const;
+
+  std::optional<uint8_t> getVRegFlagValue(StringRef Name) const override {
+    return Name == "WWM_REG" ? AMDGPU::VirtRegFlag::WWM_REG
+                             : std::optional<uint8_t>{};
+  }
+
+  SmallVector<StringLiteral>
+  getVRegFlagsOfReg(Register Reg, const MachineFunction &MF) const override;
 };
 
 namespace AMDGPU {
diff --git a/llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-no-ir.mir b/llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-no-ir.mir
index ebbb89b7816c58..51795a4fea515e 100644
--- a/llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-no-ir.mir
+++ b/llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-no-ir.mir
@@ -578,3 +578,18 @@ body:             |
     SI_RETURN
 
 ...
+---
+name: vregs
+# FULL: registers:
+# FULL-NEXT:   - { id: 0, class: vgpr_32, preferred-register: '$vgpr1', flags: [ WWM_REG ] }
+# FULL-NEXT:   - { id: 1, class: sgpr_64, preferred-register: '$sgpr0_sgpr1', flags: [  ] }
+# FULL-NEXT:   - { id: 2, class: sgpr_64, preferred-register: '', flags: [  ] }
+registers:
+  - { id: 0, class: vgpr_32, preferred-register: $vgpr1, flags: [ WWM_REG ]}
+  - { id: 1, class: sgpr_64, preferred-register: $sgpr0_sgpr1 }
+  - { id: 2, class: sgpr_64, flags: [ ] }
+body: |
+  bb.0:
+    %2:sgpr_64 = COPY %1
+    %1:sgpr_64 = COPY %0
+...

>From feaacf0d98c893c087e1f3e22f5c9ae48c25098e Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Wed, 16 Oct 2024 09:36:59 +0000
Subject: [PATCH 2/2] Update to reflect generic parser changes

---
 llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index affb8a2654c1ab..4d94faf5facf81 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -1719,14 +1719,10 @@ bool GCNTargetMachine::parseMachineFunctionInfo(
   }
 
   for (const auto &[_, Info] : PFS.VRegInfosNamed) {
-    for (uint8_t Flag : Info->Flags) {
-      MFI->setFlag(Info->VReg, Flag);
-    }
+    MFI->setFlag(Info->VReg, Info->Flags);
   }
   for (const auto &[_, Info] : PFS.VRegInfos) {
-    for (uint8_t Flag : Info->Flags) {
-      MFI->setFlag(Info->VReg, Flag);
-    }
+    MFI->setFlag(Info->VReg, Info->Flags);
   }
 
   auto parseAndCheckArgument = [&](const std::optional<yaml::SIArgument> &A,