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

Akshat Oke via llvm-branch-commits llvm-branch-commits at lists.llvm.org
Mon Oct 7 23:31:53 PDT 2024


https://github.com/Akshat-Oke updated https://github.com/llvm/llvm-project/pull/110229

>From 1cbc26fe2de38ae4e174aec128b39c899dab9136 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 27 Sep 2024 08:58:39 +0000
Subject: [PATCH 1/4] [AMDGPU] Serialize WWM_REG vreg flag

---
 llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp | 15 +++++++++++++++
 llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h |  4 ++--
 llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp      | 11 +++++++++++
 llvm/lib/Target/AMDGPU/SIRegisterInfo.h        | 10 ++++++++++
 llvm/test/CodeGen/AMDGPU/virtual-registers.mir | 16 ++++++++++++++++
 5 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/virtual-registers.mir

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index 1f2148c2922de9..28578a875c164c 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -1712,6 +1712,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 669f98dd865d61..e28c24bf8f8500 100644
--- a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
@@ -693,8 +693,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 9e1c4941dba283..84569b3f11df67 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -3839,3 +3839,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 409e5418abc8ec..2c3707e119178a 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
@@ -454,6 +454,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

>From 4c26b4f3b1ad8952767625ba949eaa750aec0652 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 4 Oct 2024 06:31:06 +0000
Subject: [PATCH 2/4] Correct TRI methods to optional<> and SmallString

---
 llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp | 6 +++---
 llvm/lib/Target/AMDGPU/SIRegisterInfo.h   | 8 ++++----
 2 files changed, 7 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
index 84569b3f11df67..af1b736714b738 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.cpp
@@ -3840,13 +3840,13 @@ SIRegisterInfo::getSubRegAlignmentNumBits(const TargetRegisterClass *RC,
   return 0;
 }
 
-SmallVector<std::string>
+SmallVector<SmallString<8>>
 SIRegisterInfo::getVRegFlagsOfReg(Register Reg,
                                   const MachineFunction &MF) const {
-  SmallVector<std::string> RegFlags;
+  SmallVector<SmallString<8>> RegFlags;
   const SIMachineFunctionInfo *FuncInfo = MF.getInfo<SIMachineFunctionInfo>();
   if (FuncInfo->checkFlag(Reg, AMDGPU::VirtRegFlag::WWM_REG)) {
-    RegFlags.push_back("WWM_REG");
+    RegFlags.push_back(SmallString<8>("WWM_REG"));
   }
   return RegFlags;
 }
diff --git a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
index 2c3707e119178a..6ea60dc53685bb 100644
--- a/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIRegisterInfo.h
@@ -455,14 +455,14 @@ class SIRegisterInfo final : public AMDGPUGenRegisterInfo {
   unsigned getSubRegAlignmentNumBits(const TargetRegisterClass *RC,
                                      unsigned SubReg) const;
 
-  std::pair<bool, uint8_t> getVRegFlagValue(StringRef Name) const override {
+  std::optional<uint8_t> getVRegFlagValue(StringRef Name) const override {
     if (Name == "WWM_REG") {
-      return {true, AMDGPU::VirtRegFlag::WWM_REG};
+      return AMDGPU::VirtRegFlag::WWM_REG;
     }
-    return {false, 0};
+    return {};
   }
 
-  SmallVector<std::string>
+  SmallVector<SmallString<8>>
   getVRegFlagsOfReg(Register Reg, const MachineFunction &MF) const override;
 };
 

>From dea012aacc9f771e2ada19c910d6f40d3ff8a03e Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Fri, 4 Oct 2024 06:31:42 +0000
Subject: [PATCH 3/4] Move test to MIR/AMDGPU/mfi-no-ir

---
 llvm/test/CodeGen/AMDGPU/virtual-registers.mir   | 16 ----------------
 .../MIR/AMDGPU/machine-function-info-no-ir.mir   | 15 +++++++++++++++
 2 files changed, 15 insertions(+), 16 deletions(-)
 delete mode 100644 llvm/test/CodeGen/AMDGPU/virtual-registers.mir

diff --git a/llvm/test/CodeGen/AMDGPU/virtual-registers.mir b/llvm/test/CodeGen/AMDGPU/virtual-registers.mir
deleted file mode 100644
index 3ea8f6eafcf10c..00000000000000
--- a/llvm/test/CodeGen/AMDGPU/virtual-registers.mir
+++ /dev/null
@@ -1,16 +0,0 @@
-# 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
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 6651bfaeb6ed2a579b8d140b1664f011625f8819 Mon Sep 17 00:00:00 2001
From: Akshat Oke <Akshat.Oke at amd.com>
Date: Tue, 8 Oct 2024 06:24:11 +0000
Subject: [PATCH 4/4] AS

---
 llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
index 28578a875c164c..41621517706bb5 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUTargetMachine.cpp
@@ -1713,18 +1713,16 @@ bool GCNTargetMachine::parseMachineFunctionInfo(
   }
 
   auto setRegisterFlags = [&](const VRegInfo &Info) {
-    for (const auto &Flag : Info.Flags) {
+    for (uint8_t Flag : Info.Flags) {
       MFI->setFlag(Info.VReg, Flag);
     }
   };
 
-  for (const auto &P : PFS.VRegInfosNamed) {
-    const VRegInfo &Info = *P.second;
-    setRegisterFlags(Info);
+  for (const auto &[_, Info] : PFS.VRegInfosNamed) {
+    setRegisterFlags(*Info);
   }
-  for (const auto &P : PFS.VRegInfos) {
-    const VRegInfo &Info = *P.second;
-    setRegisterFlags(Info);
+  for (const auto &[_, Info] : PFS.VRegInfos) {
+    setRegisterFlags(*Info);
   }
 
   auto parseAndCheckArgument = [&](const std::optional<yaml::SIArgument> &A,



More information about the llvm-branch-commits mailing list