Does Relay/TVM supports operator fission?

Hi All,

Does TVM/Relay support operator fission? Operator fission is the opposite of operator fission where it divides a given operator into parts (a kind of unrolling).

If so, where can I start to implement my own operator fission?

If not what needs to be done in order to add operator fission?

Thanks,

2 Likes

what does you mean by part. do you split a big batch into multiple smaller batch, or do you split a operator into multiple operator, each work on the same batch size, but only does a part of the job?

1 Like

I mean create two or more operators that work on smaller batch sizes for the same operator. I think this is what you asked as follows.

1 Like

tvm dont have operator fission, but it seems like you can build one using split and concat.

1 Like

Great, but in my relay program I have lots of operators that I want to apply split/concat pattern, and this is the problem. I believe relay program is a kind of computational graph? So I was hoping that I can get a graph representation of relay program and traverse the program and apply split/concate pattern to certain operators in my relay program.

So I was not able to:

  1. Get a graph representation of relay program?
  2. I used post order to traverse the relay program, but I am not able to transofrm that visited relay program into new relay program?

the python object is a graph representation.

Thank you @MarisaKirisame

Let say I have following function, and I am wondering which python object is the graph representation? My understanding is that module is consists of functions, and module is the graph representation.

def example():
    shape = (1, 10)
    dtype = 'float32'
    tp = relay.TensorType(shape, dtype)
    x = relay.var("x", tp)
    y = relay.var("y", tp)
    out = relay.var("o", tp)

    z = relay.add(x, y)

    func = relay.Function([out], z)
    mod = relay.Module({"main": func})

The outputs are here:

>>>z
CallNode(Op(add), [Var(x, ty=TensorType([1, 10], float32)), Var(y, ty=TensorType([1, 10], float32))], (nullptr), [])

>>>func
FunctionNode([Var(o, ty=TensorType([1, 10], float32))], (nullptr), CallNode(Op(add), [Var(x, ty=TensorType([1, 10], float32)), Var(y, ty=TensorType([1, 10], float32))], (nullptr), []), [], (nullptr))

>>>mod
IRModuleNode( {GlobalVar(main): FunctionNode([Var(o, ty=TensorType([1, 10], float32))], (nullptr), CallNode(Op(add), [Var(x, ty=TensorType([1, 10], float32)), Var(y, ty=TensorType([1, 10], float32))], (nullptr), []), [], (nullptr))})

functions is the graph representation. look at the C++ code to see how to access it.

Hi johnmattew490, I am just trying to understand your requirement here. When i try to execute the piece of code you shared. I get below error image

I believe you should use IRModule(). Please correct me if i misunderstood your need. If it is IRModule(), then you can simply print it. It will show you the graph like below. image

If it is not what you are looking for. Please mention in details.