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

[Akshat Oke via llvm-commits]

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

None

>From cd8df5db290f87393c80a361f18d901ce08f3009 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] 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 23ee0c3e896eb3..16e23879cd735c 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 de9cbe403ab618..20d48aa57adbdf 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -3851,3 +3851,13 @@ SIRegisterInfo::getSubRegAlignmentNumBits(const TargetRegisterClass *RC,
   }
   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 99fa632c0300be..fe0b66f75bbaa2 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
@@ -457,6 +457,14 @@ class SIRegisterInfo final : public AMDGPUGenRegisterInfo {
   // No check if the subreg is supported by the current RC is made.
   unsigned getSubRegAlignmentNumBits(const TargetRegisterClass *RC,
                                      unsigned SubReg) 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
+...