[TIR] Utility function to decide loop mapping for auto tensorization

[masahi]

Add TensorizeInfo structure and GetTensorizeLoopMapping function, that are used for determining the correspondence of loops between a target block and an intrinsic description.

Matching is based on a heuristic: It works in all cases I tested (CPU dot product for dense / conv2d, CPU / GPU matmul), but there is no guarantee that it always finds the "right" mapping. If the mapping is not correct, tensorize would fail.

The original code is https://github.com/spectrometerHBH/tvm/blob/auto-tensorization/src/tir/schedule/analysis/analysis.cc#L1175, I modified this code to support more cases and added tests. I'm sending this PR on behalf of the team, but most of the work were done by others earlier.

Co-authored-by: Siyuan Feng [email protected]
Co-authored-by: Bohan Hou [email protected]
Co-authored-by: Hongyi Jin [email protected]
Co-authored-by: Ruihang Lai [email protected]
Co-authored-by: Wuwei Lin [email protected]

@vinx13 @junrushao1994 @spectrometerHBH @Hzfengsy @MasterJH5574 @jinhongyii

[masahi]

The logic here is very different from the one in the original code https://github.com/spectrometerHBH/tvm/blob/auto-tensorization/src/tir/schedule/analysis/analysis.cc#L1246. I was not able to understand why the original code has been written that way and it didn't work for the case where matching loops in the target block are not in the innermost positions (conv2d NCHWc on CPU, a test in

tvm/tests/python/unittest/test_tir_schedule_analysis.py

Line 199 in d6ae848

def test_get_tensorize_loop_mapping_conv2d_nchwc_vnni():

).

I think my change is simple and obvious. The condition for a match is (1) divisibility of loop extent and (2) matching iterator types (reduction vs spatial). Mapping is determined starting from the innermost axis.

Please have a look at this change carefully, and let me know if I need to bring back some logic in the original code @spectrometerHBH @vinx13

[junrushao]

I would love to have @spectrometerHBH review this change before merging

[spectrometerHBH]

The goal of the original mapping is to support

for k:
  for i:
    for j:
        C[i, j] += A[i, k] * B[k, j]

where loops are not in the same order as the tensor intrinsic description function.

[spectrometerHBH]

But it also makes sense if we don't support such cases for this PR. So I approve it.

[masahi]

Thanks @spectrometerHBH, I now understand the original code and was able to integrate the original logic to support loop permutations. Please have a look at the current diff, also cc @vinx13 @Hzfengsy @MasterJH5574

The key difference between the original code and the code I submitted yesterday was that, my code was looking at only the loop nest (ForNode) to determine the mapping, while @spectrometerHBH's mapping logic is based on iter_var/value of the block (so invariant to the order of the loop nest).

[MasterJH5574]

Thanks Masa! I just caught a minor point.

[masahi]

i.e., no matter what the permutation of loop is, we should always have

i, j, k = T.axis.remap("SSR", [i0, i1, i2])

for GEMM.

I think this is a reasonable assumption. Correct me if I'm wrong @spectrometerHBH @junrushao1994 @vinx13

[vinx13]

I agree this is a reasonable assumption. Though there might be corner cases, it covers all of the current use cases