Fix pid remapping logic when GRID_MN cannot divide NUM_XCDS (#722)

Author: zhanglx13

python/perf-kernels/tools/occ.sh

@@ -39,8 +39,8 @@ get_occ_per_CU() {
 
 $1 > output.mlir 2>&1
 
-LDS_line=$(sed -n '/triton_gpu\.shared\ /p' output.mlir | tail -n 1 | grep -o 'triton_gpu.shared = [0-9]*')
-numWarps_line=$(sed -n '/triton_gpu\.num-warps/p' output.mlir | tail -n 1 | grep -o 'triton_gpu.num-warps. = [0-9]*')
+LDS_line=$(sed -n '/ttg\.shared\ /p' output.mlir | tail -n 1 | grep -o 'ttg.shared = [0-9]*')
+numWarps_line=$(sed -n '/ttg\.num-warps/p' output.mlir | tail -n 1 | grep -o 'ttg.num-warps. = [0-9]*')
 
 LDS=${LDS_line##*=}
 num_warps=${numWarps_line##*=}

python/perf-kernels/tools/tune_gemm/matmul_kernel.py

@@ -25,11 +25,23 @@ def matmul_kernel(a_ptr, b_ptr, c_ptr, bias_ptr, M, N, K, stride_am, stride_ak,
         ## pid remapping on xcds
         # Number of pids per XCD in the new arrangement
         pids_per_xcd = (GRID_MN + NUM_XCDS - 1) // NUM_XCDS
+        # When GRID_MN cannot divide NUM_XCDS, some xcds will have
+        # pids_per_xcd pids, the other will have pids_per_xcd - 1 pids.
+        # We calculate the number of xcds that have pids_per_xcd pids as
+        # tall_xcds
+        tall_xcds = GRID_MN % NUM_XCDS
+        tall_xcds = NUM_XCDS if tall_xcds == 0 else tall_xcds
         # Compute current XCD and local pid within the XCD
         xcd = pid % NUM_XCDS
         local_pid = pid // NUM_XCDS
         # Calculate new pid based on the new grouping
-        pid = xcd * pids_per_xcd + local_pid
+        # Note that we need to consider the following two cases:
+        # 1. the currnt pid is on a tall xcd
+        # 2. the current pid is on a short xcd
+        if xcd < tall_xcds:
+            pid = xcd * pids_per_xcd + local_pid
+        else:
+            pid = tall_xcds * pids_per_xcd + (xcd - tall_xcds) * (pids_per_xcd - 1) + local_pid
 
     if GROUP_SIZE_M == 1:
         pid_m = pid // num_pid_n