[llvm] [AMDGPU] Mark workitem IDs uniform in more cases (PR #152581)

[Krzysztof Drewniak via llvm-commits]

================
@@ -113,11 +113,40 @@ define amdgpu_kernel void @workitem_id_x_not_singlethreaded_dimz() !reqd_work_gr
   ret void
 }
 
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_z_uniform_len_1'
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_z_uniform_len_1(ptr %o) !reqd_work_group_size !4 {
+  %id.z = call i32 @llvm.amdgcn.workitem.id.z()
+  store i32 %id.z, ptr %o
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_div_wavefront_size'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_div_wavefront_size(ptr %o) #3 !reqd_work_group_size !5 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  %id.sg = lshr i32 %id.x, 6
+  store i32 %id.sg, ptr %o
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_y_uniform_in_subgroup'
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_y_uniform_in_subgroup(ptr %o) #3 !reqd_work_group_size !5 {
+  %id.y = call i32 @llvm.amdgcn.workitem.id.y()
----------------
krzysz00 wrote:

The required workgroup size is `!5 = !{i32 128, i32 2, i32 1}`

This means that that, within a wave, the VGPR for the X id is either vX = [0, 1, ... 63] or vX = [64, ..., 127]

Note that the Y ID _cannot_ change within a wave. All lanes of wave 0 have tid.y = 0, same for wave 1. All lanes of wave 2 have tid.Y = 1, same for wave 3.

Therefore, the Y ID is uniform (and can be either 0 or 1)

https://github.com/llvm/llvm-project/pull/152581