[EXPR] ir_operator.h->expr_operator.h Centralize const folder logic (#2719)

include/tvm/buffer.h

@@ -10,7 +10,7 @@
 
 #include "base.h"
 #include "expr.h"
-#include "ir_operator.h"
+#include "expr_operator.h"
 #include "tvm/node/container.h"
 
 namespace tvm {

include/tvm/data_layout.h

@@ -16,7 +16,7 @@
 #include <utility>
 #include <algorithm>
 
-#include "ir_operator.h"
+#include "expr_operator.h"
 
 namespace tvm {
 

include/tvm/ir_operator.h → include/tvm/expr_operator.h

 /*!
  *  Copyright (c) 2018 by Contributors
- * \file tvm/ir_operator.h
+ * \file tvm/expr_operator.h
  * \brief Common operators defined for Expr.
  *
  * \note Most of the operator defined here perform simple constant folding
  *   when the type is int32 or int64 for simplifying the index expressions.
  */
-#ifndef TVM_IR_OPERATOR_H_
-#define TVM_IR_OPERATOR_H_
+#ifndef TVM_EXPR_OPERATOR_H_
+#define TVM_EXPR_OPERATOR_H_
 
 #include <algorithm>
 #include <type_traits>
@@ -617,4 +617,4 @@ TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(operator&&);
 TVM_DEFINE_LOGICAL_OP_CONST_VAL_OVERLOAD(operator||);
 
 }  // namespace tvm
-#endif  // TVM_IR_OPERATOR_H_
+#endif  // TVM_EXPR_OPERATOR_H_

include/tvm/operation.h

@@ -10,7 +10,7 @@
 #include <vector>
 #include <unordered_map>
 #include "expr.h"
-#include "ir_operator.h"
+#include "expr_operator.h"
 #include "tensor.h"
 #include "schedule.h"
 #include "arithmetic.h"

include/tvm/tensor.h

@@ -14,7 +14,7 @@
 
 #include "base.h"
 #include "expr.h"
-#include "ir_operator.h"
+#include "expr_operator.h"
 #include "arithmetic.h"
 
 namespace tvm {

include/tvm/tvm.h

@@ -8,7 +8,7 @@
 
 #include "base.h"
 #include "expr.h"
-#include "ir_operator.h"
+#include "expr_operator.h"
 #include "tensor.h"
 #include "operation.h"
 #include "packed_func_ext.h"

src/api/api_ir.cc

@@ -5,9 +5,8 @@
  */
 #include <tvm/expr.h>
 #include <tvm/ir.h>
-#include <tvm/ir_operator.h>
 #include <tvm/api_registry.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 
 namespace tvm {
 namespace ir {

src/arithmetic/const_fold.h

+/*!
+ *  Copyright (c) 2019 by Contributors
+ * \file const_fold.h
+ * \brief Centralized location for constant folding.
+ */
+#ifndef TVM_ARITHMETIC_CONST_FOLD_H_
+#define TVM_ARITHMETIC_CONST_FOLD_H_
+
+#include <tvm/ir.h>
+#include <algorithm>
+
+namespace tvm {
+namespace arith {
+
+/*!
+ * \brief Try to run binary compute with constant folding.
+ *
+ * \param a The left operand.
+ * \param b The right operand.
+ * \tparam Op The operator type.
+ *
+ * \note a and b Must already matched data types with each other.
+ * \return nullptr if constant fold fails, otherwise return folded result.
+ */
+template<typename Op>
+inline Expr TryConstFold(Expr a, Expr b);
+
+/*!
+ * \brief Try to run unary compute with constant folding.
+ *
+ * \param a The left operand.
+ * \tparam Op The operator type.
+ *
+ * \note a and b Must already matched data types with each other.
+ * \return nullptr if constant fold fails, otherwise return folded result.
+ */
+template<typename Op>
+inline Expr TryConstFold(Expr a);
+
+/*!
+ * \brief Check whether type is used to represent index.
+ *
+ * Index types are frequently used in shape computation
+ * and need to be aggressively constant-folded.
+ *
+ * \param type The type to represent index.
+ * \return the checked result.
+ */
+inline bool IsIndexType(const Type& type) {
+  return type.is_int() && type.lanes() == 1 &&
+      (type.bits() == 32 || type.bits() == 64);
+}
+
+
+#define TVM_ARITH_CONST_PROPAGATION(BODY)                               \
+  using ir::IntImm;                                                     \
+  using ir::UIntImm;                                                    \
+  using ir::FloatImm;                                                   \
+  const IntImm* pa = a.as<IntImm>();                                    \
+  const IntImm* pb = b.as<IntImm>();                                    \
+  const FloatImm* fa = a.as<FloatImm>();                                \
+  const FloatImm* fb = b.as<FloatImm>();                                \
+  BODY;
+
+
+#define TVM_INDEX_CONST_PROPAGATION(BODY)                               \
+  using ir::IntImm;                                                     \
+  using ir::UIntImm;                                                    \
+  const IntImm* pa = a.as<IntImm>();                                    \
+  const IntImm* pb = b.as<IntImm>();                                    \
+  const Type& ta = a.type();                                            \
+  const Type& tb = b.type();                                            \
+  if (arith::IsIndexType(ta) && arith::IsIndexType(tb)) {               \
+    BODY;                                                               \
+  }                                                                     \
+
+
+// specialization of constant folders.
+template<>
+inline Expr TryConstFold<ir::Add>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      const Type& rtype = a.type();
+      if (pa && pb) return IntImm::make(rtype, pa->value + pb->value);
+      if (pa && pa->value == 0) return b;
+      if (pb && pb->value == 0) return a;
+      if (fa && fb) return FloatImm::make(rtype, fa->value + fb->value);
+      if (fa && fa->value == 0) return b;
+      if (fb && fb->value == 0) return a;
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::Sub>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      const Type& rtype = a.type();
+      if (pa && pb) return IntImm::make(rtype, pa->value - pb->value);
+      if (pb && pb->value == 0) return a;
+      if (fa && fb) return FloatImm::make(rtype, fa->value - fb->value);
+      if (fb && fb->value == 0) return a;
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::Mul>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      const Type& rtype = a.type();
+      if (pa && pb) return IntImm::make(rtype, pa->value * pb->value);
+      if (pa) {
+        if (pa->value == 1) return b;
+        if (pa->value == 0) return a;
+      }
+      if (pb) {
+        if (pb->value == 1) return a;
+        if (pb->value == 0) return b;
+      }
+      if (fa && fb) return FloatImm::make(rtype, fa->value * fb->value);
+      if (fa) {
+        if (fa->value == 1) return b;
+        if (fa->value == 0) return a;
+      }
+      if (fb) {
+        if (fb->value == 1) return a;
+        if (fb->value == 0) return b;
+      }
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::Div>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      const Type& rtype = a.type();
+      // due to division and mod can have different modes
+      // only constant fold positive number where rule is fixed.
+      if (pa && pb && pa->value >= 0 && pb->value > 0) {
+        return IntImm::make(rtype, pa->value / pb->value);
+      }
+      if (pa) {
+        if (pa->value == 0) return a;
+      }
+      if (pb) {
+        if (pb->value == 1) return a;
+        CHECK_NE(pb->value, 0) << "Divide by zero";
+      }
+      if (fa && fb && fb->value != 0) {
+        return FloatImm::make(rtype, fa->value / fb->value);
+      }
+      if (fa && fa->value == 0) return a;
+      if (fb) {
+        if (fb->value == 1) return a;
+        CHECK_NE(fb->value, 0) << "Divide by zero";
+      }
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::Mod>(Expr a, Expr b) {
+  TVM_INDEX_CONST_PROPAGATION({
+      const Type& rtype = a.type();
+      // due to division and mod can have different modes
+      // only constant fold positive number where rule is fixed.
+      if (pa && pb && pa->value >= 0 && pb->value > 0) {
+        return IntImm::make(rtype, pa->value % pb->value);
+      }
+      if (pa) {
+        if (pa->value == 0) return a;
+      }
+      if (pb) {
+        if (pb->value == 1) return make_zero(rtype);
+        CHECK_NE(pb->value, 0) << "Divide by zero";
+      }
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::Min>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      const Type& rtype = a.type();
+      if (pa && pb) return IntImm::make(rtype, std::min(pa->value, pb->value));
+      if (fa && fb) return FloatImm::make(rtype, std::min(fa->value, fb->value));
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::Max>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      const Type& rtype = a.type();
+      if (pa && pb) return IntImm::make(rtype, std::max(pa->value, pb->value));
+      if (fa && fb) return FloatImm::make(rtype, std::max(fa->value, fb->value));
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::GT>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      if (pa && pb) return UIntImm::make(UInt(1), pa->value > pb->value);
+      if (fa && fb) return UIntImm::make(UInt(1), fa->value > fb->value);
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::GE>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      if (pa && pb) return UIntImm::make(UInt(1), pa->value >= pb->value);
+      if (fa && fb) return UIntImm::make(UInt(1), fa->value >= fb->value);
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::LT>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      if (pa && pb) return UIntImm::make(UInt(1), pa->value < pb->value);
+      if (fa && fb) return UIntImm::make(UInt(1), fa->value < fb->value);
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::LE>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      if (pa && pb) return UIntImm::make(UInt(1), pa->value <= pb->value);
+      if (fa && fb) return UIntImm::make(UInt(1), fa->value <= fb->value);
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::EQ>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      if (pa && pb) return UIntImm::make(UInt(1), pa->value == pb->value);
+      if (fa && fb) return UIntImm::make(UInt(1), fa->value == fb->value);
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::NE>(Expr a, Expr b) {
+  TVM_ARITH_CONST_PROPAGATION({
+      if (pa && pb) return UIntImm::make(UInt(1), pa->value != pb->value);
+      if (fa && fb) return UIntImm::make(UInt(1), fa->value != fb->value);
+    });
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::And>(Expr a, Expr b) {
+  using ir::UIntImm;
+  const UIntImm* pa = a.as<UIntImm>();
+  const UIntImm* pb = b.as<UIntImm>();
+  if (pa && pa->value) return b;
+  if (pa && !pa->value) return a;
+  if (pb && pb->value) return a;
+  if (pb && !pb->value) return b;
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::Or>(Expr a, Expr b) {
+  using ir::UIntImm;
+  const UIntImm* pa = a.as<UIntImm>();
+  const UIntImm* pb = b.as<UIntImm>();
+  if (pa && pa->value) return a;
+  if (pa && !pa->value) return b;
+  if (pb && pb->value) return b;
+  if (pb && !pb->value) return a;
+  return Expr();
+}
+
+template<>
+inline Expr TryConstFold<ir::Not>(Expr a) {
+  using ir::UIntImm;
+  const UIntImm* pa = a.as<UIntImm>();
+  if (pa) {
+    return UIntImm::make(UInt(1), !(pa->value));
+  }
+  return Expr();
+}
+
+}  // namespace arith
+}  // namespace tvm
+#endif  // TVM_ARITHMETIC_CONST_FOLD_H_

src/arithmetic/modular_set.cc

@@ -4,7 +4,7 @@
  * \brief Modular set analysis
  */
 #include <tvm/arithmetic.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 #include <tvm/ir_functor_ext.h>
 #include <limits>
 #include "pattern_match.h"

src/lang/expr.cc

@@ -5,7 +5,7 @@
 #include <tvm/base.h>
 #include <tvm/expr.h>
 #include <tvm/ir.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 #include <ir/IRPrinter.h>
 #include <memory>
 

src/lang/ir_operator.cc → src/lang/expr_operator.cc

 /*!
  *  Copyright (c) 2017 by Contributors
- * \file ir_operator.cc
+ * \file expr_operator.cc
  */
 #include <tvm/base.h>
 #include <tvm/ir.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 #include <cmath>
+// Centralized header for constant folders.
+#include "../arithmetic/const_fold.h"
 
 namespace tvm {
 
-/*!
- * \brief Check whether type is used to represent index.
- *
- * Index types are frequently used in shape computation
- * and need to be aggressively constant-folded.
- *
- * \param type The type to represent index.
- * \return the checked result.
- */
-inline bool IsIndexType(const Type& type) {
-  return type.is_int() && type.lanes() == 1 &&
-      (type.bits() == 32 || type.bits() == 64);
-}
-
 // simple cast that only checks if type matches and cast
 inline Expr SimpleCast(const Type& t, Expr value) {
   if (value.type() == t) return value;
@@ -135,45 +123,14 @@ Expr reinterpret(const Type& t, Expr value) {
   return ir::Call::make(t, ir::Call::reinterpret, { value }, ir::Call::PureIntrinsic);
 }
 
-#define TVM_INDEX_CONST_PROPAGATION(BODY)                               \
-  using ir::IntImm;                                                     \
-  using ir::UIntImm;                                                    \
-  const IntImm* pa = a.as<IntImm>();                                    \
-  const IntImm* pb = b.as<IntImm>();                                    \
-  const Type& ta = a.type();                                            \
-  const Type& tb = b.type();                                            \
-  if (IsIndexType(ta) && IsIndexType(tb)) {                             \
-    BODY;                                                               \
-  }                                                                     \
-  BinaryOpMatchTypes(a, b);
-
-#define TVM_ARITH_CONST_PROPAGATION(BODY)                               \
-  using ir::IntImm;                                                     \
-  using ir::UIntImm;                                                    \
-  using ir::FloatImm;                                                   \
-  BinaryOpMatchTypes(a, b);                                             \
-  const IntImm* pa = a.as<IntImm>();                                    \
-  const IntImm* pb = b.as<IntImm>();                                    \
-  const FloatImm* fa = a.as<FloatImm>();                                \
-  const FloatImm* fb = b.as<FloatImm>();                                \
-  BODY;
-
-
 Expr operator+(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      const Type& ta = a.type();
-      const Type& tb = b.type();
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
-      if (pa && pb) return IntImm::make(rtype, pa->value + pb->value);
-      if (pa && pa->value == 0) return SimpleCast(rtype, b);
-      if (pb && pb->value == 0) return SimpleCast(rtype, a);
-      if (fa && fb) return FloatImm::make(rtype, fa->value + fb->value);
-      if (fa && fa->value == 0) return SimpleCast(rtype, b);
-      if (fb && fb->value == 0) return SimpleCast(rtype, a);
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::Add>(a, b);
+  if (ret.defined()) return ret;
   return ir::Add::make(a, b);
 }
 
+// negation
 Expr operator-(Expr a) {
   using ir::IntImm;
   using ir::FloatImm;
@@ -185,114 +142,44 @@ Expr operator-(Expr a) {
 }
 
 Expr operator-(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      const Type& ta = a.type();
-      const Type& tb = b.type();
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
-      if (pa && pb) return IntImm::make(rtype, pa->value - pb->value);
-      if (pb && pb->value == 0) return SimpleCast(rtype, a);
-      if (fa && fb) return FloatImm::make(rtype, fa->value - fb->value);
-      if (fb && fb->value == 0) return SimpleCast(rtype, a);
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::Sub>(a, b);
+  if (ret.defined()) return ret;
   return ir::Sub::make(a, b);
 }
 
 Expr operator*(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      const Type& ta = a.type();
-      const Type& tb = b.type();
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
-      if (pa && pb) return IntImm::make(rtype, pa->value * pb->value);
-      if (pa) {
-        if (pa->value == 1) return SimpleCast(rtype, b);
-        if (pa->value == 0) return SimpleCast(rtype, a);
-      }
-      if (pb) {
-        if (pb->value == 1) return SimpleCast(rtype, a);
-        if (pb->value == 0) return SimpleCast(rtype, b);
-      }
-      if (fa && fb) return FloatImm::make(rtype, fa->value * fb->value);
-      if (fa) {
-        if (fa->value == 1) return SimpleCast(rtype, b);
-        if (fa->value == 0) return SimpleCast(rtype, a);
-      }
-      if (fb) {
-        if (fb->value == 1) return SimpleCast(rtype, a);
-        if (fb->value == 0) return SimpleCast(rtype, b);
-      }
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::Mul>(a, b);
+  if (ret.defined()) return ret;
   return ir::Mul::make(a, b);
 }
 
 Expr operator/(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      const Type& ta = a.type();
-      const Type& tb = b.type();
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
-      // due to division and mod can have different modes
-      // only constant fold positive number where rule is fixed.
-      if (pa && pb && pa->value >= 0 && pb->value > 0) {
-        return IntImm::make(rtype, pa->value / pb->value);
-      }
-      if (pa) {
-        if (pa->value == 0) return SimpleCast(rtype, a);
-      }
-      if (pb) {
-        if (pb->value == 1) return SimpleCast(rtype, a);
-        CHECK_NE(pb->value, 0) << "Divide by zero";
-      }
-      if (fa && fb && fb->value != 0) {
-        return FloatImm::make(rtype, fa->value / fb->value);
-      }
-      if (fa && fa->value == 0) {
-        return SimpleCast(rtype, a);
-      }
-      if (fb) {
-        if (fb->value == 1) return SimpleCast(rtype, a);
-        CHECK_NE(fb->value, 0) << "Divide by zero";
-      }
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::Div>(a, b);
+  if (ret.defined()) return ret;
   return ir::Div::make(a, b);
 }
 
 Expr operator%(Expr a, Expr b) {
-  TVM_INDEX_CONST_PROPAGATION({
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
-      // due to division and mod can have different modes
-      // only constant fold positive number where rule is fixed.
-      if (pa && pb && pa->value >= 0 && pb->value > 0) {
-        return IntImm::make(rtype, pa->value % pb->value);
-      }
-      if (pa) {
-        if (pa->value == 0) return SimpleCast(rtype, a);
-      }
-      if (pb) {
-        if (pb->value == 1) return make_zero(rtype);
-        CHECK_NE(pb->value, 0) << "Divide by zero";
-      }
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::Mod>(a, b);
+  if (ret.defined()) return ret;
   return ir::Mod::make(a, b);
 }
 
 Expr min(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      const Type& ta = a.type();
-      const Type& tb = b.type();
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
-      if (pa && pb) return IntImm::make(rtype, std::min(pa->value, pb->value));
-      if (fa && fb) return FloatImm::make(rtype, std::min(fa->value, fb->value));
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::Min>(a, b);
+  if (ret.defined()) return ret;
   return ir::Min::make(a, b);
 }
 
 Expr max(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      const Type& ta = a.type();
-      const Type& tb = b.type();
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
-      if (pa && pb) return IntImm::make(rtype, std::max(pa->value, pb->value));
-      if (fa && fb) return FloatImm::make(rtype, std::max(fa->value, fb->value));
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::Max>(a, b);
+  if (ret.defined()) return ret;
   return ir::Max::make(a, b);
 }
 
@@ -328,129 +215,116 @@ Expr likely(Expr cond) {
 }
 
 Expr operator>(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      if (pa && pb) return UIntImm::make(UInt(1), pa->value > pb->value);
-      if (fa && fb) return UIntImm::make(UInt(1), fa->value > fb->value);
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::GT>(a, b);
+  if (ret.defined()) return ret;
   return ir::GT::make(a, b);
 }
 
 Expr operator>=(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      if (pa && pb) return UIntImm::make(UInt(1), pa->value >= pb->value);
-      if (fa && fb) return UIntImm::make(UInt(1), fa->value >= fb->value);
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::GE>(a, b);
+  if (ret.defined()) return ret;
   return ir::GE::make(a, b);
 }
 
 Expr operator<(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      if (pa && pb) return UIntImm::make(UInt(1), pa->value < pb->value);
-      if (fa && fb) return UIntImm::make(UInt(1), fa->value < fb->value);
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::LT>(a, b);
+  if (ret.defined()) return ret;
   return ir::LT::make(a, b);
 }
 
 Expr operator<=(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      if (pa && pb) return UIntImm::make(UInt(1), pa->value <= pb->value);
-      if (fa && fb) return UIntImm::make(UInt(1), fa->value <= fb->value);
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::LE>(a, b);
+  if (ret.defined()) return ret;
   return ir::LE::make(a, b);
 }
 
 Expr operator==(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      if (pa && pb) return UIntImm::make(UInt(1), pa->value == pb->value);
-      if (fa && fb) return UIntImm::make(UInt(1), fa->value == fb->value);
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::EQ>(a, b);
+  if (ret.defined()) return ret;
   return ir::EQ::make(a, b);
 }
 
 Expr operator!=(Expr a, Expr b) {
-  TVM_ARITH_CONST_PROPAGATION({
-      if (pa && pb) return UIntImm::make(UInt(1), pa->value != pb->value);
-      if (fa && fb) return UIntImm::make(UInt(1), fa->value != fb->value);
-    });
+  BinaryOpMatchTypes(a, b);
+  Expr ret = arith::TryConstFold<ir::NE>(a, b);
+  if (ret.defined()) return ret;
   return ir::NE::make(a, b);
 }
 
 Expr operator&&(Expr a, Expr b) {
-  using ir::UIntImm;
-  if (a.type().is_bool() && b.type().is_bool()) {
-    const UIntImm* pa = a.as<UIntImm>();
-    const UIntImm* pb = b.as<UIntImm>();
-    if (pa && pa->value) return b;
-    if (pa && !pa->value) return a;
-    if (pb && pb->value) return a;
-    if (pb && !pb->value) return b;
-  }
+  CHECK(a.type().is_bool());
+  CHECK(b.type().is_bool());
+  Expr ret = arith::TryConstFold<ir::And>(a, b);
+  if (ret.defined()) return ret;
   return ir::And::make(a, b);
 }
 
 Expr operator||(Expr a, Expr b) {
-  using ir::UIntImm;
-  if (a.type().is_bool() && b.type().is_bool()) {
-    const UIntImm* pa = a.as<UIntImm>();
-    const UIntImm* pb = b.as<UIntImm>();
-    if (pa && pa->value) return a;
-    if (pa && !pa->value) return b;
-    if (pb && pb->value) return b;
-    if (pb && !pb->value) return a;
-  }
+  CHECK(a.type().is_bool());
+  CHECK(b.type().is_bool());
+  Expr ret = arith::TryConstFold<ir::Or>(a, b);
+  if (ret.defined()) return ret;
   return ir::Or::make(a, b);
 }
 
 Expr operator!(Expr a) {
-  using ir::UIntImm;
-  const UIntImm* pa = a.as<UIntImm>();
-  if (pa) {
-    return UIntImm::make(UInt(1), !(pa->value));
-  }
+  CHECK(a.type().is_bool());
+  Expr ret = arith::TryConstFold<ir::Not>(a);
+  if (ret.defined()) return ret;
   return ir::Not::make(a);
 }
 
 Expr operator>>(Expr a, Expr b) {
+  BinaryOpMatchTypes(a, b);
   TVM_INDEX_CONST_PROPAGATION({
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
+      const Type& rtype = a.type();
       if (pa && pb) return IntImm::make(rtype, (pa->value >> pb->value));
       if (pb) {
-        if (pb->value == 0) return SimpleCast(rtype, a);
+        if (pb->value == 0) return a;
       }
     });
   return ir::Call::make(a.type(), ir::Call::shift_right, { a, b }, ir::Call::PureIntrinsic);
 }
 
 Expr operator<<(Expr a, Expr b) {
+  BinaryOpMatchTypes(a, b);
   TVM_INDEX_CONST_PROPAGATION({
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
+      const Type& rtype = a.type();
       if (pa && pb) return IntImm::make(rtype, (pa->value << pb->value));
       if (pb) {
-        if (pb->value == 0) return SimpleCast(rtype, a);
+        if (pb->value == 0) return a;
       }
     });
   return ir::Call::make(a.type(), ir::Call::shift_left, { a, b }, ir::Call::PureIntrinsic);
 }
 
 Expr operator&(Expr a, Expr b) {
+  BinaryOpMatchTypes(a, b);
   TVM_INDEX_CONST_PROPAGATION({
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
+      const Type& rtype = a.type();
       if (pa && pb) return IntImm::make(rtype, (pa->value & pb->value));
     });
   return ir::Call::make(a.type(), ir::Call::bitwise_and, { a, b }, ir::Call::PureIntrinsic);
 }
 
 Expr operator|(Expr a, Expr b) {
+  BinaryOpMatchTypes(a, b);
   TVM_INDEX_CONST_PROPAGATION({
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
+      const Type& rtype = a.type();
       if (pa && pb) return IntImm::make(rtype, (pa->value | pb->value));
     });
   return ir::Call::make(a.type(), ir::Call::bitwise_or, { a, b }, ir::Call::PureIntrinsic);
 }
 
 Expr operator^(Expr a, Expr b) {
+  BinaryOpMatchTypes(a, b);
   TVM_INDEX_CONST_PROPAGATION({
-      Type rtype = ta.bits() >= tb.bits() ? ta : tb;
+      const Type& rtype = a.type();
       if (pa && pb) return IntImm::make(rtype, (pa->value ^ pb->value));
     });
   return ir::Call::make(a.type(), ir::Call::bitwise_xor, { a, b }, ir::Call::PureIntrinsic);

src/op/hybrid_op.cc

@@ -7,8 +7,8 @@
 #include <tvm/arithmetic.h>
 #include <tvm/ir.h>
 #include <tvm/ir_mutator.h>
-#include <tvm/ir_operator.h>
 #include <tvm/ir_pass.h>
+#include <tvm/expr_operator.h>
 #include <ir/Expr.h>
 #include <unordered_set>
 #include <string>

src/pass/ir_util.h

@@ -7,7 +7,7 @@
 #define TVM_PASS_IR_UTIL_H_
 
 #include <tvm/ir.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 #include <tvm/runtime/device_api.h>
 #include <vector>
 

src/pass/storage_flatten.cc

@@ -8,7 +8,7 @@
 #include <tvm/expr.h>
 #include <tvm/operation.h>
 #include <tvm/ir_mutator.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 #include <tvm/ir_pass.h>
 #include <tvm/buffer.h>
 #include <tvm/target_info.h>

src/relay/op/nn/pad.cc

@@ -4,7 +4,7 @@
  * \brief Implementation of operator pad
  */
 #include <tvm/data_layout.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 #include <tvm/relay/op.h>
 #include <tvm/relay/attrs/nn.h>
 #include <topi/nn.h>

src/relay/op/tensor/transform.cc

@@ -5,7 +5,7 @@
  */
 #include <tvm/relay/op.h>
 #include <tvm/relay/attrs/transform.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 #include <tvm/ir.h>
 #include <tvm/data_layout.h>
 #include <topi/transform.h>

src/relay/pass/fuse_ops.cc

@@ -6,7 +6,7 @@
  * \brief This is a backend-aware optimization pass.
  *   Fuse necessary ops into a single one.
  */
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 #include <tvm/relay/pass.h>
 #include <tvm/relay/expr_functor.h>
 #include <tvm/relay/op_attr_types.h>

tests/cpp/ir_mutator_test.cc

 #include <dmlc/logging.h>
 #include <gtest/gtest.h>
 #include <tvm/ir_mutator.h>
-#include <tvm/ir_operator.h>
+#include <tvm/expr_operator.h>
 
 namespace {
 using namespace tvm::ir;