relay_ast_replace_div2mul

relay_ast_example.py

import numpy as np

from collections import Counter

import tvm
from tvm import relay
# from tvm.relay import ExprFunctor, ExprMutator, ExprVisitor
from tvm.relay.expr_functor import ExprMutator, Call


class ReplaceOP(ExprMutator):
    def visit_call(self, call):
        new_fn = self.visit(call.op)
        args = []
        for arg in call.args:
            args.append(self.visit(arg))
        if str(call.op) == "divide" and isinstance(call.args[1], relay.expr.Constant):
            print("replacing divide to multiply")
            return relay.multiply(args[0], relay.const(1.0 / call.args[1].data.numpy()) )
        return Call(new_fn, args, call.attrs)

def test_replace_op():
    x = relay.var("x", shape=[1, 10])
    out = relay.divide(x, relay.const(0.5))
    expr = relay.Function([x, ], out)
    print(expr)
    '''
    fn (%x: Tensor[(1, 10), float32]) {
        divide(%x, 0.5f)
    }
    '''
    print(ReplaceOP().visit(expr))
    '''
    fn (%x: Tensor[(1, 10), float32]) {
       multiply(%x, 2f)
    }
    '''


if __name__ == "__main__":
    test_replace_op()