[ARCH][UX] Enforce i64 index in shape/index expression
tqchen opened this issue · 2 comments
Up until now we have allowed mixed integer types in ShapeExpr, this have created problems for TIR lowering and test cases. In most cases, the integer variable created (the constant) were by default i32.
However, as we start to work with Array, i32 is not sufficient to capture size of big arrays. The low level TIR already have ability to be able to narrow i64 loops into i32 if it is decided as safe to do so. As a result, it is better to have stronger consistency in relax side to enforce shape indices to always be in i64.
This is an issue to improve our overall architecture consistency.
Action Items
- A0: Add an auxiliary function
relax.index_var()which returns tir.Var(i64)- Update testcases to use shape var creation
- A1: Update constructor of ShapeExpr, to do the following things:
- Automatically convert constant integer (IntImm) to i64 IntImm.
- If the value is general PrimExpr, do ICHECK the dtype of PrimExpr to be i64 -- this will results in errors in some of existing code that we need to continue to fix -- likely replace var decl to shape var
- A2: The above guard will uncover problematic places that we need to further fix, likely to fix the shape deduction logics to maintain the above invariance.
A1 would help us to uncover such inconsistency early and improve our overall infra consistency in the codebase
This would be easy to add to the language spec, too.