[Arith] Bound for Shape Variables

[yzhliu]

Background,

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

https://discuss.tvm.ai/t/significant-increase-in-the-amount-of-cuda-code-gen-after-migrating-indexdiv-mod-to-floordiv-mod

[tqchen]

I am not sure if we want to introduce this specialized assertion as a special Expr. It might be a better idea to have it as an intrinsic

[yzhliu]

@tqchen @icemelon9 modified per comments, please review again.

[yzhliu]

@tqchen do you have any idea why this is failing (I cannot reproduce locally, and seems to be flaky) https://ci.tvm.ai/blue/organizations/jenkins/tvm/detail/PR-4486/13/pipeline#step-112-log-519

The generated IR is different, but both should be correct. This PR:

produce T_divide {
  parallel (ax0.ax1.fused, 0, (tvm_assert_bound(n0, 0, n0)*tvm_assert_bound(n1, 0, n1))) {
    for (ax2, 0, tvm_assert_bound(n2, 0, n2)) {
      for (ax3, 0, tvm_assert_bound(n3, 0, n3)) {
        if (((ax0.ax1.fused/tvm_assert_bound(n1, 0, n1)) < tvm_assert_bound(n0, 0, n0))) {
          if ((0 <= (ax0.ax1.fused % tvm_assert_bound(n1, 0, n1)))) {
            if (((ax0.ax1.fused % tvm_assert_bound(n1, 0, n1)) < tvm_assert_bound(n1, 0, n1))) {
              T_divide[(((((ax0.ax1.fused/tvm_assert_bound(n1, 0, n1))*stride) + ((ax0.ax1.fused % tvm_assert_bound(n1, 0, n1))*stride)) + (ax2*stride)) + (ax3*stride))] = (A[(((((ax0.ax1.fused/tvm_assert_bound(n1, 0, n1))*stride) + ((ax0.ax1.fused % tvm_assert_bound(n1, 0, n1))*stride)) + (ax2*stride)) + (ax3*stride))]/float32(k))
            }
          }
        }
      }
    }
  }
}

previously:

produce T_divide {
  parallel (ax0.ax1.fused, 0, (n0*n1)) {
    for (ax2, 0, n2) {
      for (ax3, 0, n3) {
        if ((0 <= select((((n1 >= 0) && ((ax0.ax1.fused % n1) >= 0)) || ((n1 < 0) && ((ax0.ax1.fused % n1) <= 0))), (ax0.ax1.fused/n1), ((ax0.ax1.fused/n1) - 1)))) {
          if ((select((((n1 >= 0) && ((ax0.ax1.fused % n1) >= 0)) || ((n1 < 0) && ((ax0.ax1.fused % n1) <= 0))), (ax0.ax1.fused/n1), ((ax0.ax1.fused/n1) - 1)) < n0)) {
            if ((0 <= select((((n1 >= 0) && ((ax0.ax1.fused % n1) >= 0)) || ((n1 < 0) && ((ax0.ax1.fused % n1) <= 0))), (ax0.ax1.fused % n1), ((ax0.ax1.fused % n1) + n1)))) {
              if ((select((((n1 >= 0) && ((ax0.ax1.fused % n1) >= 0)) || ((n1 < 0) && ((ax0.ax1.fused % n1) <= 0))), (ax0.ax1.fused % n1), ((ax0.ax1.fused % n1) + n1)) < n1)) {
                T_divide[((((select((((n1 >= 0) && ((ax0.ax1.fused % n1) >= 0)) || ((n1 < 0) && ((ax0.ax1.fused % n1) <= 0))), (ax0.ax1.fused/n1), ((ax0.ax1.fused/n1) - 1))*stride) + (select((((n1 >= 0) && ((ax0.ax1.fused % n1) >= 0)) || ((n1 < 0) && ((ax0.ax1.fused % n1) <= 0))), (ax0.ax1.fused % n1), ((ax0.ax1.fused % n1) + n1))*stride)) + (ax2*stride)) + (ax3*stride))] = (A[((((select((((n1 >= 0) && ((ax0.ax1.fused % n1) >= 0)) || ((n1 < 0) && ((ax0.ax1.fused % n1) <= 0))), (ax0.ax1.fused/n1), ((ax0.ax1.fused/n1) - 1))*stride) + (select((((n1 >= 0) && ((ax0.ax1.fused % n1) >= 0)) || ((n1 < 0) && ((ax0.ax1.fused % n1) <= 0))), (ax0.ax1.fused % n1), ((ax0.ax1.fused % n1) + n1))*stride)) + (ax2*stride)) + (ax3*stride))]/float32(k))
              }
            }
          }
        }
      }
    }
  }
}

[yzhliu]

find the problem...

[yzhliu]

@tqchen @icemelon9 It is ready for review

[tqchen]

rename -> is_var_or_assert_bound

[tqchen]

NOTE, we are in progress of updating IRMutators to new style, would be great if we can directly change it here related #4607

[tqchen]

wrap_assert_bound_on_var

[tqchen]

After looking at the set of changes, seems we still have quite a few places where we want to try to assume the properties of a Var.

This makes me wonder if we should also revisit the option to bring more information to Variable, but instead of putting them inside Variable, we can subclass the Variable to create special vars(e.g. IterVar) that contains these information. Given that things are easier with the new Object system. See https://discuss.tvm.ai/t/discuss-embed-more-bound-information-into-var-or-expr/4079/28

Of course this does not remove the value of the additional constraint wrapper, @icemelon9 @yzhliu thoughts?

[yzhliu]

@tqchen I also feel it would be better than wrapping with a call node

[tqchen]

If we are going to use a Var instead, we will can introduce a special subclass of Variable called ShapeVarNode, which corresponds to shape can contain non-neg info, the solution mentioned in the post will depend on migrating all the Mutator to the new style(which will be done in a few days).

Note that will mean we will need to force user to construct shape_var when passing to shapes.
In the meanwhile, the intrinsic approach can co-exist with the shape_var one even if we decided to introduce that one, so it might be fine if we first proceed with this solution

[tqchen]

@yzhliu can you rebase, @icemelon9 @ZihengJiang can you help to do a round of review?

[yzhliu]

@tqchen You want me to check current approach in, or directly switch to IterVar implement instead?
I feel IterVar would be cleaner, currently we do if node is assert_bound in many places.

[tqchen]

@yzhliu I am fine either way, you can make a decision.

Since both are useful approaches and won't affect each other. We might want to think about the name a bit if we are going to sub-class Var (while it makes sense to use IterVar for iterators, it may not sounds too related to shape, perhaps ShapeVar would be a better name).

[tqchen]

close this pr as it is superceded by #4684