[AMDGPU] Hint that s_cselect_b32 source-0 and destination ideally are assigned the same physical register.

Addresses #129984.

Author: bababuck

llvm/lib/Target/AMDGPU/SIShrinkInstructions.cpp

@@ -897,6 +897,15 @@ bool SIShrinkInstructions::run(MachineFunction &MF) {
         if (!Src1->isImm() || !isKImmOperand(*Src1))
           continue;
 
+        // Hint that the source and destination register should be allocated
+        // as the same register so that we can shrink to S_CMOVK_I32 on the
+        // post-allocation SIShrinkInstructions pass.
+        if (Dest->getReg().isVirtual()) {
+          MRI->setRegAllocationHint(Dest->getReg(), 0, Src0->getReg());
+          MRI->setRegAllocationHint(Src0->getReg(), 0, Dest->getReg());
+          continue;
+        }
+
         // The first source and destination must be the same register
         if (Src0->getReg() != Dest->getReg())
           continue;

llvm/test/CodeGen/AMDGPU/fcopysign.f16.ll

@@ -1960,9 +1960,9 @@ define amdgpu_kernel void @s_copysign_out_f16_mag_f64_sign_f16(ptr addrspace(1)
 ; GFX11-NEXT:    s_movk_i32 s6, 0x7e00
 ; GFX11-NEXT:    s_cmovk_i32 s5, 0x7c00
 ; GFX11-NEXT:    s_cmp_lg_u32 s3, 0
-; GFX11-NEXT:    s_cselect_b32 s3, s6, 0x7c00
+; GFX11-NEXT:    s_cmovk_i32 s6, 0x7c00
 ; GFX11-NEXT:    s_cmpk_eq_i32 s2, 0x40f
-; GFX11-NEXT:    s_cselect_b32 s2, s3, s5
+; GFX11-NEXT:    s_cselect_b32 s2, s6, s5
 ; GFX11-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
 ; GFX11-NEXT:    v_bfi_b32 v0, 0x7fff, s2, v0
 ; GFX11-NEXT:    global_store_b16 v1, v0, s[0:1]

llvm/test/CodeGen/AMDGPU/fptrunc.ll

@@ -134,9 +134,9 @@ define amdgpu_kernel void @fptrunc_f64_to_f16(ptr addrspace(1) %out, double %in)
 ; SI-NEXT:    s_cmp_lt_i32 s0, 31
 ; SI-NEXT:    s_cmovk_i32 s6, 0x7c00
 ; SI-NEXT:    s_cmp_lg_u32 s1, 0
-; SI-NEXT:    s_cselect_b32 s1, s2, 0x7c00
+; SI-NEXT:    s_cmovk_i32 s2, 0x7c00
 ; SI-NEXT:    s_cmpk_eq_i32 s0, 0x40f
-; SI-NEXT:    s_cselect_b32 s0, s1, s6
+; SI-NEXT:    s_cselect_b32 s0, s2, s6
 ; SI-NEXT:    s_lshr_b32 s1, s7, 16
 ; SI-NEXT:    s_and_b32 s1, s1, 0x8000
 ; SI-NEXT:    s_or_b32 s6, s1, s0
@@ -314,9 +314,9 @@ define amdgpu_kernel void @fptrunc_f64_to_f16(ptr addrspace(1) %out, double %in)
 ; GFX10-SAFE-SDAG-NEXT:    s_movk_i32 s6, 0x7e00
 ; GFX10-SAFE-SDAG-NEXT:    s_cmovk_i32 s5, 0x7c00
 ; GFX10-SAFE-SDAG-NEXT:    s_cmp_lg_u32 s4, 0
-; GFX10-SAFE-SDAG-NEXT:    s_cselect_b32 s4, s6, 0x7c00
+; GFX10-SAFE-SDAG-NEXT:    s_cmovk_i32 s6, 0x7c00
 ; GFX10-SAFE-SDAG-NEXT:    s_cmpk_eq_i32 s2, 0x40f
-; GFX10-SAFE-SDAG-NEXT:    s_cselect_b32 s2, s4, s5
+; GFX10-SAFE-SDAG-NEXT:    s_cselect_b32 s2, s6, s5
 ; GFX10-SAFE-SDAG-NEXT:    s_lshr_b32 s3, s3, 16
 ; GFX10-SAFE-SDAG-NEXT:    s_and_b32 s3, s3, 0x8000
 ; GFX10-SAFE-SDAG-NEXT:    s_or_b32 s2, s3, s2
@@ -446,9 +446,9 @@ define amdgpu_kernel void @fptrunc_f64_to_f16(ptr addrspace(1) %out, double %in)
 ; GFX11-SAFE-SDAG-NEXT:    s_movk_i32 s6, 0x7e00
 ; GFX11-SAFE-SDAG-NEXT:    s_cmovk_i32 s5, 0x7c00
 ; GFX11-SAFE-SDAG-NEXT:    s_cmp_lg_u32 s4, 0
-; GFX11-SAFE-SDAG-NEXT:    s_cselect_b32 s4, s6, 0x7c00
+; GFX11-SAFE-SDAG-NEXT:    s_cmovk_i32 s6, 0x7c00
 ; GFX11-SAFE-SDAG-NEXT:    s_cmpk_eq_i32 s2, 0x40f
-; GFX11-SAFE-SDAG-NEXT:    s_cselect_b32 s2, s4, s5
+; GFX11-SAFE-SDAG-NEXT:    s_cselect_b32 s2, s6, s5
 ; GFX11-SAFE-SDAG-NEXT:    s_lshr_b32 s3, s3, 16
 ; GFX11-SAFE-SDAG-NEXT:    s_delay_alu instid0(SALU_CYCLE_1) | instskip(NEXT) | instid1(SALU_CYCLE_1)
 ; GFX11-SAFE-SDAG-NEXT:    s_and_b32 s3, s3, 0x8000

llvm/test/CodeGen/AMDGPU/shrink-select.mir

@@ -11,8 +11,8 @@ body:             |
     ; GCN-NEXT: renamable $sgpr1 = S_MOV_B32 0
     ; GCN-NEXT: renamable $sgpr2 = S_MOV_B32 0
     ; GCN-NEXT: renamable $sgpr0 = S_ADD_U32 killed renamable $sgpr0, killed renamable $sgpr1, implicit-def $scc
-    ; GCN-NEXT: renamable $sgpr1 = S_CSELECT_B32 killed renamable $sgpr2, 31744, implicit killed $scc
-    ; GCN-NEXT: S_ENDPGM 0, implicit killed renamable $sgpr1, implicit killed renamable $sgpr0
+    ; GCN-NEXT: renamable $sgpr2 = S_CMOVK_I32 31744, implicit killed $scc
+    ; GCN-NEXT: S_ENDPGM 0, implicit killed renamable $sgpr2, implicit killed renamable $sgpr0
     %0:sgpr_32 = S_MOV_B32 0
     %1:sgpr_32 = S_MOV_B32 0
     %2:sgpr_32 = S_MOV_B32 0