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.

Author: mssefat

llvm/lib/Target/AMDGPU/SIFoldOperands.cpp

@@ -1068,15 +1068,19 @@ void SIFoldOperandsImpl::foldOperand(
     if (MovOp == AMDGPU::COPY)
       return;
 
-    // Fold if the destination register class of the MOV instruction (ResRC)
-    // is a superclass of (or equal to) the destination register class of the COPY (DestRC).
-    // If this condition fails, folding would be illegal.
+    // Check for common register subclass between destination (DestRC) and MOV
+    // result (ResRC). If exists, verify this common subclass is a superclass of
+    // (or equal to) the destination register class, otherwise folding is
+    // illegal.
+
     const MCInstrDesc &MovDesc = TII->get(MovOp);
-    if (MovDesc.getNumDefs() > 0 && MovDesc.operands()[0].RegClass != -1) {
-      const TargetRegisterClass *ResRC = 
-          TRI->getRegClass(MovDesc.operands()[0].RegClass);
-      if (!DestRC -> hasSuperClassEq(ResRC)) return;
-    }
+    assert(MovDesc.getNumDefs() > 0 && MovDesc.operands()[0].RegClass != -1);
+    const TargetRegisterClass *ResRC =
+        TRI->getRegClass(MovDesc.operands()[0].RegClass);
+    const TargetRegisterClass *CommonRC = TRI->getCommonSubClass(DestRC, ResRC);
+
+    if (!CommonRC || !DestRC->hasSuperClassEq(CommonRC))
+      return;
 
     MachineInstr::mop_iterator ImpOpI = UseMI->implicit_operands().begin();
     MachineInstr::mop_iterator ImpOpE = UseMI->implicit_operands().end();

llvm/test/CodeGen/AMDGPU/fold-imm-copy.mir

@@ -208,52 +208,70 @@ 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:
+    ; GCN-LABEL: name: s_mov_b32_inlineimm_copy_s_to_av_32
+    ; GCN: [[S_MOV_B32_:%[0-9]+]]:sreg_32 = S_MOV_B32 32
+    ; GCN-NEXT: [[COPY:%[0-9]+]]:av_32 = COPY [[S_MOV_B32_]]
+    ; GCN-NEXT: $agpr0 = COPY [[COPY]]
+    ; GCN-NEXT: S_ENDPGM 0
+    %0:sreg_32 = S_MOV_B32 32
+    %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:
+    ; GCN-LABEL: name: v_mov_b32_inlineimm_copy_v_to_av_32
+    ; GCN: [[V_MOV_B32_e32_:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 32, implicit $exec
+    ; GCN-NEXT: [[COPY:%[0-9]+]]:av_32 = COPY [[V_MOV_B32_e32_]]
+    ; GCN-NEXT: $agpr0 = COPY [[COPY]]
+    ; GCN-NEXT: S_ENDPGM 0
+   %0:vgpr_32 = V_MOV_B32_e32 32, implicit $exec
+   %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: s_mov_b32_imm_literal_copy_s_to_av_32
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; GCN-LABEL: name: s_mov_b32_imm_literal_copy_s_to_av_32
+    ; GCN: [[S_MOV_B32_:%[0-9]+]]:sreg_32 = S_MOV_B32 999
+    ; GCN-NEXT: [[COPY:%[0-9]+]]:av_32 = COPY [[S_MOV_B32_]]
+    ; GCN-NEXT: $agpr0 = COPY [[COPY]]
+    ; GCN-NEXT: S_ENDPGM 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: v_mov_b32_imm_literal_copy_v_to_av_32
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; GCN-LABEL: name: v_mov_b32_imm_literal_copy_v_to_av_32
+    ; GCN: [[V_MOV_B32_e32_:%[0-9]+]]:vgpr_32 = V_MOV_B32_e32 999, implicit $exec
+    ; GCN-NEXT: [[COPY:%[0-9]+]]:av_32 = COPY [[V_MOV_B32_e32_]]
+    ; GCN-NEXT: $agpr0 = COPY [[COPY]]
+    ; GCN-NEXT: S_ENDPGM 0
+    %0:vgpr_32 = V_MOV_B32_e32 999, implicit $exec
+    %1:av_32 = COPY %0
+    $agpr0 = COPY %1
+    S_ENDPGM 0
 
-# ...
+...