Triton 矩阵乘法

对于矩阵相乘： , 矩阵大小 

考虑简单情况 M 能够被 BLOCK_SIZE_M = 3 整除，K 能够被 BLOCK_SIZE_K = 3 整除，N 能够被 BLOCK_SIZE_N = 3 整除，伪代码如下：

for m in range(0,M,BLOCK_SIZE_M):       
	for n in range(0,N,BLOCK_SIZE_N):   
    	acc = zeros(BLOCK_SIZE_M,BLOCK_SIZE_N,dtype=float32)
    	for k in range(0,K,BLOCK_SIZE_K):
		    a = A[m:m+BLOCK_SIZE_M,k:k+BLOCK_SIZE_K] # Triton 改写
		    b = B[k:k+BLOCK_SIZE_K,n:n+BLOCK_SIZE_N]
		    acc += matmul(a,b)
	    C[m:m+BLOCK_SIZE_M,n:n+BLOCK_SIZE_N] = acc

Triton 可以很好的把上面的代码进行替换

选 C 矩阵中的 (m,n) 块

两个 for 循环，表示 C 矩阵的分块数目，块与块都可以并行，一共 (M // BLOCK_SIZE_M ) x  (N // BLOCK_SIZE_N) 个，官方代码用一唯 pid 表示具体某个块 (m,n)，pid 取值范围 0 - num_pid_m x num_pid_n

for m in range(0,M,BLOCK_SIZE_M):       
	for n in range(0,N,BLOCK_SIZE_N):   

# 用 Triton 替换
pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) 
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) # pid 取值 0 - num_pid_m x num_pid_n 

从 A 和 B 矩阵取两块数据相乘得到块 (m,n)

如图1所示，现在需要从 A、B 矩阵取红色框区域的数据，根据 pid 计算所属行列 (pid_m,pid_n) = (m,n)。代码很抽象，把 C 矩阵中每一整块 BLOCK_SIZE_M x N 大小当成一个 group，这么搞的目的是为了提高 L2 Cache 命中率，但是没必要再弄出来一个 group 的概念，国外的也没说清楚

pid = tl.program_id(axis=0)
num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
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 % num_pid_in_group) % group_size_m)
pid_n = (pid % num_pid_in_group) // group_size_m

根据 (pid_m,pid_n) 计算地址

offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M
offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N
offs_k = tl.arange(0, BLOCK_SIZE_K)
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)

计算、保存数据

offs_cm[:, None] 跟 torch.Tensor 的用法一样的，Triton 不好调试，可以用 torch 代替

accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
    # Load the next block of A and B, generate a mask by checking the K dimension.
    # If it is out of bounds, set it to 0.
    a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
    b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)
    # We accumulate along the K dimension.
    accumulator = tl.dot(a, b, accumulator)
    # Advance the ptrs to the next K block.
    a_ptrs += BLOCK_SIZE_K * stride_ak
    b_ptrs += BLOCK_SIZE_K * stride_bk


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)
tl.store(c_ptrs, c, mask=c_mask)

参考：

https://isamu-website.medium.com/understanding-the-triton-tutorials-part-1-6191b59ba4c