[RFC][ARITH] Migrate to FloorDiv/Mod for Index Caculation

[tqchen]

This is a followup of #3478

Background: DIvMod Variants

Trunc div/mod

IEEE standard integer division is truncated division, directly maps to the hardware. This is the division used by C and TVM by default. The problem of trunc div is that we usually relies on knowing the sign of both operands. For example, we cannot simplify (x * 4 + 3)/4 => x if x can be negative.

Floor div/mod

Always round down. This is the division mode adopted by python. To simplify floordiv, we usually need to know the sign of divisor, but not necessarily the sign of dividend floordiv(x * 4 + 3,4) => x .

Euclidean div/mod

Always ensures that the remainder is non-negative. This is the division mode used by Halide by default. Note that when the divisor is positive, Euclidean and floordiv are equivalent. But floordiv will produce a negative remainder when the divisor is negative.

Challenges of using FloorDiv

FloorDiv/Mod has a better advantage of doing simplifications. The main challenge of using FloorDiv/Mod is when we lower the code into low level divisions, where we have to make use of trunc div/mod.

In most cases, when we can prove the positiveness of the divident and divisor, we can simply lower the results into truncdiv without problem to get the best of efficiency as well as analysis. However, we need be careful about cases where these proves fail due to lack of information, and provide these information into the IR.

Potential Bad Case that Needs Our Attention

for (i, 0, n) {
    A[floordiv(i, m) * m] = 0 
}

Consider the above code. We cannot simply translate floordiv(i, m) to truncdiv(i, m) because we do not know that m is non-negative in the above IR.

This is however due to lack of additional information provided. In many cases, m is intended to be a symbolic shape of a tensor, in which case it is supposed to be non-negative. If we have a hint in the code as in below, then we will be able to handle it correctly.

assert m >= 0
for (i, 0, n) {
    A[floordiv(i, m) * m] = 0 
}

This is something we need to keep in mind especially when we handle cases like symbolic shape operations.

Tradeoffs and Solutions

• FloorDiv is relatively easier for simplification as we do not need to know the sign of operand a
  • we can do floordiv(a+ 17, 16) -> floordiv(a + 1, 16), regardless of the sign of a
  • However, we need to prove the sign of the operands eventually when lower to the low level code, as we can only lower floordiv(a, b)->truncdiv(a, b) if a>=0 and b>=0
• TruncDiv is easier to for codegen, as it directly corresponds to low level instructions, but it is harder for simplification:
  • truncdiv(a+ 17, 16) -> truncdiv(a + 1, 16) only if a+1>=0.

So we have to make the choice of division mode for index calculations. We isolate that choice by introducing indexdiv/mod functions, that do the division for a and b that are non-negative.

• Currently, we lower indexdiv-> truncdiv, we plan to move to floordiv
• We can plugin additional hints: indexdiv(a, b)->floordiv(a, nonneg_hint(b) if b is Var else b)

Steps

• Introduce lowering of FloorDiv/Mod [ARITH] Add Lowering rule for FloorDiv/Mod #3976
• Migrate internal index calculations of split/fuse to floordiv
• Consider remove operator overload of div and mod and force developers to use truncdiv or floordiv, so they can be aware of the difference

[tqchen]

cc @dmlc/tvm-committer

[tqchen]

https://discuss.tvm.ai/t/discuss-embed-more-bound-information-into-var-or-expr/4079

[tqchen]

Migration is finished #4008, @icemelon9 we need to followup on https://discuss.tvm.ai/t/discuss-embed-more-bound-information-into-var-or-expr/4079/14 to get the best symbolic shape code