chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,195 @@
|
||||
# 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
|
||||
Reference in New Issue
Block a user