Converts the expensive non linear functions to their fast but approximate counterparts. Parameters ---------- mod: IRModule The module to be transformed
| 25 | |
| 26 | @expr_functor.mutator |
| 27 | class FastMathCodeGenerator(PyExprMutator): |
| 28 | """ |
| 29 | Converts the expensive non linear functions to their fast but approximate counterparts. |
| 30 | |
| 31 | Parameters |
| 32 | ---------- |
| 33 | mod: IRModule |
| 34 | The module to be transformed |
| 35 | """ |
| 36 | |
| 37 | def __init__(self, mod): |
| 38 | super().__init__(mod) |
| 39 | |
| 40 | def visit_call_(self, call: Call) -> Expr: |
| 41 | if call.op.name == "relax.nn.softmax": |
| 42 | return self.builder_.call_te(topi.nn.fast_softmax, call.args[0], call.attrs.axis) |
| 43 | if call.op.name == "relax.exp": |
| 44 | return self.builder_.call_te(topi.fast_exp, call.args[0]) |
| 45 | if call.op.name == "relax.erf": |
| 46 | return self.builder_.call_te(topi.fast_erf, call.args[0]) |
| 47 | if call.op.name == "relax.tanh": |
| 48 | return self.builder_.call_te(topi.fast_tanh, call.args[0]) |
| 49 | |
| 50 | return super().visit_call_(call) |
| 51 | |
| 52 | |
| 53 | @tvm.transform.module_pass(opt_level=0, name="FastMathTransform") |
no outgoing calls
no test coverage detected
searching dependent graphs…