chore: import upstream snapshot with attribution

python/tvm/s_tir/dlight/base/transform.py

@@ -0,0 +1,94 @@
+# 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.
+"""
+Apply ScheduleRules onto an IRModule to generate default schedules without tuning,
+or a space for MetaSchedule tuning
+"""
+
+from tvm import s_tir, tirx
+from tvm.ir import IRModule
+from tvm.ir.transform import PassContext, module_pass
+from tvm.target import Target
+
+from .schedule_rule import ScheduleRule
+
+
+def _is_scheduled(func: tirx.PrimFunc) -> bool:
+    if not isinstance(func, tirx.PrimFunc):
+        return False
+    if "tirx.is_scheduled" not in func.attrs:
+        return False
+    return func.attrs["tirx.is_scheduled"] == 1
+
+
+def _get_target(func: tirx.PrimFunc) -> Target:
+    target = func.attrs.get("target")
+    if target is None:
+        return Target.current(allow_none=False)
+    else:
+        return target
+
+
+@module_pass(opt_level=0, name="ApplyDefaultSchedule")
+class ApplyDefaultSchedule:  # pylint: disable=too-few-public-methods
+    """A IRModule pass that applies a list of ScheduleRules to all PrimFuncs in the module."""
+
+    def __init__(self, *rules: ScheduleRule):
+        """Construct a new ApplyDefaultSchedule pass.
+
+        Parameters
+        ----------
+        *rules : ScheduleRule
+            The ScheduleRules to apply to all PrimFuncs in the module.
+        """
+        self.rules = list(rules)
+
+    def transform_module(  # pylint: disable=missing-function-docstring
+        self,
+        mod: IRModule,
+        _: PassContext,
+    ) -> IRModule:
+        updated_functions = {}
+        for g_var, func in mod.functions_items():
+            if isinstance(func, tirx.PrimFunc) and not _is_scheduled(func):
+                target = _get_target(func)
+
+                sch = _apply_rules(func, target, self.rules, tunable=False)
+                if sch is not None:
+                    assert len(sch) == 1
+                    updated_functions[g_var] = (
+                        sch[0].mod["main"].with_attr("tirx.is_scheduled", True)
+                    )
+        for g_var, func in updated_functions.items():
+            mod[g_var] = func
+        return mod
+
+
+def _apply_rules(
+    func: tirx.PrimFunc,
+    target: Target,
+    rules: list[ScheduleRule],
+    tunable: bool,
+) -> list[s_tir.Schedule] | None:
+    for rule in rules:
+        space = rule.apply(func, target, tunable)
+        if space is None:
+            continue
+        if isinstance(space, s_tir.Schedule):
+            space = [space]
+        return space
+    return None