wehub-resource-sync
108 lines
4.0 KiB
Python
108 lines
4.0 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
|
|
"""A fallback schedule rule for GPU operators."""
|
|
|
|
from tvm import s_tir, tirx
|
|
from tvm.target import Target
|
|
|
|
from .. import base
|
|
from ..analysis import normalize_prim_func
|
|
from ..base import try_inline
|
|
from .base import GPUScheduleRule
|
|
|
|
|
|
def _has_internal_thread_env(stmt: tirx.Stmt) -> bool:
|
|
"""Check whether a statement already launches GPU threads internally,
|
|
e.g. via `T.launch_thread` (AttrStmt "thread_extent") or nested
|
|
thread-bound loops. Such blocks manage their own thread environment
|
|
and must not be wrapped in an additional thread binding."""
|
|
found = False
|
|
|
|
def _visit(node):
|
|
nonlocal found
|
|
if isinstance(node, tirx.AttrStmt) and node.attr_key in ("thread_extent", "virtual_thread"):
|
|
found = True
|
|
elif isinstance(node, tirx.For) and node.kind == tirx.ForKind.THREAD_BINDING:
|
|
found = True
|
|
|
|
tirx.stmt_functor.post_order_visit(stmt, _visit)
|
|
return found
|
|
|
|
|
|
class Fallback(GPUScheduleRule):
|
|
"""
|
|
A fallback schedule rule for all GPU operators. It will try to inline all the blocks first,
|
|
and then apply a simple block/grid mapping to the spatial loops on top of the remaining blocks.
|
|
"""
|
|
|
|
def apply( # pylint: disable=too-many-locals,missing-docstring
|
|
self,
|
|
func: tirx.PrimFunc,
|
|
target: Target,
|
|
_: bool,
|
|
) -> s_tir.Schedule:
|
|
if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
|
|
return None
|
|
max_threads_per_block = base.max_threads_per_block(target)
|
|
|
|
sch = s_tir.Schedule(func)
|
|
block_infos = normalize_prim_func(sch)
|
|
|
|
if block_infos is None:
|
|
return None
|
|
|
|
block_infos = try_inline(sch, block_infos)
|
|
reduction_blocks: list[tuple[s_tir.schedule.SBlockRV, s_tir.schedule.LoopRV]] = []
|
|
for block in block_infos:
|
|
s_loops: list[s_tir.schedule.LoopRV] = []
|
|
r_loops: list[s_tir.schedule.LoopRV] = []
|
|
o_loops: list[s_tir.schedule.LoopRV] = []
|
|
dom_kind = block.dom_kind()
|
|
block = block.block_rv
|
|
|
|
if any(
|
|
[sch.get(loop_rv).thread_binding is not None for loop_rv in sch.get_loops(block)]
|
|
):
|
|
continue
|
|
|
|
if len(sch.get_loops(block)) == 0 and _has_internal_thread_env(sch.get(block).body):
|
|
# The block (e.g. an opaque sort kernel) launches its own
|
|
# threads; binding an outer loop would conflict with them.
|
|
continue
|
|
|
|
for loop, iter_type in zip(sch.get_loops(block), dom_kind):
|
|
{"S": s_loops, "R": r_loops, "O": o_loops}[iter_type].append(loop)
|
|
|
|
if not s_loops:
|
|
s_loops.append(sch.add_unit_loop(block))
|
|
sch.reorder(*s_loops, *r_loops, *o_loops)
|
|
bx, tx = sch.split( # pylint: disable=invalid-name
|
|
sch.fuse(*s_loops),
|
|
factors=[None, max_threads_per_block],
|
|
)
|
|
sch.bind(bx, "blockIdx.x")
|
|
sch.bind(tx, "threadIdx.x")
|
|
|
|
if len(r_loops) > 0:
|
|
reduction_blocks.append((block, r_loops[0]))
|
|
|
|
for block, r_loop in reduction_blocks:
|
|
sch.decompose_reduction(block, r_loop)
|
|
|
|
return sch
|