[AMDGPU][gfx13] use s_add_pc_i64 for rank-call (#3161)

Make the rank-call more efficient.

Author: cmc-rep

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -6836,15 +6836,19 @@ SITargetLowering::EmitInstrWithCustomInserter(MachineInstr &MI,
         .addReg(WaveIDInWaveGroup)
         .addImm(Rank);
     BuildMI(*BB, &MI, DL, TII->get(AMDGPU::S_CBRANCH_SCC0)).addMBB(SplitBB);
+
     // Call inside the conditional branch.
     Register CalleeAddrReg = MI.getOperand(1).getReg();
-    BuildMI(*RankCallBB, RankCallBB->end(), DL, TII->get(AMDGPU::S_SETPC_B64))
-        .addReg(CalleeAddrReg);
+    auto CalleeAddrDef = MRI.getVRegDef(CalleeAddrReg);
+    assert(CalleeAddrDef->getOpcode() == AMDGPU::SI_PC_ADD_REL_OFFSET64);
+    // Use s_add_pc_i64, bypass the address computation.
+    BuildMI(*RankCallBB, RankCallBB->end(), DL, TII->get(AMDGPU::S_ADD_PC_I64))
+        .addGlobalAddress(CalleeAddrDef->getOperand(1).getGlobal(), 0,
+                          SIInstrInfo::MO_REL64);
+
     // Update IDX0 for the next rank-call. Use the global address of the rank
     // callee as the source. In AsmPrinter, it will be replaced with the
     // MCSymbol representing the number of VGPRs of that callee.
-    auto CalleeAddrDef =  MRI.getVRegDef(CalleeAddrReg);
-    assert(CalleeAddrDef->getOpcode() == AMDGPU::SI_PC_ADD_REL_OFFSET64);
     BuildMI(*SplitBB, SplitBB->begin(), DL, TII->get(AMDGPU::S_ADD_GPR_IDX_U32),
             AMDGPU::IDX0)
         .addGlobalAddress(CalleeAddrDef->getOperand(1).getGlobal(), 0,

llvm/test/CodeGen/AMDGPU/rank-specialization-lowered.ll

@@ -219,47 +219,45 @@ define dso_local amdgpu_kernel void @test_kernel_1() local_unnamed_addr #1 !reqd
 ; CHECK-NEXT:  .LBB5_3: ; %bb.rank_0_2_3_4_5_6_7
 ; CHECK-NEXT:    s_set_gpr_idx_u32 idx0, 0
 ; CHECK-NEXT:    s_cmp_eq_u32 s9, 0
-; CHECK-NEXT:    s_get_pc_i64 s[10:11]
-; CHECK-NEXT:    s_add_nc_u64 s[10:11], s[10:11], test_kernel_1.rank_0_2_3_4_5_6_7@rel64+4
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB5_5
 ; CHECK-NEXT:  ; %bb.4:
-; CHECK-NEXT:    s_set_pc_i64 s[10:11]
+; CHECK-NEXT:    s_add_pc_i64 test_kernel_1.rank_0_2_3_4_5_6_7@rel64
 ; CHECK-NEXT:  .LBB5_5: ; %bb.rank_0_2_3_4_5_6_7
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, max(128, dummy_store.num_vgpr)
 ; CHECK-NEXT:    s_cmp_eq_u32 s9, 2
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB5_7
 ; CHECK-NEXT:  ; %bb.6:
-; CHECK-NEXT:    s_set_pc_i64 s[10:11]
+; CHECK-NEXT:    s_add_pc_i64 test_kernel_1.rank_0_2_3_4_5_6_7@rel64
 ; CHECK-NEXT:  .LBB5_7: ; %bb.rank_0_2_3_4_5_6_7
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, max(128, dummy_store.num_vgpr)
 ; CHECK-NEXT:    s_cmp_eq_u32 s9, 3
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB5_9
 ; CHECK-NEXT:  ; %bb.8:
-; CHECK-NEXT:    s_set_pc_i64 s[10:11]
+; CHECK-NEXT:    s_add_pc_i64 test_kernel_1.rank_0_2_3_4_5_6_7@rel64
 ; CHECK-NEXT:  .LBB5_9: ; %bb.rank_0_2_3_4_5_6_7
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, max(128, dummy_store.num_vgpr)
 ; CHECK-NEXT:    s_cmp_eq_u32 s9, 4
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB5_11
 ; CHECK-NEXT:  ; %bb.10:
-; CHECK-NEXT:    s_set_pc_i64 s[10:11]
+; CHECK-NEXT:    s_add_pc_i64 test_kernel_1.rank_0_2_3_4_5_6_7@rel64
 ; CHECK-NEXT:  .LBB5_11: ; %bb.rank_0_2_3_4_5_6_7
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, max(128, dummy_store.num_vgpr)
 ; CHECK-NEXT:    s_cmp_eq_u32 s9, 5
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB5_13
 ; CHECK-NEXT:  ; %bb.12:
-; CHECK-NEXT:    s_set_pc_i64 s[10:11]
+; CHECK-NEXT:    s_add_pc_i64 test_kernel_1.rank_0_2_3_4_5_6_7@rel64
 ; CHECK-NEXT:  .LBB5_13: ; %bb.rank_0_2_3_4_5_6_7
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, max(128, dummy_store.num_vgpr)
 ; CHECK-NEXT:    s_cmp_eq_u32 s9, 6
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB5_15
 ; CHECK-NEXT:  ; %bb.14:
-; CHECK-NEXT:    s_set_pc_i64 s[10:11]
+; CHECK-NEXT:    s_add_pc_i64 test_kernel_1.rank_0_2_3_4_5_6_7@rel64
 ; CHECK-NEXT:  .LBB5_15: ; %bb.rank_0_2_3_4_5_6_7
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, max(128, dummy_store.num_vgpr)
 ; CHECK-NEXT:    s_cmp_eq_u32 s9, 7
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB5_17
 ; CHECK-NEXT:  ; %bb.16:
-; CHECK-NEXT:    s_set_pc_i64 s[10:11]
+; CHECK-NEXT:    s_add_pc_i64 test_kernel_1.rank_0_2_3_4_5_6_7@rel64
 ; CHECK-NEXT:  .LBB5_17: ; %bb.rank_0_2_3_4_5_6_7
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, max(128, dummy_store.num_vgpr)
 ; CHECK-NEXT:    s_cbranch_execnz .LBB5_2
@@ -268,10 +266,7 @@ define dso_local amdgpu_kernel void @test_kernel_1() local_unnamed_addr #1 !reqd
 ; CHECK-NEXT:    s_cmp_eq_u32 s9, 1
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB5_20
 ; CHECK-NEXT:  ; %bb.19:
-; CHECK-NEXT:    s_get_pc_i64 s[10:11]
-; CHECK-NEXT:    s_add_nc_u64 s[10:11], s[10:11], test_kernel_1.rank_1@rel64+4
-; CHECK-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
-; CHECK-NEXT:    s_set_pc_i64 s[10:11]
+; CHECK-NEXT:    s_add_pc_i64 test_kernel_1.rank_1@rel64
 ; CHECK-NEXT:  .LBB5_20: ; %bb.rank_1
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, max(128, dummy_store.num_vgpr, dummy_rank1a.num_vgpr, dummy_rank1b.num_vgpr)
 ; CHECK-NEXT:    s_endpgm

llvm/test/CodeGen/AMDGPU/wavegroup-rank-func.ll

@@ -143,28 +143,19 @@ define amdgpu_kernel void @main(ptr addrspace(1) %inbuf, ptr addrspace(1) %wbuf,
 ; CHECK-NEXT:    s_cmp_eq_u32 s3, 0
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB3_2
 ; CHECK-NEXT:  ; %bb.1:
-; CHECK-NEXT:    s_get_pc_i64 s[4:5]
-; CHECK-NEXT:    s_add_nc_u64 s[4:5], s[4:5], input@rel64+4
-; CHECK-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
-; CHECK-NEXT:    s_set_pc_i64 s[4:5]
+; CHECK-NEXT:    s_add_pc_i64 input@rel64
 ; CHECK-NEXT:  .LBB3_2: ; %entry
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, 18
 ; CHECK-NEXT:    s_cmp_eq_u32 s3, 1
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB3_4
 ; CHECK-NEXT:  ; %bb.3:
-; CHECK-NEXT:    s_get_pc_i64 s[4:5]
-; CHECK-NEXT:    s_add_nc_u64 s[4:5], s[4:5], compute@rel64+4
-; CHECK-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
-; CHECK-NEXT:    s_set_pc_i64 s[4:5]
+; CHECK-NEXT:    s_add_pc_i64 compute@rel64
 ; CHECK-NEXT:  .LBB3_4: ; %entry
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, 0
 ; CHECK-NEXT:    s_cmp_eq_u32 s3, 2
 ; CHECK-NEXT:    s_cbranch_scc0 .LBB3_6
 ; CHECK-NEXT:  ; %bb.5:
-; CHECK-NEXT:    s_get_pc_i64 s[4:5]
-; CHECK-NEXT:    s_add_nc_u64 s[4:5], s[4:5], output@rel64+4
-; CHECK-NEXT:    s_delay_alu instid0(SALU_CYCLE_1)
-; CHECK-NEXT:    s_set_pc_i64 s[4:5]
+; CHECK-NEXT:    s_add_pc_i64 output@rel64
 ; CHECK-NEXT:  .LBB3_6: ; %entry
 ; CHECK-NEXT:    s_add_gpr_idx_u32 idx0, 1
 ; CHECK-NEXT:    s_endpgm