Use pass instead of pattern rewriter

Author: yongwww

include/tvm/relax/attrs/manipulate.h

@@ -156,21 +156,6 @@ struct ScatterElementsAttrs : public tvm::AttrsNode<ScatterElementsAttrs> {
         "either \"update\", \"add\", \"mul\", \"mean\", \"min\" or \"max\".");
   }
 };  // struct ScatterElementsAttrs
-
-/*! \brief Attributes used in sort operator */
-struct SortAttrs : public tvm::AttrsNode<SortAttrs> {
-  Optional<Integer> axis;
-  Optional<Bool> is_ascend;
-
-  TVM_DECLARE_ATTRS(SortAttrs, "relax.attrs.SortAttrs") {
-    TVM_ATTR_FIELD(axis).describe(
-        "Axis along which the sort is computed."
-        "The default (None) is to compute the sort over the flattened array.");
-    TVM_ATTR_FIELD(is_ascend).describe(
-        "Whether to sort in ascending or descending order."
-        "If it is not specified, it defaults to the ascending order.");
-  }
-};  // struct SortAttrs
 }  // namespace relax
 }  // namespace tvm
 

include/tvm/relax/attrs/sort.h

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * 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
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/relax/attrs/sort.h
+ * \brief Attributes for sorting operators.
+ */
+#ifndef TVM_RELAX_ATTRS_SORT_H_
+#define TVM_RELAX_ATTRS_SORT_H_
+
+#include <tvm/relax/expr.h>
+#include <tvm/tir/index_map.h>
+
+namespace tvm {
+namespace relax {
+
+/*! \brief Attributes used in sort operator */
+struct SortAttrs : public tvm::AttrsNode<SortAttrs> {
+  Optional<Integer> axis;
+  Optional<Bool> descending;
+
+  TVM_DECLARE_ATTRS(SortAttrs, "relax.attrs.SortAttrs") {
+    TVM_ATTR_FIELD(axis).describe(
+        "Axis along which the sort is computed."
+        "The default (None) is to compute the sort over the flattened array.");
+    TVM_ATTR_FIELD(descending).describe(
+        "Whether to sort in descending order."
+        "If it is not specified, it defaults to the ascending order.");
+  }
+};  // struct SortAttrs
+}  // namespace relax
+}  // namespace tvm
+
+#endif  // TVM_RELAX_ATTRS_SORT_H_

include/tvm/relax/dataflow_pattern.h

@@ -56,7 +56,6 @@ class NotPattern;
 class ShapePattern;
 class TypePattern;
 class DataTypePattern;
-class TargetPattern;
 class AttrPattern;
 class SameShapeConstraint;
 
@@ -121,8 +120,6 @@ class DFPattern : public ObjectRef {
   TVM_DLL DataTypePattern HasDtype(const std::string& dtype) const;
   /*! \brief Syntatic Sugar for creating a ShapePattern */
   TVM_DLL ShapePattern HasShape(const Array<PrimExpr>& shape) const;
-    /*! \brief Syntatic Sugar for creating a TargetPattern with a target */
-  TVM_DLL TargetPattern HasTarget(const Target& target) const;
   /*! \brief Syntatic Sugar for creating a ShapePattern */
   TVM_DLL SameShapeConstraint HasSameShapeAs(const DFPattern& other) const;
   /*! \brief Syntatic Sugar for duplicating the current pattern */
@@ -846,34 +843,6 @@ class DataTypePattern : public DFPattern {
   TVM_DEFINE_OBJECT_REF_METHODS(DataTypePattern, DFPattern, DataTypePatternNode);
 };
 
-/*!
- * \brief A pattern that asserting a root pattern has a certain target.
- * \sa TargetPattern
- */
-class TargetPatternNode : public DFPatternNode {
- public:
-  DFPattern pattern; /*!< The root pattern to match */
-  Target target;    /*!< The target to match */
-
-  void VisitAttrs(tvm::AttrVisitor* v) {
-    v->Visit("pattern", &pattern);
-    v->Visit("target", &target);
-  }
-
-  static constexpr const char* _type_key = "relax.dpl.TargetPattern";
-  TVM_DECLARE_FINAL_OBJECT_INFO(TargetPatternNode, DFPatternNode);
-};
-
-/*!
- * \brief Managed reference to TargetPatternNode.
- * \sa TargetPatternNode
- */
-class TargetPattern : public DFPattern {
- public:
-  TVM_DLL TargetPattern(DFPattern pattern, Target target);
-  TVM_DEFINE_OBJECT_REF_METHODS(TargetPattern, DFPattern, TargetPatternNode);
-};
-
 /*!
  * \brief A pattern that asserting a root pattern has certain attributes.
  * \sa AttrPattern

include/tvm/relax/dataflow_pattern_functor.h

@@ -91,7 +91,6 @@ class DFPatternFunctor<R(const DFPattern& n, Args...)> {
   virtual R VisitDFPattern_(const ExprPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
   virtual R VisitDFPattern_(const FunctionPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
   virtual R VisitDFPattern_(const ShapePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
-  virtual R VisitDFPattern_(const TargetPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
   virtual R VisitDFPattern_(const TupleGetItemPatternNode* op,
                             Args... args) DFPATTERN_FUNCTOR_DEFAULT;
   virtual R VisitDFPattern_(const TuplePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
@@ -128,7 +127,6 @@ class DFPatternFunctor<R(const DFPattern& n, Args...)> {
     RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExprPatternNode);
     RELAX_DFPATTERN_FUNCTOR_DISPATCH(FunctionPatternNode);
     RELAX_DFPATTERN_FUNCTOR_DISPATCH(ShapePatternNode);
-    RELAX_DFPATTERN_FUNCTOR_DISPATCH(TargetPatternNode);
     RELAX_DFPATTERN_FUNCTOR_DISPATCH(TupleGetItemPatternNode);
     RELAX_DFPATTERN_FUNCTOR_DISPATCH(TuplePatternNode);
     RELAX_DFPATTERN_FUNCTOR_DISPATCH(TypePatternNode);
@@ -163,7 +161,6 @@ class DFPatternVisitor : public DFPatternFunctor<void(const DFPattern&)> {
   void VisitDFPattern_(const ExprPatternNode* op) override;
   void VisitDFPattern_(const FunctionPatternNode* op) override;
   void VisitDFPattern_(const ShapePatternNode* op) override;
-  void VisitDFPattern_(const TargetPatternNode* op) override;
   void VisitDFPattern_(const TupleGetItemPatternNode* op) override;
   void VisitDFPattern_(const TuplePatternNode* op) override;
   void VisitDFPattern_(const TypePatternNode* op) override;

python/tvm/relax/backend/__init__.py

@@ -18,4 +18,4 @@
 
 from . import contrib
 from .pattern_registry import get_pattern, get_patterns_with_prefix
-from .dispatch_ops import DispatchOps
+from .dispatch_sort_scan import DispatchSortScan