From 5e53a8dadb0019ee87936c1278fa222781257005 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault@amd.com>
Date: Fri, 13 Dec 2024 11:52:11 +0900
Subject: [PATCH] AMDGPU: Fix verifier assert with out of bounds subregister
 indexes (#119799)

The manual check for aligned VGPR classes would assert if a virtual
register used an index not supported by the register class.
---
 llvm/lib/Target/AMDGPU/SIInstrInfo.cpp        | 11 ++---
 ...ported-subreg-index-aligned-vgpr-check.mir | 41 +++++++++++++++++++
 2 files changed, 47 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/MachineVerifier/AMDGPU/unsupported-subreg-index-aligned-vgpr-check.mir

diff --git a/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp b/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
index 4a94d69029794..057412d41e7a2 100644
--- a/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIInstrInfo.cpp
@@ -4771,11 +4771,12 @@ bool SIInstrInfo::verifyInstruction(const MachineInstr &MI,
     if (ST.needsAlignedVGPRs()) {
       const TargetRegisterClass *RC = RI.getRegClassForReg(MRI, Reg);
       if (RI.hasVectorRegisters(RC) && MO.getSubReg()) {
-        const TargetRegisterClass *SubRC =
-            RI.getSubRegisterClass(RC, MO.getSubReg());
-        RC = RI.getCompatibleSubRegClass(RC, SubRC, MO.getSubReg());
-        if (RC)
-          RC = SubRC;
+        if (const TargetRegisterClass *SubRC =
+                RI.getSubRegisterClass(RC, MO.getSubReg())) {
+          RC = RI.getCompatibleSubRegClass(RC, SubRC, MO.getSubReg());
+          if (RC)
+            RC = SubRC;
+        }
       }
 
       // Check that this is the aligned version of the class.
diff --git a/llvm/test/MachineVerifier/AMDGPU/unsupported-subreg-index-aligned-vgpr-check.mir b/llvm/test/MachineVerifier/AMDGPU/unsupported-subreg-index-aligned-vgpr-check.mir
new file mode 100644
index 0000000000000..651a9d71ae320
--- /dev/null
+++ b/llvm/test/MachineVerifier/AMDGPU/unsupported-subreg-index-aligned-vgpr-check.mir
@@ -0,0 +1,41 @@
+# RUN: not --crash llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx942 -run-pass=none -filetype=null %s 2>&1 | FileCheck %s
+
+# sub16_sub17_sub18_sub19 is outside the bounds of a 512-bit
+# register. Make sure the verification doesn't assert. This was only
+# broken for targets that require even aligned VGPRs due to the manual
+# alignment check.
+
+---
+name: uses_invalid_subregister_for_regclass
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    %0:vgpr_32 = V_MOV_B32_e32 0, implicit $exec
+    S_NOP 0, implicit-def %1:vreg_512_align2
+
+
+    ; CHECK: *** Bad machine code: Invalid subregister index for virtual register ***
+    ; CHECK-NEXT: - function:    uses_invalid_subregister_for_regclass
+    ; CHECK-NEXT: - basic block: %bb.0
+    ; CHECK-NEXT: - instruction: GLOBAL_STORE_DWORDX4_SADDR %0:vgpr_32, %1.sub16_sub17_sub18_sub19:vreg_512_align2, undef $sgpr8_sgpr9, 80, 0, implicit $exec :: (store (s128), addrspace 1)
+    ; CHECK-NEXT: - operand 1:   %1.sub16_sub17_sub18_sub19:vreg_512_align2
+    ; CHECK-NEXT: Register class VReg_512_Align2 does not support subreg index 166
+
+    ; CHECK: *** Bad machine code: Subtarget requires even aligned vector registers ***
+    ; CHECK-NEXT: - function:    uses_invalid_subregister_for_regclass
+    ; CHECK-NEXT: - basic block: %bb.0
+    ; CHECK-NEXT: - instruction: GLOBAL_STORE_DWORDX4_SADDR %0:vgpr_32, %2.sub16_sub17_sub18_sub19:vreg_512, undef $sgpr8_sgpr9, 80, 0, implicit $exec :: (store (s128), addrspace 1)
+    GLOBAL_STORE_DWORDX4_SADDR %0, %1.sub16_sub17_sub18_sub19, undef $sgpr8_sgpr9, 80, 0, implicit $exec :: (store (s128), addrspace 1)
+
+    ; Test with unaligned class
+    ; CHECK: *** Bad machine code: Invalid subregister index for virtual register ***
+    ; CHECK-NEXT: - function:    uses_invalid_subregister_for_regclass
+    ; CHECK-NEXT: - basic block: %bb.0
+    ; CHECK-NEXT: - instruction: GLOBAL_STORE_DWORDX4_SADDR %0:vgpr_32, %2.sub16_sub17_sub18_sub19:vreg_512, undef $sgpr8_sgpr9, 80, 0, implicit $exec :: (store (s128), addrspace 1)
+    ; CHECK-NEXT: - operand 1:   %2.sub16_sub17_sub18_sub19:vreg_512
+    ; CHECK-NEXT: Register class VReg_512 does not support subreg index 166
+    S_NOP 0, implicit-def %2:vreg_512
+    GLOBAL_STORE_DWORDX4_SADDR %0, %2.sub16_sub17_sub18_sub19, undef $sgpr8_sgpr9, 80, 0, implicit $exec :: (store (s128), addrspace 1)
+    S_ENDPGM 0
+
+...