What does "topi.nn.conv2d_alter_layout" do?

#1

I am trying to understand the compilation process, but I couldn’t look into the code of “topi.nn.conv2d_alter_layout” that is called inside “alter_conv2d_layout” function. Please note that the function I’m asking is at the last line.

When I tried to debug, the program jumps from:

Questions

  1. What does the function topi.nn.conv2d_alter_layout do?
  2. Where to look into its code? Where does it call to? For information, I’m running with the VTA target.

For reference (tvm/nnvm/python/nnvm/top/nn.py)

@reg.register_alter_op_layout("conv2d")
def alter_conv2d_layout(attrs, inputs, tinfos):
    """Replace conv2d op with other layouts or algorithms"""
    import nnvm.symbol as sym

    # map relay op names to nnvm op names
    sym.contrib_conv2d_winograd_without_weight_transform = \
            sym.contrib.conv2d_winograd_without_weight_transform
    sym.contrib_conv2d_winograd_weight_transform = \
            sym.contrib.conv2d_winograd_weight_transform
    sym.nn = sym

    # map relay argument names to nnvm argument names
    raw_reshape = sym.reshape
    def _reshape(*args, **kwargs):
        if "newshape" in kwargs:
            kwargs['shape'] = kwargs.pop('newshape')
        return raw_reshape(*args, **kwargs)
    sym.reshape = _reshape

    return topi.nn.conv2d_alter_layout(attrs,  # <---- source code ?
                                       inputs,
                                       tinfos,
                                       sym)
#2

It depends on what target you are compiling for. For example, if your target is “llvm”, then the x86 backend will be used, and topi.nn.conv2d_alter_layout above will dispatch to this line. This dispatch mechanism was introduced in this PR.

1 Like
#3

@masahi Thank you! I’ve been looking into the decorators for quite a long time and your point to the dispatch mechanism is really helpful.

Just a few quick questions for clarification. Please correct me if I’m wrong. So in a bigger picture:

  1. The deep learning graph contains many operators. And the “topi.nn.conv2d_alter_layout” is used to “alter” the layout of the conv2d layer so that the operator can be compiled and run on the specific different backend.
  2. Then if I want to make the TVM compiler support new operators such as “Quantize”, I can somehow step into the graph compilation process. Then call to my custom “alter function” to convert the new operator to NNVM equivalent right?
#4

conv2d_alter_layout is used to rewrite the default conv2d op in NCHW layout to more efficient one offered by various backends.

I don’t understand your second point. If you want to add a new operator, you have to add a NNVM op corresponding to it. There is no “convert the new operator to NNVM”.

Your confusion is totally understandable given that we have basically no documentation around your inquiry. There has been some effort to improve our doc, see here for example. Your feedback will be appreciated.

#5

@masahi
Thank you for your sharing.
About alter_op_layout, i have two questions:

  1. For weight transform, will node layout_transform() be calcuated in compile time or runtime?
  2. For specific hardware, if i want to transform weight from NCHW to NCHW4o4i, and require layout of each 4o4i weight to be 256 bytes aligned . How can i achieve this?

Thanks

#6
  1. For weight transform, will node layout_transform() be calcuated in compile time or runtime?

For inference, compile time.

  1. For specific hardware, if i want to transform weight from NCHW to NCHW4o4i, and require layout of each 4o4i weight to be 256 bytes aligned . How can i achieve this?

layout transform only takes care of the shape. the alignment could be specified by tvm.schedule.Buffer in my understanding.

2 Likes