[llvm] bec839d - [AMDGPU] Serialize WWM_REG vreg flag (#110229)

[via llvm-commits]

Author: Akshat Oke
Date: 2024-10-14T14:37:21+05:30
New Revision: bec839d8eed9dd13fa7eaffd50b28f8f913de2e2

URL: https://github.com/llvm/llvm-project/commit/bec839d8eed9dd13fa7eaffd50b28f8f913de2e2
DIFF: https://github.com/llvm/llvm-project/commit/bec839d8eed9dd13fa7eaffd50b28f8f913de2e2.diff

LOG: [AMDGPU] Serialize WWM_REG vreg flag (#110229)

Added: 
    

Modified: 
    llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
    llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
    llvm/lib/Target/AMDGPU/SIRegisterInfo.h
    llvm/test/CodeGen/MIR/AMDGPU/machine-function-info-no-ir.mir

Removed: 
    


################################################################################
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
+...