wehub-resource-sync
39 lines
1.4 KiB
Python
39 lines
1.4 KiB
Python
"""Helper functions for checking max num thread."""
|
|
|
|
from tvm.target import Target
|
|
|
|
|
|
def get_max_num_threads_per_block(target: Target) -> int:
|
|
"""
|
|
max(max_num_threads, max_threads_per_block); if latter does not exist, return max_num_threads.
|
|
We add this method since some targets have both fields and `max_threads_per_block` is larger.
|
|
"""
|
|
max_num_threads = target.attrs.get("max_num_threads")
|
|
max_threads_per_block = target.attrs.get("max_threads_per_block", None)
|
|
if max_threads_per_block is None:
|
|
return max_num_threads
|
|
return max(max_num_threads, max_threads_per_block)
|
|
|
|
|
|
def check_thread_limits(target: Target, bdx: int, bdy: int, bdz: int, gdz: int):
|
|
"""
|
|
Check whether max num threads exceeded given a target.
|
|
|
|
Parameters
|
|
----------
|
|
bdx: threadIdx.x
|
|
bdy: threadIdx.y
|
|
bdz: threadIdx.z
|
|
gdz: blockIdx.z
|
|
"""
|
|
max_num_threads_per_block = get_max_num_threads_per_block(target)
|
|
|
|
assert bdx * bdy * bdz <= max_num_threads_per_block, (
|
|
f"{target.kind} max num threads exceeded: {bdx}*{bdy}*{bdz}>{max_num_threads_per_block}"
|
|
)
|
|
|
|
if target.kind.name == "webgpu":
|
|
# https://gpuweb.github.io/gpuweb/#dom-supported-limits-maxcomputeworkgroupsizez
|
|
assert bdz <= 64, f"webgpu's threadIdx.z cannot exceed 64, but got bdz={bdz}"
|
|
assert gdz == 1, f"webgpu's blockIdx.z should be 1, but got gdz={gdz}"
|