mirror of
https://github.com/ROCm/ROCm.git
synced 2026-02-21 03:00:39 -05:00
8351f49fc7
* Add gemm tuning script v3 * Introduce --jobs to control the number of files to generate * Switch to trans convention used by Tensile * Rerun rocprof if it crashes * update README * Remove peak perf and efficiency
53 lines
2.1 KiB
Python
53 lines
2.1 KiB
Python
import triton
|
|
import triton.language as tl
|
|
|
|
@triton.jit
|
|
def matmul_kernel(
|
|
a_ptr, b_ptr, c_ptr,
|
|
M, N, K,
|
|
stride_am, stride_ak,
|
|
stride_bk, stride_bn,
|
|
stride_cm, stride_cn,
|
|
BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr,
|
|
SPLIT_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr,
|
|
):
|
|
pid = tl.program_id(axis=0)
|
|
pid_z = tl.program_id(1)
|
|
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
|
|
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
|
|
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)
|
|
pid_m = first_pid_m + (pid % group_size_m)
|
|
pid_n = (pid % num_pid_in_group) // group_size_m
|
|
if SPLIT_K == 1:
|
|
offs_k = tl.arange(0, BLOCK_SIZE_K)
|
|
else:
|
|
offs_k = pid_z * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K)
|
|
offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M))
|
|
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N))
|
|
a_ptrs = a_ptr + offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak
|
|
b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn
|
|
acc_dtype = tl.float32 if c_ptr.type.element_ty != tl.int8 else tl.int32
|
|
accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=acc_dtype)
|
|
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K * SPLIT_K)):
|
|
a = tl.load(a_ptrs)
|
|
b = tl.load(b_ptrs)
|
|
accumulator += tl.dot(a, b)
|
|
a_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_ak
|
|
b_ptrs += BLOCK_SIZE_K * SPLIT_K * stride_bk
|
|
c = accumulator.to(c_ptr.type.element_ty)
|
|
offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
|
|
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
|
|
c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :]
|
|
c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)
|
|
if SPLIT_K == 1:
|
|
tl.store(c_ptrs, c, mask=c_mask)
|
|
else:
|
|
tl.atomic_add(c_ptrs, c, mask=c_mask)
|