Fix AutoTVM int8 vnni dense task extraction problem

[mengceng15]

This PR is a minor fix on the AutoTVM int8 vnni dense task extraction.
VNNI dense strategy is used only when the input-weight-out datatypes are u8s8s32 and the weight layout is NC16n4c.
Without altering the weight layout, for a simple dense workload:

def @main(%data: Tensor[(1, 16), uint8] /* ty=Tensor[(1, 16), uint8] /, %weight: Tensor[(32, 16), int8] / ty=Tensor[(32, 16), int8] /) -> Tensor[(1, 32), int32] {
nn.dense(%data, %weight, units=None, out_dtype="int32") / ty=Tensor[(1, 32), int32] */
}

There are two non-VNNI tasks extracted:
[Task(func_name=dense_nopack.x86, args=(('TENSOR', (1, 16), 'uint8'), ('TENSOR', (32, 16), 'int8'), None, 'int32'), kwargs={}, workload=('dense_nopack.x86', ('TENSOR', (1, 16), 'uint8'), ('TENSOR', (32, 16), 'int8'), None, 'int32')),
Task(func_name=dense_pack.x86, args=(('TENSOR', (1, 16), 'uint8'), ('TENSOR', (32, 16), 'int8'), None, 'int32'), kwargs={}, workload=('dense_pack.x86', ('TENSOR', (1, 16), 'uint8'), ('TENSOR', (32, 16), 'int8'), None, 'int32'))]

WIth this fix, one VNNI task is extracted:
[Task(func_name=dense_vnni.x86, args=(('TENSOR', (1, 16), 'uint8'), ('TENSOR', (2, 4, 16, 4), 'int8'), None, 'int32'), kwargs={}, workload=('dense_vnni.x86', ('TENSOR', (1, 16), 'uint8'), ('TENSOR', (2, 4, 16, 4), 'int8'), None, 'int32'))]

[masahi]

So autotvm doesn't run alter layout during tuning, and task extraction for an op that requires alter layout is done in a very strange way. When I added VNNI dense in #10230, I didn't want to follow the existing ad-hoc approach that supports task extraction for such ops (like https://github.com/apache/tvm/blob/main/python/tvm/topi/x86/conv2d.py#L223-L235).

For now, if you want to tune vnni dense with autotvm, you need to apply alter layout before you extract tasks.

[elvin-n]

This change have to break backward compatibility. The current ideology of AlterOpLayout and autotvm assumes that tuning does not get converted layout, just plain ones and it is responsibility of schedule to create transformed compute even for flattened data and it is responsibility of AlterOplLyout to understand that flattened data in log stands for blocked one and substitute workload to apply tuned statistics to blocked configuration.

For example here is how conv2d compute get 4d tensor, figure out that it is autotvm session and override them by 5d tensors

And here is a place where AlterOpLayout understnabds that it should correct workload to handle flattened initial tensors to packed one when AlterOpLayout is executed during network compilation.

Approach looks hacky but how it is implemented now and change in the tuning now will make previous tuned kernels not applicable for newer versions

BTW, I tried your change and it broke even tuning in cpu fp32 mode

[mengceng15]

@elvin-n Thanks for the detailed explanation. I just realize that this change could introduce serious compatibility problems due to the autotvm tuning design.

Initially, I followed @masahi 's method to apply alter op layout before the task extraction. Everything worked well. However, it could be a bit confusing to the user. Seems there is no document or hint mentioning that this workaround should be used. My intention to raise this PR was to clean up the confusion.

[masahi]

The "right" fix is to add autotvm.GLOBAL_SCOPE.in_tuning thing in the VNNI dense op. I didn't want to do that in my PR, but you are welcome to do so if this would solve your problem.

[mengceng15]

@masahi Thank you for the suggestion.

[elvin-n]

When I added VNNI dense in #10230, I didn't want to follow the existing ad-hoc approach that supports task extraction for such ops

BTW, why did not you want to follow this standard autotvm rules?
@masahi

[masahi]

BTW, why did not you want to follow this standard autotvm rules?

Just like you said "Approach looks hacky but how it is implemented now", I don't like how things are done currently. I know VNNI op is much faster than generic ones that don't require alter layout, so it makes no sense that I have to add a new compute that takes unaltered weight as input just for task extraction.

I know @tkonolige loves this topic.