[TIR] Fix buffer shape and IndexMap indices dtype mismatch

[masahi]

After the PR #13327, I'm getting a dtype-mismatch error at

tvm/src/tir/schedule/primitive/layout_transformation.cc

Line 310 in b6fae9b

IterVar(Range::FromMinExtent(make_zero(dim.dtype()), dim), virtual_var, kDataPar));

, because dim.dtype() is now int64 while virtual_var.dtype() is int32. The dtypes of initial_indices in IndexMap are fixed to int32 (see below), this is in conflict with the above PR which made the dtypes of buffer shapes int64.

tvm/src/tir/ir/index_map.cc

Line 51 in 458ca81

initial_indices.push_back(Var("i" + std::to_string(i), DataType::Int(32)));

tvm/python/tvm/tir/function.py

Line 395 in 78b5322

default_index_dtype = "int32"

Since initial_indices is used everywhere to construct loop_var / iter_var / iter_values of For and Block etc, I'm adding a dtype legalization at the beginning of TransformLayout, when the dtypes of the input buffer and initial_indices do not match.

@vinx13 @Lunderberg

[tvm-bot]

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

• cc @Hzfengsy, @junrushao _{See #10317 for details}
• Built docs for commit aa7f08e can be found here.

_{Generated by tvm-bot}

[Lunderberg]

Good catch, and the change looks reasonable. The MapIndices method already uses SubstituteWithDataTypeLegalization, so it makes sense for the transform layout to use it as well.

[Lunderberg]

Regarding the readability of the unit test, I did some poking around, and the following PrimFunc triggers the same error on main when the layout is transformed using lambda h: [h//8, h%8].

@T.prim_func
def func(A: T.Buffer[T.int64(58), "int32"]):
    for i in T.serial(T.int64(58)):
        with T.block("block"):
            vi = T.axis.remap("S", [i])
            A[vi] = 0

[masahi]

@Lunderberg Thanks for finding the simple test case! I hit this error when running this Hexagon test

tvm/tests/python/contrib/test_hexagon/metaschedule_e2e/test_resnet50_int8.py

Line 397 in b29ab5c

def test_packed_8x8x32_resnet50(hexagon_launcher):

. The original test case in this PR, with its cache_read, index map etc, was directly taken from this Hexagon test.

The bug is triggered when we hit the code path

tvm/src/tir/schedule/primitive/layout_transformation.cc

Line 282 in b6fae9b

// If we got to this point, all indices used to access the

I was not sure what makes this code path hit and why existing tests didn't hit it.

[masahi]

@Lunderberg @vinx13 PTAL, thanks.

[Lunderberg]

Thank you for making the changes, and LGTM!