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,88 @@
+# 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=invalid-name, unused-argument, redefined-argument-from-local
+"""Relax Optimize Layout Transform pass."""
+
+import tvm_ffi
+
+from tvm.ir.module import IRModule
+from tvm.ir.transform import PassContext
+from tvm.relax import Expr
+from tvm.relax.dpl import TuplePattern, is_op, rewrite_call, wildcard
+
+from . import function_pass
+
+
+@function_pass(opt_level=0)
+class OptimizeLayoutTransform:
+    """
+    Pass to remove redundant transform layout operators
+    introduced by AlterOpImpl pass.
+    """
+
+    def __init__(self):
+        self.input = wildcard()
+        pattern_transform_layout = is_op("relax.layout_transform")(self.input)
+        pattern_1 = is_op("relax.layout_transform")(pattern_transform_layout)
+
+        self.gv_ = wildcard()
+        args = TuplePattern([pattern_transform_layout])
+        pattern_2 = is_op("relax.call_tir")(self.gv_, args)
+        self.pattern_2 = is_op("relax.layout_transform")(pattern_2)
+
+        self.pattern = pattern_1 | self.pattern_2
+
+    def transform_function(self, func: Expr, mod: IRModule, ctx: PassContext) -> IRModule:
+        """
+        Tranformation function to pattern match layout_transform -> layout_transform
+        pattern
+
+        Parameters
+        ----------
+        func: Expr
+            The relax function to be optimized
+
+        mod: IRModule
+            The ir module
+
+        ctx: PassContext
+            Relax pass context
+        """
+
+        self.mod = mod
+        updated_func = func
+
+        # Skip primitive functions
+        if "Primitive" in func.attrs.keys() and func.attrs["Primitive"] != 0:
+            return updated_func
+
+        def rewriter(expr, matches):
+            arg1 = matches[self.pattern]
+            if self.pattern_2 not in matches.keys():
+                arg2 = matches[self.input]
+            else:
+                arg2 = matches[self.gv_]
+                if "remove_pad" == self.mod[arg2].attrs["operator_name"]:
+                    arg2 = matches[self.input]
+            if hasattr(arg1.ty, "shape") and hasattr(arg2.ty, "shape"):
+                if tvm_ffi.structural_equal(arg1.ty.shape, arg2.ty.shape):
+                    return arg2
+            return expr
+
+        updated_func = rewrite_call(self.pattern, rewriter, func)
+
+        return updated_func