[llvm] [AMDGPU] Optimize block count calculations to the new ABI (PR #174112)

[Matt Arsenault via llvm-commits]

================
@@ -0,0 +1,219 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -mtriple=amdgcn-amd-amdhsa -S -passes=amdgpu-lower-kernel-attributes,instcombine,simplifycfg %s | FileCheck %s
+
+define i32 @num_blocks(i32 noundef %dim) {
+; CHECK-LABEL: define i32 @num_blocks(
+; CHECK-SAME: i32 noundef [[DIM:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*]]:
+; CHECK-NEXT:    [[TMP1:%.*]] = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+; CHECK-NEXT:    switch i32 [[DIM]], label %[[DEFAULT:.*]] [
+; CHECK-NEXT:      i32 0, label %[[EXIT:.*]]
+; CHECK-NEXT:      i32 1, label %[[DIM_Y:.*]]
+; CHECK-NEXT:      i32 2, label %[[DIM_Z:.*]]
+; CHECK-NEXT:    ]
+; CHECK:       [[DIM_Y]]:
+; CHECK-NEXT:    [[TMP2:%.*]] = getelementptr i8, ptr addrspace(4) [[TMP1]], i64 4
+; CHECK-NEXT:    br label %[[EXIT]]
+; CHECK:       [[DIM_Z]]:
+; CHECK-NEXT:    [[TMP3:%.*]] = getelementptr i8, ptr addrspace(4) [[TMP1]], i64 8
+; CHECK-NEXT:    br label %[[EXIT]]
+; CHECK:       [[DEFAULT]]:
+; CHECK-NEXT:    unreachable
+; CHECK:       [[EXIT]]:
+; CHECK-NEXT:    [[RETVAL_IN:%.*]] = phi ptr addrspace(4) [ [[TMP2]], %[[DIM_Y]] ], [ [[TMP3]], %[[DIM_Z]] ], [ [[TMP1]], %[[ENTRY]] ]
+; CHECK-NEXT:    [[RETVAL_0_I:%.*]] = load i32, ptr addrspace(4) [[RETVAL_IN]], align 4, !invariant.load [[META0:![0-9]+]], !noundef [[META0]]
+; CHECK-NEXT:    ret i32 [[RETVAL_0_I]]
+;
+entry:
+  %dispatch = call ptr addrspace(4) @llvm.amdgcn.dispatch.ptr()
+  %implicitarg = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+
+  switch i32 %dim, label %default [
+  i32 0, label %dim_x
+  i32 1, label %dim_y
+  i32 2, label %dim_z
+  ]
+
+dim_x:
+  %d_gep_x = getelementptr i8, ptr addrspace(4) %dispatch, i32 12
+  %grid_size_x = load i32, ptr addrspace(4) %d_gep_x, align 4
+  %i_gep_x = getelementptr i8, ptr addrspace(4) %implicitarg, i32 12
+  %wg_size_x = load i16, ptr addrspace(4) %i_gep_x, align 2
+  %conv_x = zext i16 %wg_size_x to i32
+  %count_x = udiv i32 %grid_size_x, %conv_x
+  br label %exit
+
+dim_y:
+  %d_gep_y = getelementptr i8, ptr addrspace(4) %dispatch, i32 16
+  %grid_size_y = load i32, ptr addrspace(4) %d_gep_y, align 4
+  %i_gep_y = getelementptr i8, ptr addrspace(4) %implicitarg, i32 14
+  %wg_size_y = load i16, ptr addrspace(4) %i_gep_y, align 2
+  %conv_y = zext i16 %wg_size_y to i32
+  %count_y = udiv i32 %grid_size_y, %conv_y
+  br label %exit
+
+dim_z:
+  %d_gep_z = getelementptr i8, ptr addrspace(4) %dispatch, i32 20
+  %grid_size_z = load i32, ptr addrspace(4) %d_gep_z, align 4
+  %i_gep_z = getelementptr i8, ptr addrspace(4) %implicitarg, i32 16
+  %wg_size_z = load i16, ptr addrspace(4) %i_gep_z, align 2
+  %conv_z = zext i16 %wg_size_z to i32
+  %count_z = udiv i32 %grid_size_z, %conv_z
+  br label %exit
+
+default:
+  unreachable
+
+exit:
+  %retval = phi i32 [ %count_x, %dim_x ], [ %count_y, %dim_y ], [ %count_z, %dim_z ]
+  ret i32 %retval
+}
+
+ at dim_const = constant i32 1, align 4
+
+define i32 @constant() {
+; CHECK-LABEL: define i32 @constant() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[IMPLICITARG:%.*]] = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+; CHECK-NEXT:    [[TMP0:%.*]] = getelementptr i8, ptr addrspace(4) [[IMPLICITARG]], i64 4
+; CHECK-NEXT:    [[TMP1:%.*]] = load i32, ptr addrspace(4) [[TMP0]], align 4, !invariant.load [[META0]], !noundef [[META0]]
+; CHECK-NEXT:    ret i32 [[TMP1]]
+;
+entry:
+  %dispatch = call ptr addrspace(4) @llvm.amdgcn.dispatch.ptr()
+  %implicitarg = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %dim = load i32, ptr @dim_const, align 4
+
+  switch i32 %dim, label %default [
+  i32 0, label %dim_x
+  i32 1, label %dim_y
+  i32 2, label %dim_z
+  ]
+
+dim_x:
+  %d_gep_x = getelementptr i8, ptr addrspace(4) %dispatch, i32 12
+  %grid_size_x = load i32, ptr addrspace(4) %d_gep_x, align 4
+  %i_gep_x = getelementptr i8, ptr addrspace(4) %implicitarg, i32 12
+  %wg_size_x = load i16, ptr addrspace(4) %i_gep_x, align 2
+  %conv_x = zext i16 %wg_size_x to i32
+  %count_x = udiv i32 %grid_size_x, %conv_x
+  br label %exit
+
+dim_y:
+  %d_gep_y = getelementptr i8, ptr addrspace(4) %dispatch, i32 16
+  %grid_size_y = load i32, ptr addrspace(4) %d_gep_y, align 4
+  %i_gep_y = getelementptr i8, ptr addrspace(4) %implicitarg, i32 14
+  %wg_size_y = load i16, ptr addrspace(4) %i_gep_y, align 2
+  %conv_y = zext i16 %wg_size_y to i32
+  %count_y = udiv i32 %grid_size_y, %conv_y
+  br label %exit
+
+dim_z:
+  %d_gep_z = getelementptr i8, ptr addrspace(4) %dispatch, i32 20
+  %grid_size_z = load i32, ptr addrspace(4) %d_gep_z, align 4
+  %i_gep_z = getelementptr i8, ptr addrspace(4) %implicitarg, i32 16
+  %wg_size_z = load i16, ptr addrspace(4) %i_gep_z, align 2
+  %conv_z = zext i16 %wg_size_z to i32
+  %count_z = udiv i32 %grid_size_z, %conv_z
+  br label %exit
+
+default:
+  unreachable
+
+exit:
+  %retval = phi i32 [ %count_x, %dim_x ], [ %count_y, %dim_y ], [ %count_z, %dim_z ]
+  ret i32 %retval
+}
+
+define i64 @larger() {
+; CHECK-LABEL: define i64 @larger() {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[IMPLICITARG:%.*]] = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+; CHECK-NEXT:    [[GRID_SIZE_X:%.*]] = load i32, ptr addrspace(4) [[IMPLICITARG]], align 4, !invariant.load [[META0]], !noundef [[META0]]
+; CHECK-NEXT:    [[CONV_GRID_X:%.*]] = zext i32 [[GRID_SIZE_X]] to i64
+; CHECK-NEXT:    ret i64 [[CONV_GRID_X]]
+;
+entry:
+  %dispatch = call ptr addrspace(4) @llvm.amdgcn.dispatch.ptr()
+  %d_gep_x = getelementptr i8, ptr addrspace(4) %dispatch, i32 12
+  %grid_size_x = load i32, ptr addrspace(4) %d_gep_x, align 4
+  %implicitarg = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %i_gep_x = getelementptr i8, ptr addrspace(4) %implicitarg, i32 12
+  %wg_size_x = load i16, ptr addrspace(4) %i_gep_x, align 2
----------------
arsenm wrote:

Check load of wrong type and wrong cast type 

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