64 bits integer simplification in TIR

elfswallow · October 6, 2022, 4:55am

I encountered a bug when I try to simplify a 64 bits integer TIR expression: floordiv(uint32(floordiv((int64(floordiv(iii, (uint32)1))*(int64)536870912), (int64)4294967296)), (uint32)1) Using arith::Analyzer to “Simplify” this expression will result in floormod 0 error. The cause is that the definition of floormod is PrimExpr(PrimExpr, int) which truncates int64_t to int at here:

github.com

apache/tvm/blob/3d5e440fff581e30b614037997feb81bf587cefe/include/tvm/tir/op.h#L1018


  inline PrimExpr Name(bool a, const PrimExpr& b) { return Name(PrimExpr(a), b); }


#define TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD_SPANNED(Name)          \
  inline PrimExpr Name(const PrimExpr& a, bool b, Span span = Span()) { \
    return Name(a, PrimExpr(b), span);                                  \
  }                                                                     \
  inline PrimExpr Name(bool a, const PrimExpr& b, Span span = Span()) { \
    return Name(PrimExpr(a), b, span);                                  \
  }


#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD(Name) \
  inline PrimExpr Name(const PrimExpr& a, int b) { \
    return Name(a, tir::make_const(a.dtype(), b)); \
  }                                                \
  inline PrimExpr Name(int a, const PrimExpr& b) { return Name(tir::make_const(b.dtype(), a), b); }


#define TVM_DEFINE_INT_OP_CONST_VAL_OVERLOAD_SPANNED(Name)             \
  inline PrimExpr Name(const PrimExpr& a, int b, Span span = Span()) { \
    return Name(a, tir::make_const(a.dtype(), b), span);               \
  }                                                                    \
  inline PrimExpr Name(int a, const PrimExpr& b, Span span = Span()) { \

Adding int64_t overloading version will solve the problem, but a potential issue is that any existing invocation of this function with uint64_t type will raise complaining of “call of overloadded … is ambiguous”.

vinx13 · October 6, 2022, 8:38pm

ideally we should create IntImm explicitly with specific dtype. where did your expression come from? is it possible to fix such usage by creating IntImm?

elfswallow · October 7, 2022, 1:58am

@vinx13, that should fix this bug too. In my case, the direct caller to here is

github.com

apache/tvm/blob/bb2cec1941db6f5a67b85550faa9279c4c14e2a3/src/arith/rewrite_simplify.cc#L799


                   c1.Eval()->value > 0 && c3.Eval()->value > 0);


if (floordiv(x * c1 + y, c2).Match(ret) || floordiv(x * c1, c2).Match(ret) ||
    floordiv(y + x * c1, c2).Match(ret)) {
  int64_t c1val = c1.Eval()->value;
  int64_t c2val = c2.Eval()->value;
  PrimExpr yval = y.EvalOr(Integer(0));
  if (c2val == 0) return ret;


  // try eliminate residue part
  PrimExpr residue =
      floordiv(x.Eval() * floormod(c1.Eval(), c2val) + floormod(yval, c2val), c2val);
  PrimExpr y_div = CanProveEqual(floordiv(yval, c2val), 0) ? 0 : floordiv(yval, c2val);
  auto bound = analyzer_->const_int_bound(residue);
  if (bound.defined() && bound->max_value == bound->min_value) {
    return x.Eval() * floordiv(c1val, c2.Eval()) + (y_div + Integer(bound->max_value));
  }


  // try simplify divisor
  if (c1val > 0 && c2val > 0 && c2val % c1val == 0 &&
      CanProveLess(floormod(yval, c2val), c1val)) {

I assume that there are many places doing the same way, i.e., passing int64_t or uint64_t directly to this function.

vinx13 · October 7, 2022, 5:22pm

In this case, it should use c1.Eval() / c2.Eval() instead, which is the same way as other rewriting rules and it respects the original dtypes. Actually, it’s not safe to directly create IntImm with i64, or adding i64 overload here because that will always promote the whole expression to i64.

elfswallow · October 8, 2022, 2:13am

I think the intention of the original code in op.h is to prioritize the type of PrimExpr operand if some of the operands are not “TIR typed”. So even if 64bit overload is added, the eventual expression type is determined by the PrimExpr operand. So if the expression is finally a 64bit computation, it should result from the type of the PrimExpr operand, and I think that should also be the intention of the caller. On the other hand, if PrimExpr is 32 bit, it won’t escalate the compute precision to 64 bits.

vinx13 · October 10, 2022, 11:27pm

That makes sense. So either solution should work