wehub-resource-sync
196 lines
8.1 KiB
Python
196 lines
8.1 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# pylint: disable=invalid-name, too-many-nested-blocks
|
|
"""Backend kernels for cumsum operator."""
|
|
|
|
import math
|
|
|
|
from tvm.script import tirx as T
|
|
from tvm.tirx import PrimFunc
|
|
|
|
|
|
def _is_power_of_two(n: int):
|
|
"""Check if n is a power of 2."""
|
|
return n > 0 and (n & (n - 1)) == 0
|
|
|
|
|
|
def gpu_2d_continuous_cumsum(
|
|
ty_len: int = 4,
|
|
tx_len: int = 32,
|
|
thread_elem: int = 4,
|
|
in_dtype: str = "int32",
|
|
out_dtype: str | None = None,
|
|
) -> PrimFunc:
|
|
"""Generate GPU kernel for 2D continuous cumsum, i.e. The cumsum axis is -1
|
|
|
|
Parameters
|
|
----------
|
|
ty_len : int
|
|
The length of `threadIdx.y`
|
|
|
|
tx_len : int
|
|
The length of `threadIdx.x`
|
|
|
|
thread_elem : int
|
|
The number of elements processed by single thread
|
|
|
|
in_dtype : str
|
|
The input data type
|
|
|
|
out_dtype : Optional[str]
|
|
The output data type, if None, it will be the same as in_dtype
|
|
|
|
Returns
|
|
-------
|
|
cumsum : PrimFunc
|
|
The generated cumsum kernel
|
|
"""
|
|
|
|
out_dtype = out_dtype or in_dtype
|
|
|
|
# Configuration for GPU kernel
|
|
TX = T.int64(tx_len) # threadIdx.x
|
|
TY = T.int64(ty_len) # threadIdx.y
|
|
N = T.int64(thread_elem) # number of elements in single thread
|
|
|
|
if not _is_power_of_two(TX) or not _is_power_of_two(TY) or not _is_power_of_two(N):
|
|
raise ValueError("Configuration of TX, TY, N must be power of 2")
|
|
|
|
# number of elements to be processed by single warp
|
|
warp_elem = T.int64(tx_len * thread_elem)
|
|
# number of elements to be processed by single block(SM)
|
|
block_elem = T.int64(tx_len * ty_len * thread_elem)
|
|
|
|
LOG_TX = T.int64(int(math.log2(tx_len)))
|
|
LOG_BLOCK_N = T.int64(int(math.log2(tx_len * ty_len * thread_elem)))
|
|
|
|
@T.macro
|
|
def block_inclusive_inside_block(
|
|
batch: T.int64,
|
|
cur_len: T.int64,
|
|
source: T.Buffer,
|
|
output: T.Buffer,
|
|
tmp_buf: T.Buffer,
|
|
src_offset: T.int64,
|
|
tmp_offset: T.int64,
|
|
):
|
|
for by in T.thread_binding(batch, thread="blockIdx.y"):
|
|
for bx in T.thread_binding(T.ceildiv(cur_len, block_elem), thread="blockIdx.x"):
|
|
with T.sblock():
|
|
local_buf = T.sblock_alloc_buffer((thread_elem,), out_dtype, scope="local")
|
|
shared_buf = T.sblock_alloc_buffer((block_elem,), out_dtype, scope="shared")
|
|
for ty in T.thread_binding(TY, thread="threadIdx.y"):
|
|
for tx in T.thread_binding(TX, thread="threadIdx.x"):
|
|
tx_idx: T.let[T.int64] = (
|
|
bx * block_elem + ty * warp_elem + tx * thread_elem
|
|
)
|
|
# Load data from global memory
|
|
for i in T.vectorized(N):
|
|
local_buf[i] = T.if_then_else(
|
|
tx_idx + i < cur_len,
|
|
T.Cast(out_dtype, source[by, src_offset + tx_idx + i]),
|
|
T.Cast(out_dtype, 0),
|
|
)
|
|
# Inclusive scan inside thread
|
|
for i in T.unroll(1, N):
|
|
local_buf[i] += local_buf[i - 1]
|
|
# Store data to shared memory
|
|
for i in T.vectorized(N):
|
|
shared_buf[ty * warp_elem + tx * thread_elem + i] = local_buf[i]
|
|
# Inclusive scan inside warp
|
|
for i in T.unroll(LOG_TX):
|
|
for j in T.vectorized(N):
|
|
idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem
|
|
if tx >= (1 << i):
|
|
shared_buf[idx + j] += shared_buf[
|
|
idx - (1 << i) * thread_elem + N - 1
|
|
]
|
|
# Inclusive scan inside block
|
|
for i in T.unroll(1, TY):
|
|
for j in T.vectorized(N):
|
|
if ty == 0:
|
|
idx: T.let[T.int64] = i * warp_elem + tx * thread_elem
|
|
shared_buf[idx + j] += shared_buf[i * warp_elem - 1]
|
|
# Write sum of block to global memory
|
|
for i in T.vectorized(N):
|
|
idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem + i
|
|
if bx * block_elem + idx < cur_len:
|
|
output[by, src_offset + bx * block_elem + idx] = shared_buf[idx]
|
|
if tx == 0 and ty == 0:
|
|
for i in T.vectorized(N):
|
|
tmp_buf[by, tmp_offset + bx] = shared_buf[block_elem - 1]
|
|
|
|
@T.macro
|
|
def update_cross_block(
|
|
batch: T.int64,
|
|
cur_len: T.int64,
|
|
source: T.Buffer,
|
|
output: T.Buffer,
|
|
src_offset: T.int64,
|
|
out_offset: T.int64,
|
|
):
|
|
for by in T.thread_binding(batch, thread="blockIdx.y"):
|
|
for bx in T.thread_binding(T.ceildiv(cur_len, block_elem), thread="blockIdx.x"):
|
|
for ty in T.thread_binding(TY, thread="threadIdx.y"):
|
|
for tx in T.thread_binding(TX, thread="threadIdx.x"):
|
|
for i in T.serial(N):
|
|
idx: T.let[T.int64] = bx * block_elem + ty * warp_elem + i * TX + tx
|
|
if idx < cur_len:
|
|
output[by, out_offset + idx] += T.if_then_else(
|
|
bx > 0, source[by, src_offset + bx - 1], 0
|
|
)
|
|
|
|
@T.prim_func(private=True, s_tir=True)
|
|
def cumsum(var_a: T.handle, var_out: T.handle):
|
|
T.func_attr({"tirx.is_scheduled": True}) # prevent further scheduling
|
|
m, n = T.int64(), T.int64()
|
|
A = T.match_buffer(var_a, [m, n], dtype=in_dtype)
|
|
Out = T.match_buffer(var_out, [m, n], dtype=out_dtype)
|
|
Tmp = T.alloc_buffer([m, n], dtype=out_dtype)
|
|
total_rounds: T.let[T.int64] = (
|
|
T.Cast("int64", T.ceil(T.log2(T.Cast("float32", n)))) // LOG_BLOCK_N
|
|
)
|
|
|
|
block_inclusive_inside_block(
|
|
m, n, A, Out, Tmp, src_offset=T.int64(0), tmp_offset=T.int64(0)
|
|
)
|
|
for i in range(total_rounds):
|
|
cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (i + 1)))
|
|
block_inclusive_inside_block(
|
|
m,
|
|
cur_len,
|
|
Tmp,
|
|
Tmp,
|
|
Tmp,
|
|
src_offset=i * T.ceildiv(n, block_elem),
|
|
tmp_offset=(i + 1) * T.ceildiv(n, block_elem),
|
|
)
|
|
for i in range(total_rounds - 1):
|
|
real_idx: T.let[T.int64] = total_rounds - 1 - i - 1
|
|
cur_len: T.let[T.int64] = T.ceildiv(n, 1 << (LOG_BLOCK_N * (real_idx + 1)))
|
|
update_cross_block(
|
|
m,
|
|
cur_len,
|
|
Tmp,
|
|
Tmp,
|
|
src_offset=(real_idx + 1) * T.ceildiv(n, block_elem),
|
|
out_offset=real_idx * T.ceildiv(n, block_elem),
|
|
)
|
|
update_cross_block(m, n, Tmp, Out, src_offset=0, out_offset=0)
|
|
|
|
return cumsum
|