wehub-resource-sync
84 lines
2.8 KiB
Python
84 lines
2.8 KiB
Python
# 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
|