wehub-resource-sync
745 lines
27 KiB
Python
745 lines
27 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.
|
|
# 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,
|
|
)
|