[llvm] [AMDGPU] Fix register class constraints for si-fold-operands pass when folding immediate into copies (PR #131387)

[via llvm-commits]

https://github.com/mssefat updated https://github.com/llvm/llvm-project/pull/131387

>From 30124cf82bfb3123f4ce04d4b11a755e710a0d21 Mon Sep 17 00:00:00 2001
From: mssefat <syadus.sefat at gmail.com>
Date: Fri, 14 Mar 2025 16:14:19 -0400
Subject: [PATCH] [AMDGPU] Fix register class constraints for si-fold-operands
 pass

Fixes #130020

This fixes an issue where the si-fold-operands pass would incorrectly
fold immediate values into COPY instructions targeting av_32 registers,
which is illegal.

The pass now properly checks register class constraints before attempting to
fold the immediates.
---
 llvm/lib/Target/AMDGPU/SIFoldOperands.cpp  |  5 ++
 llvm/test/CodeGen/AMDGPU/fold-imm-copy.mir | 95 +++++++++++-----------
 2 files changed, 51 insertions(+), 49 deletions(-)

diff --git a/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp b/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
index 91df516b80857..5a7fefaafd768 100644
--- a/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
+++ b/llvm/lib/Target/AMDGPU/SIFoldOperands.cpp
@@ -1047,6 +1047,11 @@ void SIFoldOperandsImpl::foldOperand(
     if (MovOp == AMDGPU::COPY)
       return;
 
+    // Check if the destination register of the MOV operation belongs
+    // to a vector superclass. Folding would be illegal.
+    if (TRI->isVectorSuperClass(DestRC))
+      return;
+
     MachineInstr::mop_iterator ImpOpI = UseMI->implicit_operands().begin();
     MachineInstr::mop_iterator ImpOpE = UseMI->implicit_operands().end();
     while (ImpOpI != ImpOpE) {
diff --git a/llvm/test/CodeGen/AMDGPU/fold-imm-copy.mir b/llvm/test/CodeGen/AMDGPU/fold-imm-copy.mir
index dce11209d0514..4709bc947a380 100644
--- a/llvm/test/CodeGen/AMDGPU/fold-imm-copy.mir
+++ b/llvm/test/CodeGen/AMDGPU/fold-imm-copy.mir
@@ -185,52 +185,49 @@ body:             |
 
 ...
 
-# FIXME: Register class restrictions of av register not respected,
-# issue 130020
-
-# ---
-# name: s_mov_b32_inlineimm_copy_s_to_av_32
-# tracksRegLiveness: true
-# body:             |
-#   bb.0:
-#     %0:sreg_32 = S_MOV_B32 32
-#     %1:av_32 = COPY %0
-#     $agpr0 = COPY %1
-#     S_ENDPGM 0
-
-# ...
-
-# ---
-# name: v_mov_b32_inlineimm_copy_v_to_av_32
-# tracksRegLiveness: true
-# body:             |
-#  bb.0:
-#    %0:vgpr_32 = V_MOV_B32_e32 32, implicit $exec
-#    %1:av_32 = COPY %0
-#    $agpr0 = COPY %1
-#    S_ENDPGM 0
-# ...
-
-# ---
-# name: s_mov_b32_imm_literal_copy_s_to_av_32
-# tracksRegLiveness: true
-# body:             |
-#   bb.0:
-#     %0:sreg_32 = S_MOV_B32 999
-#     %1:av_32 = COPY %0
-#     $agpr0 = COPY %1
-#     S_ENDPGM 0
-
-# ...
-
-# ---
-# name: v_mov_b32_imm_literal_copy_v_to_av_32
-# tracksRegLiveness: true
-# body:             |
-#   bb.0:
-#     %0:vgpr_32 = V_MOV_B32_e32 999, implicit $exec
-#     %1:av_32 = COPY %0
-#     $agpr0 = COPY %1
-#     S_ENDPGM 0
-
-# ...
+---
+name: s_mov_b32_inlineimm_copy_s_to_av_32
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    %0:sreg_32 = S_MOV_B32 32
+    %1:av_32 = COPY %0
+    $agpr0 = COPY %1
+    S_ENDPGM 0
+
+...
+
+---
+name: v_mov_b32_inlineimm_copy_v_to_av_32
+tracksRegLiveness: true
+body:             |
+ bb.0:
+   %0:vgpr_32 = V_MOV_B32_e32 32, implicit $exec
+   %1:av_32 = COPY %0
+   $agpr0 = COPY %1
+   S_ENDPGM 0
+...
+
+---
+name: s_mov_b32_imm_literal_copy_s_to_av_32
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    %0:sreg_32 = S_MOV_B32 999
+    %1:av_32 = COPY %0
+    $agpr0 = COPY %1
+    S_ENDPGM 0
+
+...
+
+---
+name: v_mov_b32_imm_literal_copy_v_to_av_32
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    %0:vgpr_32 = V_MOV_B32_e32 999, implicit $exec
+    %1:av_32 = COPY %0
+    $agpr0 = COPY %1
+    S_ENDPGM 0
+
+...