[TOPI] Implementation Guideline

[tqchen]

There has been a few discussions on this and I am creating this issue to consolidate what we have so far. The specific question we have is that how should we create APIs for data flow declaration and Schedule interface in TOPI, here are a two key guidelines

Tensor in,Tensor out in dataflow declaration

The qoute(tensor in/tensor out) comes from google brain team. This is a general principle for compositional APIs design. Imagine we want to create a dataflow declaration for conv-relu, instead of create a declaration function for conv_relu, we want to create two functions (conv, and relu), and compose them

def conv(input, weight):
     out =  tvm.comput(...)
     return out

def relu(input):
     return tvm.compute(input.shape, lambda i*:tvm.max(0, input(*)))

input = tvm.placeholder((n, c, h, w))
weight = tvm.placeholder((c1, c2, hh, ww))
net= conv(input, weight)
net = relu(net)
# Now that net contains data flow declaration of conv-relu

Seperate data flow declaration from schedule

While it is usually convenient to put schedule logic together with the data flow declaration, it is also somewhat harmful. Imagine we have a schedule for conv-relu, what if we want to schedule conv-sigmoid? They contains the essentially the same pattern, but in the old style we need to create schedule for each of them. So ideally, we want to create a generic schedule function for a class of dataflows, without directly touching the dataflow part.

To get the tensors needed for schedule, we can recover them by traversing the dataflow DAG. Here is an possible skeleton to schedule schedule conv-ewise

# generic schedule convolution ewise
def schedule_conv_map(op):  
      # find the conv part
     s = create_schedule(op)
     conv_args = []
     def schedule_conv(data, filter, conv):
         # schedule conv here
         pass
     # visit, maybe need deduplicate
     def  visit(op):
        if is_ewise(op):
          if not is_output(op):  
              s[op].compute_inline()
           for t in op.input_tensors:
               visit(t.op)
        if is_conv(op):
           conv = op.output(0)
           data = op.input_tensors[0]
           filter = op.input_tensors[1]
           schedule_conv(conv,  data, filter)

[tqchen]

@Huyuwei @icemelon9 @ZihengJiang Please feel free to include more detailed examples, in comments, eventually we should put the guideline in document

[ZihengJiang]

SGTM, about the second part, maybe we can create an API like dispatch_schedule([(lambda op: is_conv(op), lambda op: conv_schedule(op)), ...])

A more generic method requires us to partition the description into small parts.

[tqchen]

For now, it is a bit hard to have a very generic api for all schedules, mainly because the schedule is context dependent (on its inputs or parent), so we will start from specific ones, like schedule_conv_ewise, which can take all conv ewise patterns, eventually when we see enough of them, we can get to a generic dispatching mechanism.

For now, creating util(like detect ewise, traverse) is good for us to start the first step