[ARITH] Tight bound for floormod

[hzfan]

Estimate the range of floormod(a, b) when b < 0. Fix #6691

[giuseros]

Just a suggestion: why don't you move this (very well done) explanation at the beginning of the function? It seems to cover more than just the else branch

[hzfan]

Sure. Thanks for the suggestion.

[tqchen]

@hzfan can you also comment why negative b appears in the case of #6691, is it because we did not capture bound of b? Usually negative bound should not appear

[hzfan]

@tqchen The ir (before narrowing) is like

for (ax0.ax1.fused.ax2.outer.fused: int64, 0i64, 42i64) "parallel" {
  for (n.oc_chunk.fused.oh.outer.fused: int64, 0i64, a_very_long_upper_bound) {
    ...
    ...
    for (oh.inner: int32, 0, 2) {
      for (ow.inner: int32, 0, 14) {
        for (oc_block: int32, 0, 16) "vectorized" {
          ...floormod(n.oc_chunk.fused.oh.outer.fused, cast(int64, floordiv(((cast(int32, ((floormod(ax0.ax1.fused.ax2.outer.fused, 7i64)*2i64) + 1i64)) + 1) - cast(int32, (floormod(ax0.ax1.fused.ax2.outer.fused, 7i64)*2i64))), 2)))...
        }
      }
    }
  }
}

Let a = n.oc_chunk.fused.oh.outer.fused (the first operand of floormod), b = cast(int64, floordiv(((cast(int32, ((floormod(ax0.ax1.fused.ax2.outer.fused, 7i64)*2i64) + 1i64)) + 1) - cast(int32, (floormod(ax0.ax1.fused.ax2.outer.fused, 7i64)*2i64))), 2)) (second operand of floormod)
const_int_bound gives 0 <= a <= 6 and -5 <= b <= 7.

In a simplified form,

b = floordiv( cast(i32, c * 2 + 1) + 1 - cast(i32, c * 2), 2)
c = floormod(ax0.ax1.fused.ax2.outer.fused, 7)

I guess cast(i32, c * 2 + 1) + 1 - cast(i32, c * 2) is not simplified as 2, instead, its bound is analyzed directly, so the min becomes negative due to the minus operation.

[hzfan]

It sort of make sense, because at the narrowing pass, we still don't know c * 2 + 1 fits in i32 and the cast cannot be removed yet.

I added the above example as a test. The FPS improves to 400 with INDEX_DEFAULT_I64=ON

[tqchen]

I see, in this particular case, perhaps it makes sense to optimize such pattern and make sure cast(i32, c * 2 + 1) + 1 - cast(i32, c * 2) get simplified as well. One way to do so is to first change i32 casts to i64 before narrow.

It would also be useful to find out why cast(value, i32) is inserted since we preferred i64 in most cases.

We can do that in another PR. This PR's improvement is certainly useful

[tqchen]

THanks @hzfan @giuseros . It would also be great if we can followup further on the above case:

• Find out why cast i32 is inserted (ideally we should be all in i64 in compute mode)
• Consider insert a cast removing pass if cast i32 persists, or update rewrite simplify to handle cast