chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,28 @@
|
||||
# isort: skip_file
|
||||
# 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.
|
||||
"""The Relax Metal backend compilation pipeline and other passes."""
|
||||
|
||||
from .cumsum import gpu_2d_continuous_cumsum
|
||||
from .pipeline import (
|
||||
dataflow_lower_passes,
|
||||
finalize_passes,
|
||||
get_default_pipeline,
|
||||
legalize_passes,
|
||||
library_dispatch_passes,
|
||||
)
|
||||
from .sampling import generic_get_sample_index, gpu_multinomial_from_uniform
|
||||
@@ -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
|
||||
@@ -0,0 +1,89 @@
|
||||
# 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.
|
||||
"""The Relax generic GPU backend compilation pipeline and other passes."""
|
||||
|
||||
import tvm
|
||||
from tvm import relax
|
||||
|
||||
|
||||
def library_dispatch_passes(target: tvm.target.Target): # pylint: disable=unused-argument
|
||||
"""The default library dispatch passes for generic GPU backend."""
|
||||
return [
|
||||
relax.backend.DispatchSampling(),
|
||||
relax.backend.DispatchSortScan(),
|
||||
]
|
||||
|
||||
|
||||
def legalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
|
||||
"""The default legalization passes for generic GPU backend."""
|
||||
from tvm.s_tir import dlight as dl # pylint: disable=import-outside-toplevel
|
||||
|
||||
return [
|
||||
tvm.relax.transform.LegalizeOps(),
|
||||
tvm.relax.transform.AnnotateTIROpPattern(),
|
||||
tvm.relax.transform.FoldConstant(),
|
||||
tvm.relax.transform.FuseOps(),
|
||||
tvm.relax.transform.FuseTIR(),
|
||||
dl.ApplyDefaultSchedule(
|
||||
dl.gpu.Matmul(),
|
||||
dl.gpu.GEMV(),
|
||||
dl.gpu.Reduction(),
|
||||
dl.gpu.GeneralReduction(),
|
||||
dl.gpu.Fallback(),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def dataflow_lower_passes(target: tvm.target.Target): # pylint: disable=unused-argument
|
||||
"""The default dataflow lowering passes for generic GPU backend."""
|
||||
return [
|
||||
relax.transform.RewriteDataflowReshape(),
|
||||
relax.transform.ToNonDataflow(),
|
||||
relax.transform.RemovePurityChecking(),
|
||||
relax.transform.CallTIRRewrite(),
|
||||
]
|
||||
|
||||
|
||||
def finalize_passes(target: tvm.target.Target): # pylint: disable=unused-argument
|
||||
"""The default finalization passes for generic GPU backend."""
|
||||
return [
|
||||
relax.transform.StaticPlanBlockMemory(),
|
||||
relax.transform.LowerAllocTensor(),
|
||||
relax.transform.KillAfterLastUse(),
|
||||
relax.transform.LowerRuntimeBuiltin(),
|
||||
relax.transform.ComputePrimValue(),
|
||||
relax.transform.VMShapeLower(),
|
||||
relax.transform.AttachGlobalSymbol(),
|
||||
]
|
||||
|
||||
|
||||
def get_default_pipeline(target: tvm.target.Target):
|
||||
"""Return the default compilation pipeline for generic GPU."""
|
||||
|
||||
@tvm.transform.module_pass(opt_level=0)
|
||||
def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext):
|
||||
with target:
|
||||
seq = tvm.transform.Sequential(
|
||||
library_dispatch_passes(target)
|
||||
+ legalize_passes(target)
|
||||
+ dataflow_lower_passes(target)
|
||||
+ finalize_passes(target)
|
||||
)
|
||||
mod = seq(mod)
|
||||
return mod
|
||||
|
||||
return _pipeline
|
||||
@@ -0,0 +1,345 @@
|
||||
# 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 sampling operator."""
|
||||
|
||||
import math
|
||||
from collections.abc import Callable
|
||||
|
||||
import tvm
|
||||
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_multinomial_from_uniform(
|
||||
prob_dtype: str = "float32",
|
||||
sample_dtype: str = "float32",
|
||||
sample_indices_dtype: str = "int64",
|
||||
dtype: str = "int64",
|
||||
ty_len: int = 4,
|
||||
tx_len: int = 32,
|
||||
thread_elem: int = 4,
|
||||
eps: float = 1e-6,
|
||||
) -> PrimFunc:
|
||||
"""Generate GPU kernel for multinomial_from_uniform operator.
|
||||
|
||||
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
|
||||
|
||||
prob_dtype : str
|
||||
The probability data type
|
||||
|
||||
sample_dtype : str
|
||||
The sample data type
|
||||
|
||||
sample_indices_dtype : str
|
||||
The sample indices data type
|
||||
|
||||
dtype : str
|
||||
The output data type
|
||||
|
||||
Returns
|
||||
-------
|
||||
func : PrimFunc
|
||||
The generated function
|
||||
"""
|
||||
|
||||
target = tvm.target.Target.current()
|
||||
target_dtype = "int32" if "webgpu" in str(target) else "int64"
|
||||
|
||||
TX = T.int64(tx_len) # threadIdx.x
|
||||
TY = T.int64(ty_len) # threadIdx.y
|
||||
|
||||
# number of elements to be processed by single thread
|
||||
thread_elem = T.int64(thread_elem)
|
||||
# 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_TY = T.int64(int(math.log2(ty_len)))
|
||||
|
||||
if (
|
||||
not _is_power_of_two(tx_len)
|
||||
or not _is_power_of_two(ty_len)
|
||||
or not _is_power_of_two(thread_elem)
|
||||
):
|
||||
raise ValueError(
|
||||
"Configuration of tx_len, ty_len, thread_elem must be power of 2,"
|
||||
f"but got {tx_len}, {ty_len}, {thread_elem}"
|
||||
)
|
||||
|
||||
@T.macro
|
||||
def block_cumsum(
|
||||
ty: T.int64,
|
||||
tx: T.int64,
|
||||
source_local: T.Buffer,
|
||||
output_shared: T.Buffer,
|
||||
):
|
||||
"""cumsum inside block (SM)"""
|
||||
# Inclusive scan inside thread
|
||||
for i in T.unroll(1, thread_elem):
|
||||
source_local[i] += source_local[i - 1]
|
||||
# Store data to shared memory
|
||||
for i in T.vectorized(thread_elem):
|
||||
output_shared[ty * warp_elem + tx * thread_elem + i] = source_local[i]
|
||||
# Inclusive scan inside warp
|
||||
for i in T.unroll(LOG_TX):
|
||||
for j in T.vectorized(thread_elem):
|
||||
idx: T.let[T.int64] = ty * warp_elem + tx * thread_elem
|
||||
if tx >= (1 << i):
|
||||
output_shared[idx + j] += output_shared[
|
||||
idx - (1 << i) * thread_elem + thread_elem - 1
|
||||
]
|
||||
# Inclusive scan inside block
|
||||
for i in T.unroll(1, TY):
|
||||
for j in T.vectorized(thread_elem):
|
||||
if ty == 0:
|
||||
idx: T.let[T.int64] = i * warp_elem + tx * thread_elem
|
||||
output_shared[idx + j] += output_shared[i * warp_elem - 1]
|
||||
|
||||
def compare_bool_not_equal(a: T.bool, b: T.bool) -> T.bool:
|
||||
# Vulkan does not support compare two bool value direct
|
||||
# return a != b
|
||||
return T.Cast("int8", a) != T.Cast("int8", b)
|
||||
|
||||
@T.macro
|
||||
def block_adjacent_difference_left(
|
||||
ty: T.int64,
|
||||
tx: T.int64,
|
||||
source_local: T.Buffer,
|
||||
output_local: T.Buffer,
|
||||
):
|
||||
with T.sblock():
|
||||
shared_buf = T.sblock_alloc_buffer((TX * TY,), "bool", scope="shared")
|
||||
tx_idx: T.let[T.int64] = ty * TX + tx
|
||||
shared_buf[tx_idx] = source_local[thread_elem - 1]
|
||||
output_local[0] = T.if_then_else(
|
||||
tx_idx != 0,
|
||||
compare_bool_not_equal(source_local[0], shared_buf[tx_idx - 1]),
|
||||
source_local[0],
|
||||
)
|
||||
for i in T.unroll(1, thread_elem):
|
||||
output_local[i] = compare_bool_not_equal(source_local[i], source_local[i - 1])
|
||||
|
||||
def op_reduce_min(a, b):
|
||||
return T.min(a, b)
|
||||
|
||||
def op_reduce_sum(a, b):
|
||||
return a + b
|
||||
|
||||
@T.macro
|
||||
def block_reduce_with_mask(
|
||||
ty: T.int64,
|
||||
tx: T.int64,
|
||||
init_value,
|
||||
data_local: T.Buffer,
|
||||
output_local: T.Buffer,
|
||||
dtype: str,
|
||||
reduce_op: Callable, # T.macro
|
||||
mask_local: T.Buffer | None = None,
|
||||
):
|
||||
with T.sblock():
|
||||
local_sum = T.sblock_alloc_buffer((), dtype, scope="local")
|
||||
shared_buf = T.sblock_alloc_buffer((TX * TY,), dtype, scope="shared")
|
||||
idx: T.let[T.int64] = ty * TX + tx
|
||||
|
||||
local_sum[()] = T.Cast(dtype, init_value)
|
||||
for i in T.unroll(thread_elem):
|
||||
if mask_local is not None:
|
||||
if mask_local[i]:
|
||||
local_sum[()] = reduce_op(local_sum[()], data_local[i])
|
||||
else:
|
||||
local_sum[()] = reduce_op(local_sum[()], data_local[i])
|
||||
shared_buf[idx] = local_sum[()]
|
||||
|
||||
for i in T.unroll(LOG_TX + LOG_TY):
|
||||
if idx % (1 << (i + 1)) == 0:
|
||||
shared_buf[idx] = reduce_op(shared_buf[idx], shared_buf[idx + (1 << i)])
|
||||
output_local[()] = shared_buf[0]
|
||||
|
||||
@T.macro
|
||||
def single_batch_sampling(
|
||||
prob,
|
||||
row_idx,
|
||||
vocab_size,
|
||||
ty,
|
||||
tx,
|
||||
step_iter,
|
||||
threshold,
|
||||
aggregate,
|
||||
uniform_sample,
|
||||
sample_id_local,
|
||||
):
|
||||
with T.sblock():
|
||||
prob_gt_threshold = T.sblock_alloc_buffer((thread_elem,), prob_dtype, scope="local")
|
||||
cumsum = T.sblock_alloc_buffer((block_elem,), prob_dtype, scope="shared")
|
||||
greater_than_u = T.sblock_alloc_buffer((thread_elem,), "bool", scope="local")
|
||||
mask = T.sblock_alloc_buffer((thread_elem,), "bool", scope="local")
|
||||
valid = T.sblock_alloc_buffer((thread_elem,), "bool", scope="local")
|
||||
indices = T.sblock_alloc_buffer((thread_elem), dtype, scope="local")
|
||||
step_aggregate = T.sblock_alloc_buffer((), prob_dtype, scope="local")
|
||||
# Load prob data from global memory to local memory
|
||||
for v in T.unroll(thread_elem):
|
||||
idx: T.let[T.int64] = step_iter * block_elem + ty * warp_elem + tx * thread_elem + v
|
||||
prob_local: T.let = T.if_then_else(
|
||||
idx < vocab_size,
|
||||
prob[row_idx, idx],
|
||||
T.Cast(prob_dtype, 0),
|
||||
)
|
||||
prob_gt_threshold[v] = T.if_then_else(
|
||||
prob_local > threshold, prob_local, T.Cast(prob_dtype, 0)
|
||||
)
|
||||
valid[v] = prob_local > threshold and idx < vocab_size
|
||||
|
||||
block_reduce_with_mask(
|
||||
ty,
|
||||
tx,
|
||||
init_value=0,
|
||||
data_local=prob_gt_threshold,
|
||||
output_local=step_aggregate,
|
||||
dtype=prob_dtype,
|
||||
reduce_op=op_reduce_sum,
|
||||
mask_local=None,
|
||||
)
|
||||
if T.tvm_thread_invariant(aggregate[()] + step_aggregate[()] >= uniform_sample - eps):
|
||||
block_cumsum(ty, tx, prob_gt_threshold, cumsum)
|
||||
# Note: it should be `T.vectorized` instead of `T.unroll`
|
||||
# However, it will cause vulkan codegen error
|
||||
for v in T.unroll(thread_elem):
|
||||
greater_than_u[v] = (
|
||||
cumsum[ty * warp_elem + tx * thread_elem + v] + aggregate[()]
|
||||
>= uniform_sample - eps
|
||||
)
|
||||
|
||||
block_adjacent_difference_left(ty, tx, greater_than_u, mask)
|
||||
# Same as above, it should be `T.vectorized`
|
||||
for v in T.unroll(thread_elem):
|
||||
mask[v] = mask[v] and valid[v]
|
||||
indices[v] = step_iter * block_elem + ty * warp_elem + tx * thread_elem + v
|
||||
block_reduce_with_mask(
|
||||
ty,
|
||||
tx,
|
||||
init_value=vocab_size - 1,
|
||||
data_local=indices,
|
||||
output_local=sample_id_local,
|
||||
dtype=dtype,
|
||||
reduce_op=op_reduce_min,
|
||||
mask_local=mask,
|
||||
)
|
||||
|
||||
aggregate[()] += step_aggregate[()]
|
||||
|
||||
@T.prim_func(s_tir=True)
|
||||
def parallel_sampling_from_prob(
|
||||
var_prob: T.handle,
|
||||
var_uniform_samples: T.handle,
|
||||
var_row_indices: T.handle,
|
||||
var_sampled_token_ids: T.handle,
|
||||
):
|
||||
T.func_attr({"tirx.is_scheduled": True})
|
||||
n, vocab_size, batch_size = T.int64(), T.int64(), T.int64()
|
||||
# match buffers
|
||||
prob = T.match_buffer(var_prob, (n, vocab_size), prob_dtype)
|
||||
uniform_samples = T.match_buffer(var_uniform_samples, (batch_size, 1), sample_dtype)
|
||||
row_indices = T.match_buffer(var_row_indices, (batch_size, 1), sample_indices_dtype)
|
||||
token_ids = T.match_buffer(var_sampled_token_ids, (batch_size, 1), dtype)
|
||||
# local buffers
|
||||
aggregate = T.sblock_alloc_buffer((), prob_dtype, scope="local")
|
||||
sample_id_local = T.sblock_alloc_buffer((), dtype, scope="local")
|
||||
step_iter = T.sblock_alloc_buffer((), "int32", scope="local")
|
||||
|
||||
for bx in T.thread_binding(batch_size, thread="blockIdx.x"):
|
||||
row_idx: T.let[T.int64] = T.Cast("int64", row_indices[bx, 0])
|
||||
for ty in T.thread_binding(TY, thread="threadIdx.y"):
|
||||
for tx in T.thread_binding(TX, thread="threadIdx.x"):
|
||||
u: T.let[T.float32] = uniform_samples[bx, 0]
|
||||
aggregate[()] = T.Cast(prob_dtype, 0)
|
||||
step_iter[()] = T.int32(0)
|
||||
# at least one iteration
|
||||
while T.tvm_thread_invariant(
|
||||
(step_iter[()] == 0 or aggregate[()] < u - eps)
|
||||
and T.Cast(target_dtype, step_iter[()])
|
||||
< T.Cast(target_dtype, T.ceildiv(vocab_size, block_elem))
|
||||
):
|
||||
single_batch_sampling(
|
||||
prob,
|
||||
row_idx,
|
||||
vocab_size,
|
||||
ty,
|
||||
tx,
|
||||
T.Cast(target_dtype, step_iter[()]),
|
||||
0.0,
|
||||
aggregate,
|
||||
u,
|
||||
sample_id_local,
|
||||
)
|
||||
step_iter[()] += 1
|
||||
if tx == 0 and ty == 0:
|
||||
token_ids[bx, 0] = sample_id_local[()]
|
||||
|
||||
return parallel_sampling_from_prob
|
||||
|
||||
|
||||
def generic_get_sample_index(
|
||||
prob_dtype: str = "float32",
|
||||
sample_dtype: str = "float32",
|
||||
sample_indices_dtype: str = "int64",
|
||||
dtype: str = "int64",
|
||||
):
|
||||
"""Generate a generic get_sample_index kernel."""
|
||||
|
||||
@T.prim_func(private=True, s_tir=True)
|
||||
def _get_sample_index(A: T.handle, B: T.handle, C: T.handle, D: T.handle):
|
||||
batch, vocab_size = T.int64(), T.int64()
|
||||
prob = T.match_buffer(A, (batch, vocab_size), prob_dtype)
|
||||
out_batch = T.int64()
|
||||
usample = T.match_buffer(B, (out_batch, 1), sample_dtype)
|
||||
sample_indices = T.match_buffer(C, (out_batch, 1), sample_indices_dtype)
|
||||
output_index = T.match_buffer(D, (out_batch, 1), dtype)
|
||||
|
||||
for ax0, ax1 in T.grid(out_batch, vocab_size):
|
||||
with T.sblock("T_get_sample_index"):
|
||||
v_ax0, v_ax1 = T.axis.remap("SS", [ax0, ax1])
|
||||
T.writes(output_index[v_ax0, 0])
|
||||
if (
|
||||
usample[v_ax0, T.int64(0)] < prob[sample_indices[v_ax0, T.int64(0)], v_ax1]
|
||||
or v_ax1 + 1 == vocab_size
|
||||
):
|
||||
if v_ax1 == 0:
|
||||
output_index[v_ax0, 0] = 0
|
||||
elif (
|
||||
usample[v_ax0, T.int64(0)]
|
||||
>= prob[sample_indices[v_ax0, T.int64(0)], v_ax1 - 1]
|
||||
):
|
||||
output_index[v_ax0, 0] = v_ax1
|
||||
|
||||
return _get_sample_index
|
||||
Reference in New Issue
Block a user