[ARITH] migrate indexdiv/mod to floordiv/mod

python/tvm/expr.py

@@ -92,16 +92,13 @@ def __rtruediv__(self, other):
         return _generic.divide(other, self)
 
     def __floordiv__(self, other):
-        # return _generic.floordiv(self, other)
-        return _generic.divide(self, other)
+        return _generic.floordiv(self, other)
 
     def __rfloordiv__(self, other):
-        # return _generic.floordiv(other, self)
-        return _generic.divide(other, self)
+        return _generic.floordiv(other, self)
 
     def __mod__(self, other):
-        raise div_ambiguity_error()
-        # return _make._OpMod(self, other)
+        return _make._OpFloorMod(self, other)
 
     def __neg__(self):
         neg_one = _api_internal._const(-1, self.dtype)

src/lang/attr_functor.h

@@ -6,9 +6,9 @@
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
- * 
+ *
  *   http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
@@ -87,6 +87,8 @@ class AttrFunctor<R(const NodeRef& n, Args...)> {
   virtual R VisitAttr_(const ir::Mul* op, Args... args) ATTR_FUNCTOR_DEFAULT;
   virtual R VisitAttr_(const ir::Div* op, Args... args) ATTR_FUNCTOR_DEFAULT;
   virtual R VisitAttr_(const ir::Mod* op, Args... args) ATTR_FUNCTOR_DEFAULT;
+  virtual R VisitAttr_(const ir::FloorDiv* op, Args... args) ATTR_FUNCTOR_DEFAULT;
+  virtual R VisitAttr_(const ir::FloorMod* op, Args... args) ATTR_FUNCTOR_DEFAULT;
   virtual R VisitAttr_(const ir::Min* op, Args... args) ATTR_FUNCTOR_DEFAULT;
   virtual R VisitAttr_(const ir::Max* op, Args... args) ATTR_FUNCTOR_DEFAULT;
   virtual R VisitAttr_(const ir::GE* op, Args... args) ATTR_FUNCTOR_DEFAULT;
@@ -119,6 +121,9 @@ class AttrFunctor<R(const NodeRef& n, Args...)> {
     ATTR_FUNCTOR_DISPATCH(Sub);
     ATTR_FUNCTOR_DISPATCH(Mul);
     ATTR_FUNCTOR_DISPATCH(Div);
+    ATTR_FUNCTOR_DISPATCH(Mod);
+    ATTR_FUNCTOR_DISPATCH(FloorDiv);
+    ATTR_FUNCTOR_DISPATCH(FloorMod);
     ATTR_FUNCTOR_DISPATCH(Min);
     ATTR_FUNCTOR_DISPATCH(Max);
     ATTR_FUNCTOR_DISPATCH(GE);
@@ -160,6 +165,8 @@ class AttrsEqualHandler :
   bool VisitAttr_(const ir::Mul* lhs, const NodeRef& other) final;
   bool VisitAttr_(const ir::Div* lhs, const NodeRef& other) final;
   bool VisitAttr_(const ir::Mod* lhs, const NodeRef& other) final;
+  bool VisitAttr_(const ir::FloorDiv* lhs, const NodeRef& other) final;
+  bool VisitAttr_(const ir::FloorMod* lhs, const NodeRef& other) final;
   bool VisitAttr_(const ir::Min* lhs, const NodeRef& other) final;
   bool VisitAttr_(const ir::Max* lhs, const NodeRef& other) final;
   bool VisitAttr_(const ir::GE* lhs, const NodeRef& other) final;
@@ -201,6 +208,8 @@ class AttrsHashHandler :
   size_t VisitAttr_(const ir::Mul* op) final;
   size_t VisitAttr_(const ir::Div* op) final;
   size_t VisitAttr_(const ir::Mod* op) final;
+  size_t VisitAttr_(const ir::FloorDiv* op) final;
+  size_t VisitAttr_(const ir::FloorMod* op) final;
   size_t VisitAttr_(const ir::Min* op) final;
   size_t VisitAttr_(const ir::Max* op) final;
   size_t VisitAttr_(const ir::GE* op) final;

src/lang/attrs.cc

@@ -6,9 +6,9 @@
  * to you under the Apache License, Version 2.0 (the
  * "License"); you may not use this file except in compliance
  * with the License.  You may obtain a copy of the License at
- * 
+ *
  *   http://www.apache.org/licenses/LICENSE-2.0
- * 
+ *
  * Unless required by applicable law or agreed to in writing,
  * software distributed under the License is distributed on an
  * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
@@ -154,6 +154,8 @@ TVM_DEFINE_ATTRS_BINOP_EQUAL(Sub);
 TVM_DEFINE_ATTRS_BINOP_EQUAL(Mul);
 TVM_DEFINE_ATTRS_BINOP_EQUAL(Div);
 TVM_DEFINE_ATTRS_BINOP_EQUAL(Mod);
+TVM_DEFINE_ATTRS_BINOP_EQUAL(FloorDiv);
+TVM_DEFINE_ATTRS_BINOP_EQUAL(FloorMod);
 TVM_DEFINE_ATTRS_BINOP_EQUAL(Max);
 TVM_DEFINE_ATTRS_BINOP_EQUAL(Min);
 TVM_DEFINE_ATTRS_BINOP_EQUAL(GE);
@@ -266,6 +268,8 @@ TVM_DEFINE_ATTRS_BINOP_HASH(Sub);
 TVM_DEFINE_ATTRS_BINOP_HASH(Mul);
 TVM_DEFINE_ATTRS_BINOP_HASH(Div);
 TVM_DEFINE_ATTRS_BINOP_HASH(Mod);
+TVM_DEFINE_ATTRS_BINOP_HASH(FloorDiv);
+TVM_DEFINE_ATTRS_BINOP_HASH(FloorMod);
 TVM_DEFINE_ATTRS_BINOP_HASH(Max);
 TVM_DEFINE_ATTRS_BINOP_HASH(Min);
 TVM_DEFINE_ATTRS_BINOP_HASH(GE);

src/lang/buffer.cc

@@ -32,8 +32,8 @@
 namespace tvm {
 
 // TODO(tqchen): change to floormod/div
-using IndexMod = ir::Mod;
-using IndexDiv = ir::Div;
+using IndexMod = ir::FloorMod;
+using IndexDiv = ir::FloorDiv;
 
 Array<Expr> SimplifyArray(Array<Expr> array) {
   for (size_t i = 0; i < array.size(); ++i) {

src/lang/expr_operator.cc

@@ -208,11 +208,11 @@ Expr operator%(Expr a, Expr b) {
 
 // TODO(tqchen): switch to floordiv
 Expr indexdiv(Expr a, Expr b) {
-  return truncdiv(a, b);
+  return floordiv(a, b);
 }
 
 Expr indexmod(Expr a, Expr b) {
-  return truncmod(a, b);
+  return floormod(a, b);
 }
 
 Expr floordiv(Expr a, Expr b) {

tests/python/unittest/test_codegen_device.py

@@ -48,6 +48,8 @@ def test_add_pipeline():
     stmt = tvm.ir_pass.Simplify(stmt)
     fapi = tvm.ir_pass.MakeAPI(stmt, "myadd", [Ab, Bb, Db], 0, True)
     fsplits = [x for x in tvm.ir_pass.SplitHostDevice(fapi)]
+    # lower the floordiv
+    fsplits = [tvm.ir_pass.LowerIntrin(x, "generic") for x in fsplits]
     fsplits[0] = tvm.ir_pass.LowerTVMBuiltin(fsplits[0])
 
     def check_target(device, host="stackvm"):

tests/python/unittest/test_codegen_vm_basic.py

@@ -37,6 +37,7 @@ def tvm_call_back_get_shape(shape0):
     stmt = tvm.make.Evaluate(tvm.call_packed("tvm_call_back_get_shape", Ab.shape[0]))
     fapi = tvm.ir_pass.MakeAPI(stmt, "print_shape", [Ab], 0, True)
     fapi = tvm.ir_pass.LowerTVMBuiltin(fapi)
+    fapi = tvm.ir_pass.LowerIntrin(fapi, "generic")
     run_jit(fapi, lambda f: f(a))