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,83 @@
+# 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, missing-function-docstring, abstract-method
+"""Relax Remove Redundant Reshape ops"""
+
+import tvm_ffi
+
+from tvm import IRModule, relax
+from tvm.ir.transform import PassContext
+from tvm.relax import Expr
+from tvm.relax.dpl import is_op, rewrite_call, wildcard
+
+from . import function_pass
+
+
+@function_pass(opt_level=0)
+class RemoveRedundantReshape:
+    """
+    Transformation pass to remove redundant reshape operator
+    """
+
+    def __init__(self):
+        self.input1 = wildcard()
+        shape1 = wildcard()
+        pattern_redundant_reshape = is_op("relax.reshape")(self.input1, shape1)
+        self.no_op_reshape = pattern_redundant_reshape
+        shape2 = wildcard()
+        self.repeated_reshape = is_op("relax.reshape")(pattern_redundant_reshape, shape2)
+        self.pattern = self.repeated_reshape | self.no_op_reshape
+
+    def transform_function(self, func: Expr, mod: IRModule, ctx: PassContext) -> IRModule:
+        """
+        Tarnsformation function to remove redundant reshape
+        where tensors before and after reshape are of same dimentions.
+
+        Parameters
+        --------------
+        func: Expr
+            The relax function to be optimized
+
+        mod: IRModule
+            The IR module
+
+        ctx: PassContext
+            Relax pass context
+        """
+
+        updated_func = func
+
+        # Skip primitive functions
+        if "Primitive" in func.attrs.keys() and func.attrs["Primitive"] != 0:
+            return updated_func
+
+        def rewriter(expr, matches):
+            arg = matches[self.input1]
+
+            if self.repeated_reshape in matches:
+                output_shape = matches[self.repeated_reshape].args[1]
+                return relax.op.reshape(arg, output_shape)
+
+            elif self.no_op_reshape in matches:
+                output_shape = matches[self.no_op_reshape].args[1]
+                if arg.ty.shape and tvm_ffi.structural_equal(arg.ty.shape, output_shape):
+                    return arg
+            return expr
+
+        updated_func = rewrite_call(self.pattern, rewriter, func)
+
+        return updated_func