File size: 3,156 Bytes
2976eec | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | # SPDX-License-Identifier: MIT
# Copyright (C) 2024-2026, Advanced Micro Devices, Inc. All rights reserved.
import triton
import triton.language as tl
@triton.jit
def remap_xcd_chunked(
pid, GRID_MN, NUM_XCDS: tl.constexpr = 8, CHUNK_SIZE: tl.constexpr = 2
):
# Compute current XCD and local PID
xcd = pid % NUM_XCDS
# distribute the modulo pids in round robin
if pid > (GRID_MN // (NUM_XCDS * CHUNK_SIZE)) * (NUM_XCDS * CHUNK_SIZE):
return pid
local_pid = pid // NUM_XCDS
# Calculate chunk index and position within chunk
chunk_idx = local_pid // CHUNK_SIZE
pos_in_chunk = local_pid % CHUNK_SIZE
# Calculate new PID
new_pid = chunk_idx * NUM_XCDS * CHUNK_SIZE + xcd * CHUNK_SIZE + pos_in_chunk
return new_pid
@triton.jit
def remap_xcd(pid, GRID_MN, NUM_XCDS: tl.constexpr = 8):
## 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
if tall_xcds == 0:
tall_xcds = tl.cast(NUM_XCDS, tall_xcds.type)
# 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
# Note that we need to consider the following two cases:
# 1. the current 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
)
return pid
@triton.jit
def pid_grid(pid: int, num_pid_m: int, num_pid_n: int, GROUP_SIZE_M: tl.constexpr = 1):
"""
Maps 1D pid to 2D grid coords (pid_m, pid_n).
Args:
- pid: 1D pid
- num_pid_m: grid m size
- num_pid_n: grid n size
- GROUP_SIZE_M: tl.constexpr: default is 1
"""
if GROUP_SIZE_M == 1:
pid_m = pid // num_pid_n
pid_n = pid % num_pid_n
else:
num_pid_in_group = GROUP_SIZE_M * num_pid_n
group_id = pid // num_pid_in_group
first_pid_m = group_id * GROUP_SIZE_M
group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
tl.assume(group_size_m >= 0)
pid_m = first_pid_m + (pid % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m
return pid_m, pid_n
@triton.jit
def pid_grid_3d(pid: int, num_pid_m: int, num_pid_n: int, num_pid_k):
"""
Maps 1D pid to 3D grid coords (pid_m, pid_n, pid_k).
Args:
- pid: 1D pid
- num_pid_m: grid m size
- num_pid_n: grid n size
- num_pid_k: grid k size
Returns:
- pid_m, pid_n, pid_k: 3D grid coordinates
"""
pid_m = pid % num_pid_m
pid_n = (pid // num_pid_m) % num_pid_n
pid_k = pid // (num_pid_m * num_pid_n) % num_pid_k
return pid_m, pid_n, pid_k
|