wehub-resource-sync
89 lines
3.0 KiB
Python
89 lines
3.0 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, 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
|