[FIX][IR] Recover rewrite pattern hot paths

Author: tqchen

src/arith/pattern_match.h

@@ -199,18 +199,15 @@ class PVar : public Pattern<PVar<T>> {
   // Store PVars by reference in the expression.
   using Nested = const PVar<T>&;
 
-  void InitMatch_() const {
-    value_ = nullptr;
-    filled_ = false;
-  }
+  void InitMatch_() const { filled_ = false; }
 
   bool Match_(const T& value) const {
     if (!filled_) {
       value_ = value;
       filled_ = true;
       return true;
     } else {
-      return PEqualChecker<T>()(value_.value(), value);
+      return PEqualChecker<T>()(value_, value);
     }
   }
 
@@ -226,14 +223,14 @@ class PVar : public Pattern<PVar<T>> {
 
   T Eval() const {
     TVM_FFI_ICHECK(filled_);
-    return value_.value();
+    return value_;
   }
 
-  T EvalOr(const T& default_value) const { return filled_ ? value_.value() : default_value; }
+  T EvalOr(const T& default_value) const { return filled_ ? value_ : default_value; }
 
  protected:
   /*! \brief The matched value */
-  mutable ffi::Optional<T> value_;
+  mutable T value_;
   /*! \brief whether the variable has been filled */
   mutable bool filled_{false};
 };
@@ -285,7 +282,7 @@ class PVarWithDataType : public PVarWithCheck<PVarWithDataType<T, DType>, T> {
  public:
   explicit PVarWithDataType(const DType& dtype) : dtype_(dtype) {}
 
-  bool Match_(const T& value) const { return dtype_.Match_(value.ty()); }
+  bool Match_(const T& value) const { return dtype_.Match_(value.ty()->dtype); }
 
  protected:
   typename DType::Nested dtype_;
@@ -294,15 +291,15 @@ class PVarWithDataType : public PVarWithCheck<PVarWithDataType<T, DType>, T> {
 /*!
  * \brief Pattern variable container for data type with lanes.
  */
-class PVecDataType : public PVarWithCheck<PVecDataType, PrimType> {
+class PVecDataType : public PVarWithCheck<PVecDataType, DLDataType> {
  public:
   /*! \brief construct vector dtype placeholder with element type check */
-  explicit PVecDataType(PrimType elem_dtype) : elem_dtype_(elem_dtype) {}
+  explicit PVecDataType(DLDataType elem_dtype) : elem_dtype_(elem_dtype) {}
 
-  bool Match_(PrimType dtype) const { return dtype.code() == elem_dtype_.code(); }
+  bool Match_(DLDataType dtype) const { return dtype.code == elem_dtype_.code; }
 
  protected:
-  PrimType elem_dtype_;
+  DLDataType elem_dtype_;
 };
 
 /*!
@@ -543,15 +540,15 @@ class PCastExpr : public Pattern<PCastExpr<DType, TA>> {
 
   bool Match_(const ffi::ObjectRef& node) const {
     if (const tirx::CastNode* ptr = node.as<tirx::CastNode>()) {
-      if (!dtype_.Match_(ptr->ty())) return false;
+      if (!dtype_.Match_(ptr->ty()->dtype)) return false;
       if (!value_.Match_(ptr->value)) return false;
       return true;
     } else {
       return false;
     }
   }
 
-  PrimExpr Eval() const { return tirx::Cast(dtype_.Eval(), value_.Eval()); }
+  PrimExpr Eval() const { return tirx::Cast(PrimType(dtype_.Eval()), value_.Eval()); }
 
  private:
   typename DType::Nested dtype_;
@@ -561,7 +558,7 @@ class PCastExpr : public Pattern<PCastExpr<DType, TA>> {
 /*!
  * \brief Construct a cast pattern.
  *
- * \param dtype The target data type, can be PVar<PrimType> or PConst<PrimType>.
+ * \param dtype The target data type, can be PVar<DLDataType> or PConst<DLDataType>.
  * \param value The input type.
  *
  * \return The result pattern.

src/ir/type.cc

@@ -90,11 +90,11 @@ PrimType::PrimType(DLDataTypeCode code, int bits, int lanes)
 PrimType PrimType::Int(int bits, int lanes) {
   if (lanes == 1) {
     if (bits == 32) {
-      static const PrimType i32_ty(DLDataType{kDLInt, 32, 1});
+      thread_local PrimType i32_ty(DLDataType{kDLInt, 32, 1});
       return i32_ty;
     }
     if (bits == 64) {
-      static const PrimType i64_ty(DLDataType{kDLInt, 64, 1});
+      thread_local PrimType i64_ty(DLDataType{kDLInt, 64, 1});
       return i64_ty;
     }
   }
@@ -107,7 +107,7 @@ PrimType PrimType::UInt(int bits, int lanes) {
 
 PrimType PrimType::Float(int bits, int lanes) {
   if (bits == 32 && lanes == 1) {
-    static const PrimType f32_ty(DLDataType{kDLFloat, 32, 1});
+    thread_local PrimType f32_ty(DLDataType{kDLFloat, 32, 1});
     return f32_ty;
   }
   return PrimType(DLDataType{kDLFloat, static_cast<uint8_t>(bits), static_cast<uint16_t>(lanes)});
@@ -119,7 +119,7 @@ PrimType PrimType::BFloat(int bits, int lanes) {
 
 PrimType PrimType::Bool(int lanes) {
   if (lanes == 1) {
-    static const PrimType bool_ty(DLDataType{kDLBool, 8, 1});
+    thread_local PrimType bool_ty(DLDataType{kDLBool, 8, 1});
     return bool_ty;
   }
   return PrimType(DLDataType{kDLBool, 8, static_cast<uint16_t>(lanes)});