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,107 @@
+# 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
+"""A fallback schedule rule for GPU operators."""
+
+from tvm import s_tir, tirx
+from tvm.target import Target
+
+from .. import base
+from ..analysis import normalize_prim_func
+from ..base import try_inline
+from .base import GPUScheduleRule
+
+
+def _has_internal_thread_env(stmt: tirx.Stmt) -> bool:
+    """Check whether a statement already launches GPU threads internally,
+    e.g. via `T.launch_thread` (AttrStmt "thread_extent") or nested
+    thread-bound loops. Such blocks manage their own thread environment
+    and must not be wrapped in an additional thread binding."""
+    found = False
+
+    def _visit(node):
+        nonlocal found
+        if isinstance(node, tirx.AttrStmt) and node.attr_key in ("thread_extent", "virtual_thread"):
+            found = True
+        elif isinstance(node, tirx.For) and node.kind == tirx.ForKind.THREAD_BINDING:
+            found = True
+
+    tirx.stmt_functor.post_order_visit(stmt, _visit)
+    return found
+
+
+class Fallback(GPUScheduleRule):
+    """
+    A fallback schedule rule for all GPU operators. It will try to inline all the blocks first,
+    and then apply a simple block/grid mapping to the spatial loops on top of the remaining blocks.
+    """
+
+    def apply(  # pylint: disable=too-many-locals,missing-docstring
+        self,
+        func: tirx.PrimFunc,
+        target: Target,
+        _: bool,
+    ) -> s_tir.Schedule:
+        if not isinstance(func, tirx.PrimFunc) or not self.is_target_available(target):
+            return None
+        max_threads_per_block = base.max_threads_per_block(target)
+
+        sch = s_tir.Schedule(func)
+        block_infos = normalize_prim_func(sch)
+
+        if block_infos is None:
+            return None
+
+        block_infos = try_inline(sch, block_infos)
+        reduction_blocks: list[tuple[s_tir.schedule.SBlockRV, s_tir.schedule.LoopRV]] = []
+        for block in block_infos:
+            s_loops: list[s_tir.schedule.LoopRV] = []
+            r_loops: list[s_tir.schedule.LoopRV] = []
+            o_loops: list[s_tir.schedule.LoopRV] = []
+            dom_kind = block.dom_kind()
+            block = block.block_rv
+
+            if any(
+                [sch.get(loop_rv).thread_binding is not None for loop_rv in sch.get_loops(block)]
+            ):
+                continue
+
+            if len(sch.get_loops(block)) == 0 and _has_internal_thread_env(sch.get(block).body):
+                # The block (e.g. an opaque sort kernel) launches its own
+                # threads; binding an outer loop would conflict with them.
+                continue
+
+            for loop, iter_type in zip(sch.get_loops(block), dom_kind):
+                {"S": s_loops, "R": r_loops, "O": o_loops}[iter_type].append(loop)
+
+            if not s_loops:
+                s_loops.append(sch.add_unit_loop(block))
+            sch.reorder(*s_loops, *r_loops, *o_loops)
+            bx, tx = sch.split(  # pylint: disable=invalid-name
+                sch.fuse(*s_loops),
+                factors=[None, max_threads_per_block],
+            )
+            sch.bind(bx, "blockIdx.x")
+            sch.bind(tx, "threadIdx.x")
+
+            if len(r_loops) > 0:
+                reduction_blocks.append((block, r_loops[0]))
+
+        for block, r_loop in reduction_blocks:
+            sch.decompose_reduction(block, r_loop)
+
+        return sch