Migrate IntImm & FloatImm ObjectRef to not-null

include/tvm/ir/expr.h

@@ -266,7 +266,7 @@ class IntImm : public PrimExpr {
    */
   TVM_DLL IntImm(DataType dtype, int64_t value);
 
-  TVM_DEFINE_OBJECT_REF_METHODS(IntImm, PrimExpr, IntImmNode);
+  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(IntImm, PrimExpr, IntImmNode);
 };
 
 /*!
@@ -310,7 +310,7 @@ class FloatImm : public PrimExpr {
    */
   TVM_DLL FloatImm(DataType dtype, double value);
 
-  TVM_DEFINE_OBJECT_REF_METHODS(FloatImm, PrimExpr, FloatImmNode);
+  TVM_DEFINE_NOTNULLABLE_OBJECT_REF_METHODS(FloatImm, PrimExpr, FloatImmNode);
 };
 
 /*!
@@ -350,7 +350,6 @@ inline Bool operator&&(const Bool& a, const Bool& b) {
  */
 class Integer : public IntImm {
  public:
-  Integer() {}
   /*!
    * \brief constructor from node.
    */

include/tvm/relay/attrs/transform.h

@@ -94,20 +94,20 @@ struct ReshapeAttrs : public tvm::AttrsNode<ReshapeAttrs> {
 };  // struct ReshapeAttrs
 
 struct ScatterAttrs : public tvm::AttrsNode<ScatterAttrs> {
-  Integer axis;
+  Integer axis = Integer(0);
 
   TVM_DECLARE_ATTRS(ScatterAttrs, "relay.attrs.ScatterAttrs") {
     TVM_ATTR_FIELD(axis).set_default(0).describe("The axis over which to select values.");
   }
 };
 
 struct TakeAttrs : public tvm::AttrsNode<TakeAttrs> {
-  Integer axis;
+  Optional<Integer> axis;
   std::string mode;
 
   TVM_DECLARE_ATTRS(TakeAttrs, "relay.attrs.TakeAttrs") {
     TVM_ATTR_FIELD(axis)
-        .set_default(NullValue<Integer>())
+        .set_default(NullValue<Optional<Integer>>())
         .describe("The axis over which to select values.");
     TVM_ATTR_FIELD(mode).set_default("clip").describe(
         "Specify how out-of-bound indices will behave."
@@ -145,7 +145,7 @@ struct ArangeAttrs : public tvm::AttrsNode<ArangeAttrs> {
 
 /*! \brief Attributes used in stack operators */
 struct StackAttrs : public tvm::AttrsNode<StackAttrs> {
-  Integer axis;
+  Integer axis = Integer(0);
   TVM_DECLARE_ATTRS(StackAttrs, "relay.attrs.StackAttrs") {
     TVM_ATTR_FIELD(axis).set_default(0).describe(
         "The axis in the result array along which the input arrays are stacked.");
@@ -154,12 +154,12 @@ struct StackAttrs : public tvm::AttrsNode<StackAttrs> {
 
 /*! \brief Attributes used in repeat operators */
 struct RepeatAttrs : public tvm::AttrsNode<RepeatAttrs> {
-  Integer repeats;
-  Integer axis;
+  Integer repeats = Integer(0);
+  Optional<Integer> axis;
   TVM_DECLARE_ATTRS(RepeatAttrs, "relay.attrs.RepeatAttrs") {
     TVM_ATTR_FIELD(repeats).describe("The number of repetitions for each element.");
     TVM_ATTR_FIELD(axis)
-        .set_default(NullValue<Integer>())
+        .set_default(NullValue<Optional<Integer>>())
         .describe(" The axis along which to repeat values.");
   }
 };  // struct RepeatAttrs
@@ -176,10 +176,10 @@ struct TileAttrs : public tvm::AttrsNode<TileAttrs> {
 
 /*! \brief Attributes used in reverse operators */
 struct ReverseAttrs : public tvm::AttrsNode<ReverseAttrs> {
-  Integer axis;
+  Optional<Integer> axis;
   TVM_DECLARE_ATTRS(ReverseAttrs, "relay.attrs.ReverseAttrs") {
     TVM_ATTR_FIELD(axis)
-        .set_default(NullValue<Integer>())
+        .set_default(NullValue<Optional<Integer>>())
         .describe("The axis along which to reverse elements.");
   }
 };  // struct ReverseAttrs

include/tvm/relay/attrs/vision.h

@@ -172,7 +172,7 @@ struct ROIPoolAttrs : public tvm::AttrsNode<ROIPoolAttrs> {
 
 /*! \brief Attributes used in yolo reorg operators */
 struct YoloReorgAttrs : public tvm::AttrsNode<YoloReorgAttrs> {
-  Integer stride;
+  Integer stride = Integer(1);
 
   TVM_DECLARE_ATTRS(YoloReorgAttrs, "relay.attrs.YoloReorgAttrs") {
     TVM_ATTR_FIELD(stride).set_default(1).describe("Stride value for yolo reorg");

include/tvm/runtime/container.h

@@ -1452,6 +1452,7 @@ template <typename T>
 class Optional : public ObjectRef {
  public:
   using ContainerType = typename T::ContainerType;
+  using RefType = T;
   static_assert(std::is_base_of<ObjectRef, T>::value, "Optional is only defined for ObjectRef.");
   // default constructors.
   Optional() = default;
@@ -1474,6 +1475,7 @@ class Optional : public ObjectRef {
     data_ = nullptr;
     return *this;
   }
+
   // normal value handling.
   Optional(T other)  // NOLINT(*)
       : ObjectRef(std::move(other)) {}
@@ -1546,6 +1548,12 @@ class Optional : public ObjectRef {
     if (*this == nullptr) return RetType(true);
     return value() != other;
   }
+
+  const ContainerType* operator->() const {
+    CHECK(data_ != nullptr);
+    return static_cast<const ContainerType*>(data_.get());
+  }
+
   static constexpr bool _type_is_nullable = true;
 };
 
@@ -1561,6 +1569,7 @@ struct PackedFuncValueConverter<Optional<T>> {
   }
 };
 
+
 }  // namespace runtime
 
 // expose the functions to the root namespace.

src/arith/canonical_simplify.cc

@@ -706,7 +706,7 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const DivNode* op) {
   // const folding
   PrimExpr const_res = TryConstFold<Div>(a, b);
   if (const_res.defined()) return const_res;
-  PVar<IntImm> c1;
+  PVarOpt<Optional<IntImm>> c1;
   // x / c1
   if (c1.Match(b) && c1.Eval()->value > 0) {
     int64_t cval = c1.Eval()->value;
@@ -764,7 +764,7 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const FloorDivNode* op) {
   // const folding
   PrimExpr const_res = TryConstFold<FloorDiv>(a, b);
   if (const_res.defined()) return const_res;
-  PVar<IntImm> c1;
+  PVarOpt<Optional<IntImm>> c1;
   // x / c1
   if (c1.Match(b) && c1.Eval()->value > 0) {
     int64_t cval = c1.Eval()->value;
@@ -868,7 +868,7 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const ModNode* op) {
   PrimExpr const_res = TryConstFold<Mod>(a, b);
   if (const_res.defined()) return const_res;
 
-  PVar<IntImm> c1;
+  PVarOpt<Optional<IntImm>> c1;
   // x % c1
   if (c1.Match(b) && c1.Eval()->value > 0) {
     int64_t cval = c1.Eval()->value;
@@ -936,7 +936,7 @@ PrimExpr CanonicalSimplifier::Impl::VisitExpr_(const FloorModNode* op) {
   PrimExpr const_res = TryConstFold<FloorMod>(a, b);
   if (const_res.defined()) return const_res;
 
-  PVar<IntImm> c1;
+  PVarOpt<Optional<IntImm>> c1;
   // x % c1
   if (c1.Match(b) && c1.Eval()->value > 0) {
     int64_t cval = c1.Eval()->value;

src/arith/const_int_bound.cc

@@ -520,7 +520,7 @@ class ConstIntBoundAnalyzer::Impl
    */
   static std::vector<BoundInfo> DetectBoundInfo(const PrimExpr& cond) {
     PVar<PrimExpr> x, y;
-    PVar<IntImm> c;
+    PVarOpt<Optional<IntImm>> c;
     // NOTE: canonical form always use <= or <
     if ((c <= x).Match(cond)) {
       return {BoundInfo(x.Eval(), MakeBound(c.Eval()->value, kPosInf))};

src/arith/int_set.cc

@@ -414,7 +414,7 @@ class IntervalSetEvaluator : public ExprFunctor<IntervalSet(const PrimExpr&)> {
   IntervalSet VisitExpr_(const RampNode* op) final {
     CHECK(eval_vec_);
     IntervalSet base = Eval(op->base);
-    PVar<IntImm> stride;
+    PVarOpt<Optional<IntImm>> stride;
     if (stride.Match(op->stride)) {
       DataType t = op->base.dtype();
       int64_t vstride = stride.Eval()->value;

src/arith/modular_set.cc

@@ -104,7 +104,7 @@ class ModularSetAnalyzer::Impl : public ExprFunctor<ModularSetAnalyzer::Entry(co
   // Detect useful constraints and use them in the analysis scope.
   std::function<void()> EnterConstraint(const PrimExpr& constraint) {
     PVar<Var> var;
-    PVar<IntImm> coeff, base;
+    PVarOpt<Optional<IntImm>> coeff, base;
     // pattern match interesting constraints
     if ((truncmod(var, coeff) == base).Match(constraint) ||
         (floormod(var, coeff) == base).Match(constraint)) {

src/arith/pattern_match.h

@@ -200,6 +200,57 @@ class PVar : public Pattern<PVar<T>> {
   mutable bool filled_{false};
 };
 
+/*!
+ * \brief Pattern variable container.
+ *
+ * PVarOpt is a variant of PVar to incorporate optional objectref.
+ *
+ * \tparam T the type of the hole.
+ *
+ * \note PVarOpt is not thread safe.
+ *       Do not use the same PVarOpt in multiple threads.
+ */
+template <typename T>
+class PVarOpt : public Pattern<PVarOpt<T>> {
+ public:
+  // Store PVars by reference in the expression.
+  using Nested = const PVarOpt<T>&;
+  using RefType = typename T::RefType;
+
+  void InitMatch_() const { filled_ = false; }
+
+  bool Match_(const RefType& value) const {
+    if (!filled_) {
+      value_ = value;
+      filled_ = true;
+      return true;
+    } else {
+      return PEqualChecker<RefType>()(value_.value(), value);
+    }
+  }
+
+  template <typename NodeRefType,
+            typename = typename std::enable_if<std::is_base_of<NodeRefType, RefType>::value>::type>
+  bool Match_(const NodeRefType& value) const {
+    if (const auto* ptr = value.template as<typename T::ContainerType>()) {
+      return Match_(GetRef<RefType>(ptr));
+    } else {
+      return false;
+    }
+  }
+
+  RefType Eval() const {
+    CHECK(filled_);
+    return value_.value();
+  }
+
+ protected:
+  /*! \brief The matched value */
+  mutable T value_;
+  /*! \brief whether the variable has been filled */
+  mutable bool filled_{false};
+};
+
 /*!
  * \brief Constant Pattern variable container.
  *

src/arith/rewrite_simplify.cc

@@ -123,7 +123,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const AddNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, b1, b2, s1, s2;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2, c3;
+  PVarOpt<Optional<IntImm>> c1, c2, c3;
   // Pattern var for lanes in broadcast and ramp
   PVar<int> lanes;
   // Vector rules
@@ -230,7 +230,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const SubNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, b1, b2, s1, s2;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2, c3;
+  PVarOpt<Optional<IntImm>> c1, c2, c3;
   // Pattern var for lanes in broadcast and ramp
   PVar<int> lanes;
   // Vector rules
@@ -413,7 +413,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const MulNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, b1, b2, s1, s2;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2;
+  PVarOpt<Optional<IntImm>> c1, c2;
   // Pattern var for lanes in broadcast and ramp
   PVar<int> lanes;
   // Vector rules
@@ -446,7 +446,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const DivNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, b1;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2, c3;
+  PVarOpt<Optional<IntImm>> c1, c2, c3;
   // Pattern var for lanes in broadcast and ramp
   PVar<int> lanes;
 
@@ -621,7 +621,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const ModNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, b1;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2;
+  PVarOpt<Optional<IntImm>> c1, c2;
   // Pattern var for lanes in broadcast and ramp
   PVar<int> lanes;
 
@@ -701,7 +701,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorDivNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, b1;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2, c3;
+  PVarOpt<Optional<IntImm>> c1, c2, c3;
   // Pattern var for lanes in broadcast and ramp
   PVar<int> lanes;
 
@@ -819,7 +819,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const FloorModNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, b1;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2;
+  PVarOpt<Optional<IntImm>> c1, c2;
   // Pattern var for lanes in broadcast and ramp
   PVar<int> lanes;
 
@@ -884,7 +884,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const MinNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, s1, s2;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2;
+  PVarOpt<Optional<IntImm>> c1, c2;
   PVar<int> lanes;
 
   // vector rule
@@ -1056,7 +1056,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const MaxNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, s1, s2;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2;
+  PVarOpt<Optional<IntImm>> c1, c2;
   PVar<int> lanes;
 
   // vector rule
@@ -1219,7 +1219,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const EQNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y;
   // Pattern var match IntImm
-  PVar<IntImm> c1;
+  PVarOpt<Optional<IntImm>> c1;
   PVar<int> lanes;
 
   // vector rule
@@ -1267,7 +1267,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const LTNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y, z, s1, s2;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2;
+  PVarOpt<Optional<IntImm>> c1, c2;
   PVar<int> lanes;
 
   // vector rule
@@ -1422,7 +1422,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const AndNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2;
+  PVarOpt<Optional<IntImm>> c1, c2;
   PVar<int> lanes;
 
   if (op->dtype.lanes() != 1) {
@@ -1461,7 +1461,7 @@ PrimExpr RewriteSimplifier::Impl::VisitExpr_(const OrNode* op) {
   // Pattern var to match any expression
   PVar<PrimExpr> x, y;
   // Pattern var match IntImm
-  PVar<IntImm> c1, c2;
+  PVarOpt<Optional<IntImm>> c1, c2;
   PVar<int> lanes;
 
   if (op->dtype.lanes() != 1) {