forked from ccf-ai-infra/Intro-ops
52 lines
2.0 KiB
Python
52 lines
2.0 KiB
Python
from __future__ import annotations
|
||
|
||
import tilelang
|
||
import tilelang.language as T
|
||
|
||
|
||
@tilelang.jit(
|
||
# * 关闭 warp-specialized 相关优化 pass，避免编译器使用更复杂的 warp-specialization 调度
|
||
pass_configs={
|
||
tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
|
||
},
|
||
)
|
||
def reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int):
|
||
'''
|
||
Args:
|
||
- src: 输入 tensor
|
||
- BLOCK_N: int 类型，tile 处理的行数
|
||
- BLOCK_M: int 类型，tile 处理的列数
|
||
'''
|
||
N, M = T.const("N, M")
|
||
dtype = T.float32
|
||
src: T.Tensor((N, M), dtype)
|
||
out = T.empty((N,), dtype)
|
||
|
||
# TODO: implement a tiled row-wise reduce_sum kernel.
|
||
#
|
||
# Suggested steps:
|
||
# 1. Launch one TileLang kernel over row tiles.
|
||
# 2. Allocate fragments for the input tile and running output tile.
|
||
# 3. Clear the running output fragment before accumulation.
|
||
# 4. Iterate over column tiles with T.Serial(M // BLOCK_M).
|
||
# 5. T.copy each input tile into the fragment.
|
||
# 6. Call T.reduce_sum on the fragment and accumulate into the output fragment.
|
||
# 7. Copy the final output fragment back to global memory.
|
||
with T.Kernel(T.ceildiv(N, BLOCK_N), threads=256) as bx:
|
||
# fragment：当前 program 内的小块临时数据，通常更靠近寄存器/计算单元
|
||
x_frag = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) # 当前 kernel 的临时输入块
|
||
acc_frag = T.alloc_fragment((BLOCK_N,), dtype) # kernel 内部临时存储局部段的向量
|
||
partial_frag = T.alloc_fragment((BLOCK_N,), dtype)
|
||
T.clear(acc_frag)
|
||
|
||
row_base = bx * BLOCK_N
|
||
for k in T.Serial(T.ceildiv(M, BLOCK_M)):
|
||
col_base = k * BLOCK_M
|
||
T.copy(src[row_base : row_base + BLOCK_N, col_base : col_base + BLOCK_M], x_frag)
|
||
T.reduce_sum(x_frag, partial_frag, dim=1)
|
||
|
||
for i in T.Parallel(BLOCK_N):
|
||
acc_frag[i] += partial_frag[i]
|
||
T.copy(acc_frag, out[row_base : row_base + BLOCK_N])
|
||
return out
|