forked from ccf-ai-infra/Intro-ops
27 lines
782 B
Python
27 lines
782 B
Python
from __future__ import annotations
|
|
|
|
import tilelang
|
|
import tilelang.language as T
|
|
|
|
|
|
@tilelang.jit
|
|
def vector_add_kernel(a, b, BLOCK_N: int, dtype):
|
|
N = T.const("N")
|
|
a: T.Tensor((N,), dtype)
|
|
b: T.Tensor((N,), dtype)
|
|
out = T.empty((N,), dtype)
|
|
|
|
# TODO: implement a tile-wise vector add kernel.
|
|
#
|
|
# Suggested steps:
|
|
# 1. Launch one TileLang kernel over the N // BLOCK_N tiles.
|
|
# 2. Compute the tile base offset.
|
|
# 3. Use T.Parallel(BLOCK_N) to fill out[base + i] = a[base + i] + b[base + i].
|
|
with T.Kernel(T.ceildiv(N, BLOCK_N), threads=128) as bx:
|
|
base = bx * BLOCK_N
|
|
for i in T.Parallel(BLOCK_N):
|
|
idx = base + i
|
|
if (idx < N):
|
|
out[base + i] = a[base + i] + b[base + i]
|
|
return out
|