chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,744 @@
|
||||
# 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.
|
||||
# ruff: noqa: E741, F821
|
||||
"""A rule for low-batch GEMM / decode-GEMM using GEMV schedule."""
|
||||
|
||||
from functools import reduce
|
||||
from typing import Literal
|
||||
|
||||
import tvm_ffi
|
||||
|
||||
from tvm import arith, s_tir, tirx
|
||||
from tvm.target import Target
|
||||
|
||||
from ..analysis import (
|
||||
SBlockInfo,
|
||||
collect_block_iter_vars_used_in_access_region,
|
||||
collect_vars_used_in_prim_expr,
|
||||
get_max_shared_memory_per_block,
|
||||
is_broadcast_epilogue,
|
||||
normalize_prim_func,
|
||||
)
|
||||
from ..base import auto_vectorize, get_bytes, get_extent, try_inline_contiguous_spatial
|
||||
from .base import GPUScheduleRule
|
||||
|
||||
|
||||
def _get_reduction_expr(block: tirx.SBlock) -> tirx.Expr | None:
|
||||
# Detect and return `Y` in `X[...] = X[...] + Y`
|
||||
buffer_store = block.body
|
||||
if not isinstance(buffer_store, tirx.BufferStore):
|
||||
return None
|
||||
if not isinstance(buffer_store.value, tirx.Add):
|
||||
return None
|
||||
if not tvm_ffi.structural_equal(
|
||||
buffer_store.value.a,
|
||||
tirx.BufferLoad(buffer_store.buffer, block.body.indices),
|
||||
map_free_vars=True,
|
||||
):
|
||||
return None
|
||||
return buffer_store.value.b
|
||||
|
||||
|
||||
def is_gemv(sch: s_tir.Schedule, block_info: SBlockInfo) -> list[tirx.Buffer] | None:
|
||||
"""Check if the block is a low batch GEMM.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
|
||||
sch : s_tir.Schedule
|
||||
The schedule
|
||||
|
||||
block_info : SBlockInfo
|
||||
The block info to be checked
|
||||
|
||||
|
||||
Returns
|
||||
-------
|
||||
ret : Optional[List[tirx.Buffer]]
|
||||
The vector-like buffers used in the low batch GEMM if it is a low batch GEMM,
|
||||
otherwise None.
|
||||
"""
|
||||
block = block_info.block_rv
|
||||
block_stmt = sch.get(block)
|
||||
conditions = []
|
||||
conditions.append(block_info.is_reduction())
|
||||
conditions.append(len(block_stmt.reads) >= 2)
|
||||
conditions.append(len(block_stmt.writes) == 1)
|
||||
conditions.append(_get_reduction_expr(block_stmt) is not None)
|
||||
conditions.append(
|
||||
len(collect_block_iter_vars_used_in_access_region(block_stmt, block_stmt.writes[0].region))
|
||||
> 0
|
||||
)
|
||||
if not all(conditions):
|
||||
return None
|
||||
const_iter_vars = set(
|
||||
iter_var.var
|
||||
for iter_var in block_stmt.iter_vars
|
||||
if isinstance(iter_var.dom.extent, tirx.IntImm)
|
||||
)
|
||||
if len(block_stmt.iter_vars) - len(const_iter_vars) != 1:
|
||||
return None
|
||||
symbolic_iter_var = next(
|
||||
iter_var
|
||||
for iter_var in block_stmt.iter_vars
|
||||
if not isinstance(iter_var.dom.extent, tirx.IntImm)
|
||||
)
|
||||
if symbolic_iter_var.iter_type != tirx.stmt.IterVar.DataPar:
|
||||
return None
|
||||
ret = [
|
||||
read.buffer
|
||||
for read in block_stmt.reads
|
||||
if len(
|
||||
collect_block_iter_vars_used_in_access_region(block_stmt, read.region) & const_iter_vars
|
||||
)
|
||||
< len(const_iter_vars)
|
||||
and len(
|
||||
collect_block_iter_vars_used_in_access_region(block_stmt, read.region) & const_iter_vars
|
||||
)
|
||||
> 0
|
||||
]
|
||||
return ret if 0 < len(ret) < len(block_stmt.reads) else None
|
||||
|
||||
|
||||
def detect_dominant_read(block: tirx.SBlock, const_iter_vars: set[tirx.Var]) -> tirx.Expr:
|
||||
"""Detect the dominant read indices in the block."""
|
||||
dominant_read = None
|
||||
num_read_iters = -1
|
||||
for buffer_region in block.reads:
|
||||
tir_vars = (
|
||||
collect_block_iter_vars_used_in_access_region(block, buffer_region.region)
|
||||
& const_iter_vars
|
||||
)
|
||||
if num_read_iters < len(tir_vars):
|
||||
num_read_iters = len(tir_vars)
|
||||
dominant_read = buffer_region
|
||||
assert dominant_read is not None
|
||||
(result,) = dominant_read.buffer.offset_of([e.min for e in dominant_read.region])
|
||||
return result
|
||||
|
||||
|
||||
def normalize(
|
||||
sch: s_tir.Schedule,
|
||||
block_info: SBlockInfo,
|
||||
) -> bool | None:
|
||||
"""Normalize the main block."""
|
||||
block_stmt: tirx.SBlock = sch.get(block_info.block_rv)
|
||||
const_iter_vars = set(
|
||||
iter_var.var
|
||||
for iter_var in block_stmt.iter_vars
|
||||
if isinstance(iter_var.dom.extent, tirx.IntImm)
|
||||
)
|
||||
dynamic_iter_vars = set(
|
||||
iter_var.var for iter_var in block_stmt.iter_vars if iter_var.var not in const_iter_vars
|
||||
)
|
||||
access = arith.normalize_to_iter_sum(
|
||||
detect_dominant_read(block_stmt, const_iter_vars),
|
||||
input_iters={i.var: i.dom for i in block_stmt.iter_vars},
|
||||
)
|
||||
buffers_use_vars = [
|
||||
collect_block_iter_vars_used_in_access_region(block_stmt, buf.region)
|
||||
for buf in block_stmt.writes
|
||||
]
|
||||
buffers_use_vars.extend(
|
||||
[
|
||||
collect_block_iter_vars_used_in_access_region(block_stmt, buf.region)
|
||||
for buf in block_stmt.reads
|
||||
]
|
||||
)
|
||||
if collect_vars_used_in_prim_expr(access.base) & set(
|
||||
iter_var.var for iter_var in block_stmt.iter_vars
|
||||
):
|
||||
return None
|
||||
iter_to_info = {i.var: i for i in block_info.iters}
|
||||
batch_loops, s_loops, r_loops = [], [], []
|
||||
inner_axis = access.args[-1].source.source
|
||||
is_inner_reduction = iter_to_info[inner_axis].kind == "R"
|
||||
|
||||
for split_expr in access.args:
|
||||
var = split_expr.source.source
|
||||
info = iter_to_info.get(var)
|
||||
loop = info.loop_rv
|
||||
is_reduction = info.kind == "R"
|
||||
# No C loops as we do not compute_inline weights into main block
|
||||
if is_reduction:
|
||||
r_loops.append(loop)
|
||||
elif all([var in buf_vars for buf_vars in buffers_use_vars]):
|
||||
batch_loops.append(loop)
|
||||
else:
|
||||
s_loops.append(loop)
|
||||
|
||||
assert s_loops
|
||||
assert r_loops
|
||||
dynamic_loops = [iter_to_info[var].loop_rv for var in dynamic_iter_vars]
|
||||
assert len(dynamic_loops) == 1
|
||||
sch.reorder(*dynamic_loops, *s_loops, *r_loops)
|
||||
sch.fuse(*s_loops)
|
||||
sch.fuse(*r_loops)
|
||||
return is_inner_reduction
|
||||
|
||||
|
||||
class LowBatchGEMV(GPUScheduleRule):
|
||||
"""A rule for low batch GEMM / decode-GEMM."""
|
||||
|
||||
def __init__(self, bucket=4):
|
||||
self.bucket = bucket
|
||||
|
||||
def apply( # pylint: disable=too-many-locals,too-many-branches,too-many-return-statements
|
||||
self,
|
||||
func: tirx.PrimFunc,
|
||||
target: Target,
|
||||
_: bool,
|
||||
) -> None | s_tir.Schedule | list[s_tir.Schedule]:
|
||||
if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
|
||||
return None
|
||||
sch = s_tir.Schedule(func)
|
||||
block_infos = normalize_prim_func(sch)
|
||||
if block_infos is None:
|
||||
return None
|
||||
reduction_block_infos = [
|
||||
block_info for block_info in block_infos if block_info.is_reduction()
|
||||
]
|
||||
if len(reduction_block_infos) != 1:
|
||||
return None
|
||||
reduction_block_info = reduction_block_infos[0]
|
||||
vector_input_buffers = is_gemv(sch, reduction_block_info)
|
||||
if vector_input_buffers is None:
|
||||
return None
|
||||
batch_pad = self.bucket
|
||||
pad_value = [
|
||||
iter.dom if isinstance(iter.dom, int) else batch_pad
|
||||
for iter in reduction_block_info.iters
|
||||
]
|
||||
sch.pad_einsum(reduction_block_info.block_rv, pad_value)
|
||||
block_infos = normalize_prim_func(sch)
|
||||
dequantize_block = None
|
||||
pad_input_block = None
|
||||
for block_info in block_infos:
|
||||
if "dequantize" in block_info.name:
|
||||
dequantize_block = block_info.block_rv
|
||||
elif "pad" in block_info.name and len(sch.get_producers(block_info.block_rv)) == 0:
|
||||
pad_input_block = block_info.block_rv
|
||||
block_infos = [
|
||||
block_info
|
||||
for block_info in block_infos
|
||||
if "pad" not in block_info.name and "dequantize" not in block_info.name
|
||||
]
|
||||
block_infos = try_inline_contiguous_spatial(sch, block_infos)
|
||||
if len(block_infos) == 1:
|
||||
epilogue = None
|
||||
elif len(block_infos) == 2:
|
||||
epilogue = block_infos[1]
|
||||
if not epilogue.is_injective():
|
||||
return None
|
||||
else:
|
||||
return None
|
||||
|
||||
block_info = block_infos[0]
|
||||
if len(block_info.iters) not in [2, 3]:
|
||||
# either [B, S, R] = [B, S, R] * [B, R]
|
||||
# or [S, R] = [S, R] * [R]
|
||||
return None
|
||||
block = block_info.block_rv
|
||||
vector_input_buffers = is_gemv(sch, block_info)
|
||||
if vector_input_buffers is None:
|
||||
return None
|
||||
|
||||
# Step 1. Normalize the block, merge spatial and reduction iters
|
||||
is_inner_reduction = normalize(sch, block_info)
|
||||
# Step 2. Do the scheduling
|
||||
if is_inner_reduction is None:
|
||||
return None
|
||||
elif is_inner_reduction:
|
||||
self.sch_inner_reduction(
|
||||
sch,
|
||||
target,
|
||||
block,
|
||||
dequantize_block,
|
||||
pad_input_block,
|
||||
vector_input_buffers,
|
||||
epilogue,
|
||||
batch_pad,
|
||||
)
|
||||
return sch
|
||||
elif self.bucket <= 4:
|
||||
self.sch_outer_reduction(
|
||||
sch,
|
||||
target,
|
||||
block,
|
||||
dequantize_block,
|
||||
pad_input_block,
|
||||
vector_input_buffers,
|
||||
epilogue,
|
||||
batch_pad,
|
||||
)
|
||||
return sch
|
||||
else:
|
||||
return None
|
||||
|
||||
def sch_inner_reduction( # pylint: disable=too-many-arguments, invalid-name, unused-argument
|
||||
self,
|
||||
sch: s_tir.Schedule,
|
||||
target: Target,
|
||||
block: s_tir.schedule.SBlockRV,
|
||||
dequantize_block: s_tir.schedule.SBlockRV | None,
|
||||
pad_input_block: s_tir.schedule.SBlockRV | None,
|
||||
vector_input_buffers: list[tirx.Buffer],
|
||||
epilogue_info: SBlockInfo | None,
|
||||
batch_pad: int,
|
||||
):
|
||||
"""Schedule the inner reduction block."""
|
||||
|
||||
def get_max_factor(n, factors):
|
||||
factors = sorted(factors, reverse=True)
|
||||
for factor in factors:
|
||||
if n % factor == 0:
|
||||
return factor
|
||||
return 1
|
||||
|
||||
def apply(
|
||||
sch: s_tir.Schedule,
|
||||
gemv,
|
||||
TAG_S,
|
||||
TAG_R,
|
||||
TS,
|
||||
TR,
|
||||
TILE_S,
|
||||
TILE_R,
|
||||
VEC_LOAD,
|
||||
VEC_C,
|
||||
LOAD_V_SHARED,
|
||||
LOAD_V_VEC,
|
||||
UNROLL,
|
||||
):
|
||||
# rfactor: reduce to tx * vec_c
|
||||
|
||||
_, s, r = sch.get_loops(block=gemv)
|
||||
bx, ts, tile_s = sch.split(s, factors=[None, TS, TILE_S], preserve_unit_iters=True)
|
||||
r, tr, tile_r_vec_n, vec_c = sch.split(
|
||||
r, factors=[None, TR, TILE_R // VEC_C, VEC_C], preserve_unit_iters=True
|
||||
)
|
||||
sch.reorder(r, tile_r_vec_n, tr, vec_c)
|
||||
tr_vec_c = sch.fuse(tr, vec_c)
|
||||
rf = sch.rfactor(tr_vec_c, 0)
|
||||
|
||||
# rfactor: reduce to tx
|
||||
_, bx, ts, tile_s, tr_vec_c = sch.get_loops(block=gemv)
|
||||
tr, vec_c = sch.split(tr_vec_c, factors=[TR, None], preserve_unit_iters=True)
|
||||
rf2 = sch.rfactor(tr, 0)
|
||||
# bind, vectorize compute
|
||||
batch_loop, bx, ts, tile_s, r, tile_r_vec_n, tr_vec_c = sch.get_loops(block=rf)
|
||||
tr, vec_c = sch.split(tr_vec_c, factors=[TR, None], preserve_unit_iters=True)
|
||||
sch.reorder(bx, ts, tr, r, tile_s, tile_r_vec_n, vec_c)
|
||||
sch.bind(bx, "blockIdx.x")
|
||||
sch.bind(ts, TAG_S)
|
||||
sch.bind(tr, TAG_R)
|
||||
sch.vectorize(vec_c)
|
||||
by, batch = sch.split(batch_loop, factors=[None, batch_pad])
|
||||
sch.bind(by, "blockIdx.y")
|
||||
sch.reorder(bx, ts, tr, r, batch)
|
||||
|
||||
shared_mem_usage = 0
|
||||
for buf in vector_input_buffers:
|
||||
buf_size = reduce(
|
||||
lambda x, y: x * y, buf.shape, tirx.IntImm(buf.shape[0].ty, 1)
|
||||
) * get_bytes(buf.dtype)
|
||||
shared_mem_usage += buf_size
|
||||
max_smem = get_max_shared_memory_per_block(target)
|
||||
LOAD_V_SHARED = (
|
||||
LOAD_V_SHARED
|
||||
and isinstance(shared_mem_usage, tirx.IntImm)
|
||||
and shared_mem_usage.value <= max_smem
|
||||
)
|
||||
|
||||
# vectorize load A
|
||||
# (TODO) this is now actually problematic since the number of loops is dependent on the
|
||||
# number of dimensions of A_q
|
||||
if dequantize_block is not None:
|
||||
sch.compute_at(dequantize_block, r, preserve_unit_loops=True)
|
||||
sch.set_scope(dequantize_block, 0, "local")
|
||||
|
||||
s_local, r_local = sch.get_loops(block=dequantize_block)[-2:]
|
||||
s_local, vec_load = sch.split(
|
||||
s_local, factors=[None, VEC_LOAD], preserve_unit_iters=True
|
||||
)
|
||||
sch.reorder(s_local, r_local, vec_load) # either s_local or r_local should be 1
|
||||
sch.vectorize(vec_load)
|
||||
|
||||
# load vector into shared memory, shape should be the whole vector
|
||||
if LOAD_V_SHARED:
|
||||
assert len(vector_input_buffers) == 1
|
||||
V_shared = sch.cache_read(rf, read_buffer_index=0, storage_scope="shared")
|
||||
sch.compute_at(V_shared, tr, preserve_unit_loops=True)
|
||||
l = sch.get_loops(block=V_shared)[-1]
|
||||
loop: tirx.For = sch.get(l)
|
||||
if isinstance(loop.extent, tirx.IntImm):
|
||||
# avoid introducing predicates when vector length is too large
|
||||
vec_length = max(
|
||||
min(
|
||||
get_max_factor(
|
||||
(int)(loop.extent),
|
||||
[TS * TR * 1, TS * TR * 2, TS * TR * 4, TS * TR * 8],
|
||||
)
|
||||
// TS
|
||||
// TR,
|
||||
LOAD_V_VEC,
|
||||
),
|
||||
1,
|
||||
)
|
||||
else:
|
||||
vec_length = LOAD_V_VEC
|
||||
if TAG_R == "threadIdx.x":
|
||||
_, ty, tx, vec = sch.split(
|
||||
l, factors=[None, TS, TR, vec_length], preserve_unit_iters=True
|
||||
)
|
||||
else:
|
||||
_, ty, tx, vec = sch.split(
|
||||
l, factors=[None, TR, TS, vec_length], preserve_unit_iters=True
|
||||
)
|
||||
sch.bind(ty, "threadIdx.y")
|
||||
sch.bind(tx, "threadIdx.x")
|
||||
sch.vectorize(vec)
|
||||
if pad_input_block is not None:
|
||||
sch.compute_inline(pad_input_block)
|
||||
|
||||
# reduce tile_s * tr * vec to tile_s * tr
|
||||
sch.reverse_compute_at(rf2, loop=bx, preserve_unit_loops=True)
|
||||
tr, vec_c, batch_loop, *ts_tile_s = sch.get_loops(block=rf2)[2:]
|
||||
ts_tile_s = sch.fuse(*ts_tile_s)
|
||||
ts_o, ts_i, tile_s = sch.split(
|
||||
ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
|
||||
)
|
||||
tile_s, vec_s = sch.split(
|
||||
tile_s,
|
||||
factors=[None, get_max_factor(TILE_S, [1, 2, 4, 8])],
|
||||
preserve_unit_iters=True,
|
||||
)
|
||||
assert sch.get(ts_o).extent.value == 1
|
||||
ts = sch.fuse(ts_o, ts_i)
|
||||
sch.reorder(ts, tr, tile_s, batch_loop, vec_s, vec_c)
|
||||
sch.bind(ts, TAG_S)
|
||||
sch.bind(tr, TAG_R)
|
||||
sch.vectorize(vec_s)
|
||||
|
||||
# reduce tile_s * tr to tile_s
|
||||
sch.reverse_compute_at(gemv, loop=bx, preserve_unit_loops=True)
|
||||
|
||||
tr, batch_loop, *ts_tile_s = sch.get_loops(block=gemv)[2:]
|
||||
ts_tile_s = sch.fuse(*ts_tile_s)
|
||||
ts_o, ts_i, tile_s = sch.split(
|
||||
ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
|
||||
)
|
||||
assert sch.get(ts_o).extent.value == 1
|
||||
ts = sch.fuse(ts_o, ts_i)
|
||||
sch.reorder(tile_s, batch_loop, ts, tr)
|
||||
sch.bind(ts, TAG_S)
|
||||
sch.bind(tr, TAG_R)
|
||||
|
||||
sch.decompose_reduction(rf, loop=sch.get_loops(block=rf)[4])
|
||||
sch.decompose_reduction(rf2, loop=sch.get_loops(block=rf2)[-1])
|
||||
|
||||
sch.set_scope(rf, buffer_index=0, storage_scope="local")
|
||||
sch.set_scope(rf2, buffer_index=0, storage_scope="local")
|
||||
|
||||
unroll_factor = UNROLL
|
||||
|
||||
sch.annotate(
|
||||
block_or_loop=sch.get_loops(rf)[4],
|
||||
ann_key="pragma_auto_unroll_max_step",
|
||||
ann_val=unroll_factor,
|
||||
)
|
||||
sch.annotate(
|
||||
block_or_loop=sch.get_loops(rf)[4], ann_key="pragma_unroll_explicit", ann_val=1
|
||||
)
|
||||
|
||||
sch.annotate(
|
||||
block_or_loop=sch.get_loops(rf2)[4],
|
||||
ann_key="pragma_auto_unroll_max_step",
|
||||
ann_val=unroll_factor,
|
||||
)
|
||||
sch.annotate(
|
||||
block_or_loop=sch.get_loops(rf2)[4], ann_key="pragma_unroll_explicit", ann_val=1
|
||||
)
|
||||
|
||||
if LOAD_V_SHARED:
|
||||
sch.annotate(
|
||||
block_or_loop=sch.get_loops(V_shared)[-4],
|
||||
ann_key="pragma_unroll_explicit",
|
||||
ann_val=unroll_factor,
|
||||
)
|
||||
sch.annotate(
|
||||
block_or_loop=sch.get_loops(V_shared)[-4], ann_key="pragma_vectorize", ann_val=1
|
||||
)
|
||||
|
||||
epilogue = sch.get_consumers(gemv)
|
||||
# Schedule epilogue
|
||||
if epilogue:
|
||||
epilogue = epilogue[0]
|
||||
if is_broadcast_epilogue(sch, block, epilogue):
|
||||
sch.reverse_compute_at(epilogue, bx)
|
||||
sch.set_scope(block, 0, "shared")
|
||||
_, _, _, *s = sch.get_loops(epilogue) # pylint: disable=invalid-name
|
||||
_, tx = sch.split(sch.fuse(*s), factors=[None, TX])
|
||||
sch.bind(tx, TAG_S)
|
||||
else:
|
||||
sch.reverse_compute_at(epilogue, bx, preserve_unit_loops=True)
|
||||
ts_tile_s = sch.fuse(*sch.get_loops(epilogue)[3:])
|
||||
ts_tile_s = sch.get_loops(epilogue)[-1]
|
||||
ts_o, ts_i, tile_s = sch.split(
|
||||
ts_tile_s, factors=[None, TS, TILE_S], preserve_unit_iters=True
|
||||
)
|
||||
assert sch.get(ts_o).extent.value == 1
|
||||
ts = sch.fuse(ts_o, ts_i)
|
||||
sch.bind(ts, TAG_S)
|
||||
sch.set_scope(block, 0, "local")
|
||||
|
||||
return sch
|
||||
|
||||
# Specify the `len_tx` and `len_ty` according to the loop extent
|
||||
_, s, r = sch.get_loops(block=block)
|
||||
len_s, len_r = get_extent(sch, s), get_extent(sch, r)
|
||||
|
||||
TAG_S, TAG_R = "threadIdx.y", "threadIdx.x"
|
||||
if target.kind.name == "cuda":
|
||||
VEC_C = 4
|
||||
LOAD_V_SHARED = True
|
||||
LOAD_V_VEC = 8
|
||||
UNROLL = 256
|
||||
if isinstance(len_s, int):
|
||||
if len_s > len_r:
|
||||
TS, TR = 4, 64
|
||||
else:
|
||||
TS, TR = 16, 32
|
||||
elif target.kind.name == "metal":
|
||||
VEC_C = 4
|
||||
LOAD_V_SHARED = False
|
||||
LOAD_V_VEC = -1
|
||||
UNROLL = 8
|
||||
if isinstance(len_s, int):
|
||||
if len_s > len_r:
|
||||
TS, TR = 8, 32
|
||||
else:
|
||||
TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
|
||||
TS, TR = 8, 32
|
||||
elif target.kind.name == "rocm":
|
||||
VEC_C = 4
|
||||
LOAD_V_SHARED = True
|
||||
LOAD_V_VEC = 8
|
||||
UNROLL = 256
|
||||
if isinstance(len_s, int):
|
||||
if len_s > len_r:
|
||||
TS, TR = 1, 128
|
||||
else:
|
||||
TS, TR = 8, 64
|
||||
elif target.kind.name == "opencl" and "android" in str(target.host):
|
||||
TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
|
||||
VEC_C = 8
|
||||
LOAD_V_SHARED = False
|
||||
LOAD_V_VEC = -1
|
||||
UNROLL = 8
|
||||
TS, TR = 2, 32
|
||||
elif target.kind.name == "vulkan":
|
||||
VEC_C = 4
|
||||
LOAD_V_SHARED = True
|
||||
LOAD_V_VEC = 4
|
||||
UNROLL = 256
|
||||
if isinstance(len_s, int):
|
||||
if len_s > len_r:
|
||||
TS, TR = 4, 32
|
||||
else:
|
||||
TS, TR = 16, 32
|
||||
elif target.kind.name == "opencl" and "mali" in str(target.attrs):
|
||||
VEC_C = 8
|
||||
LOAD_V_SHARED = False
|
||||
LOAD_V_VEC = -1
|
||||
UNROLL = 64
|
||||
TS, TR = 1, 64
|
||||
else:
|
||||
VEC_C = 1
|
||||
LOAD_V_SHARED = False
|
||||
LOAD_V_VEC = -1
|
||||
UNROLL = 64
|
||||
TS, TR = 1, 64
|
||||
|
||||
if not isinstance(len_s, int):
|
||||
TS, TR = 1, 64
|
||||
|
||||
while TS * TR > int(target.attrs["max_num_threads"]):
|
||||
if TS > 1:
|
||||
TS //= 2
|
||||
else:
|
||||
TR //= 2
|
||||
|
||||
TILE_S, TILE_R = 2, max(get_max_factor(len_r, [TR * 1, TR * 2, TR * 4, TR * 8]) // TR, 1)
|
||||
VEC_C = min(get_max_factor(TILE_R, [1, 2, 4, 8]), VEC_C)
|
||||
VEC_LOAD = 1
|
||||
return apply(
|
||||
sch,
|
||||
gemv=block,
|
||||
TAG_S=TAG_S,
|
||||
TAG_R=TAG_R,
|
||||
TS=TS,
|
||||
TR=TR,
|
||||
TILE_S=TILE_S,
|
||||
TILE_R=TILE_R,
|
||||
VEC_LOAD=VEC_LOAD,
|
||||
VEC_C=VEC_C,
|
||||
LOAD_V_SHARED=LOAD_V_SHARED,
|
||||
LOAD_V_VEC=LOAD_V_VEC,
|
||||
UNROLL=UNROLL,
|
||||
)
|
||||
|
||||
def sch_outer_reduction( # pylint: disable=too-many-arguments, invalid-name, unused-argument
|
||||
self,
|
||||
sch: s_tir.Schedule,
|
||||
target: Target,
|
||||
block: s_tir.schedule.SBlockRV,
|
||||
dequantize_block: s_tir.schedule.SBlockRV | None,
|
||||
pad_input_block: s_tir.schedule.SBlockRV | None,
|
||||
vector_input_buffers: list[tirx.Buffer],
|
||||
epilogue_info: SBlockInfo | None,
|
||||
batch_pad: int,
|
||||
):
|
||||
"""Schedule the outer reduction block."""
|
||||
|
||||
# Need to detect from the block
|
||||
DEC_PACK = 8
|
||||
SCALE_PACK = 4
|
||||
|
||||
def apply(
|
||||
sch: s_tir.Schedule,
|
||||
main_block: s_tir.schedule.SBlockRV,
|
||||
TAG_S: Literal["threadIdx.x", "threadIdx.y"],
|
||||
TAG_R: Literal["threadIdx.x", "threadIdx.y"],
|
||||
TS: int,
|
||||
TR: int,
|
||||
VEC: int,
|
||||
UNROLL: int,
|
||||
):
|
||||
# rfactor: reduce to tx * vec_c
|
||||
b, s, r = sch.get_loops(main_block)
|
||||
by, batch = sch.split(b, [None, batch_pad], preserve_unit_iters=True)
|
||||
bx, ts = sch.split(s, [None, TS], preserve_unit_iters=True)
|
||||
r, tr, scale_c, vec_c = sch.split(
|
||||
r, [None, TR, SCALE_PACK, DEC_PACK], preserve_unit_iters=True
|
||||
)
|
||||
sch.reorder(by, bx, ts, r, batch, scale_c, tr, vec_c)
|
||||
tr_vec_c = sch.fuse(tr, vec_c)
|
||||
rf = sch.rfactor(tr_vec_c, 0)
|
||||
|
||||
# rfactor: reduce to tx
|
||||
by, bx, ts, batch, tr_vec_c = sch.get_loops(block=main_block)
|
||||
tr, vec_c = sch.split(tr_vec_c, [TR, DEC_PACK], preserve_unit_iters=True)
|
||||
rf2 = sch.rfactor(tr, 0)
|
||||
|
||||
# bind, vectorize compute
|
||||
by, bx, ts, r, batch, scale_c, tr_vec_c = sch.get_loops(block=rf)
|
||||
tr, vec_c = sch.split(tr_vec_c, [TR, DEC_PACK], preserve_unit_iters=True)
|
||||
sch.reorder(by, bx, ts, tr, r, scale_c, batch, vec_c)
|
||||
sch.bind(by, "blockIdx.y")
|
||||
sch.bind(bx, "blockIdx.x")
|
||||
sch.bind(ts, TAG_S)
|
||||
sch.bind(tr, TAG_R)
|
||||
auto_vectorize(sch, vec_c, VEC)
|
||||
|
||||
if dequantize_block is not None:
|
||||
sch.compute_at(dequantize_block, scale_c, preserve_unit_loops=True)
|
||||
sch.set_scope(dequantize_block, 0, "local")
|
||||
auto_vectorize(sch, sch.fuse(*sch.get_loops(dequantize_block)[6:]), VEC)
|
||||
|
||||
B0_local = sch.cache_read(dequantize_block, 0, "local")
|
||||
sch.compute_at(B0_local, r, preserve_unit_loops=True)
|
||||
auto_vectorize(sch, sch.fuse(*sch.get_loops(B0_local)[5:]), VEC)
|
||||
|
||||
B1_local = sch.cache_read(dequantize_block, 1, "local")
|
||||
sch.compute_at(B1_local, r, preserve_unit_loops=True)
|
||||
auto_vectorize(sch, sch.fuse(*sch.get_loops(B1_local)[5:]), VEC)
|
||||
else:
|
||||
# Only support quantized workloads for now
|
||||
sch = None
|
||||
return
|
||||
|
||||
if LOAD_V_SHARED:
|
||||
sch.set_scope(pad_input_block, 0, "shared")
|
||||
sch.compute_at(pad_input_block, r, preserve_unit_loops=True)
|
||||
sch.storage_align(pad_input_block, 0, axis=-2, factor=8, offset=1)
|
||||
tr, ts, v = sch.split(sch.fuse(*sch.get_loops(pad_input_block)[5:]), [TR, TS, None])
|
||||
sch.bind(tr, TAG_R)
|
||||
sch.bind(ts, TAG_S)
|
||||
auto_vectorize(sch, v, VEC)
|
||||
else:
|
||||
sch.compute_inline(pad_input_block)
|
||||
|
||||
# reduce tile_s * tr * vec to tile_s * tr
|
||||
sch.reverse_compute_at(rf2, bx, preserve_unit_loops=True)
|
||||
tr, vec_c, batch, ts = sch.get_loops(rf2)[2:]
|
||||
sch.reorder(ts, tr, batch, vec_c)
|
||||
sch.bind(ts, TAG_S)
|
||||
sch.bind(tr, TAG_R)
|
||||
|
||||
# reduce tile_s * tr to tile_s
|
||||
sch.reverse_compute_at(main_block, bx, preserve_unit_loops=True)
|
||||
tr, batch, ts = sch.get_loops(main_block)[2:]
|
||||
sch.reorder(batch, ts, tr)
|
||||
sch.bind(ts, TAG_S)
|
||||
sch.bind(tr, TAG_R)
|
||||
# unroll(batch, 1)
|
||||
|
||||
sch.decompose_reduction(rf, loop=sch.get_loops(block=rf)[4])
|
||||
sch.decompose_reduction(rf2, loop=sch.get_loops(block=rf2)[4])
|
||||
|
||||
sch.set_scope(rf, buffer_index=0, storage_scope="local")
|
||||
sch.set_scope(rf2, buffer_index=0, storage_scope="local")
|
||||
|
||||
epilogue = sch.get_consumers(main_block)
|
||||
# Schedule epilogue
|
||||
if epilogue:
|
||||
epilogue = epilogue[0]
|
||||
if is_broadcast_epilogue( # pylint: disable=no-else-raise
|
||||
sch, main_block, epilogue
|
||||
):
|
||||
raise NotImplementedError
|
||||
else:
|
||||
sch.reverse_compute_at(epilogue, bx, preserve_unit_loops=True)
|
||||
batch, ts = sch.get_loops(epilogue)[2:]
|
||||
sch.bind(ts, TAG_S)
|
||||
sch.set_scope(main_block, 0, "local")
|
||||
|
||||
if target.kind.name == "metal":
|
||||
TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
|
||||
TS, TR = 64, 4
|
||||
LOAD_V_SHARED = True
|
||||
VEC = 4
|
||||
UNROLL = 8
|
||||
else:
|
||||
# fallback configuration
|
||||
TAG_S, TAG_R = "threadIdx.x", "threadIdx.y"
|
||||
TS, TR = 32, 4
|
||||
LOAD_V_SHARED = False
|
||||
VEC = 1
|
||||
UNROLL = 64
|
||||
|
||||
return apply(
|
||||
sch,
|
||||
block,
|
||||
TAG_S,
|
||||
TAG_R,
|
||||
TS,
|
||||
TR,
|
||||
VEC,
|
||||
UNROLL,
|
||||
)
|
||||
Reference in New Issue
Block a user