Upstreaming New Op to TOPI


#1

I was wondering about the process for adding a new operator to TOPI. In a previous post, I remarked on the absence of a bias_add operator in TOPI when it seemed to already exist in Relay. This absence came up as a result of trying to prototype operators/schedules for a new hardware backend.

Since bias_add did not exist in TOPI, I implemented the op, and was wondering about how to go about upstreaming that op. Is there a suite of tests where this op should be added to verify functionality/stability? And are there any specific format guidelines for defining new operators?

I’ve attached the bias_add operator below:

@tvm.tag_scope(tag=tag.INJECTIVE)
def bias_add(x, b, layout=None):
    """Add biases to each feature map along channels dim.

    Parameters
    ----------
    x : tvm.Tensor
        Input tensor. 4D
        
    b : tvm.Tensor
        Biases tensor. 1D = Channels of x
        
    layout: string
        layout of data

    Returns
    -------
    y : tvm.Tensor
        The result.
    """
    # Infer layout from input layout
    if layout is None:
        layout = 'none'
        if 'nchw' in x.op.tag: layout = 'NCHW'
        if 'nhwc' in x.op.tag: layout = 'NHWC'
    
    if layout == "NCHW":
        batches, channels, height, width = x.shape
        func = lambda bb, cc, hh, ww: tvm.sum(
                    x[bb, cc, hh, ww] + b[cc],
                    tvm.const(0, x.dtype))
    elif layout == "NHWC":
        batches, height, width, channels = x.shape
        func = lambda bb, hh, ww, cc: tvm.sum(
                    x[bb, hh, ww, cc] + b[cc],
                    tvm.const(0, x.dtype))
    else:
        raise ValueError("bias_add does not support this layout {} yet".format(layout))
    
    # Ensure we have enough biases
    biases = b.shape
    assert len(biases) == 1
    biases = biases[0]
    assert int(biases) == int(channels)
    
    return tvm.compute(
        x.shape, func)

#2

bias_add is an op that will get canonicalized directly into add, so as long as add is supported, we are fine


#3

Maybe this will help? https://github.com/dmlc/tvm/pull/3027


#4

That solves my problems perfectly! TOPI is given direct access to a standard bias_add op used across other modules and that’s exactly what I was concerned about.

Thanks for updating me.