InferCorrectLayout

Questions
#1

Hello everyone,

i have been implementing my version of the Resampler OP (from TF Frontend) to our TVM Stack. Now (to my understanding) by adding the “InferCorrectLayout” Attribute to the RelayCall Node i should be able to also automatically change the Layout of my Custom OP’s Inputs/Outputs when the layout is changed for nn.conv2d (see here. So my problem is: The Resampler OP has a data tensor of a certain layout (NHWC, NCHW etc.) which i want to adapt depending on the desired_layout. But Resampler also takes a so called warp tensor with n Dimensions. The same applies for the output shape. In theory i want a layout_transform insertion ONLY on the first input Tensor and leave the other input tensor and the output tensor untouched with regards to its layout. From my intuition

Array<Array<Layout>> ResamplerInferCorrectLayout(...){
   /*similar code to UpsamplingInferCorrectLayout*/
   return Array<Array<Layout>>{{inferred_layout, Layout::Undef()}, {Layout::Undef()}};
}

Sadly this does not work and i dont yield any layout_transform near my Resampler OP. If i understand correctly that has to do with the fact, that a Layout::Undef() tells TVM to overall not insert any layout_transform. Do you have any idea how i can achieve the “passthrough” of the 2nd Input tensor and outputtensor with regards to layout_transform?

#2

My impression is actually opposite to yours. My impression is if you return Undef or incompatible layout, then layout_transform will be inserted to guarantee the correctness. For example:

conv2d(NCHW, OIHW) -(NCHW)-> transpose(axis=[0, 2, 3, 1]) -(NHWC)->

The output layout of conv2d is NCHW, and the output of transpose transposes to NHWC. If we now want to convert conv2d to NHWC, then we have

conv2d(NHWC, HWIO)  -(NHWC)-> transpose(?)  -->

In this case, if transpose has no InferCorrectLayout registered, or InferCorrectLayout returns Undef, then the ConvertLayout pass will keep the original transpose op and insert layout_transform to keep NCHW in and NHWC out:

conv2d(NHWC, HWIO) -(NHWC)-> layout_transform  -(NCHW)-> transpose(axis=[0, 2, 3, 1])  - (NHWC) ->

On the other hand, if transpose's InferCorrectLayout is powerful enough, it will adjust its attributes to accept the new input layout and return the new layout:

conv2d(NHWC, HWIO)  -(NHWC)-> transpose(axis=[0, 1, 2, 3]) - (NHWC) ->

In this case, layout_transfrom won’t be inserted.

1 Like
#3

Thank you for your response @comaniac !

When testing your statement by uncommenting the InferCorrectLayout property of nn.resampler and feeding a simple conv2D->Resampler model (applying a LayoutTransform) i actually get the desired behaviour (it updated the layout-attribute of my resampler node and having the output tensor of conv2d in NCHW layout instead of NHWC).

Now when trying to run an even simpler model containing just a Resampler OP (without Conv2D) i try to convert the layout from NHWC to NCHW aswell. I have implemented a convert_function in python/tvm/relay/op/nn/_nn.py and defined a InferCorrectLayout for the call node.

But when trying to convert the layouts by using the following code lines, i dont get any triggers and yet no layout_transform:

desired_layouts={"nn.resampler": ["NCHW"]} 
seq = tvm.transform.Sequential([relay.transform.RemoveUnusedFunctions(),
                                relay.transform.ConvertLayout(desired_layouts)])
with tvm.transform.PassContext(opt_level=3):
    irmodule = seq(irmodule)

I compared my changes to the ones mentioned in the “official” Layout Pass guide (this one).

Are there additional hooks that need to be set or strategies (or something similar) that need to be defined in order for transform_layout to start running the layout pass?

#4

It doesn’t make sense to convert a model with only one Resampler op. Think about it. What’s the layout of an input variable? An input variable only has shape and type but no layout. As a result, InferCorrectLayout basically does nothing.

Based on that, ConvertLayout actually propagates the known layouts through the graph. When it meets a conv2d, it clearly knows the input and output layout, so it could pass the output layout to the next op so on so forth. The ops that don’t have layout attribute rely on the propagated layout information to infer the correct layout.

1 Like
#5

Okay thanks alot for you fast and helpful answers :slight_smile:

#6

Hello @comaniac,

I met a similar problem about ConvertLayout pass and tried to follow what you explained. From the previous comment, the ConvertLayout pass could work on the Ops which input has layout information.

I tried an example as following. The out_layout of the nn.conv2d will be NHWC, and I want to convert the data layout to make the input of nn.avg_pool2d becomes NCHW.

x = relay.var("x", shape=(1, 56, 56, 64))
y = relay.var("y", shape=(3, 3, 64, 32))
conv2d = relay.nn.conv2d(
            x,
            y,
            channels=32,
            kernel_size=(3, 3),
            padding=(1, 1),
            data_layout="NHWC",
            kernel_layout="HWIO")
avg_pool2d = relay.nn.avg_pool2d(conv2d, pool_size=(3, 3), layout="NHWC")
func = relay.Function([x, y], avg_pool2d)
mod = tvm.IRModule.from_expr(func)

seq = tvm.transform.Sequential([relay.transform.InferType(),
                                relay.transform.ConvertLayout({"nn.avg_pool2d": ["NCHW"]})])
with tvm.transform.PassContext(opt_level=3):
    mod = seq(mod)

The result is

def @main(%x: Tensor[(1, 56, 56, 64), float32], %y: Tensor[(3, 3, 64, 32), float32]) -> Tensor[(1, 54, 54, 32), float32] {
  %0 = nn.conv2d(%x, %y, padding=[1, 1, 1, 1], channels=32, kernel_size=[3, 3], data_layout="NHWC", kernel_layout="HWIO") /* ty=Tensor[(1, 56, 56, 32), float32] */;
  nn.avg_pool2d(%0, pool_size=[3, 3], padding=[0, 0, 0, 0], layout="NHWC") /* ty=Tensor[(1, 54, 54, 32), float32] */
}

Seems the ConvertLayout pass doesn’t change anything in the graph.

According to the result, I traced this pass and found it will extract the user-defined layout at here.

As the new layout is saved in the new_call->attrs. It should update the new_in and then we can invoke InferCorrectLayouts in line357 to get the correct infer_out (based on the PoolInferCorrectLayout).

But it doesn’t update it, so it’ll apply the original layout (NHWC for this case).

It seems to be the reason that why this pass can’t work as expected.