wehub-resource-sync
116 lines
3.6 KiB
Python
116 lines
3.6 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
|
|
"""Utility methods for generic GPU."""
|
|
|
|
from tvm import DataType, s_tir, tirx
|
|
from tvm.ir import PrimType
|
|
from tvm.target import Target
|
|
|
|
|
|
def get_bytes(dtype: DataType | PrimType | str) -> int:
|
|
if isinstance(dtype, PrimType):
|
|
dtype = dtype.dtype
|
|
if isinstance(dtype, str):
|
|
dtype = DataType(dtype)
|
|
return dtype.itemsize
|
|
|
|
|
|
def get_extent(sch: s_tir.Schedule, loop_rv: s_tir.schedule.LoopRV):
|
|
loop: tirx.For = sch.get(loop_rv)
|
|
return loop.extent.value if isinstance(loop.extent, tirx.IntImm) else loop.extent
|
|
|
|
|
|
def auto_vectorize(sch: s_tir.Schedule, loop: s_tir.schedule.LoopRV, max_vec: int):
|
|
"""Auto vectorize the loop."""
|
|
extent = get_extent(sch, loop)
|
|
if not isinstance(extent, int):
|
|
return
|
|
v = loop if extent <= max_vec else sch.split(loop, factors=[None, max_vec])[-1]
|
|
sch.vectorize(v)
|
|
|
|
|
|
def max_threads_per_block(target: Target) -> int:
|
|
"""Get the maximum number of threads per block for a given target.
|
|
|
|
Parameters
|
|
----------
|
|
target : Target
|
|
The target to get the maximum number of threads per block for.
|
|
|
|
Returns
|
|
-------
|
|
max_threads_per_block : int
|
|
The maximum number of threads per block for the given target.
|
|
"""
|
|
for name in ["max_threads_per_block", "max_num_threads"]:
|
|
result = target.attrs.get(name, None)
|
|
if result is not None:
|
|
return result
|
|
if target.kind.name == "cuda":
|
|
return 1024
|
|
return 256
|
|
|
|
|
|
def suggest_threads_per_block(
|
|
target: Target,
|
|
loops: list[tirx.For],
|
|
max_threads_for_dynamic_loop: int = 32,
|
|
) -> list[int]:
|
|
if target.kind.name == "cuda":
|
|
threads = 1024
|
|
elif target.kind.name == "rocm":
|
|
threads = 256
|
|
elif target.kind.name == "metal":
|
|
threads = 256
|
|
elif target.kind.name == "opencl":
|
|
threads = 256
|
|
else:
|
|
threads = 64
|
|
results: list[int | None] = []
|
|
dynamic: list[int] = []
|
|
for i, loop in enumerate(loops):
|
|
loop_extent = loop.extent
|
|
if isinstance(loop_extent, tirx.IntImm):
|
|
loop_extent = loop_extent.value
|
|
extent = 1
|
|
while extent <= loop_extent and extent <= threads:
|
|
extent *= 2
|
|
extent //= 2
|
|
assert extent >= 1
|
|
assert threads % extent == 0
|
|
threads //= extent
|
|
results.append(extent)
|
|
else:
|
|
results.append(None)
|
|
dynamic.append(i)
|
|
|
|
for i in dynamic:
|
|
extent = 1
|
|
while extent <= max_threads_for_dynamic_loop and extent <= threads:
|
|
extent *= 2
|
|
extent //= 2
|
|
assert extent >= 1
|
|
assert threads % extent == 0
|
|
threads //= extent
|
|
results[i] = extent
|
|
|
|
if dynamic:
|
|
results[dynamic[0]] *= threads
|
|
|
|
return results
|