chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,183 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
# pylint: disable=invalid-name
|
||||
"""Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
|
||||
|
||||
from tvm import arith, s_tir, tirx
|
||||
from tvm.target import Target
|
||||
|
||||
from ..analysis import normalize_prim_func
|
||||
from ..base import try_inline_contiguous_spatial
|
||||
from .base import GPUScheduleRule
|
||||
|
||||
|
||||
class GeneralReduction(GPUScheduleRule):
|
||||
"""General Reduction rule for operators including softmax, layer norm, RMS norm, etc"""
|
||||
|
||||
def apply( # pylint: disable=too-many-locals
|
||||
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
|
||||
|
||||
if target.kind.name == "cuda":
|
||||
len_tx = 256
|
||||
unroll_depth = 256
|
||||
elif target.kind.name == "opencl":
|
||||
len_tx = 256
|
||||
unroll_depth = 64
|
||||
else:
|
||||
len_tx = 64
|
||||
unroll_depth = 64
|
||||
|
||||
sch = s_tir.Schedule(func)
|
||||
block_infos = normalize_prim_func(sch)
|
||||
block_infos = try_inline_contiguous_spatial(sch, block_infos)
|
||||
if block_infos is None or len(block_infos) == 0:
|
||||
return None
|
||||
|
||||
dom_kind = block_infos[0].dom_kind()
|
||||
num_leading_s = len(dom_kind) - len(dom_kind.lstrip("S"))
|
||||
num_trailing_r = len(dom_kind) - len(dom_kind.rstrip("R"))
|
||||
|
||||
# Align the number of block iters of the last block.
|
||||
num_last_block_iter = len(block_infos[-1].dom_kind())
|
||||
if num_last_block_iter < len(dom_kind):
|
||||
# If the last block is a scalar value, there is nothing left to
|
||||
# tile/parallelise, and `iters` is an empty tuple.
|
||||
# Add a unit thread loop so the final write happens inside a valid
|
||||
# GPU thread environment.
|
||||
if num_last_block_iter == 0:
|
||||
# Put every block (both the running reductions and the final
|
||||
# scalar write) inside a trivial GPU thread. The very first block
|
||||
# gets a `blockIdx.x` wrapper so that kernels still have a unique
|
||||
# block scope.
|
||||
for i, info in enumerate(block_infos):
|
||||
loop_rv = sch.add_unit_loop(info.block_rv)
|
||||
if i == 0:
|
||||
sch.bind(loop_rv, "blockIdx.x")
|
||||
else:
|
||||
sch.bind(loop_rv, "threadIdx.x")
|
||||
|
||||
return sch
|
||||
|
||||
def f_layout_mapping(*iters):
|
||||
analyzer = arith.Analyzer()
|
||||
# Try to match the iters of last block to the iters of the first block.
|
||||
# For matched positions, use the iter from the input `iters`.
|
||||
# For unmatched positions, use a new iter which is constant 0.
|
||||
num_matched = 0
|
||||
target_layout_iters = []
|
||||
for block_iter in block_infos[0].iters:
|
||||
if num_matched < len(iters) and analyzer.can_prove_equal(
|
||||
block_iter.dom, block_infos[-1].iters[num_matched].dom
|
||||
):
|
||||
target_layout_iters.append(iters[num_matched])
|
||||
num_matched += 1
|
||||
else:
|
||||
target_layout_iters.append(tirx.const(0, iters[0].ty))
|
||||
|
||||
# If all the iters of the last block can match, return the new layout.
|
||||
if num_matched == len(iters):
|
||||
return target_layout_iters
|
||||
# Otherwise, fallback to appending zeros in the beginning.
|
||||
return [tirx.const(0, iters[0].ty)] * (len(dom_kind) - num_last_block_iter) + list(
|
||||
iters
|
||||
)
|
||||
|
||||
index_map = tirx.IndexMap.from_func(f_layout_mapping, ndim=num_last_block_iter)
|
||||
sch.transform_block_layout(block_infos[-1].block_rv, index_map)
|
||||
|
||||
try:
|
||||
# TODO: fix num_leading_s = 0 case
|
||||
assert num_trailing_r > 0
|
||||
for block in block_infos[1:-1]:
|
||||
assert block.dom_kind() == dom_kind
|
||||
assert block_infos[-1].is_injective()
|
||||
assert len(block_infos[-1].dom_kind()) <= len(dom_kind)
|
||||
except AssertionError:
|
||||
return None
|
||||
|
||||
if "R" not in block_infos[-1].dom_kind():
|
||||
# The final block is a spatial block.
|
||||
# It is possible that the loop order of the last block is not the same as
|
||||
# previous blocks.
|
||||
# Thus we reorder spatial loops to align with reduction loops for followup schedule.
|
||||
# We first collect all the buffers written by reduction blocks,
|
||||
# then in the final block, any index of those buffers are spatial.
|
||||
reduced_buffers = []
|
||||
for block_info in block_infos[:-1]:
|
||||
for buffer_write in sch.get(block_info.block_rv).writes:
|
||||
reduced_buffers.append(buffer_write.buffer)
|
||||
|
||||
spatial_block = sch.get(block_infos[-1].block_rv)
|
||||
spatial_loops = set()
|
||||
block_var_to_loop_var = {}
|
||||
loops = sch.get_loops(block_infos[-1].block_rv)
|
||||
for block_iter, loop_rv in zip(spatial_block.iter_vars, loops):
|
||||
block_var_to_loop_var[block_iter.var] = sch.get(loop_rv).loop_var
|
||||
|
||||
def _visit_expr(e: tirx.Expr):
|
||||
if isinstance(e, tirx.Var) and e in block_var_to_loop_var:
|
||||
spatial_loops.add(block_var_to_loop_var[e])
|
||||
|
||||
for buffer_read in spatial_block.reads:
|
||||
buffer = buffer_read.buffer
|
||||
if buffer in reduced_buffers:
|
||||
for read_range in buffer_read.region:
|
||||
tirx.stmt_functor.post_order_visit(read_range.min, _visit_expr)
|
||||
tirx.stmt_functor.post_order_visit(read_range.extent, _visit_expr)
|
||||
|
||||
s_loops = []
|
||||
other_loops = []
|
||||
for loop_rv in loops:
|
||||
loop = sch.get(loop_rv)
|
||||
if loop.loop_var in spatial_loops or loop.extent == 1:
|
||||
s_loops.append(loop_rv)
|
||||
else:
|
||||
other_loops.append(loop_rv)
|
||||
sch.reorder(*s_loops, *other_loops)
|
||||
|
||||
loops = sch.get_loops(block_infos[-1].block_rv)
|
||||
bx = sch.fuse(*loops[:num_leading_s])
|
||||
r_loop, tx = sch.split(loops[-1], [None, len_tx])
|
||||
sch.reorder(tx, r_loop)
|
||||
sch.bind(bx, "blockIdx.x")
|
||||
sch.bind(tx, "threadIdx.x")
|
||||
sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
|
||||
sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)
|
||||
|
||||
for block in reversed(block_infos[:-1]):
|
||||
block = block.block_rv
|
||||
for i, _ in enumerate(sch.get(block).writes):
|
||||
sch.set_scope(block, buffer_index=i, storage_scope="shared")
|
||||
sch.compute_at(block, bx, preserve_unit_loops=True)
|
||||
r_loop = sch.fuse(*sch.get_loops(block)[-num_trailing_r:])
|
||||
r_loop, tx = sch.split(r_loop, [None, len_tx])
|
||||
sch.reorder(tx, r_loop)
|
||||
sch.bind(tx, "threadIdx.x")
|
||||
sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
|
||||
sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)
|
||||
|
||||
# TODO: It's just a workaround to avoid unroll spatial loops, because of the bug of
|
||||
# the pass lower-thread-allreduce. We should fix it in the future.
|
||||
# sch.annotate(bx, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
|
||||
# sch.annotate(bx, ann_key="pragma_unroll_explicit", ann_val=1)
|
||||
return sch
|
||||
Reference in New Issue
Block a user