AMDGPU/GlobalISel: Legalize workgroup ID intrinsics

[tstellar]

LGTM.

[arsenm]

r364834

lib/Target/AMDGPU/AMDGPUCallLowering.cpp

static void allocateSpecialEntryInputVGPRs(CCState &CCInfo,
   }
 }
 
 static void allocateSystemSGPRs(CCState &CCInfo,
                                 MachineFunction &MF,
                                 SIMachineFunctionInfo &Info,
                                 CallingConv::ID CallConv,
                                 bool IsShader) {
+  const LLT S32 = LLT::scalar(32);
+  MachineRegisterInfo &MRI = MF.getRegInfo();
+
+  if (Info.hasWorkGroupIDX()) {
+    Register Reg = Info.addWorkGroupIDX();
+    MRI.setType(MF.addLiveIn(Reg, &AMDGPU::SReg_32_XM0RegClass), S32);
+    CCInfo.AllocateReg(Reg);
+  }
+
+  if (Info.hasWorkGroupIDY()) {
+    Register Reg = Info.addWorkGroupIDY();
+    MRI.setType(MF.addLiveIn(Reg, &AMDGPU::SReg_32_XM0RegClass), S32);
+    CCInfo.AllocateReg(Reg);
+  }
+
+  if (Info.hasWorkGroupIDZ()) {
+    unsigned Reg = Info.addWorkGroupIDZ();
+    MRI.setType(MF.addLiveIn(Reg, &AMDGPU::SReg_32_XM0RegClass), S32);
+    CCInfo.AllocateReg(Reg);
+  }
+
+  if (Info.hasWorkGroupInfo()) {
+    unsigned Reg = Info.addWorkGroupInfo();
+    MRI.setType(MF.addLiveIn(Reg, &AMDGPU::SReg_32_XM0RegClass), S32);
+    CCInfo.AllocateReg(Reg);
+  }
+
   if (Info.hasPrivateSegmentWaveByteOffset()) {
     // Scratch wave offset passed in system SGPR.
     unsigned PrivateSegmentWaveByteOffsetReg;
 
     if (IsShader) {
       PrivateSegmentWaveByteOffsetReg =
         Info.getPrivateSegmentWaveByteOffsetSystemSGPR();
 

lib/Target/AMDGPU/AMDGPULegalizerInfo.cpp

case Intrinsic::amdgcn_workitem_id_x:
     return legalizePreloadedArgIntrin(MI, MRI, B,
                                       AMDGPUFunctionArgInfo::WORKITEM_ID_X);
   case Intrinsic::amdgcn_workitem_id_y:
     return legalizePreloadedArgIntrin(MI, MRI, B,
                                       AMDGPUFunctionArgInfo::WORKITEM_ID_Y);
   case Intrinsic::amdgcn_workitem_id_z:
     return legalizePreloadedArgIntrin(MI, MRI, B,
                                       AMDGPUFunctionArgInfo::WORKITEM_ID_Z);
+  case Intrinsic::amdgcn_workgroup_id_x:
+    return legalizePreloadedArgIntrin(MI, MRI, B,
+                                      AMDGPUFunctionArgInfo::WORKGROUP_ID_X);
+  case Intrinsic::amdgcn_workgroup_id_y:
+    return legalizePreloadedArgIntrin(MI, MRI, B,
+                                      AMDGPUFunctionArgInfo::WORKGROUP_ID_Y);
+  case Intrinsic::amdgcn_workgroup_id_z:
+    return legalizePreloadedArgIntrin(MI, MRI, B,
+                                      AMDGPUFunctionArgInfo::WORKGROUP_ID_Z);
   default:
     return true;
   }
 
   return true;
 }

test/CodeGen/AMDGPU/GlobalISel/irtranslator-amdgpu_kernel-system-sgprs.ll

 ; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=fiji -O0 -amdgpu-ir-lower-kernel-arguments=0 -stop-after=irtranslator -global-isel %s -o - | FileCheck -check-prefix=HSA %s
 
 ; HSA-LABEL: name: default_kernel
 ; HSA: liveins:
 ; HSA-NEXT: - { reg: '$sgpr0_sgpr1_sgpr2_sgpr3', virtual-reg: '%0' }
 ; HSA-NEXT: - { reg: '$vgpr0', virtual-reg: '%1' }
 ; HSA-NEXT: - { reg: '$sgpr4', virtual-reg: '%2' }
+; HSA-NEXT: - { reg: '$sgpr5', virtual-reg: '%3' }
 ; HSA-NEXT: frameInfo:
 define amdgpu_kernel void @default_kernel() {
   ret void
 }
+
+
+; HSA-LABEL: name: workgroup_id_x{{$}}
+; HSA: liveins:
+; HSA-NEXT: - { reg: '$sgpr0_sgpr1_sgpr2_sgpr3', virtual-reg: '%0' }
+; HSA-NEXT: - { reg: '$vgpr0', virtual-reg: '%1' }
+; HSA-NEXT: - { reg: '$sgpr4', virtual-reg: '%2' }
+; HSA-NEXT: - { reg: '$sgpr5', virtual-reg: '%3' }
+; HSA-NEXT: frameInfo:
+define amdgpu_kernel void @workgroup_id_x() {
+  %id = call i32 @llvm.amdgcn.workgroup.id.x()
+  store volatile i32 %id, i32 addrspace(1)* undef
+  ret void
+}
+
+; HSA-LABEL: name: workgroup_id_y{{$}}
+; HSA: liveins:
+; HSA-NEXT: - { reg: '$sgpr0_sgpr1_sgpr2_sgpr3', virtual-reg: '%0' }
+; HSA-NEXT: - { reg: '$vgpr0', virtual-reg: '%1' }
+; HSA-NEXT: - { reg: '$sgpr4', virtual-reg: '%2' }
+; HSA-NEXT: - { reg: '$sgpr5', virtual-reg: '%3' }
+; HSA-NEXT: - { reg: '$sgpr6', virtual-reg: '%4' }
+; HSA-NEXT: frameInfo:
+define amdgpu_kernel void @workgroup_id_y() {
+  %id = call i32 @llvm.amdgcn.workgroup.id.y()
+  store volatile i32 %id, i32 addrspace(1)* undef
+  ret void
+}
+
+; HSA-LABEL: name: workgroup_id_z{{$}}
+; HSA: liveins:
+; HSA-NEXT: - { reg: '$sgpr0_sgpr1_sgpr2_sgpr3', virtual-reg: '%0' }
+; HSA-NEXT: - { reg: '$vgpr0', virtual-reg: '%1' }
+; HSA-NEXT: - { reg: '$sgpr4', virtual-reg: '%2' }
+; HSA-NEXT: - { reg: '$sgpr5', virtual-reg: '%3' }
+; HSA-NEXT: - { reg: '$sgpr6', virtual-reg: '%4' }
+; HSA-NEXT: frameInfo:
+define amdgpu_kernel void @workgroup_id_z() {
+  %id = call i32 @llvm.amdgcn.workgroup.id.z()
+  store volatile i32 %id, i32 addrspace(1)* undef
+  ret void
+}
+
+; HSA-LABEL: name: workgroup_id_xy{{$}}
+; HSA: liveins:
+; HSA-NEXT: - { reg: '$sgpr0_sgpr1_sgpr2_sgpr3', virtual-reg: '%0' }
+; HSA-NEXT: - { reg: '$vgpr0', virtual-reg: '%1' }
+; HSA-NEXT: - { reg: '$sgpr4', virtual-reg: '%2' }
+; HSA-NEXT: - { reg: '$sgpr5', virtual-reg: '%3' }
+; HSA-NEXT: - { reg: '$sgpr6', virtual-reg: '%4' }
+; HSA-NEXT: frameInfo:
+define amdgpu_kernel void @workgroup_id_xy() {
+  %id0 = call i32 @llvm.amdgcn.workgroup.id.x()
+  store volatile i32 %id0, i32 addrspace(1)* undef
+  %id1 = call i32 @llvm.amdgcn.workgroup.id.y()
+  store volatile i32 %id1, i32 addrspace(1)* undef
+  ret void
+}
+
+; HSA-LABEL: name: workgroup_id_xyz{{$}}
+; HSA: liveins:
+; HSA-NEXT: - { reg: '$sgpr0_sgpr1_sgpr2_sgpr3', virtual-reg: '%0' }
+; HSA-NEXT: - { reg: '$vgpr0', virtual-reg: '%1' }
+; HSA-NEXT: - { reg: '$sgpr4', virtual-reg: '%2' }
+; HSA-NEXT: - { reg: '$sgpr5', virtual-reg: '%3' }
+; HSA-NEXT: - { reg: '$sgpr6', virtual-reg: '%4' }
+; HSA-NEXT: frameInfo:
+define amdgpu_kernel void @workgroup_id_xyz() {
+  %id0 = call i32 @llvm.amdgcn.workgroup.id.x()
+  store volatile i32 %id0, i32 addrspace(1)* undef
+  %id1 = call i32 @llvm.amdgcn.workgroup.id.y()
+  store volatile i32 %id1, i32 addrspace(1)* undef
+  %id2 = call i32 @llvm.amdgcn.workgroup.id.y()
+  store volatile i32 %id2, i32 addrspace(1)* undef
+  ret void
+}
+
+; HSA-LABEL: name: workgroup_id_yz{{$}}
+; HSA: liveins:
+; HSA-NEXT: - { reg: '$sgpr0_sgpr1_sgpr2_sgpr3', virtual-reg: '%0' }
+; HSA-NEXT: - { reg: '$vgpr0', virtual-reg: '%1' }
+; HSA-NEXT: - { reg: '$sgpr4', virtual-reg: '%2' }
+; HSA-NEXT: - { reg: '$sgpr5', virtual-reg: '%3' }
+; HSA-NEXT: - { reg: '$sgpr6', virtual-reg: '%4' }
+; HSA-NEXT: frameInfo:
+define amdgpu_kernel void @workgroup_id_yz() {
+  %id0 = call i32 @llvm.amdgcn.workgroup.id.x()
+  store volatile i32 %id0, i32 addrspace(1)* undef
+  %id1 = call i32 @llvm.amdgcn.workgroup.id.y()
+  store volatile i32 %id1, i32 addrspace(1)* undef
+  ret void
+}
+
+; HSA-LABEL: name: workgroup_id_xz{{$}}
+; HSA: liveins:
+; HSA-NEXT: - { reg: '$sgpr0_sgpr1_sgpr2_sgpr3', virtual-reg: '%0' }
+; HSA-NEXT: - { reg: '$vgpr0', virtual-reg: '%1' }
+; HSA-NEXT: - { reg: '$sgpr4', virtual-reg: '%2' }
+; HSA-NEXT: - { reg: '$sgpr5', virtual-reg: '%3' }
+; HSA-NEXT: - { reg: '$sgpr6', virtual-reg: '%4' }
+; HSA-NEXT: frameInfo:
+define amdgpu_kernel void @workgroup_id_xz() {
+  %id0 = call i32 @llvm.amdgcn.workgroup.id.x()
+  store volatile i32 %id0, i32 addrspace(1)* undef
+  %id1 = call i32 @llvm.amdgcn.workgroup.id.z()
+  store volatile i32 %id1, i32 addrspace(1)* undef
+  ret void
+}
+
+declare i32 @llvm.amdgcn.workgroup.id.x() #0
+declare i32 @llvm.amdgcn.workgroup.id.y() #0
+declare i32 @llvm.amdgcn.workgroup.id.z() #0
+
+attributes #0 = { nounwind readnone speculatable }

test/CodeGen/AMDGPU/GlobalISel/llvm.amdgcn.workgroup.id.ll

+; RUN: llc -global-isel -mtriple=amdgcn-unknown-amdhsa -mattr=-code-object-v3 -mcpu=kaveri -verify-machineinstrs < %s | FileCheck -check-prefix=ALL -check-prefix=CO-V2 -check-prefix=CI-HSA  %s
+; RUN: llc -global-isel -mtriple=amdgcn-unknown-amdhsa -mattr=-code-object-v3 -mcpu=carrizo -verify-machineinstrs < %s | FileCheck -check-prefix=ALL -check-prefix=CO-V2 -check-prefix=VI-HSA  %s
+; RUN: llc -global-isel -mtriple=amdgcn-- -mcpu=tahiti -verify-machineinstrs < %s | FileCheck -check-prefix=ALL -check-prefix=UNKNOWN-OS -check-prefix=SI-MESA %s
+; RUN: llc -global-isel -mtriple=amdgcn-- -mcpu=tonga -verify-machineinstrs < %s | FileCheck -check-prefix=ALL -check-prefix=UNKNOWN-OS -check-prefix=VI-MESA %s
+; RUN: llc -global-isel -mtriple=amdgcn-unknown-mesa3d -mattr=-code-object-v3 -mcpu=tahiti -verify-machineinstrs < %s | FileCheck -check-prefixes=ALL,CO-V2,SI-MESA %s
+; RUN: llc -global-isel -mtriple=amdgcn-unknown-mesa3d -mattr=-code-object-v3 -mcpu=tonga -verify-machineinstrs < %s | FileCheck -check-prefixes=ALL,CO-V2,VI-MESA %s
+
+declare i32 @llvm.amdgcn.workgroup.id.x() #0
+declare i32 @llvm.amdgcn.workgroup.id.y() #0
+declare i32 @llvm.amdgcn.workgroup.id.z() #0
+
+; ALL-LABEL {{^}}test_workgroup_id_x:
+
+; CO-V2: .amd_kernel_code_t
+; CO-V2: user_sgpr_count = 6
+; CO-V2: enable_sgpr_workgroup_id_x = 1
+; CO-V2: enable_sgpr_workgroup_id_y = 0
+; CO-V2: enable_sgpr_workgroup_id_z = 0
+; CO-V2: enable_sgpr_workgroup_info = 0
+; CO-V2: enable_vgpr_workitem_id = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_x = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_y = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_z = 0
+; CO-V2: .end_amd_kernel_code_t
+
+; UNKNOWN-OS: v_mov_b32_e32 [[VCOPY:v[0-9]+]], s2{{$}}
+; CO-V2: v_mov_b32_e32 [[VCOPY:v[0-9]+]], s6{{$}}
+
+; ALL: {{buffer|flat}}_store_dword {{.*}}[[VCOPY]]
+
+; CO-V2: COMPUTE_PGM_RSRC2:USER_SGPR: 6
+; ALL-NOCO-V2: COMPUTE_PGM_RSRC2:USER_SGPR: 2
+; ALL: COMPUTE_PGM_RSRC2:TGID_X_EN: 1
+; ALL: COMPUTE_PGM_RSRC2:TGID_Y_EN: 0
+; ALL: COMPUTE_PGM_RSRC2:TGID_Z_EN: 0
+; ALL: COMPUTE_PGM_RSRC2:TIDIG_COMP_CNT: 0
+define amdgpu_kernel void @test_workgroup_id_x(i32 addrspace(1)* %out) #1 {
+  %id = call i32 @llvm.amdgcn.workgroup.id.x()
+  store i32 %id, i32 addrspace(1)* %out
+  ret void
+}
+
+; ALL-LABEL {{^}}test_workgroup_id_y:
+; CO-V2: user_sgpr_count = 6
+; CO-V2: enable_sgpr_workgroup_id_x = 1
+; CO-V2: enable_sgpr_workgroup_id_y = 1
+; CO-V2: enable_sgpr_workgroup_id_z = 0
+; CO-V2: enable_sgpr_workgroup_info = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_x = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_y = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_z = 0
+
+; UNKNOWN-OS: v_mov_b32_e32 [[VCOPY:v[0-9]+]], s3{{$}}
+; HSA: v_mov_b32_e32 [[VCOPY:v[0-9]+]], s7{{$}}
+
+; ALL: {{buffer|flat}}_store_dword {{.*}}[[VCOPY]]
+
+; CO-V2: COMPUTE_PGM_RSRC2:USER_SGPR: 6
+; ALL-NOCO-V2: COMPUTE_PGM_RSRC2:USER_SGPR: 2
+; ALL: COMPUTE_PGM_RSRC2:TGID_X_EN: 1
+; ALL: COMPUTE_PGM_RSRC2:TGID_Y_EN: 1
+; ALL: COMPUTE_PGM_RSRC2:TGID_Z_EN: 0
+; ALL: COMPUTE_PGM_RSRC2:TIDIG_COMP_CNT: 0
+define amdgpu_kernel void @test_workgroup_id_y(i32 addrspace(1)* %out) #1 {
+  %id = call i32 @llvm.amdgcn.workgroup.id.y()
+  store i32 %id, i32 addrspace(1)* %out
+  ret void
+}
+
+; ALL-LABEL {{^}}test_workgroup_id_z:
+; CO-V2: user_sgpr_count = 6
+; CO-V2: enable_sgpr_workgroup_id_x = 1
+; CO-V2: enable_sgpr_workgroup_id_y = 0
+; CO-V2: enable_sgpr_workgroup_id_z = 1
+; CO-V2: enable_sgpr_workgroup_info = 0
+; CO-V2: enable_vgpr_workitem_id = 0
+; CO-V2: enable_sgpr_private_segment_buffer = 1
+; CO-V2: enable_sgpr_dispatch_ptr = 0
+; CO-V2: enable_sgpr_queue_ptr = 0
+; CO-V2: enable_sgpr_kernarg_segment_ptr = 1
+; CO-V2: enable_sgpr_dispatch_id = 0
+; CO-V2: enable_sgpr_flat_scratch_init = 0
+; CO-V2: enable_sgpr_private_segment_size = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_x = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_y = 0
+; CO-V2: enable_sgpr_grid_workgroup_count_z = 0
+
+; UNKNOWN-OS: v_mov_b32_e32 [[VCOPY:v[0-9]+]], s3{{$}}
+; HSA: v_mov_b32_e32 [[VCOPY:v[0-9]+]], s7{{$}}
+
+; ALL: {{buffer|flat}}_store_dword {{.*}}[[VCOPY]]
+
+; CO-V2: COMPUTE_PGM_RSRC2:USER_SGPR: 6
+; ALL-NOCO-V2: COMPUTE_PGM_RSRC2:USER_SGPR: 2
+; ALL: COMPUTE_PGM_RSRC2:TGID_X_EN: 1
+; ALL: COMPUTE_PGM_RSRC2:TGID_Y_EN: 0
+; ALL: COMPUTE_PGM_RSRC2:TGID_Z_EN: 1
+; ALL: COMPUTE_PGM_RSRC2:TIDIG_COMP_CNT: 0
+define amdgpu_kernel void @test_workgroup_id_z(i32 addrspace(1)* %out) #1 {
+  %id = call i32 @llvm.amdgcn.workgroup.id.z()
+  store i32 %id, i32 addrspace(1)* %out
+  ret void
+}
+
+attributes #0 = { nounwind readnone }
+attributes #1 = { nounwind }