[BYOC][Layout conversion] How to use Convert layout pass with BYOC flow?

[masahi]

@comaniac

As a result, I would suggest embedding layout transform from NCHW to NCHWc inside the external codegen

Yes this is what is happening right now inside ex_fun1, ex_fun2 right now. Inside my extern function, I operate entirely on NCHWc layout.

The motivation is very simple. It is for removing unnecessary layout transform at the beginning and end of each of my extern function. If TVM (host side) also operates on NCHWc layout, it should be possible to keep NCHWc layout end to end.

The fallback ops are supposed to be cheap, so I don’t really care if NCHWc layout would improve their perf or what AutoTVM does to fall back ops. You can assume that I know the right inner c dim, so it is not a responsibility of TVM to figure out this value.

If you think about how layout transform works in TVM (InferCorrectLayout + way to turn NCHW op to NCHWc op (conv2d → conv2d_NCHWc etc) and look at my module below again,

@masahi

def main(...)
    %1 = ex_fun1(%0) // some complicated extern func
    %2 = max_pooling(%1) // falls back to CPU
    %3 = ex_fun2(%2) // another extern_func
    %4 = global_pool(%3) // falls back to CPU
    %5 = dense(%4)

I think the only new requirement is for TVM to know what the src and dst layout of callee functions. The layout transform pass can treat these extern func calls as if they were just like any other ops, as long as src and dst layout are specified. Then InferCorrectLayout should be able to infer the right layout constraint between ops (and extern functions), and nn.layout_transform would be inserted at the right place.

So it is really a simple generalization of existing passes. The new requirement is to handle the case when call->op is not Op, but Function. Does this make sense? cc @anijain2305

I hope by now my possible-solution below also makes sense.

@masahi

• Add src and dst layouts as new attributes to Function
• Extend InferCorrectLayouts that currently assumes call->op to be Op , to take care of Function too
• Similarly, extend AlterOpLayout to handle Function . This requires passing a user defined Function → Function rewriter that decorates Function with src and dst layouts, in place of falter_layout used for op conversion.