wehub-resource-sync
132 lines
5.0 KiB
Python
132 lines
5.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.
|
|
"""A rule for GEMV and DecodeGEMV."""
|
|
|
|
from tvm import s_tir, tirx
|
|
from tvm.target import Target
|
|
|
|
from ..analysis import SBlockInfo, normalize_prim_func
|
|
from ..analysis.gemv import is_gemv, normalize
|
|
from ..base import get_extent, try_inline_contiguous_spatial
|
|
from .base import CPUScheduleRule
|
|
|
|
|
|
class GEMV(CPUScheduleRule):
|
|
"""A rule for GEMV and DecodeGEMV."""
|
|
|
|
def apply( # pylint: disable=too-many-locals,too-many-branches,too-many-return-statements, no-else-return
|
|
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)
|
|
block_infos = try_inline_contiguous_spatial(sch, block_infos)
|
|
if block_infos is None:
|
|
return None
|
|
if len(block_infos) == 1:
|
|
epilogue = None
|
|
elif len(block_infos) == 2:
|
|
epilogue = block_infos[1]
|
|
if not epilogue.is_injective():
|
|
return None
|
|
else:
|
|
return None
|
|
|
|
block_info = block_infos[0]
|
|
if len(block_info.iters) not in [2, 3]:
|
|
# either [B, S, R] = [B, S, R] * [B, R]
|
|
# or [S, R] = [S, R] * [R]
|
|
return None
|
|
block = block_info.block_rv
|
|
vector_input_buffers = is_gemv(sch, block_info)
|
|
if vector_input_buffers is None:
|
|
return None
|
|
|
|
# Step 1. Normalize the block, merge spatial and reduction iters
|
|
is_inner_reduction = normalize(sch, block_info)
|
|
|
|
# Step 2. Do the scheduling
|
|
if is_inner_reduction is None:
|
|
return None
|
|
elif is_inner_reduction:
|
|
return self.sch_inner_reduction(sch, target, block, vector_input_buffers, epilogue)
|
|
else:
|
|
# sch_outer reduction
|
|
return None
|
|
|
|
def sch_inner_reduction( # pylint: disable=too-many-arguments, too-many-positional-arguments, invalid-name, unused-argument
|
|
self,
|
|
sch: s_tir.Schedule,
|
|
target: Target,
|
|
block: s_tir.schedule.SBlockRV,
|
|
vector_input_buffers: list[tirx.Buffer],
|
|
epilogue_info: SBlockInfo | None,
|
|
):
|
|
"""Schedule the inner reduction block."""
|
|
|
|
def apply( # pylint: disable=unused-variable, too-many-locals
|
|
sch: s_tir.Schedule,
|
|
gemv,
|
|
vector_width: int = 8,
|
|
parallel_threads: int = 8,
|
|
unroll_factor: int = 256,
|
|
):
|
|
batch, s, r, c = sch.get_loops(block)
|
|
len_batch, len_s, len_r, len_c = (
|
|
get_extent(sch, batch),
|
|
get_extent(sch, s),
|
|
get_extent(sch, r),
|
|
get_extent(sch, c),
|
|
)
|
|
len_S = len_batch * len_s
|
|
len_R = len_r * len_c
|
|
|
|
if isinstance(len_S, int) and isinstance(len_R, int):
|
|
if len_S > len_R:
|
|
tile_s, tile_r = 128, 64 # Larger tiling for s-axis when len_S is larger
|
|
else:
|
|
tile_s, tile_r = 64, 128 # Larger tiling for r-axis when len_R is larger
|
|
else:
|
|
tile_s, tile_r = 64, 64 # Default tile sizes for unknown extents
|
|
|
|
tile_c = min(vector_width, len_c) # Ensure c-axis tiling aligns with SIMD vector width
|
|
|
|
# Apply loop tiling (improves cache locality)
|
|
s_outer, s_inner = sch.split(s, factors=[None, tile_s])
|
|
r_outer, r_inner = sch.split(r, factors=[None, tile_r])
|
|
c_outer, c_inner = sch.split(c, factors=[None, tile_c])
|
|
|
|
# Apply vectorization (SIMD optimization)
|
|
sch.vectorize(s_inner) # Vectorize computation along c-axis for AVX/NEON
|
|
|
|
# Enable parallel execution
|
|
sch.parallel(s_outer) # Parallelize along the s-axis (major computation loop)
|
|
|
|
# Apply loop unrolling for better CPU performance
|
|
sch.annotate(r_outer, "pragma_auto_unroll_max_step", unroll_factor)
|
|
sch.annotate(r_outer, "pragma_unroll_explicit", 1)
|
|
return sch
|
|
|
|
return apply(
|
|
sch,
|
|
gemv=block,
|
|
)
|