[llvm-branch-commits] [llvm] [AMDGPU] Serialize WWM_REG vreg flag (PR #110229)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Akshat Oke (Akshat-Oke)

<details>
<summary>Changes</summary>



---
Full diff: https://github.com/llvm/llvm-project/pull/110229.diff


5 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp (+15) 
- (modified) llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h (+2-2) 
- (modified) llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp (+11) 
- (modified) llvm/lib/Target/AMDGPU/SIRegisterInfo.h (+10) 
- (added) llvm/test/CodeGen/AMDGPU/virtual-registers.mir (+16) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index abd50748f2cc05..c42b443230acc9 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -1628,6 +1628,21 @@ bool GCNTargetMachine::parseMachineFunctionInfo(
     MFI->reserveWWMRegister(ParsedReg);
   }
 
+  auto setRegisterFlags = [&](const VRegInfo &Info) {
+    for (const auto &Flag : Info.Flags) {
+      MFI->setFlag(Info.VReg, Flag);
+    }
+  };
+
+  for (const auto &P : PFS.VRegInfosNamed) {
+    const VRegInfo &Info = *P.second;
+    setRegisterFlags(Info);
+  }
+  for (const auto &P : PFS.VRegInfos) {
+    const VRegInfo &Info = *P.second;
+    setRegisterFlags(Info);
+  }
+
   auto parseAndCheckArgument = [&](const std::optional<yaml::SIArgument> &A,
                                    const TargetRegisterClass &RC,
                                    ArgDescriptor &Arg, unsigned UserSGPRs,
diff --git a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
index aff0b34947d688..7fffc706c4b027 100644
--- a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
@@ -684,8 +684,8 @@ class SIMachineFunctionInfo final : public AMDGPUMachineFunction,
 
   void setFlag(Register Reg, uint8_t Flag) {
     assert(Reg.isVirtual());
-    if (VRegFlags.inBounds(Reg))
-      VRegFlags[Reg] |= Flag;
+    VRegFlags.grow(Reg);
+    VRegFlags[Reg] |= Flag;
   }
 
   bool checkFlag(Register Reg, uint8_t Flag) const {
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
index 2d1cd1bda3afe1..ad3e592127de80 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -3614,3 +3614,14 @@ SIRegisterInfo::getSubRegAlignmentNumBits(const TargetRegisterClass *RC,
   }
   return 0;
 }
+
+SmallVector<std::string>
+SIRegisterInfo::getVRegFlagsOfReg(Register Reg,
+                                  const MachineFunction &MF) const {
+  SmallVector<std::string> 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 88d5686720985e..bce5a2aa792bd3 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
@@ -449,6 +449,16 @@ 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::pair<bool, uint8_t> getVRegFlagValue(StringRef Name) const override {
+    if (Name == "WWM_REG") {
+      return {true, AMDGPU::VirtRegFlag::WWM_REG};
+    }
+    return {false, 0};
+  }
+
+  SmallVector<std::string>
+  getVRegFlagsOfReg(Register Reg, const MachineFunction &MF) const override;
 };
 
 namespace AMDGPU {
diff --git a/llvm/test/CodeGen/AMDGPU/virtual-registers.mir b/llvm/test/CodeGen/AMDGPU/virtual-registers.mir
new file mode 100644
index 00000000000000..3ea8f6eafcf10c
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/virtual-registers.mir
@@ -0,0 +1,16 @@
+# RUN: llc -mtriple=amdgcn -run-pass=none -o - %s | FileCheck %s
+# This test ensures that the MIR parser parses virtual register flags correctly
+
+---
+name: vregs
+# CHECK: registers:
+# CHECK-NEXT:   - { id: 0, class: vgpr_32, preferred-register: '$vgpr1', flags: [ WWM_REG ] }
+# CHECK-NEXT:   - { id: 1, class: sgpr_64, preferred-register: '$sgpr0_sgpr1', flags: [  ] }
+# CHECK-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 }
+body: |
+  bb.0:
+    %2:sgpr_64 = COPY %1
+    %1:sgpr_64 = COPY %0

``````````

</details>


https://github.com/llvm/llvm-project/pull/110229