How does TOPI interact with other modules in the TVM stack (ie. Relay)?

Questions
#1

The TOPI library seems to have a disconnect from graph oriented modules such as Relay.

While trying to create the schedule for a conv2d+bias_add+relu, I discovered a lack of a bias_add compute definition in the TOPI library. Before implementing my own simple bias_add with tvm.compute(… tvm.sum(…)), I tried exploring the Relay library to see how it implemented bias_add for any graph it might receive that contained that operator.

This is where my question came up; why does Relay have its own implementation of bias_add, whereas TOPI does not? I was under the impression that TOPI used compute and schedule definitions to support conversions between high-level graph ops and low-level IR. There seems to be a disconnect between these two and I’m trying to figure out why.

Is it because these two libraries are geared towards inherently different tasks and only have some overlap in functionality? Or is there just a lack of integration with each other?

Upstreaming New Op to TOPI
#2

First some clarifying definitions:

  • Relay/NNVM-V2: is the graph IR doing graph-level optimizations.
  • TVM: is a tensor DSL providing primitive for defining operators and how to schedule them.
  • TOPI: defines operators and scheduling mechanism using DSL provided by TVM. It’s primarily used for AutoTVM. In fact, NNVM (old) and new Relay call into TOPI.

tvm-stack

You may find more details in

  1. Relay RFC
  2. Original NNVM design before merging into TVM
2 Likes
#3

Okay this makes a lot more sense after going back to the Relay code with all of that in mind.

I’m still a little confused why the bias_add implementation below exists in Relay instead of in TOPI.

@reg.register_compute("nn.bias_add")
def compute_bias_add(attrs, inputs, out_dtype, target):
    """Compute definition of conv2d_transpose"""
    axis = attrs.axis
    bias = inputs[1]
    data_ndim = len(inputs[0].shape)
    if axis < 0:
        axis = axis + data_ndim
    num_newaxis = data_ndim - axis - 1

    if num_newaxis:
        bias = topi.expand_dims(bias, axis=1, num_newaxis=num_newaxis)
    return [topi.add(inputs[0], bias)]

This seems like something that should be handled by TOPI with a specific bias_add operator so that prototyping with bias_add can be done outside of Relay?

#4

I think the main reason is for Relay frontend functionality to support different frameworks such as mxnet, onnx, tf, etc. so it needs to abstract over some common ops including bias_add, and it’s calling into topi anyway. I agree, depending on use cases it seems to be a good option to have it directly in topi.

Contributions are welcome in any case!

2 Likes