chore: import upstream snapshot with attribution

python/mlc_llm/support/max_thread_check.py

@@ -0,0 +1,38 @@
+"""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}"