wehub-resource-sync
130 lines
4.9 KiB
Python
130 lines
4.9 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.
|
|
"""Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
|
|
|
|
from tvm import arith, s_tir, tirx
|
|
from tvm.s_tir import Schedule
|
|
from tvm.s_tir.schedule import SBlockRV
|
|
from tvm.target import Target
|
|
|
|
from ..analysis import detect_dominant_read, normalize_prim_func
|
|
from ..base import try_inline_contiguous_spatial
|
|
from .base import GPUScheduleRule
|
|
|
|
|
|
class Transpose(GPUScheduleRule):
|
|
"""Schedule rule for transpose"""
|
|
|
|
def is_transpose(self, sch: Schedule, block_rv: SBlockRV):
|
|
block = sch.get(block_rv)
|
|
if isinstance(block.body, tirx.BufferStore):
|
|
rhs = block.body.value
|
|
if isinstance(rhs, tirx.BufferLoad):
|
|
lhs_indices = block.body.indices
|
|
rhs_indices = rhs.indices
|
|
if list(lhs_indices) != list(rhs_indices) and set(lhs_indices) == set(rhs_indices):
|
|
return True
|
|
return False
|
|
|
|
def apply( # pylint: disable=too-many-locals
|
|
self,
|
|
func: tirx.PrimFunc,
|
|
target: Target,
|
|
_: bool,
|
|
) -> None | s_tir.Schedule | list[s_tir.Schedule]:
|
|
# pylint: disable=invalid-name
|
|
if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
|
|
return None
|
|
if target.kind.name == "cuda":
|
|
len_tx = 16
|
|
len_ty = 8
|
|
unroll_depth = 256
|
|
elif target.kind.name == "opencl":
|
|
len_tx = 16
|
|
len_ty = 8
|
|
unroll_depth = 64
|
|
else:
|
|
len_tx = 8
|
|
len_ty = 4
|
|
unroll_depth = 64
|
|
len_vec = 4
|
|
|
|
sch = s_tir.Schedule(func)
|
|
blocks = normalize_prim_func(sch)
|
|
transpose_block_idx = -1
|
|
for idx, block in reversed(list(enumerate(blocks))):
|
|
if self.is_transpose(sch, block.block_rv):
|
|
transpose_block_idx = idx
|
|
break
|
|
if not block.is_injective():
|
|
return None
|
|
if transpose_block_idx == -1:
|
|
return None
|
|
transpose_block = blocks[transpose_block_idx].block_rv
|
|
|
|
prologue = None # the optional decoding block
|
|
if transpose_block_idx > 0:
|
|
spatials = try_inline_contiguous_spatial(sch, blocks[: transpose_block_idx - 1])
|
|
assert len(spatials) == 0
|
|
prologue = blocks[transpose_block_idx - 1].block_rv
|
|
|
|
loops = sch.get_loops(transpose_block)
|
|
if len(loops) != 2:
|
|
# transpose with more than 2 axes is not supported
|
|
return None
|
|
|
|
c_factor = 1
|
|
if prologue is not None:
|
|
block_stmt = sch.get(prologue)
|
|
result = arith.normalize_to_iter_sum(
|
|
detect_dominant_read(block_stmt),
|
|
input_iters={i.var: i.dom for i in block_stmt.iter_vars},
|
|
)
|
|
if len(result.args) > 0:
|
|
c_factor = int(result.args[0].lower_factor)
|
|
|
|
i, j = loops
|
|
i, vi = sch.split(i, factors=[None, c_factor], preserve_unit_iters=True)
|
|
bi, ti = sch.split(i, factors=[None, len_ty], preserve_unit_iters=True)
|
|
bj, tj = sch.split(j, factors=[None, len_tx], preserve_unit_iters=True)
|
|
sch.reorder(bi, bj, ti, tj, vi)
|
|
sch.bind(bi, "blockIdx.y")
|
|
sch.bind(bj, "blockIdx.x")
|
|
sch.bind(ti, "threadIdx.y")
|
|
sch.bind(tj, "threadIdx.x")
|
|
len_vec = min(len_vec, c_factor)
|
|
_, vi = sch.split(vi, factors=[None, len_vec])
|
|
if len_vec > 1:
|
|
sch.vectorize(vi)
|
|
|
|
cache_read = sch.cache_read(transpose_block, read_buffer_index=0, storage_scope="shared")
|
|
sch.compute_at(cache_read, bj)
|
|
loops = sch.get_loops(cache_read)[2:]
|
|
fused = sch.fuse(*loops)
|
|
_, ty, tx, v = sch.split(fused, factors=[None, len_ty, len_tx, c_factor])
|
|
sch.bind(ty, "threadIdx.y")
|
|
sch.bind(tx, "threadIdx.x")
|
|
sch.unroll(v)
|
|
sch.storage_align(block=cache_read, buffer_index=0, axis=0, factor=32, offset=1)
|
|
|
|
sch.annotate(bi, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
|
|
sch.annotate(bi, ann_key="pragma_unroll_explicit", ann_val=1)
|
|
|
|
if prologue is not None:
|
|
sch.compute_inline(prologue)
|
|
return sch
|