Kernels:

kernels-community
/

aiter-kernels

Trusted publisher

Kernel card Files Files and versions

aiter-kernels / build /torch-rocm /utils /_triton /pid_preprocessing.py

kernels-bot's picture

Uploaded using `kernel-builder`.

89181fc verified 25 days ago

3.16 kB

	# 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