[TIR] Unify index data type when creating prim func

[vinx13]

• Added data type pass unification pass to by default promote data types of all indices and shapes to int64 when creating prim func.
• Added some fixes for lowering passes to make it compatible with int64 data type.

[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 c6db9f5 can be found here.

_{Generated by tvm-bot}

[junrushao]

Happy to review and let's fix the CI :-)

[junrushao]

LGTM!

[ibsidorenko]

Hello, @vinx13 !

I have one question about this PR... Is it possible to give more information or motivation about why do we need to convert indexes into "int64" data type?

A few words about why I am asking:
I am working on MetaScheduler for Hexagon target. And found that this PR dramatically reduce performance for some operations.

Example: Average Pooling 2D
For this operator we use indexes in its compute function and pool2d divisor.
Before IndexDataTypeNormalizer:
pool_avg[ax0, ax1, ax2, ax3, ax4] = (pool_sum[ax0, ax1, ax2, ax3, ax4] / max((((min(1, (34 - ax2)) + 2) - max((1 - ax2), 0))*((min(1, (34 - ax3)) + 2) - max((1 - ax3), 0))), 1))
After IndexDataTypeNormalizer:
pool_avg[ax0, ax1, ax2, ax3, ax4] = cast(int32, (cast(int64, pool_sum[ax0, ax1, ax2, ax3, ax4]) / max((((min(1i64, (34i64 - ax2)) + 2i64) - max((1i64 - ax2), 0i64))*((min(1i64, (34i64 - ax3)) + 2i64) - max((1i64 - ax3), 0i64))), 1i64)))

As you can see we get extra cast("int64"). Unfortunately, Hexagon does not support vectorization of "int64" data types and performance became very very poor.

P.S. Just for experiment I have reverted conversion of indexes into int64 and get performance gain +40% (!!!).

So, I would like to fix it somehow but I would like to know motivation for this PR.

Thank you in advance!

Just FYI cc @masahi

[vinx13]

Sometimes the model contains mixed indices type (e.g both int32 and int64). It causes dtype mismatch error during scheduling. It is expected that this pass doesn't hurt performance since there is another pass NarrowDataType that should convert it back to int32 (or other smaller types) if possible.

In your case, the result of conversion seems more complicated than expected, it contains some unnecessary cast. I'd expect something like

pool_avg[ax0, ax1, ax2, ax3, ax4] = (pool_sum[ax0, ax1, ax2, ax3, ax4] / max((((min(1, (34 - ax2)) + 2) - max((1 - ax2), 0))*((min(1, (34 - ax3)) + 2) - max((1 - ax3), 0))), 1))

that looks the same as before, except that variables and constant ints have int64 dtypes