[llvm] [AMDGPU] Fix SIPreAllocateWWMRegs to reserve AV-class WWM defs (PR #211560)

[via llvm-commits]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Arseniy Obolenskiy (aobolensk)

<details>
<summary>Changes</summary>

isVGPR() rejects the unified VGPR+AGPR register class used on gfx90A+, so strict-WWM defs allocated to an AV-class register were left out of WWMReservedRegs and could be clobbered by the post-WWM allocator

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


2 Files Affected:

- (modified) llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp (+1-1) 
- (modified) llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs.mir (+39) 


``````````diff
diff --git a/llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp b/llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp
index 2afd573c94f3f..bf484cef98da4 100644
--- a/llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp
+++ b/llvm/lib/Target/AMDGPU/SIPreAllocateWWMRegs.cpp
@@ -109,7 +109,7 @@ bool SIPreAllocateWWMRegs::processDef(MachineOperand &MO) {
   if (Reg.isPhysical())
     return false;
 
-  if (!TRI->isVGPR(*MRI, Reg))
+  if (!SIRegisterInfo::hasVGPRs(MRI->getRegClass(Reg)))
     return false;
 
   if (VRM->hasPhys(Reg))
diff --git a/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs.mir b/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs.mir
index bde9ea336fae9..eb825bf80b3e0 100644
--- a/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs.mir
+++ b/llvm/test/CodeGen/AMDGPU/si-pre-allocate-wwm-regs.mir
@@ -44,6 +44,45 @@ body: |
 ...
 ---
 
+name: pre_allocate_wwm_regs_strict_av_class
+tracksRegLiveness: true
+body: |
+  bb.0:
+    liveins: $sgpr1
+    ; CHECK-LABEL: name: pre_allocate_wwm_regs_strict_av_class
+    ; CHECK: wwmReservedRegs:
+    ; CHECK-NEXT: - '$vgpr0'
+    ; CHECK-NEXT: - '$vgpr1'
+    ; CHECK: liveins: $sgpr1
+    ; CHECK-NEXT: {{  $}}
+    ; CHECK-NEXT: [[DEF:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
+    ; CHECK-NEXT: renamable $sgpr4_sgpr5 = ENTER_STRICT_WWM -1, implicit-def $exec, implicit-def $scc, implicit $exec
+    ; CHECK-NEXT: $vgpr0 = AV_MOV_B32_IMM_PSEUDO 0, implicit $exec
+    ; CHECK-NEXT: dead $vgpr1 = V_MOV_B32_dpp [[DEF]], [[DEF]], 323, 12, 15, 0, implicit $exec
+    ; CHECK-NEXT: $exec = EXIT_STRICT_WWM killed renamable $sgpr4_sgpr5
+    ; CHECK-NEXT: dead [[COPY:%[0-9]+]]:av_32 = COPY $vgpr0
+    ;
+    ; CHECK2-LABEL: name: pre_allocate_wwm_regs_strict_av_class
+    ; CHECK2: wwmReservedRegs:
+    ; CHECK2-NEXT: - '$vgpr0'
+    ; CHECK2-NEXT: - '$vgpr1'
+    ; CHECK2: liveins: $sgpr1
+    ; CHECK2-NEXT: {{  $}}
+    ; CHECK2-NEXT: [[DEF:%[0-9]+]]:vgpr_32 = IMPLICIT_DEF
+    ; CHECK2-NEXT: renamable $sgpr4_sgpr5 = ENTER_STRICT_WWM -1, implicit-def $exec, implicit-def $scc, implicit $exec
+    ; CHECK2-NEXT: $vgpr0 = AV_MOV_B32_IMM_PSEUDO 0, implicit $exec
+    ; CHECK2-NEXT: dead $vgpr1 = V_MOV_B32_dpp [[DEF]], [[DEF]], 323, 12, 15, 0, implicit $exec
+    ; CHECK2-NEXT: $exec = EXIT_STRICT_WWM killed renamable $sgpr4_sgpr5
+    ; CHECK2-NEXT: dead [[COPY:%[0-9]+]]:av_32 = COPY $vgpr0
+    %0:vgpr_32 = IMPLICIT_DEF
+    renamable $sgpr4_sgpr5 = ENTER_STRICT_WWM -1, implicit-def $exec, implicit-def $scc, implicit $exec
+    %1:av_32 = AV_MOV_B32_IMM_PSEUDO 0, implicit $exec
+    %2:vgpr_32 = V_MOV_B32_dpp %0, %0, 323, 12, 15, 0, implicit $exec
+    $exec = EXIT_STRICT_WWM killed renamable $sgpr4_sgpr5
+    %3:av_32 = COPY %1
+...
+---
+
 name: pre_allocate_wwm_spill_to_vgpr
 tracksRegLiveness: true
 body: |

``````````

</details>


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