wehub-resource-sync
1244 lines
45 KiB
Python
1244 lines
45 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=missing-docstring, invalid-name
|
|
# ruff: noqa: E722
|
|
"""A GEMM schedule rule for GPU operators."""
|
|
|
|
from dataclasses import dataclass
|
|
from enum import Enum
|
|
|
|
from tvm import s_tir, tirx
|
|
from tvm.ir import Range
|
|
from tvm.s_tir.schedule.schedule import SBlockRV
|
|
from tvm.script import tirx as T
|
|
from tvm.target import Target
|
|
from tvm.tirx import Expr, IterVar, Var
|
|
from tvm.tirx.analysis import undefined_vars
|
|
|
|
from ..analysis import IterInfo, SBlockInfo, get_root_block
|
|
from .base import GPUScheduleRule
|
|
|
|
|
|
def _collect_producers(sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV):
|
|
result = []
|
|
for producer in sch.get_producers(block):
|
|
result.append(producer)
|
|
result.extend(_collect_producers(sch, producer))
|
|
return result
|
|
|
|
|
|
def _collect_consumers(sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV):
|
|
result = []
|
|
for consumer in sch.get_consumers(block):
|
|
result.append(consumer)
|
|
result.extend(_collect_consumers(sch, consumer))
|
|
return result
|
|
|
|
|
|
def auto_inline_producers(
|
|
sch: s_tir.Schedule,
|
|
block: s_tir.schedule.SBlockRV,
|
|
):
|
|
while True:
|
|
inlined_cnt = 0
|
|
producers = _collect_producers(sch, block)
|
|
for producer in producers:
|
|
try:
|
|
sch.compute_inline(producer)
|
|
inlined_cnt += 1
|
|
except: # pylint: disable=bare-except
|
|
continue
|
|
if inlined_cnt == 0:
|
|
return
|
|
|
|
|
|
def auto_inline_consumers(
|
|
sch: s_tir.Schedule,
|
|
block: s_tir.schedule.SBlockRV,
|
|
):
|
|
while True:
|
|
inlined_cnt = 0
|
|
consumers = _collect_consumers(sch, block)
|
|
for consumer in consumers:
|
|
try:
|
|
sch.compute_inline(consumer)
|
|
inlined_cnt += 1
|
|
except: # pylint: disable=bare-except
|
|
continue
|
|
for consumer in consumers:
|
|
try:
|
|
sch.reverse_compute_inline(consumer)
|
|
inlined_cnt += 1
|
|
except: # pylint: disable=bare-except
|
|
continue
|
|
if inlined_cnt == 0:
|
|
return
|
|
|
|
|
|
def auto_inline_consumer_chain(
|
|
sch: s_tir.Schedule,
|
|
block: s_tir.schedule.SBlockRV,
|
|
):
|
|
auto_inline_consumers(sch, block)
|
|
remaining_consumers = sch.get_consumers(block)
|
|
|
|
if len(remaining_consumers) != 0:
|
|
# Some blocks have failed to be inlined to the producer cache-write stage.
|
|
# This could be due to another producer block that has not been scheduled.
|
|
for c in remaining_consumers:
|
|
for p in sch.get_producers(c):
|
|
if sch.get(p) != sch.get(block):
|
|
auto_inline_producers(sch, p)
|
|
sch.compute_inline(p)
|
|
|
|
# Try inlining into the cache-write stage again, this time it should succeed.
|
|
auto_inline_consumers(sch, block)
|
|
|
|
msg = "There are some consumers of the cache-write stage that are not properly inlined."
|
|
assert len(sch.get_consumers(block)) == 0, msg
|
|
|
|
|
|
class IterKind(Enum):
|
|
"""Iter kinds for GEMM-liked programs.
|
|
We can simplify the computation to C[S, I, J] += A[S, I, K] * B[S, J, K],
|
|
where `I, J, K` are fundamental axes for gemm and `S` represents all
|
|
other spatial axes (e.g. batches)
|
|
kIter_S: spatial axes
|
|
kIter_I: I axes
|
|
kIter_J: J axes
|
|
kIter_K: K axes
|
|
kIter_T: trivial axes (i.e. with extent 1)
|
|
"""
|
|
|
|
kIter_S = 0
|
|
kIter_I = 1
|
|
kIter_J = 2
|
|
kIter_K = 3
|
|
kIter_T = 4
|
|
|
|
|
|
@dataclass
|
|
class IterTrait:
|
|
kind: IterKind
|
|
extent: Expr
|
|
|
|
|
|
def _is_one(x: Expr) -> bool:
|
|
return isinstance(x, tirx.IntImm) and x.value == 1
|
|
|
|
|
|
def make_iter_fusion_index_map(
|
|
traits: list[IterTrait],
|
|
kind_order: list[IterKind],
|
|
) -> tirx.IndexMap:
|
|
fused_iters: dict[IterKind, Expr] = {}
|
|
input_iters: list[tirx.Var] = []
|
|
for i, trait in enumerate(traits):
|
|
v_i = tirx.Var(f"i{i}", trait.extent.ty)
|
|
input_iters.append(v_i)
|
|
if trait.kind == IterKind.kIter_T:
|
|
continue
|
|
if trait.kind not in kind_order:
|
|
raise ValueError(f"Unknown iter kind {trait.kind}")
|
|
if trait.kind in fused_iters:
|
|
fused_iters[trait.kind] = fused_iters[trait.kind] * trait.extent + v_i
|
|
else:
|
|
fused_iters[trait.kind] = v_i
|
|
|
|
final_indices: list[tirx.Expr] = [
|
|
fused_iters.get(kind, tirx.IntImm(traits[0].extent.ty, 0)) for kind in kind_order
|
|
]
|
|
|
|
return tirx.IndexMap(input_iters, final_indices, None)
|
|
|
|
|
|
def detect_iter_traits(block: tirx.SBlock) -> tuple[list[IterTrait]] | None:
|
|
"""Detect iter traits based on the pattern C[S, I, J] += A[S, I, K] * B[S, J, K]
|
|
|
|
Parameters
|
|
----------
|
|
block : tirx.SBlock
|
|
The block to be analyzed
|
|
|
|
Returns
|
|
-------
|
|
traits : Optional[Tuple[List[IterTrait]]]
|
|
The detected iter traits for axes in A, B and C. None if the block
|
|
does not match the pattern.
|
|
|
|
"""
|
|
|
|
if len(block.reads) != 2 or len(block.writes) != 1:
|
|
return None
|
|
|
|
def get_access_axes(region: list[Range]) -> set[Var]:
|
|
axes: set[Var] = set()
|
|
for r in region:
|
|
if not _is_one(r.extent):
|
|
raise ValueError("Expect elemwise block access")
|
|
axes = axes.union(set(undefined_vars(r.min)))
|
|
return axes
|
|
|
|
try:
|
|
A_axes = get_access_axes(block.reads[0].region)
|
|
B_axes = get_access_axes(block.reads[1].region)
|
|
C_axes = get_access_axes(block.writes[0].region)
|
|
except ValueError:
|
|
return None
|
|
|
|
traits: dict[Var, IterTrait] = {}
|
|
for iter_var in block.iter_vars:
|
|
var = iter_var.var
|
|
kind: IterKind
|
|
if _is_one(iter_var.dom.extent):
|
|
kind = IterKind.kIter_T
|
|
elif iter_var.iter_type == iter_var.DataPar:
|
|
if var in A_axes and var in B_axes and var in C_axes:
|
|
kind = IterKind.kIter_S
|
|
elif var in A_axes and var in C_axes:
|
|
kind = IterKind.kIter_I
|
|
elif var in B_axes and var in C_axes:
|
|
kind = IterKind.kIter_J
|
|
else:
|
|
return None
|
|
elif iter_var.iter_type == tirx.IterVar.CommReduce:
|
|
if var in A_axes and var in B_axes and var not in C_axes:
|
|
kind = IterKind.kIter_K
|
|
else:
|
|
return None
|
|
else:
|
|
return None
|
|
traits[var] = IterTrait(kind, iter_var.dom.extent)
|
|
|
|
# A Gemm-kernel requires have I, J and K axes
|
|
gemm_traits = {IterKind.kIter_I, IterKind.kIter_J, IterKind.kIter_K}
|
|
if {x.kind for x in traits.values()}.intersection(gemm_traits) != gemm_traits:
|
|
return None
|
|
|
|
A_traits = [traits[iter_var.var] for iter_var in block.iter_vars if iter_var.var in A_axes]
|
|
B_traits = [traits[iter_var.var] for iter_var in block.iter_vars if iter_var.var in B_axes]
|
|
C_traits = [traits[iter_var.var] for iter_var in block.iter_vars if iter_var.var in C_axes]
|
|
block_traits = [traits[i.var] for i in block.iter_vars]
|
|
return A_traits, B_traits, C_traits, block_traits
|
|
|
|
|
|
def get_index_map(block: tirx.SBlock) -> tuple[tirx.IndexMap, ...] | None:
|
|
"""Get index maps for the block
|
|
|
|
Parameters
|
|
----------
|
|
block : tirx.SBlock
|
|
The block to be analyzed
|
|
|
|
Returns
|
|
-------
|
|
index_maps : Optional[Tuple[tirx.IndexMap]]
|
|
The index maps for the block, or None if the block is not a gemm-liked kernel
|
|
"""
|
|
traits = detect_iter_traits(block)
|
|
if traits is None:
|
|
return None
|
|
A_traits, B_traits, C_traits, block_traits = traits
|
|
|
|
A_index_map = make_iter_fusion_index_map(
|
|
A_traits, [IterKind.kIter_S, IterKind.kIter_I, IterKind.kIter_K]
|
|
)
|
|
B_index_map = make_iter_fusion_index_map(
|
|
B_traits, [IterKind.kIter_S, IterKind.kIter_J, IterKind.kIter_K]
|
|
)
|
|
C_index_map = make_iter_fusion_index_map(
|
|
C_traits, [IterKind.kIter_S, IterKind.kIter_I, IterKind.kIter_J]
|
|
)
|
|
matmul_index_map = make_iter_fusion_index_map(
|
|
block_traits, [IterKind.kIter_S, IterKind.kIter_I, IterKind.kIter_J, IterKind.kIter_K]
|
|
)
|
|
|
|
return (
|
|
matmul_index_map,
|
|
A_index_map,
|
|
B_index_map,
|
|
C_index_map,
|
|
)
|
|
|
|
|
|
def get_sblock_info(sch: s_tir.Schedule, block: s_tir.schedule.SBlockRV) -> SBlockInfo:
|
|
def _iter_kind(loop: tirx.IterVar) -> str:
|
|
return {tirx.IterVar.DataPar: "S", tirx.IterVar.CommReduce: "R"}.get(loop.iter_type, "O")
|
|
|
|
def _is_reduction_block(block: s_tir.schedule.SBlockRV):
|
|
for iter_var in sch.get(block).iter_vars:
|
|
if _iter_kind(iter_var) == "R":
|
|
return True
|
|
return False
|
|
|
|
return SBlockInfo(
|
|
name=sch.get(block).name_hint,
|
|
iters=[
|
|
IterInfo(
|
|
kind=_iter_kind(iter_var),
|
|
var=iter_var.var,
|
|
dom=iter_var.dom.extent,
|
|
loop_rv=loop_rv,
|
|
)
|
|
for loop_rv, iter_var in zip(sch.get_loops(block), sch.get(block).iter_vars)
|
|
],
|
|
block_rv=block,
|
|
reduction_block=_is_reduction_block(block),
|
|
)
|
|
|
|
|
|
def get_reduction_blocks(sch, blocks) -> bool:
|
|
# Get the main computation block
|
|
def is_reduction(block: SBlockRV) -> bool:
|
|
block_stmt = sch.get(block)
|
|
iter_types = {iter_var.iter_type for iter_var in block_stmt.iter_vars}
|
|
return iter_types == {IterVar.CommReduce, IterVar.DataPar}
|
|
|
|
def is_spatial(block: SBlockRV) -> bool:
|
|
block_stmt = sch.get(block)
|
|
iter_types = {iter_var.iter_type for iter_var in block_stmt.iter_vars}
|
|
return iter_types == {IterVar.DataPar}
|
|
|
|
# NOTE: We assume there is only one reduction block in the function
|
|
# all blocks are required to be spatial or reduction
|
|
if not all([is_reduction(block) or is_spatial(block) for block in blocks]):
|
|
return None
|
|
|
|
# There is only one reduction block
|
|
reduction_blocks = [block for block in blocks if is_reduction(block)]
|
|
if len(reduction_blocks) != 1:
|
|
return None
|
|
|
|
return reduction_blocks
|
|
|
|
|
|
def get_in_out_dtypes(block: tirx.SBlock) -> tuple[str]:
|
|
"""
|
|
Detect In/Out data types for the given block based on the analysis if read/write buffers.
|
|
"""
|
|
assert len(block.reads) > 0 and len(block.writes) > 0
|
|
in_dtype = block.reads[0].buffer.dtype
|
|
out_dtype = block.writes[0].buffer.dtype
|
|
return (in_dtype, out_dtype)
|
|
|
|
|
|
def check_sm_version(arch: str) -> int:
|
|
sm_version = arch.replace("sm_", "")
|
|
return int(sm_version) if sm_version.isdigit() else -1
|
|
|
|
|
|
class MetalMatmul(GPUScheduleRule):
|
|
"""
|
|
The schedule rule for Metal matmul computation.
|
|
"""
|
|
|
|
def apply( # pylint: disable=too-many-locals,missing-docstring
|
|
self,
|
|
func: tirx.PrimFunc,
|
|
target: Target,
|
|
_: bool,
|
|
) -> s_tir.Schedule | None:
|
|
from tvm.s_tir.tensor_intrin.metal import ( # pylint: disable=import-outside-toplevel
|
|
get_simdgroup_intrin_group,
|
|
)
|
|
|
|
if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
|
|
return None
|
|
sch = s_tir.Schedule(func)
|
|
root_block = get_root_block(sch)
|
|
blocks = sch.get_child_blocks(root_block)
|
|
|
|
reduction_blocks = get_reduction_blocks(sch, blocks)
|
|
if reduction_blocks is None:
|
|
return None
|
|
|
|
# Step 0. Configs
|
|
block_size_x: int = 16
|
|
block_size_y: int = 16
|
|
block_size_k: int = 32
|
|
micro_size: int = 8
|
|
warp_size: int = 32
|
|
ty_len: int = 1
|
|
tz_len: int = 4
|
|
vector_size: int = 4
|
|
|
|
# Step 1. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K]
|
|
# Reindex first and than analyze the index map
|
|
main_block = reduction_blocks[0]
|
|
reindex_a = sch.reindex(main_block, ("read", 0))
|
|
reindex_b = sch.reindex(main_block, ("read", 1))
|
|
reindex_c = sch.reindex(main_block, ("write", 0))
|
|
|
|
index_maps = get_index_map(sch.get(main_block))
|
|
assert index_maps is not None
|
|
matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps
|
|
|
|
sch.transform_layout(reindex_a, ("write", 0), a_index_map)
|
|
sch.transform_layout(reindex_b, ("write", 0), b_index_map)
|
|
sch.transform_layout(reindex_c, ("read", 0), c_index_map)
|
|
sch.transform_block_layout(main_block, matmul_index_map)
|
|
|
|
# Step 2. Padding for dynamic shape kernels
|
|
sch.pad_einsum(
|
|
main_block,
|
|
[
|
|
1,
|
|
ty_len * block_size_x,
|
|
tz_len * block_size_y,
|
|
block_size_k,
|
|
],
|
|
)
|
|
|
|
# Step 3. Schedule matmul to use simdgroup intrinsics
|
|
batch, i, j, k = sch.get_loops(main_block)
|
|
bx, ty, i0, i1 = sch.split(i, [None, ty_len, block_size_x // micro_size, micro_size])
|
|
by, tz, j0, j1 = sch.split(j, [None, tz_len, block_size_y // micro_size, micro_size])
|
|
k0, k1, k2 = sch.split(k, [None, block_size_k // micro_size, micro_size])
|
|
sch.reorder(bx, by, ty, tz, k0, k1, i0, j0, i1, j1, k2)
|
|
sch.bind(bx, "blockIdx.x")
|
|
sch.bind(by, "blockIdx.y")
|
|
sch.bind(batch, "blockIdx.z")
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tz, "threadIdx.z")
|
|
|
|
def fetch_to_shared(block, idx):
|
|
block_read = sch.cache_read(block, idx, "shared")
|
|
sch.compute_at(block_read, k0, preserve_unit_loops=True)
|
|
fused = sch.fuse(*sch.get_loops(block_read)[-2:])
|
|
_, _tz, _ty, _tx, vec = sch.split(fused, [None, tz_len, ty_len, warp_size, vector_size])
|
|
|
|
sch.bind(_tz, "threadIdx.z")
|
|
sch.bind(_ty, "threadIdx.y")
|
|
sch.bind(_tx, "threadIdx.x")
|
|
sch.vectorize(vec)
|
|
|
|
return block_read
|
|
|
|
a_g2s = fetch_to_shared(main_block, 0)
|
|
b_g2s = fetch_to_shared(main_block, 1)
|
|
|
|
auto_inline_producers(sch, a_g2s)
|
|
auto_inline_producers(sch, b_g2s)
|
|
|
|
# create read cache to load matrix from shared memory to wmma fragments
|
|
A_simdgroup = sch.cache_read(main_block, 0, "metal.simdgroup")
|
|
B_simdgroup = sch.cache_read(main_block, 1, "metal.simdgroup")
|
|
sch.compute_at(A_simdgroup, k1)
|
|
sch.compute_at(B_simdgroup, k1)
|
|
|
|
C_simd2s = sch.cache_write(main_block, 0, "metal.simdgroup")
|
|
C_s2g = sch.cache_write(C_simd2s, 0, "shared")
|
|
sch.reverse_compute_at(C_simd2s, tz, preserve_unit_loops=True)
|
|
sch.reverse_compute_at(C_s2g, by, preserve_unit_loops=True)
|
|
|
|
intrin_group = get_simdgroup_intrin_group(
|
|
load_scope="shared",
|
|
store_scope="shared",
|
|
dtype="float16",
|
|
trans_a=False,
|
|
trans_b=True,
|
|
)
|
|
sch.transform_layout(B_simdgroup, ("write", 0), lambda s, i, j: (s, j, i))
|
|
|
|
def tensorize_block(block: s_tir.schedule.SBlockRV, intrin: str):
|
|
*_, i, j = sch.get_loops(block)
|
|
io, ii = sch.split(i, [None, micro_size])
|
|
jo, ji = sch.split(j, [None, micro_size])
|
|
sch.reorder(io, jo, ii, ji)
|
|
sch.tensorize(ii, intrin)
|
|
|
|
C_init = sch.decompose_reduction(main_block, k0)
|
|
tensorize_block(A_simdgroup, intrin_group["load_a"])
|
|
tensorize_block(B_simdgroup, intrin_group["load_b"])
|
|
tensorize_block(C_simd2s, intrin_group["store"])
|
|
tensorize_block(C_init, intrin_group["init"])
|
|
|
|
*_, i, j, k = sch.get_loops(main_block)
|
|
sch.tensorize(i, intrin_group["compute"])
|
|
|
|
auto_inline_consumer_chain(sch, C_s2g)
|
|
fused = sch.fuse(*sch.get_loops(C_s2g)[-2:])
|
|
_, _tz, _ty, _tx, vec = sch.split(fused, [None, tz_len, ty_len, warp_size, vector_size])
|
|
sch.bind(_tz, "threadIdx.z")
|
|
sch.bind(_ty, "threadIdx.y")
|
|
sch.bind(_tx, "threadIdx.x")
|
|
sch.vectorize(vec)
|
|
|
|
return sch
|
|
|
|
|
|
class MatmulTensorization(GPUScheduleRule):
|
|
"""
|
|
The schedule rule for float16 tensor core matmul computation.
|
|
func with attr 'dlight.do_not_tensorize' will not be tensorized.
|
|
"""
|
|
|
|
def apply( # pylint: disable=too-many-locals,missing-docstring
|
|
self,
|
|
func: tirx.PrimFunc,
|
|
target: Target,
|
|
_: bool,
|
|
) -> s_tir.Schedule | None:
|
|
from tvm.s_tir.tensor_intrin.cuda import ( # pylint: disable=import-outside-toplevel
|
|
get_wmma_intrin_group,
|
|
)
|
|
|
|
if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
|
|
return None
|
|
sch = s_tir.Schedule(func)
|
|
root_block = get_root_block(sch)
|
|
blocks = sch.get_child_blocks(root_block)
|
|
|
|
if "dlight.do_not_tensorize" in func.attrs.keys():
|
|
return None
|
|
|
|
reduction_blocks = get_reduction_blocks(sch, blocks)
|
|
if reduction_blocks is None:
|
|
return None
|
|
|
|
# Start Schedule
|
|
# Step 0. Get schedule config.
|
|
# NOTE: we can analyze the config by the hardware spec in the future
|
|
|
|
# tensor core intrinsic size
|
|
micro_size_x = 16
|
|
micro_size_y = 16
|
|
micro_size_k = 16
|
|
|
|
warp_size = 32
|
|
vector_size = 4
|
|
|
|
i_factors, j_factors, k_factors = (
|
|
[None, 1, 4, 2],
|
|
[1, None, 4, 2],
|
|
[None, 4],
|
|
)
|
|
|
|
num_ty = i_factors[2] * j_factors[2]
|
|
x_pad_factor = i_factors[2] * i_factors[3]
|
|
y_pad_factor = j_factors[2] * j_factors[3]
|
|
k_pad_factor = k_factors[1]
|
|
|
|
# Step 1. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K]
|
|
# Reindex first and than analyze the index map
|
|
main_block = reduction_blocks[0]
|
|
reindex_a = sch.reindex(main_block, ("read", 0))
|
|
reindex_b = sch.reindex(main_block, ("read", 1))
|
|
reindex_c = sch.reindex(main_block, ("write", 0))
|
|
|
|
index_maps = get_index_map(sch.get(main_block))
|
|
assert index_maps is not None
|
|
matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps
|
|
|
|
sch.transform_layout(reindex_a, ("write", 0), a_index_map)
|
|
sch.transform_layout(reindex_b, ("write", 0), b_index_map)
|
|
sch.transform_layout(reindex_c, ("read", 0), c_index_map)
|
|
sch.transform_block_layout(main_block, matmul_index_map)
|
|
|
|
# Step 2. Padding for dynamic shape kernels
|
|
sch.pad_einsum(
|
|
main_block,
|
|
[
|
|
1,
|
|
micro_size_x * x_pad_factor,
|
|
micro_size_y * y_pad_factor,
|
|
micro_size_k * k_pad_factor,
|
|
],
|
|
)
|
|
|
|
# Step 3. Schedule matmul to use tensor core
|
|
block = main_block
|
|
|
|
batch, i, j, k = sch.get_loops(block)
|
|
|
|
# inner loops for tensor core computation
|
|
i, i_inner = sch.split(i, factors=[None, micro_size_x])
|
|
j, j_inner = sch.split(j, factors=[None, micro_size_y])
|
|
k, k_inner = sch.split(k, factors=[None, micro_size_k])
|
|
|
|
sch.reorder(i, j, k, i_inner, j_inner, k_inner)
|
|
|
|
block_inner = block
|
|
block_outer = sch.blockize(i_inner)
|
|
|
|
i0, i1, i2, i3 = sch.split(i, factors=i_factors)
|
|
j0, j1, j2, j3 = sch.split(j, factors=j_factors)
|
|
k0, k1 = sch.split(k, k_factors)
|
|
sch.annotate(k0, "software_pipeline_order", [0, 3, 1, 4, 5, 2, 6])
|
|
sch.annotate(k0, "software_pipeline_stage", [0, 0, 0, 0, 0, 1, 1])
|
|
sch.annotate(k1, "software_pipeline_order", [0, 1, 2])
|
|
sch.annotate(k1, "software_pipeline_stage", [0, 0, 1])
|
|
|
|
sch.reorder(i0, j0, i1, j1, j2, i2, k0, k1, i3, j3)
|
|
|
|
block_idx = sch.fuse(i0, j0)
|
|
block_idy = sch.fuse(i1, j1)
|
|
thread_idy = sch.fuse(j2, i2)
|
|
sch.bind(batch, "blockIdx.z")
|
|
sch.bind(block_idx, "blockIdx.x")
|
|
sch.bind(block_idy, "blockIdx.y")
|
|
sch.bind(thread_idy, "threadIdx.y")
|
|
|
|
def fetch_to_shared(block, idx, ndim):
|
|
block_read = sch.cache_read(block, idx, "shared.dyn")
|
|
sch.compute_at(block_read, k0)
|
|
fused = sch.fuse(*sch.get_loops(block_read)[-ndim:])
|
|
|
|
_, f_1, f_2, f_3 = sch.split(fused, factors=[None, num_ty, warp_size, vector_size])
|
|
|
|
sch.bind(f_2, "threadIdx.x")
|
|
sch.bind(f_1, "threadIdx.y")
|
|
sch.vectorize(f_3)
|
|
|
|
sch.storage_align(block_read, 0, axis=-2, factor=16, offset=8)
|
|
sch.annotate(block_read, "tirx.manifest_shared_memory_local_stage", 1)
|
|
sch.annotate(block_read, "double_buffer_scope", 0)
|
|
return block_read
|
|
|
|
a_g2s = fetch_to_shared(block_outer, 0, 2)
|
|
b_g2s = fetch_to_shared(block_outer, 1, 2)
|
|
|
|
auto_inline_producers(sch, a_g2s)
|
|
auto_inline_producers(sch, b_g2s)
|
|
|
|
# create read cache to load matrix from shared memory to wmma fragments
|
|
A_mat = sch.cache_read(block_outer, 0, "wmma.matrix_a")
|
|
B_mat = sch.cache_read(block_outer, 1, "wmma.matrix_b")
|
|
sch.compute_at(A_mat, k1)
|
|
sch.compute_at(B_mat, k1)
|
|
|
|
# create write cache to store matrix from wmma fragments to shared memory and global memory
|
|
accumulator_shared_to_global = sch.cache_write(block_outer, 0, "shared.dyn")
|
|
sch.storage_align(accumulator_shared_to_global, 0, -2, 16, 4)
|
|
|
|
store = sch.cache_write(block_outer, 0, "wmma.accumulator")
|
|
sch.reverse_compute_at(store, thread_idy)
|
|
sch.reverse_compute_at(accumulator_shared_to_global, thread_idy)
|
|
|
|
# split the store loop to match hardware intrinsic pattern
|
|
i, j = sch.get_loops(store)[-2:]
|
|
i0, i1 = sch.split(i, factors=[None, 16])
|
|
j0, j1 = sch.split(j, factors=[None, 16])
|
|
sch.reorder(i0, j0, i1, j1)
|
|
|
|
block_init_c = sch.decompose_reduction(block_outer, k0)
|
|
block_init_c_inner = sch.get_child_blocks(block_init_c)[0]
|
|
|
|
# Tensorization by hardware intrinsics
|
|
intrin_group = get_wmma_intrin_group(
|
|
load_scope="shared.dyn",
|
|
store_scope="shared.dyn",
|
|
in_dtype="float16",
|
|
out_dtype="float32",
|
|
trans_b=True,
|
|
)
|
|
|
|
try:
|
|
i, j = sch.get_loops(A_mat)[-2:]
|
|
i0, i1 = sch.split(i, factors=[None, 16])
|
|
j0, j1 = sch.split(j, factors=[None, 16])
|
|
sch.reorder(i0, j0, i1, j1)
|
|
sch.unroll(i0)
|
|
sch.unroll(j0)
|
|
sch.tensorize(i1, intrin_group["load_a"])
|
|
|
|
i, j = sch.get_loops(B_mat)[-2:]
|
|
i0, i1 = sch.split(i, factors=[None, 16])
|
|
j0, j1 = sch.split(j, factors=[None, 16])
|
|
sch.reorder(i0, j0, i1, j1)
|
|
sch.unroll(i0)
|
|
sch.unroll(j0)
|
|
sch.tensorize(i1, intrin_group["load_b"])
|
|
except: # pylint: disable=bare-except
|
|
return None
|
|
|
|
# Try to tensorize the init, store and compute block with f16 or f32 intrinsics
|
|
tensorize_success: bool = False
|
|
|
|
def tensorize_init_store_compute():
|
|
sch.tensorize(sch.get_loops(block_init_c_inner)[-2], intrin_group["init"])
|
|
sch.tensorize(sch.get_loops(store)[-2], intrin_group["store"])
|
|
sch.tensorize(sch.get_loops(block_inner)[-3], intrin_group["compute"])
|
|
|
|
try:
|
|
tensorize_init_store_compute()
|
|
tensorize_success = True
|
|
except: # pylint: disable=bare-except
|
|
intrin_group = get_wmma_intrin_group(
|
|
load_scope="shared.dyn",
|
|
store_scope="shared.dyn",
|
|
in_dtype="float16",
|
|
out_dtype="float16",
|
|
trans_b=True,
|
|
)
|
|
|
|
if not tensorize_success:
|
|
try:
|
|
tensorize_init_store_compute()
|
|
tensorize_success = True
|
|
except: # pylint: disable=bare-except
|
|
return None
|
|
auto_inline_consumer_chain(sch, accumulator_shared_to_global)
|
|
|
|
fused = sch.fuse(*sch.get_loops(accumulator_shared_to_global)[-2:])
|
|
_, f1, f2 = sch.split(fused, factors=[None, warp_size, vector_size])
|
|
sch.bind(f1, "threadIdx.x")
|
|
sch.vectorize(f2)
|
|
|
|
return sch if tensorize_success else None
|
|
|
|
|
|
class MatmulInt8Tensorization(GPUScheduleRule):
|
|
"""
|
|
The schedule rule for int8 tensor core matmul computation.
|
|
func with attr 'dlight.do_not_tensorize' will not be tensorized.
|
|
"""
|
|
|
|
def apply( # pylint: disable=too-many-locals,missing-docstring
|
|
self,
|
|
func: tirx.PrimFunc,
|
|
target: Target,
|
|
_: bool,
|
|
) -> s_tir.Schedule | None:
|
|
from tvm.s_tir.tensor_intrin.cuda import ( # pylint: disable=import-outside-toplevel
|
|
get_wmma_intrin_group,
|
|
)
|
|
|
|
if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
|
|
return None
|
|
sch = s_tir.Schedule(func)
|
|
root_block = get_root_block(sch)
|
|
blocks = sch.get_child_blocks(root_block)
|
|
|
|
if "dlight.do_not_tensorize" in func.attrs.keys():
|
|
return None
|
|
|
|
reduction_blocks = get_reduction_blocks(sch, blocks)
|
|
if reduction_blocks is None:
|
|
return None
|
|
|
|
# Start Schedule
|
|
# Step 0. Get schedule config.
|
|
# NOTE: we can analyze the config by the hardware spec in the future
|
|
|
|
# tensor core intrinsic size
|
|
micro_size_x = 16
|
|
micro_size_y = 16
|
|
micro_size_k = 16
|
|
|
|
warp_size = 32
|
|
vector_size = 4
|
|
|
|
i_factors, j_factors, k_factors = (
|
|
[None, 1, 4, 2],
|
|
[1, None, 4, 2],
|
|
[None, 1],
|
|
)
|
|
|
|
num_ty = i_factors[2] * j_factors[2]
|
|
x_pad_factor = i_factors[2] * i_factors[3]
|
|
y_pad_factor = j_factors[2] * j_factors[3]
|
|
k_pad_factor = k_factors[1]
|
|
|
|
# Step 1. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K]
|
|
# Reindex first and than analyze the index map
|
|
main_block = reduction_blocks[0]
|
|
reindex_a = sch.reindex(main_block, ("read", 0))
|
|
reindex_b = sch.reindex(main_block, ("read", 1))
|
|
reindex_c = sch.reindex(main_block, ("write", 0))
|
|
|
|
index_maps = get_index_map(sch.get(main_block))
|
|
assert index_maps is not None
|
|
matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps
|
|
|
|
sch.transform_layout(reindex_a, ("write", 0), a_index_map)
|
|
sch.transform_layout(reindex_b, ("write", 0), b_index_map)
|
|
sch.transform_layout(reindex_c, ("read", 0), c_index_map)
|
|
sch.transform_block_layout(main_block, matmul_index_map)
|
|
|
|
# Step 2. Padding for dynamic shape kernels
|
|
sch.pad_einsum(
|
|
main_block,
|
|
[
|
|
1,
|
|
micro_size_x * x_pad_factor,
|
|
micro_size_y * y_pad_factor,
|
|
micro_size_k * k_pad_factor,
|
|
],
|
|
)
|
|
|
|
# Step 3. Schedule matmul to use tensor core
|
|
block = main_block
|
|
|
|
batch, i, j, k = sch.get_loops(block)
|
|
|
|
# inner loops for tensor core computation
|
|
i, i_inner = sch.split(i, factors=[None, micro_size_x])
|
|
j, j_inner = sch.split(j, factors=[None, micro_size_y])
|
|
k, k_inner = sch.split(k, factors=[None, micro_size_k])
|
|
|
|
sch.reorder(i, j, k, i_inner, j_inner, k_inner)
|
|
|
|
block_inner = block
|
|
block_outer = sch.blockize(i_inner)
|
|
|
|
i0, i1, i2, i3 = sch.split(i, factors=i_factors)
|
|
j0, j1, j2, j3 = sch.split(j, factors=j_factors)
|
|
k0, k1 = sch.split(k, k_factors)
|
|
sch.annotate(k0, "software_pipeline_order", [0, 3, 1, 4, 5, 2, 6])
|
|
sch.annotate(k0, "software_pipeline_stage", [0, 0, 0, 0, 0, 1, 1])
|
|
sch.annotate(k1, "software_pipeline_order", [0, 1, 2])
|
|
sch.annotate(k1, "software_pipeline_stage", [0, 0, 1])
|
|
|
|
sch.reorder(i0, j0, i1, j1, j2, i2, k0, k1, i3, j3)
|
|
|
|
block_idx = sch.fuse(i0, j0)
|
|
block_idy = sch.fuse(i1, j1)
|
|
thread_idy = sch.fuse(j2, i2)
|
|
sch.bind(batch, "blockIdx.z")
|
|
sch.bind(block_idx, "blockIdx.x")
|
|
sch.bind(block_idy, "blockIdx.y")
|
|
sch.bind(thread_idy, "threadIdx.y")
|
|
|
|
def fetch_to_shared(block, idx, ndim):
|
|
block_read = sch.cache_read(block, idx, "shared.dyn")
|
|
sch.compute_at(block_read, k0)
|
|
fused = sch.fuse(*sch.get_loops(block_read)[-ndim:])
|
|
|
|
_, f_1, f_2, f_3 = sch.split(fused, factors=[None, num_ty, warp_size, vector_size])
|
|
|
|
sch.bind(f_2, "threadIdx.x")
|
|
sch.bind(f_1, "threadIdx.y")
|
|
sch.vectorize(f_3)
|
|
|
|
sch.storage_align(block_read, 0, axis=-2, factor=32, offset=16)
|
|
sch.annotate(block_read, "tirx.manifest_shared_memory_local_stage", 1)
|
|
sch.annotate(block_read, "double_buffer_scope", 0)
|
|
return block_read
|
|
|
|
a_g2s = fetch_to_shared(block_outer, 0, 2)
|
|
b_g2s = fetch_to_shared(block_outer, 1, 2)
|
|
|
|
auto_inline_producers(sch, a_g2s)
|
|
auto_inline_producers(sch, b_g2s)
|
|
|
|
# create read cache to load matrix from shared memory to wmma fragments
|
|
A_mat = sch.cache_read(block_outer, 0, "wmma.matrix_a")
|
|
B_mat = sch.cache_read(block_outer, 1, "wmma.matrix_b")
|
|
sch.compute_at(A_mat, k1)
|
|
sch.compute_at(B_mat, k1)
|
|
|
|
# create write cache to store matrix from wmma fragments to shared memory and global memory
|
|
accumulator_shared_to_global = sch.cache_write(block_outer, 0, "shared.dyn")
|
|
sch.storage_align(accumulator_shared_to_global, 0, -2, 16, 4)
|
|
|
|
store = sch.cache_write(block_outer, 0, "wmma.accumulator")
|
|
sch.reverse_compute_at(store, thread_idy)
|
|
sch.reverse_compute_at(accumulator_shared_to_global, thread_idy)
|
|
|
|
# split the store loop to match hardware intrinsic pattern
|
|
i, j = sch.get_loops(store)[-2:]
|
|
i0, i1 = sch.split(i, factors=[None, 16])
|
|
j0, j1 = sch.split(j, factors=[None, 16])
|
|
sch.reorder(i0, j0, i1, j1)
|
|
|
|
block_init_c = sch.decompose_reduction(block_outer, k0)
|
|
block_init_c_inner = sch.get_child_blocks(block_init_c)[0]
|
|
|
|
# Tensorization by hardware intrinsics
|
|
intrin_group = get_wmma_intrin_group(
|
|
load_scope="shared.dyn",
|
|
store_scope="shared.dyn",
|
|
in_dtype="int8",
|
|
out_dtype="int32",
|
|
trans_b=True,
|
|
)
|
|
|
|
try:
|
|
i, j = sch.get_loops(A_mat)[-2:]
|
|
i0, i1 = sch.split(i, factors=[None, 16])
|
|
j0, j1 = sch.split(j, factors=[None, 16])
|
|
sch.reorder(i0, j0, i1, j1)
|
|
sch.unroll(i0)
|
|
sch.unroll(j0)
|
|
sch.tensorize(i1, intrin_group["load_a"])
|
|
|
|
i, j = sch.get_loops(B_mat)[-2:]
|
|
i0, i1 = sch.split(i, factors=[None, 16])
|
|
j0, j1 = sch.split(j, factors=[None, 16])
|
|
sch.reorder(i0, j0, i1, j1)
|
|
sch.unroll(i0)
|
|
sch.unroll(j0)
|
|
sch.tensorize(i1, intrin_group["load_b"])
|
|
except: # pylint: disable=bare-except
|
|
return None
|
|
|
|
def tensorize_init_store_compute():
|
|
sch.tensorize(sch.get_loops(block_init_c_inner)[-2], intrin_group["init"])
|
|
sch.tensorize(sch.get_loops(store)[-2], intrin_group["store"])
|
|
sch.tensorize(sch.get_loops(block_inner)[-3], intrin_group["compute"])
|
|
|
|
try:
|
|
tensorize_init_store_compute()
|
|
except: # pylint: disable=bare-except
|
|
return None
|
|
|
|
auto_inline_consumer_chain(sch, accumulator_shared_to_global)
|
|
|
|
fused = sch.fuse(*sch.get_loops(accumulator_shared_to_global)[-2:])
|
|
_, f1, f2 = sch.split(fused, factors=[None, warp_size, vector_size])
|
|
sch.bind(f1, "threadIdx.x")
|
|
sch.vectorize(f2)
|
|
|
|
return sch
|
|
|
|
|
|
class Matmul(GPUScheduleRule):
|
|
"""The schedule rule for matmul-like computation"""
|
|
|
|
@dataclass
|
|
class Config:
|
|
block_size_x: int = 8
|
|
block_size_y: int = 8
|
|
vthread_x: int = 1
|
|
vthread_y: int = 1
|
|
micro_size_x: int = 4
|
|
micro_size_y: int = 4
|
|
micro_size_k: int = 8
|
|
vector_size: int = 1
|
|
unroll: int = 256 # 0 means no unroll
|
|
use_shared: bool = True
|
|
storage_align: bool = False
|
|
inner_x: bool = False
|
|
|
|
def get_configs(self, target: Target) -> Config:
|
|
"""Get the schedule config for the target"""
|
|
if target.kind.name == "cuda" or target.kind.name == "rocm":
|
|
return Matmul.Config(
|
|
block_size_x=8,
|
|
block_size_y=16,
|
|
vthread_x=1,
|
|
vthread_y=1,
|
|
micro_size_x=4,
|
|
micro_size_y=4,
|
|
micro_size_k=16,
|
|
vector_size=2,
|
|
unroll=256,
|
|
use_shared=True,
|
|
storage_align=True,
|
|
inner_x=False,
|
|
)
|
|
elif target.kind.name == "opencl" and (
|
|
("android" in str(target.host)) or ("adreno" in str(target.attrs))
|
|
):
|
|
return Matmul.Config(
|
|
block_size_x=32,
|
|
block_size_y=4,
|
|
vthread_x=1,
|
|
vthread_y=1,
|
|
micro_size_x=8,
|
|
micro_size_y=2,
|
|
micro_size_k=16,
|
|
vector_size=8,
|
|
unroll=16,
|
|
use_shared=False,
|
|
storage_align=False,
|
|
inner_x=True,
|
|
)
|
|
else:
|
|
return Matmul.Config()
|
|
|
|
def apply( # pylint: disable=too-many-locals,missing-docstring
|
|
self,
|
|
func: tirx.PrimFunc,
|
|
target: Target,
|
|
_: bool,
|
|
) -> s_tir.Schedule | None:
|
|
if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
|
|
return None
|
|
sch = s_tir.Schedule(func)
|
|
config = self.get_configs(target)
|
|
root_block = get_root_block(sch)
|
|
blocks = sch.get_child_blocks(root_block)
|
|
|
|
reduction_blocks = get_reduction_blocks(sch, blocks)
|
|
if reduction_blocks is None:
|
|
return None
|
|
|
|
main_block = reduction_blocks[0]
|
|
block_stmt = sch.get(main_block)
|
|
|
|
main_block_info = get_sblock_info(sch, main_block)
|
|
iter_infos = main_block_info.iters
|
|
if not get_index_map(block_stmt):
|
|
return None
|
|
|
|
# Checks if it's a inner reduction by getting the last matrix's inner Index
|
|
def is_inner_reduction(block_stmt, iter_infos):
|
|
end_it = block_stmt.reads[-1].region[-1].min
|
|
return {it.var: it.kind for it in iter_infos}.get(end_it, "O") == "R"
|
|
|
|
if (
|
|
target.kind.name == "opencl"
|
|
and (("android" in str(target.host)) or ("adreno" in str(target.attrs)))
|
|
) and not is_inner_reduction(block_stmt, iter_infos):
|
|
ret = self.sch_outer_reduction(sch, config, main_block, blocks)
|
|
if ret is not None:
|
|
return ret
|
|
|
|
# Step 0. Normalize generic matmul to C[S, I, J] += A[S, I, K] * B[S, J, K]
|
|
# Reindex first and than analyze the index map
|
|
reindex_a = sch.reindex(main_block, ("read", 0))
|
|
reindex_b = sch.reindex(main_block, ("read", 1))
|
|
reindex_c = sch.reindex(main_block, ("write", 0))
|
|
|
|
index_maps = get_index_map(sch.get(main_block))
|
|
assert index_maps is not None
|
|
matmul_index_map, a_index_map, b_index_map, c_index_map = index_maps
|
|
|
|
sch.transform_layout(reindex_a, ("write", 0), a_index_map)
|
|
sch.transform_layout(reindex_b, ("write", 0), b_index_map)
|
|
sch.transform_layout(reindex_c, ("read", 0), c_index_map)
|
|
sch.transform_block_layout(main_block, matmul_index_map)
|
|
|
|
# Step 1. Check Tensor Core support
|
|
# Tensorization config:
|
|
# If any value of I, J, K is fixed and less than this threshold,
|
|
# tensorization rule will not be applied.
|
|
minimal_tensorize_threshold = 64
|
|
block_stmt = sch.get(main_block)
|
|
if target.kind.name == "cuda" and check_sm_version(str(target.attrs.get("arch", ""))) >= 70:
|
|
apply_tensorization: bool = True
|
|
# the batch dimension is not taken into consideration.
|
|
for item_var in block_stmt.iter_vars[1:]:
|
|
extent = item_var.dom.extent
|
|
if isinstance(extent, tirx.expr.IntImm):
|
|
if extent.value <= minimal_tensorize_threshold:
|
|
apply_tensorization = False
|
|
if apply_tensorization:
|
|
# Analyze read/write buffers and choose correct tensorizer: int8 or fp16.
|
|
in_dtype, out_dtype = get_in_out_dtypes(block_stmt)
|
|
tensorize_sch = None
|
|
if in_dtype == "int8" and out_dtype == "int32":
|
|
tensorize_sch = MatmulInt8Tensorization().apply(func, target, _)
|
|
elif in_dtype == "float16" and out_dtype in ["float16", "float32"]:
|
|
tensorize_sch = MatmulTensorization().apply(func, target, _)
|
|
if tensorize_sch is not None:
|
|
return tensorize_sch
|
|
elif target.kind.name == "metal":
|
|
try:
|
|
return MetalMatmul().apply(func, target, _)
|
|
except: # pylint: disable=bare-except
|
|
pass
|
|
|
|
# Step 2. Schedule matmul
|
|
y_kernel_size = config.vthread_y * config.block_size_y * config.micro_size_y
|
|
x_kernel_size = config.vthread_x * config.block_size_x * config.micro_size_x
|
|
if config.inner_x:
|
|
sch.pad_einsum(
|
|
main_block,
|
|
[1, y_kernel_size, x_kernel_size, config.micro_size_k],
|
|
)
|
|
batch, y, x, k = sch.get_loops(main_block)
|
|
else:
|
|
sch.pad_einsum(
|
|
main_block,
|
|
[1, x_kernel_size, y_kernel_size, config.micro_size_k],
|
|
)
|
|
batch, x, y, k = sch.get_loops(main_block)
|
|
by, vy, ty, yi = sch.split(
|
|
y, [None, config.vthread_y, config.block_size_y, config.micro_size_y]
|
|
)
|
|
bx, vx, tx, xi = sch.split(
|
|
x, [None, config.vthread_x, config.block_size_x, config.micro_size_x]
|
|
)
|
|
ko, ki = sch.split(k, factors=[None, config.micro_size_k])
|
|
reordered_loops = [by, bx, vy, vx, ty, tx, ko, ki] + (
|
|
[yi, xi] if config.inner_x else [xi, yi]
|
|
)
|
|
sch.reorder(*reordered_loops)
|
|
by = sch.fuse(batch, by)
|
|
sch.bind(bx, "blockIdx.x")
|
|
sch.bind(by, "blockIdx.y")
|
|
sch.bind(vy, "vthread.y")
|
|
sch.bind(vx, "vthread.x")
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
inner_loop = config.micro_size_x if config.inner_x else config.micro_size_y
|
|
if inner_loop % config.vector_size == 0:
|
|
_, v = sch.split(reordered_loops[-1], [None, config.vector_size])
|
|
sch.vectorize(v)
|
|
|
|
if config.unroll > 0:
|
|
sch.annotate(tx, ann_key="pragma_auto_unroll_max_step", ann_val=config.unroll)
|
|
sch.annotate(tx, ann_key="pragma_unroll_explicit", ann_val=1)
|
|
|
|
l2g = sch.cache_write(main_block, 0, "local")
|
|
sch.reverse_compute_at(l2g, tx, preserve_unit_loops=True)
|
|
if config.micro_size_x % config.vector_size == 0:
|
|
_, v = sch.split(sch.get_loops(l2g)[-1], [None, config.vector_size])
|
|
sch.vectorize(v)
|
|
|
|
if config.use_shared:
|
|
|
|
def _cooperative_fetch(index, vec_len):
|
|
block = sch.cache_read(main_block, index, "shared")
|
|
num_loops = len(sch.get_loops(block))
|
|
sch.compute_at(block, ko, preserve_unit_loops=True)
|
|
loops = sch.get_loops(block)[-num_loops:]
|
|
ty, tx, _, vec = sch.split(
|
|
sch.fuse(*loops),
|
|
factors=[config.block_size_y, config.block_size_x, None, vec_len],
|
|
)
|
|
sch.vectorize(vec)
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
if config.storage_align:
|
|
sch.storage_align(block, 0, axis=1, factor=8, offset=vec_len)
|
|
return block
|
|
|
|
a_g2s = _cooperative_fetch(0, vec_len=config.vector_size)
|
|
b_g2s = _cooperative_fetch(1, vec_len=config.vector_size)
|
|
|
|
auto_inline_producers(sch, a_g2s)
|
|
auto_inline_producers(sch, b_g2s)
|
|
else:
|
|
auto_inline_producers(sch, main_block)
|
|
|
|
auto_inline_consumer_chain(sch, l2g)
|
|
|
|
sch.decompose_reduction(main_block, ko)
|
|
return sch
|
|
|
|
def sch_outer_reduction(
|
|
self,
|
|
sch: s_tir.Schedule,
|
|
config: Config,
|
|
reduction_block: s_tir.schedule.SBlockRV,
|
|
blocks: list[s_tir.schedule.SBlockRV],
|
|
) -> s_tir.Schedule | None:
|
|
"""Get vectorization factor"""
|
|
|
|
def get_max_factor(n, factors):
|
|
factors = sorted(factors, reverse=True)
|
|
for factor in factors:
|
|
if n % factor == 0:
|
|
return factor
|
|
return 1
|
|
|
|
reduction_loops = sch.get_loops(reduction_block)
|
|
if not len(reduction_loops) == 4:
|
|
return None
|
|
|
|
mb, ms, n, k = reduction_loops
|
|
if not (
|
|
isinstance(sch.get(n).extent, tirx.IntImm)
|
|
and isinstance(sch.get(mb).extent, tirx.IntImm)
|
|
and not isinstance(sch.get(ms).extent, tirx.IntImm)
|
|
):
|
|
return None
|
|
|
|
Threads_X, Threads_Y, VecSize, Unroll_M = (
|
|
config.block_size_x,
|
|
config.block_size_y,
|
|
config.vector_size,
|
|
config.unroll,
|
|
)
|
|
|
|
VecSize = min(get_max_factor(sch.get(n).extent // Threads_X, [1, 2, 4, 8]), VecSize)
|
|
dequant_block = None
|
|
matmul_block = reduction_block
|
|
epilogue_block = None
|
|
if blocks[-1] is not matmul_block:
|
|
epilogue_block = blocks[-1]
|
|
for blk in blocks[:-1]:
|
|
if "dequantize" in sch.get(blk).name_hint:
|
|
dequant_block = blk
|
|
elif blk is not matmul_block:
|
|
sch.compute_inline(blk)
|
|
|
|
block = sch.reindex(reduction_block, ("read", 0))
|
|
sch.pad_einsum(reduction_block, [1, Unroll_M, 1, 1])
|
|
sch.compute_inline(block)
|
|
trans_block, matmul_reindex = (
|
|
sch.get_producers(matmul_block)[0],
|
|
sch.get_consumers(matmul_block)[0],
|
|
)
|
|
|
|
if epilogue_block is not None:
|
|
sch.compute_inline(matmul_reindex)
|
|
matmul_reindex = epilogue_block
|
|
|
|
sch.transform_layout(
|
|
trans_block,
|
|
("write", 0),
|
|
T.index_map(lambda i0, i1, i2: (i0, i1 // Unroll_M, i2, i1 % Unroll_M)),
|
|
)
|
|
|
|
# transpose block schedules
|
|
# sch.set_scope(trans_block, 0, "global.texture-1d")
|
|
tb, tn, tk = sch.get_loops(trans_block)
|
|
tbx, ttx = sch.split(tk, [None, Threads_X])
|
|
tby, tty, tc = sch.split(tn, [None, Threads_Y, Unroll_M])
|
|
sch.bind(tb, "blockIdx.z")
|
|
sch.bind(tby, "blockIdx.y")
|
|
sch.bind(tbx, "blockIdx.x")
|
|
sch.bind(tty, "threadIdx.y")
|
|
sch.bind(ttx, "threadIdx.x")
|
|
sch.reorder(tb, tby, tbx, tty, ttx, tc)
|
|
sch.vectorize(tc)
|
|
|
|
mb, ms, n, k = sch.get_loops(matmul_block)
|
|
m = sch.fuse(mb, ms)
|
|
bx, tx, vec = sch.split(n, [None, Threads_X, VecSize])
|
|
by, ty, unr = sch.split(m, [None, Threads_Y, Unroll_M])
|
|
k1, k2, k3 = sch.split(k, [None, 4, 8])
|
|
sch.reorder(bx, by, tx, ty, k1, k2, k3, unr, vec)
|
|
sch.set_scope(matmul_block, 0, "local")
|
|
if dequant_block is not None:
|
|
sch.compute_at(dequant_block, k3)
|
|
sch.set_scope(dequant_block, 0, "local")
|
|
sch.bind(by, "blockIdx.y")
|
|
sch.bind(bx, "blockIdx.x")
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
sch.vectorize(vec)
|
|
|
|
inp = sch.cache_read(matmul_block, read_buffer_index=0, storage_scope="local")
|
|
sch.compute_at(inp, k3, preserve_unit_loops=True)
|
|
sch.vectorize(sch.get_loops(inp)[-1])
|
|
|
|
sch.unroll(unr)
|
|
sch.unroll(k3)
|
|
|
|
if dequant_block is not None:
|
|
Aq_local = sch.cache_read(dequant_block, read_buffer_index=0, storage_scope="local")
|
|
sch.compute_at(Aq_local, k2, preserve_unit_loops=True)
|
|
sch.vectorize(sch.get_loops(Aq_local)[-1])
|
|
As_local = sch.cache_read(dequant_block, read_buffer_index=1, storage_scope="local")
|
|
sch.compute_at(As_local, k1, preserve_unit_loops=True)
|
|
sch.vectorize(sch.get_loops(As_local)[-1])
|
|
sch.vectorize(sch.get_loops(dequant_block)[-1])
|
|
|
|
sch.reverse_compute_at(matmul_reindex, ty)
|
|
o_ur, o_vec = sch.get_loops(matmul_reindex)[-2:]
|
|
sch.vectorize(o_vec)
|
|
sch.unroll(o_ur)
|
|
sch.decompose_reduction(matmul_block, k1)
|
|
|
|
return sch
|