diff --git a/ops/copy/tilelang/copy_tl.py b/ops/copy/tilelang/copy_tl.py index bcd9c78..a10355f 100644 --- a/ops/copy/tilelang/copy_tl.py +++ b/ops/copy/tilelang/copy_tl.py @@ -3,10 +3,11 @@ from __future__ import annotations from functools import lru_cache from dataclasses import dataclass -import tilelang import tilelang.language as T import torch +from ops.copy.tilelang.kernel import copy_kernel + def _tl_dtype(dtype: torch.dtype): if dtype is torch.float16: @@ -15,29 +16,13 @@ def _tl_dtype(dtype: torch.dtype): return T.float32 raise TypeError(f"unsupported TileLang dtype: {dtype}") - -@tilelang.jit -def _copy_kernel(src, BLOCK_N: int, dtype): - N = T.const("N") - src: T.Tensor((N,), dtype) - out = T.empty((N,), dtype) - - with T.Kernel(N // BLOCK_N, threads=256) as pid_n: - T.copy( - src[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N], - out[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N], - ) - - return out - - def _block_n(n: int) -> int: return 1024 if n % 1024 == 0 else n @lru_cache(maxsize=32) def _compiled_copy(n: int, block_n: int, dtype: torch.dtype): - return _copy_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype)) + return copy_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype)) @dataclass diff --git a/ops/copy/tilelang/kernel.py b/ops/copy/tilelang/kernel.py new file mode 100644 index 0000000..38e1333 --- /dev/null +++ b/ops/copy/tilelang/kernel.py @@ -0,0 +1,19 @@ +from __future__ import annotations + +import tilelang +import tilelang.language as T + + +@tilelang.jit +def copy_kernel(src, BLOCK_N: int, dtype): + N = T.const("N") + src: T.Tensor((N,), dtype) + out = T.empty((N,), dtype) + + with T.Kernel(N // BLOCK_N, threads=256) as pid_n: + T.copy( + src[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N], + out[pid_n * BLOCK_N : (pid_n + 1) * BLOCK_N], + ) + + return out diff --git a/ops/reduce_sum/tilelang/kernel.py b/ops/reduce_sum/tilelang/kernel.py new file mode 100644 index 0000000..d53c1c8 --- /dev/null +++ b/ops/reduce_sum/tilelang/kernel.py @@ -0,0 +1,30 @@ +from __future__ import annotations + +import tilelang +import tilelang.language as T + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + }, +) +def reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int): + N, M = T.const("N, M") + dtype = T.float32 + src: T.Tensor((N, M), dtype) + out = T.empty((N,), dtype) + + with T.Kernel(N // BLOCK_N, threads=256) as pid_n: + src_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) + out_local = T.alloc_fragment((BLOCK_N,), dtype) + T.clear(out_local) + + for m_blk in T.Serial(M // BLOCK_M): + T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local) + T.reduce_sum(src_local, out_local, dim=1, clear=False) + + T.copy(out_local, out[pid_n * BLOCK_N]) + + return out diff --git a/ops/reduce_sum/tilelang/reduce_sum_tl.py b/ops/reduce_sum/tilelang/reduce_sum_tl.py index eb8e796..8b75aa5 100644 --- a/ops/reduce_sum/tilelang/reduce_sum_tl.py +++ b/ops/reduce_sum/tilelang/reduce_sum_tl.py @@ -3,40 +3,13 @@ from __future__ import annotations from dataclasses import dataclass from functools import lru_cache -import tilelang -import tilelang.language as T import torch - -@tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - }, -) -def _reduce_sum_kernel(src, BLOCK_N: int, BLOCK_M: int): - N, M = T.const("N, M") - dtype = T.float32 - src: T.Tensor((N, M), dtype) - out = T.empty((N,), dtype) - - with T.Kernel(N // BLOCK_N, threads=256) as pid_n: - src_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) - out_local = T.alloc_fragment((BLOCK_N,), dtype) - T.clear(out_local) - - for m_blk in T.Serial(M // BLOCK_M): - T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local) - T.reduce_sum(src_local, out_local, dim=1, clear=False) - - T.copy(out_local, out[pid_n * BLOCK_N]) - - return out - +from ops.reduce_sum.tilelang.kernel import reduce_sum_kernel @lru_cache(maxsize=32) def _compiled_reduce_sum(n: int, m: int, block_n: int, block_m: int): - return _reduce_sum_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m) + return reduce_sum_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m) def _blocks(n: int, m: int) -> tuple[int, int]: diff --git a/ops/softmax/tilelang/kernel.py b/ops/softmax/tilelang/kernel.py new file mode 100644 index 0000000..fc280aa --- /dev/null +++ b/ops/softmax/tilelang/kernel.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +import tilelang +import tilelang.language as T + + +@tilelang.jit( + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + }, +) +def softmax_kernel(src, BLOCK_N: int, BLOCK_M: int): + log2_e = 1.44269504 + N, M = T.const("N, M") + dtype = T.float32 + src: T.Tensor((N, M), dtype) + out = T.empty((N, M), dtype) + + with T.Kernel(N // BLOCK_N, threads=256) as pid_n: + src_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) + out_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) + cur_exp = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) + cur_max = T.alloc_fragment((BLOCK_N,), dtype) + cur_sum = T.alloc_fragment((BLOCK_N,), dtype) + lse = T.alloc_fragment((BLOCK_N,), dtype) + + T.fill(lse, -T.infinity(dtype)) + + for m_blk in T.Serial(M // BLOCK_M): + T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local) + T.reduce_max(src_local, cur_max, dim=1, clear=True) + + for i, j in T.Parallel(BLOCK_N, BLOCK_M): + cur_exp[i, j] = T.exp2(src_local[i, j] * log2_e - cur_max[i] * log2_e) + + T.reduce_sum(cur_exp, cur_sum, dim=1, clear=True) + + for i in T.Parallel(BLOCK_N): + lse[i] = cur_max[i] * log2_e + T.log2( + T.exp2(lse[i] - cur_max[i] * log2_e) + cur_sum[i] + ) + + for m_blk in T.Serial(M // BLOCK_M): + T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local) + for i, j in T.Parallel(BLOCK_N, BLOCK_M): + out_local[i, j] = T.exp2(src_local[i, j] * log2_e - lse[i]) + T.copy(out_local, out[pid_n * BLOCK_N, m_blk * BLOCK_M]) + + return out diff --git a/ops/softmax/tilelang/softmax_tl.py b/ops/softmax/tilelang/softmax_tl.py index 3a347c6..6c07ad0 100644 --- a/ops/softmax/tilelang/softmax_tl.py +++ b/ops/softmax/tilelang/softmax_tl.py @@ -3,60 +3,13 @@ from __future__ import annotations from dataclasses import dataclass from functools import lru_cache -import tilelang -import tilelang.language as T import torch - -@tilelang.jit( - pass_configs={ - tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, - tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, - }, -) -def _softmax_kernel(src, BLOCK_N: int, BLOCK_M: int): - log2_e = 1.44269504 - N, M = T.const("N, M") - dtype = T.float32 - src: T.Tensor((N, M), dtype) - out = T.empty((N, M), dtype) - - with T.Kernel(N // BLOCK_N, threads=256) as pid_n: - src_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) - out_local = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) - cur_exp = T.alloc_fragment((BLOCK_N, BLOCK_M), dtype) - cur_max = T.alloc_fragment((BLOCK_N,), dtype) - cur_sum = T.alloc_fragment((BLOCK_N,), dtype) - lse = T.alloc_fragment((BLOCK_N,), dtype) - - T.fill(lse, -T.infinity(dtype)) - - for m_blk in T.Serial(M // BLOCK_M): - T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local) - T.reduce_max(src_local, cur_max, dim=1, clear=True) - - for i, j in T.Parallel(BLOCK_N, BLOCK_M): - cur_exp[i, j] = T.exp2(src_local[i, j] * log2_e - cur_max[i] * log2_e) - - T.reduce_sum(cur_exp, cur_sum, dim=1, clear=True) - - for i in T.Parallel(BLOCK_N): - lse[i] = cur_max[i] * log2_e + T.log2( - T.exp2(lse[i] - cur_max[i] * log2_e) + cur_sum[i] - ) - - for m_blk in T.Serial(M // BLOCK_M): - T.copy(src[pid_n * BLOCK_N, m_blk * BLOCK_M], src_local) - for i, j in T.Parallel(BLOCK_N, BLOCK_M): - out_local[i, j] = T.exp2(src_local[i, j] * log2_e - lse[i]) - T.copy(out_local, out[pid_n * BLOCK_N, m_blk * BLOCK_M]) - - return out - +from ops.softmax.tilelang.kernel import softmax_kernel @lru_cache(maxsize=32) def _compiled_softmax(n: int, m: int, block_n: int, block_m: int): - return _softmax_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m) + return softmax_kernel.compile(N=n, M=m, BLOCK_N=block_n, BLOCK_M=block_m) def _blocks(n: int, m: int) -> tuple[int, int]: diff --git a/ops/vector_add/tilelang/kernel.py b/ops/vector_add/tilelang/kernel.py new file mode 100644 index 0000000..3fb9fd4 --- /dev/null +++ b/ops/vector_add/tilelang/kernel.py @@ -0,0 +1,19 @@ +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) + + with T.Kernel(N // BLOCK_N, threads=256) as pid_n: + base = pid_n * BLOCK_N + for i in T.Parallel(BLOCK_N): + out[base + i] = a[base + i] + b[base + i] + + return out diff --git a/ops/vector_add/tilelang/vector_add_tl.py b/ops/vector_add/tilelang/vector_add_tl.py index 272798f..0d9442b 100644 --- a/ops/vector_add/tilelang/vector_add_tl.py +++ b/ops/vector_add/tilelang/vector_add_tl.py @@ -3,10 +3,11 @@ from __future__ import annotations from dataclasses import dataclass from functools import lru_cache -import tilelang import tilelang.language as T import torch +from ops.vector_add.tilelang.kernel import vector_add_kernel + def _tl_dtype(dtype: torch.dtype): if dtype is torch.float16: @@ -15,29 +16,13 @@ def _tl_dtype(dtype: torch.dtype): return T.float32 raise TypeError(f"unsupported TileLang dtype: {dtype}") - -@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) - - with T.Kernel(N // BLOCK_N, threads=256) as pid_n: - base = pid_n * BLOCK_N - for i in T.Parallel(BLOCK_N): - out[base + i] = a[base + i] + b[base + i] - - return out - - def _block_n(n: int) -> int: return 1024 if n % 1024 == 0 else n @lru_cache(maxsize=32) def _compiled_vector_add(n: int, block_n: int, dtype: torch.dtype): - return _vector_add_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype)) + return vector_add_kernel.compile(N=n, BLOCK_N=block_n, dtype=_tl_dtype(dtype)) @dataclass