[llvm] [AMDGPU] Add more cases to the "this is a wave ID" recognizer (PR #177713)

[Matt Arsenault via llvm-commits]

================
@@ -131,6 +131,54 @@ define amdgpu_kernel void @workitem_id_x_div_wavefront_size(ptr %o) #3 !reqd_wor
   ret void
 }
 
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_masked'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_masked(ptr %o) #3 !reqd_work_group_size !5 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  %id.sg.shl.6 = and i32 %id.x, 192
+  store i32 %id.sg.shl.6, ptr %o
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_zext_masked'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK: DIVERGENT: %id.x.zext = zext i32 %id.x to i64
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_zext_masked(ptr %o) #3 !reqd_work_group_size !5 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  %id.x.zext = zext i32 %id.x to i64
+  %id.sg.shl.6 = and i64 %id.x.zext, 192
+  store i64 %id.sg.shl.6, ptr %o
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_div_wavefront_size_masked'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK: DIVERGENT: %id.x.masked = and i32 %id.x, 127
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_div_wavefront_size_masked(ptr %o) #3 !reqd_work_group_size !5 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  %id.x.masked = and i32 %id.x, 127
+  %id.sg = lshr i32 %id.x.masked, 6
+  store i32 %id.sg, ptr %o
+  ret void
+}
+
+; CHECK-LABEL: UniformityInfo for function 'workitem_id_x_div_wavefront_size_trunc_masked'
+; CHECK: DIVERGENT: %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+; CHECK: DIVERGENT: %id.x.trunc = trunc i32 %id.x to i16
+; CHECK: DIVERGENT: %id.x.masked = and i16 %id.x.trunc, 127
+; CHECK-NOT: DIVERGENT
+define amdgpu_kernel void @workitem_id_x_div_wavefront_size_trunc_masked(ptr %o) #3 !reqd_work_group_size !5 {
+  %id.x = call i32 @llvm.amdgcn.workitem.id.x()
+  %id.x.trunc = trunc i32 %id.x to i16
+  %id.x.masked = and i16 %id.x.trunc, 127
+  %id.sg = lshr i16 %id.x.masked, 6
+  store i16 %id.sg, ptr %o
----------------
arsenm wrote:

Instcombine also changes this pattern, make sure you're only handling canonical IR 

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