chore: import upstream snapshot with attribution

2026-07-13 13:36:25 +08:00
commit 26446540fa
3151 changed files with 974126 additions and 0 deletions
@@ -0,0 +1,115 @@
+# 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