wehub-resource-sync
99 lines
3.6 KiB
Python
99 lines
3.6 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=missing-docstring, invalid-name, unnecessary-comprehension, unused-argument
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
from tvm.contrib.hexagon import hexagon_unary_ops
|
|
|
|
|
|
def op_replace(call_node, func) -> bool:
|
|
if not isinstance(call_node, relax.Call):
|
|
return False
|
|
call_tir_op = tvm.ir.Op.get("relax.call_tir")
|
|
if call_node.op != call_tir_op:
|
|
return False
|
|
ops = [
|
|
"qnn.tanh",
|
|
"qnn.sqrt",
|
|
"qnn.rsqrt",
|
|
"qnn.exp",
|
|
"qnn.erf",
|
|
"qnn.sigmoid",
|
|
"qnn.hardswish",
|
|
"qnn.log",
|
|
"qnn.abs",
|
|
]
|
|
if func.attrs["op_attrs"]["op_name"] in ops:
|
|
return True
|
|
return False
|
|
|
|
|
|
@relax.expr_functor.mutator
|
|
class Tanh2TakeReplace(tvm.relax.PyExprMutator):
|
|
def __init__(self, mod: tvm.IRModule) -> None:
|
|
super().__init__(mod)
|
|
self.mod_ = mod
|
|
|
|
def transform(self) -> tvm.IRModule:
|
|
# Iterate over all the nodes to check for the node replaceable
|
|
for global_var, func in self.mod_.functions.items():
|
|
# Skip non-relax functions
|
|
if not isinstance(func, relax.Function):
|
|
continue
|
|
updated_func = self.visit_expr(func)
|
|
self.builder_.normalize(updated_func)
|
|
self.builder_.update_func(global_var, updated_func)
|
|
# At the end of the transformation we return the updated IRModule from the BlockBuilder.
|
|
return self.builder_.get()
|
|
|
|
def visit_call_(self, call_node: relax.Call) -> relax.Call:
|
|
call_tir_op = tvm.ir.Op.get("relax.call_tir")
|
|
if call_node.op != call_tir_op:
|
|
return call_node
|
|
|
|
var = call_node.args[0]
|
|
func = self.mod_[var]
|
|
|
|
if call_node.args[1][0].ty.dtype == "uint8":
|
|
if op_replace(call_node, func):
|
|
inp, inp_scale, inp_zp, out_scale, out_zp = [x for x in call_node.args[1]]
|
|
# LUT node creation
|
|
LUT = hexagon_unary_ops.LUT_generation(
|
|
inp_scale, inp_zp, out_scale, out_zp, call_node.args[0].name_hint
|
|
)
|
|
# Take operation node creation
|
|
take_func = hexagon_unary_ops.generate_take_primfunc(inp, call_node.ty)
|
|
take_func = take_func.without_attr("global_symbol")
|
|
take_func_gv = self.builder_.add_func(take_func, "take")
|
|
take_node = relax.call_tir(
|
|
take_func_gv,
|
|
relax.expr.Tuple(
|
|
[call_node.args[1][0], relax.expr.Constant(tvm.runtime.tensor(LUT))]
|
|
),
|
|
call_node.ty,
|
|
)
|
|
return take_node
|
|
return call_node
|
|
|
|
|
|
@tvm.ir.transform.module_pass(opt_level=2, name="replace_tanh_take")
|
|
class PassReplaceWithTakeOpPrimFuncs:
|
|
def transform_module(self, mod, ctx):
|
|
return Tanh2TakeReplace(mod).transform()
|