wehub-resource-sync
154 lines
6.2 KiB
Python
154 lines
6.2 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.
|
|
"""CPU reduction rule for operators including softmax, layer norm, RMS norm, etc."""
|
|
|
|
from tvm import s_tir, tirx
|
|
from tvm.target import Target
|
|
from tvm.target.codegen import llvm_get_vector_width
|
|
|
|
from ..analysis import normalize_prim_func
|
|
from ..base import get_extent
|
|
from .base import CPUScheduleRule
|
|
|
|
|
|
def _get_num_leading_s(dom_kind: str) -> int:
|
|
"""Count leading spatial ('S') axes in a dom_kind string."""
|
|
return len(dom_kind) - len(dom_kind.lstrip("S"))
|
|
|
|
|
|
class Reduction(CPUScheduleRule):
|
|
"""CPU reduction rule for softmax, layer norm, RMS norm, and similar operators.
|
|
|
|
Targets patterns with a mix of reduction (SR) and injective (SS) blocks,
|
|
where all blocks share the same leading spatial axes.
|
|
Example: softmax = maxelem(SR) -> exp(SS) -> expsum(SR) -> norm(SS).
|
|
|
|
Schedule strategy:
|
|
1. Parallelize leading spatial axes (batch dimension).
|
|
2. Move all blocks under the spatial loop via compute_at.
|
|
3. Vectorize injective blocks (exp, delta, norm) on their inner axis.
|
|
4. Split reduction inner axis to VLEN-sized chunks and annotate for
|
|
LLVM unrolling, preventing harmful full-unroll by the backend.
|
|
|
|
Note: vectorized reduction via rfactor is not used here because TVM's
|
|
rfactor primitive requires the reduction block to be the first child of
|
|
its enclosing loop, which is incompatible with compute_at when multiple
|
|
blocks share the same spatial loop. A follow-up using RVV reduction
|
|
intrinsics (vfredmax/vfredusum) via tensorize can address this.
|
|
"""
|
|
|
|
def apply( # pylint: disable=too-many-locals,too-many-return-statements,too-many-branches
|
|
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 or len(block_infos) < 2:
|
|
return None
|
|
|
|
# Must have at least one reduction block and last block must be injective.
|
|
if not any(not bi.is_injective() for bi in block_infos):
|
|
return None
|
|
if not block_infos[-1].is_injective():
|
|
return None
|
|
|
|
# Every block must start with at least one spatial axis, and all blocks
|
|
# must agree on the minimum number of leading spatial axes.
|
|
num_leading_s = None
|
|
for bi in block_infos:
|
|
dk = bi.dom_kind()
|
|
if not dk or dk[0] != "S":
|
|
return None
|
|
n = _get_num_leading_s(dk)
|
|
num_leading_s = n if num_leading_s is None else min(num_leading_s, n)
|
|
if not num_leading_s:
|
|
return None
|
|
|
|
# Infer dtype from the last block's write buffer.
|
|
last_block_stmt = sch.get(block_infos[-1].block_rv)
|
|
dtype_bits = (
|
|
last_block_stmt.writes[0].buffer.dtype.dtype.bits if last_block_stmt.writes else 32
|
|
)
|
|
|
|
# Determine vector lanes from target VLEN.
|
|
vlen_bits = llvm_get_vector_width(target)
|
|
if vlen_bits <= 0:
|
|
vlen_bits = 128
|
|
vec_lanes = max(vlen_bits // dtype_bits, 2)
|
|
|
|
# --- Phase 1: Parallelize spatial on the last block ---
|
|
last_block = block_infos[-1]
|
|
loops = sch.get_loops(last_block.block_rv)
|
|
if num_leading_s > 1:
|
|
spatial = sch.fuse(*loops[:num_leading_s])
|
|
else:
|
|
spatial = loops[0]
|
|
sch.parallel(spatial)
|
|
|
|
# --- Phase 2: Vectorize the last (injective) block ---
|
|
self._vectorize_inner(sch, last_block.block_rv, vec_lanes)
|
|
|
|
# --- Phase 3: compute_at all preceding blocks under spatial ---
|
|
for block_info in reversed(block_infos[:-1]):
|
|
sch.compute_at(block_info.block_rv, spatial, preserve_unit_loops=True)
|
|
|
|
# --- Phase 4: Vectorize injective, split+unroll reduction blocks ---
|
|
for block_info in block_infos[:-1]:
|
|
if block_info.is_injective():
|
|
self._vectorize_inner(sch, block_info.block_rv, vec_lanes)
|
|
else:
|
|
self._unroll_reduction_inner(sch, block_info.block_rv, vec_lanes)
|
|
|
|
return sch
|
|
|
|
@staticmethod
|
|
def _vectorize_inner(sch, block_rv, vec_lanes):
|
|
"""Split the innermost loop to vec_lanes and vectorize."""
|
|
block_loops = sch.get_loops(block_rv)
|
|
if len(block_loops) <= 1:
|
|
return
|
|
inner = block_loops[-1]
|
|
extent = get_extent(sch, inner)
|
|
if isinstance(extent, int):
|
|
if extent > vec_lanes:
|
|
_, vec_loop = sch.split(inner, factors=[None, vec_lanes])
|
|
sch.vectorize(vec_loop)
|
|
elif extent >= 2:
|
|
sch.vectorize(inner)
|
|
else:
|
|
_, vec_loop = sch.split(inner, factors=[None, vec_lanes])
|
|
sch.vectorize(vec_loop)
|
|
|
|
@staticmethod
|
|
def _unroll_reduction_inner(sch, block_rv, vec_lanes):
|
|
"""Split the reduction inner loop and annotate for unrolling."""
|
|
block_loops = sch.get_loops(block_rv)
|
|
if len(block_loops) <= 1:
|
|
return
|
|
inner = block_loops[-1]
|
|
extent = get_extent(sch, inner)
|
|
if isinstance(extent, int) and extent <= vec_lanes:
|
|
return
|
|
_, inner_loop = sch.split(inner, factors=[None, vec_lanes])
|
|
sch.annotate(inner_loop, ann_key="pragma_auto_unroll_max_step", ann_val=vec_lanes)
|
|
sch.annotate(inner_loop, ann_key="pragma_unroll_explicit", ann_val=1)
|