You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@tvm.apache.org by tq...@apache.org on 2023/02/19 00:21:07 UTC
[tvm] 41/42: [Unity] Relax dataflow pattern language (matching) (#14041)
This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch unity-staging
in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 413329dfa21fed3aac7fcce8d8c64fbe7d36ac83
Author: Jiawei Liu <ja...@gmail.com>
AuthorDate: Sat Feb 18 15:57:02 2023 -0600
[Unity] Relax dataflow pattern language (matching) (#14041)
The dataflow pattern language for Relax (originally from https://github.com/tlc-pack/relax/pull/163).
The implementation splits patterns into two parts:
- Match an Expression: match an expression syntactically (MatchExprPattern, i.e., DFPatternMatcher);
- Match a Graph: match a graph (cross multiple VarBinding) topologically (MatchGraphPattern);
Credit also given to relay pattern lang authors and reviewers of https://github.com/tlc-pack/relax/pull/163.
---
include/tvm/relax/analysis.h | 32 +
include/tvm/relax/dataflow_matcher.h | 80 ++
include/tvm/relax/dataflow_pattern.h | 810 +++++++++++++++++++
include/tvm/relax/dataflow_pattern_functor.h | 183 +++++
python/tvm/relax/analysis/analysis.py | 38 +-
python/tvm/relax/dpl/__init__.py | 21 +
python/tvm/relax/dpl/_ffi.py | 20 +
python/tvm/relax/dpl/context.py | 86 ++
python/tvm/relax/dpl/pattern.py | 1095 ++++++++++++++++++++++++++
src/relax/analysis/udchain.cc | 102 +++
src/relax/analysis/var2value.cc | 91 +++
src/relax/ir/dataflow_matcher.cc | 768 ++++++++++++++++++
src/relax/ir/dataflow_matcher_impl.h | 87 ++
src/relax/ir/dataflow_pattern.cc | 607 ++++++++++++++
src/relax/ir/dataflow_pattern_functor.cc | 111 +++
tests/python/relax/test_analysis.py | 23 +-
tests/python/relax/test_dataflow_pattern.py | 867 ++++++++++++++++++++
17 files changed, 5018 insertions(+), 3 deletions(-)
diff --git a/include/tvm/relax/analysis.h b/include/tvm/relax/analysis.h
index f9896efdf2..32e1582134 100644
--- a/include/tvm/relax/analysis.h
+++ b/include/tvm/relax/analysis.h
@@ -317,6 +317,38 @@ TVM_DLL tvm::Array<GlobalVar> CalledGlobalVars(const Expr& expr);
*/
TVM_DLL tvm::Array<GlobalVar> AllGlobalVars(const Expr& expr);
+/*!
+ * \brief Analyze var -> value mapping from VarBindings.
+ *
+ * \param m The IRModule to check.
+ * \return Var -> Value (Expr)
+ */
+TVM_DLL Map<Var, Expr> AnalyzeVar2Value(const IRModule& m);
+
+/*!
+ * \brief Analyze var -> value mapping from VarBindings.
+ *
+ * \param expr The expression to check.
+ * \return Var -> Value (Expr)
+ */
+TVM_DLL Map<Var, Expr> AnalyzeVar2Value(const Expr& expr);
+
+/*!
+ * \brief Analyze var -> value mapping from VarBindings.
+ *
+ * \param dfb The dataflow block to check.
+ * \return Var -> Value (Expr)
+ */
+TVM_DLL Map<Var, Expr> AnalyzeVar2Value(const DataflowBlock& dfb);
+
+/*!
+ * \brief Get the use-def chain of variables inside a dataflow block.
+ *
+ * \param dfb The dataflow block to be analyzed.
+ * \return A map mapping variable definitions to a set of uses.
+ */
+TVM_DLL Map<Var, Array<Var>> DataflowBlockUseDef(const DataflowBlock& dfb);
+
/*!
* \brief Annotate Op Pattern Kind for PrimFunc, which is used in relax FuseOps.
*
diff --git a/include/tvm/relax/dataflow_matcher.h b/include/tvm/relax/dataflow_matcher.h
new file mode 100644
index 0000000000..fa58308fac
--- /dev/null
+++ b/include/tvm/relax/dataflow_matcher.h
@@ -0,0 +1,80 @@
+/*
+ * 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/dataflow_matcher.h
+ * \brief A pattern matcher for matching dataflow properties.
+ */
+#ifndef TVM_RELAX_DATAFLOW_MATCHER_H_
+#define TVM_RELAX_DATAFLOW_MATCHER_H_
+
+#include <tvm/relax/dataflow_pattern.h>
+#include <tvm/runtime/container/optional.h>
+
+#include <memory>
+
+namespace tvm {
+namespace relax {
+
+/**
+ * \brief Determine if a pattern matches an expression.
+ * \note The behavior of MatchExpr is to match a relax.Expr (`expr`) syntactically through
+ * one given pattern (`pattern`).
+ *
+ * \param pattern The pattern to match
+ * \param expr The expression to match
+ * \param bindings The mapping from relax.Var to relax.Expr
+ * \return true if matched
+ * \return false if unmatched
+ */
+bool MatchExpr(DFPattern pattern, Expr expr, Optional<runtime::Map<Var, Expr>> bindings = NullOpt);
+
+/* \brief Similar to above, but return pairs of a matching pattern and an expression. */
+Optional<Map<DFPattern, Expr>> ExtractMatchedExpr(
+ DFPattern pattern, Expr expr, Optional<runtime::Map<Var, Expr>> bindings = NullOpt);
+
+/**
+ * \brief Match a sub-graph in a DataflowBlock with a graph of patterns and return the mapping.
+ * \note This algorithm returns the first matched sub-graph. Use `start_hint` to specify the
+ * starting point of the matching so that we can distinguish multiple matches.
+ *
+ * \param ctx The graph-wise patterns.
+ * \param dfb The function to match.
+ * \param start_hint The starting point expression to match to distinguish multiple matches.
+ * \param must_include_hint If start_hint is given, the return pattern must include start_hint.
+ * \return tvm::runtime::Map<DFPattern, Var>
+ */
+TVM_DLL tvm::runtime::Map<DFPattern, Var> MatchGraph(const PatternContext& ctx,
+ const DataflowBlock& dfb,
+ Optional<Var> start_hint = NullOpt,
+ bool must_include_hint = false);
+
+/**
+ * \brief Match a graph-wise pattern with the current context (PatternContext::Current()).
+ */
+inline tvm::runtime::Map<DFPattern, Var> MatchGraphDefault(const DataflowBlock& dfb,
+ Optional<Var> start_hint = NullOpt,
+ bool must_include_hint = false) {
+ return MatchGraph(PatternContext::Current(), dfb, start_hint, must_include_hint);
+}
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_DATAFLOW_MATCHER_H_
diff --git a/include/tvm/relax/dataflow_pattern.h b/include/tvm/relax/dataflow_pattern.h
new file mode 100644
index 0000000000..701879745e
--- /dev/null
+++ b/include/tvm/relax/dataflow_pattern.h
@@ -0,0 +1,810 @@
+/*
+ * 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/dataflow_pattern.h
+ * \brief A pattern language for matching dataflow properties.
+ */
+#ifndef TVM_RELAX_DATAFLOW_PATTERN_H_
+#define TVM_RELAX_DATAFLOW_PATTERN_H_
+
+#include <tvm/ir/expr.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/type.h>
+#include <tvm/runtime/container/array.h>
+#include <tvm/runtime/container/optional.h>
+#include <tvm/support/with.h>
+
+#include <cstdint>
+#include <map>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace relax {
+
+class PatternSeq;
+class CallPattern;
+class OrPattern;
+class AndPattern;
+class NotPattern;
+class ShapePattern;
+class TypePattern;
+class DataTypePattern;
+class AttrPattern;
+
+/*!
+ * \brief Create used-by relationship between lhs[-1] and rhs[0], with [*lhs, *rhs] returned.
+ *
+ * \param lhs Left hand side of the used-by relationship.
+ * \param rhs Right hand side of the used-by relationship.
+ * \param index lhs[-1] is used as the index'th argument of rhs[0].
+ * \return PatternSeq The concatenated sequence of [*lhs, *rhs].
+ */
+TVM_DLL PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index = -1);
+/*! \brief Syntax sugar of UsedBy(lhs, rhs, -1). */
+TVM_DLL PatternSeq operator^(const PatternSeq& lhs, const PatternSeq& rhs);
+
+/*!
+ * \brief Create only-used-by relationship between lhs[-1] and rhs[0], with [*lhs, *rhs] returned.
+ *
+ * \param lhs Left hand side of the used-by relationship.
+ * \param rhs Right hand side of the used-by relationship.
+ * \param index lhs[-1] is used as the index'th argument of rhs[0].
+ * \return PatternSeq The concatenated sequence of [*lhs, *rhs].
+ */
+TVM_DLL PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index = -1);
+/*! \brief Syntax sugar of OnlyUsedBy(lhs, rhs, -1). */
+TVM_DLL PatternSeq operator>>(const PatternSeq& lhs, const PatternSeq& rhs);
+
+/*!
+ * \brief Base type of all dataflow patterns.
+ * \sa DFPattern
+ */
+class DFPatternNode : public Object {
+ public:
+ static constexpr const char* _type_key = "DFPatternNode";
+ TVM_DECLARE_BASE_OBJECT_INFO(DFPatternNode, Object);
+};
+
+/*!
+ * \brief Managed reference to dataflow patterns.
+ * \sa DFPatternNode
+ */
+class DFPattern : public ObjectRef {
+ public:
+ /*! \brief Syntatic Sugar for creating a CallPattern */
+ template <typename... Args>
+ CallPattern operator()(Args&&... args) const;
+ /*! \brief Syntatic Sugar for creating a CallPattern */
+ TVM_DLL CallPattern operator()(const std::vector<DFPattern>& args) const;
+ /*! \brief Syntatic Sugar for creating an OrPattern */
+ TVM_DLL OrPattern operator|(const DFPattern& other) const;
+ /*! \brief Syntatic Sugar for creating an AndPattern */
+ TVM_DLL AndPattern operator&(const DFPattern& other) const;
+ /*! \brief Syntatic Sugar for creating a NotPattern */
+ TVM_DLL NotPattern operator~() const;
+ /*! \brief Syntatic Sugar for creating an AttrPattern */
+ TVM_DLL AttrPattern HasAttr(const Map<String, ObjectRef>& attrs) const;
+ /*! \brief Syntatic Sugar for creating a TypePattern */
+ TVM_DLL TypePattern HasType(const Type& type) const;
+ /*! \brief Syntatic Sugar for creating a DataTypePattern with a DataType */
+ TVM_DLL DataTypePattern HasDtype(const DataType& dtype) const;
+ /*! \brief Syntatic Sugar for creating a DataTypePattern with a data type's name */
+ 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 duplicating the current pattern */
+ TVM_DLL DFPattern dup() const;
+
+ /*! \brief Implicit conversion from DFPattern to PatternSeq */
+ TVM_DLL operator PatternSeq() const;
+
+ TVM_DEFINE_OBJECT_REF_METHODS(DFPattern, ObjectRef, DFPatternNode);
+};
+
+/*! \brief Constraint of a DFPattern edge (producer -> consumer) in graph-level matching */
+struct PairCons {
+ /*! \brief Constraint types of the edge */
+ enum Type {
+ kUsedBy, /*!< producer ^ consumer */
+ kOnlyUsedBy, /*!< producer >> consumer */
+ } type = kUsedBy;
+ int index = -1; /*!< The argument index of the producer in the consumer caller site */
+
+ /*!
+ * \brief Construct a new PairCons object
+ *
+ * \param t The constraint type
+ * \param index The producer is called as the index'th argument of the consumer function.
+ */
+ TVM_DLL explicit PairCons(Type t, int index = -1) : type(t), index(index) {}
+
+ bool operator==(const PairCons& other) const {
+ return type == other.type && index == other.index;
+ }
+};
+
+/*!
+ * \brief A sequence of DFPatterns that the previous DFPattern is connected to the next one.
+ * \sa PatternSeq
+ */
+class PatternSeqNode final : public Object {
+ public:
+ tvm::Array<DFPattern> patterns; /*!< The sequence of DFPatterns */
+ std::vector<PairCons> pair_constraints; /*!< Constraints between the previous and next patterns */
+
+ void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("patterns", &patterns); }
+ static constexpr const char* _type_key = "relax.dpl.PatternSeq";
+ TVM_DECLARE_BASE_OBJECT_INFO(PatternSeqNode, Object);
+};
+
+/*!
+ * \brief Managed reference to pattern sequences.
+ * \sa PatternSeqNode
+ */
+class PatternSeq final : public ObjectRef {
+ public:
+ TVM_DLL explicit PatternSeq(DFPattern init_pattern);
+ TVM_DLL explicit PatternSeq(tvm::Array<DFPattern> patterns, bool only_used_by = false);
+
+ PatternSeq UsedBy(PatternSeq other, int index = -1) const;
+ PatternSeq OnlyUsedBy(PatternSeq other, int index = -1) const;
+
+ /*! \brief Syntatic Sugar for duplicating the current pattern sequence */
+ PatternSeq dup() const;
+
+ // friend functions
+ friend PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index);
+ friend PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index);
+
+ TVM_DEFINE_OBJECT_REF_METHODS(PatternSeq, ObjectRef, PatternSeqNode);
+};
+
+/*!
+ * \brief A context to manage the graph-level pattern matching.
+ * \sa PatternContext
+ */
+class PatternContextNode : public Object {
+ public:
+ /*! \brief Constrainting matched graph with assertion to external uses */
+ enum ExternUse {
+ kMay, /*!< No constraints */
+ kMustNot, /*!< All nodes except outputs only have internal depedencies in the matched graph. */
+ } allow_extern_use = kMay;
+ // src node -> <dst node, constraint type> constraints.
+ std::map<DFPattern, std::map<DFPattern, std::vector<PairCons>>> constraints;
+
+ static constexpr const char* _type_key = "relax.dpl.PatternContext";
+ TVM_DECLARE_FINAL_OBJECT_INFO(PatternContextNode, Object);
+};
+
+/*!
+ * \brief Managed reference to a pattern context.
+ * \sa PatternContextNode
+ */
+class PatternContext : public ObjectRef {
+ public:
+ TVM_DLL explicit PatternContext(ObjectPtr<Object> n) : ObjectRef(n) {}
+ TVM_DLL explicit PatternContext(bool incremental = false);
+
+ const PatternContextNode* operator->() const {
+ ICHECK(get() != nullptr);
+ return static_cast<const PatternContextNode*>(get());
+ }
+
+ PatternContextNode* operator->() {
+ ICHECK(get() != nullptr);
+ return static_cast<PatternContextNode*>(get_mutable());
+ }
+
+ /*!
+ * \brief Build an edge constraint between two patterns (producer and consumer).
+ *
+ * \param producer The pattern corresponding to the producer node.
+ * \param consumer The pattern corresponding to the consumer node.
+ * \param cons The constraint type. \sa PairCons
+ */
+ void add_constraint(DFPattern producer, DFPattern consumer, PairCons cons) {
+ auto& vec = (*this)->constraints[producer][consumer];
+ ICHECK(std::find(vec.cbegin(), vec.cend(), cons) == vec.cend()) << "Constraint already exists";
+ vec.push_back(cons);
+ }
+
+ /*! \brief Get the pass context object on the top of the stack */
+ TVM_DLL static PatternContext Current();
+
+ class Internal;
+
+ private:
+ /*! \brief The RAII-like entry of a pass context scope */
+ TVM_DLL void EnterWithScope();
+ /*! \brief The RAII-like exit of a pass context scope */
+ TVM_DLL void ExitWithScope();
+ friend class Internal;
+ friend class With<PatternContext>;
+};
+
+/*!
+ * \brief Pattern for Relax Expression.
+ * \sa ExprPattern
+ */
+class ExprPatternNode : public DFPatternNode {
+ public:
+ Expr expr; /*!< The expression to match */
+
+ void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("expr", &expr); }
+
+ static constexpr const char* _type_key = "relax.dpl.ExprPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(ExprPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to an ExprPattern.
+ * \sa ExprPatternNode
+ */
+class ExprPattern : public DFPattern {
+ public:
+ TVM_DLL explicit ExprPattern(Expr expr);
+ TVM_DEFINE_OBJECT_REF_METHODS(ExprPattern, DFPattern, ExprPatternNode);
+};
+
+/*!
+ * \brief A Pattern to Match a Relax Variable.
+ * \note The name field matches any string if it is empty.
+ * \sa VarPattern
+ */
+class VarPatternNode : public DFPatternNode {
+ public:
+ String name;
+ const String& name_hint() const { return name; }
+ void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("name", &name); }
+
+ static constexpr const char* _type_key = "relax.dpl.VarPattern";
+ TVM_DECLARE_BASE_OBJECT_INFO(VarPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to a VarPattern.
+ * \sa VarPatternNode
+ */
+class VarPattern : public DFPattern {
+ public:
+ /*!
+ * \brief Create a pattern matching by variable name.
+ *
+ * \param name_hint Variable name to match. Any if empty ("").
+ */
+ TVM_DLL VarPattern(String name_hint);
+ TVM_DEFINE_OBJECT_REF_METHODS(VarPattern, DFPattern, VarPatternNode);
+};
+
+/*!
+ * \brief A Pattern to Match a Relax Dataflow Variable
+ * \sa DataflowVarPattern
+ */
+class DataflowVarPatternNode : public VarPatternNode {
+ public:
+ static constexpr const char* _type_key = "relax.dpl.DataflowVarPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(DataflowVarPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to a DataflowVarPattern.
+ * \sa DataflowVarPatternNode
+ */
+class DataflowVarPattern : public DFPattern {
+ public:
+ /*! \sa VarPattern::VarPattern */
+ TVM_DLL DataflowVarPattern(String name_hint);
+ TVM_DEFINE_OBJECT_REF_METHODS(DataflowVarPattern, DFPattern, DataflowVarPatternNode);
+};
+
+/*!
+ * \brief A Pattern to Match a Relax Global Variable
+ * \sa GlobalVarPattern
+ */
+class GlobalVarPatternNode : public VarPatternNode {
+ public:
+ static constexpr const char* _type_key = "relax.dpl.GlobalVarPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(GlobalVarPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to a GlobalVarPattern.
+ * \sa GlobalVarPatternNode
+ */
+class GlobalVarPattern : public DFPattern {
+ public:
+ TVM_DLL GlobalVarPattern(String name_hint);
+ TVM_DEFINE_OBJECT_REF_METHODS(GlobalVarPattern, DFPattern, GlobalVarPatternNode);
+};
+
+/*!
+ * \brief A Pattern to Match a Relax Constant.
+ * \sa ConstantPattern
+ */
+class ConstantPatternNode : public DFPatternNode {
+ public:
+ void VisitAttrs(tvm::AttrVisitor* v) {}
+
+ static constexpr const char* _type_key = "relax.dpl.ConstantPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(ConstantPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to a ConstantPattern.
+ * \sa ConstantPatternNode
+ */
+class ConstantPattern : public DFPattern {
+ public:
+ TVM_DEFINE_OBJECT_REF_METHODS(ConstantPattern, DFPattern, ConstantPatternNode);
+};
+
+/*!
+ * \brief A pattern to match a callable node in Relax.
+ * \sa CallPattern
+ */
+class CallPatternNode : public DFPatternNode {
+ public:
+ /*!
+ * \note The op field can be:
+ * - relay::Op which corresponds to the primitive operators.
+ * - user defined functions (Function, GlobalVar, Var).
+ */
+ DFPattern op; /*!< The operator (function) being invoked */
+ tvm::Array<DFPattern> args; /*!< The arguments of the function call */
+ /*!
+ * \note If varg_default_wildcard is true. Given args of [pA, pB], when matching a call whose
+ * arguments are [A, B, ...], the pattern will still match despite N(args) < N(call.args). That
+ * said, with varg_default_wildcard set to true, we match the args in the order we have, and
+ * regard the rest of the arguments as wildcards.
+ */
+ bool varg_default_wildcard; /*!< N(args) can be < N(real args) by the padding of Wildcard */
+
+ // Todo(relax-team): Dataflow pattern for StructInfo, and match sinfo_args
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("op", &op);
+ v->Visit("args", &args);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.CallPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(CallPatternNode, DFPatternNode);
+};
+
+class CallPattern : public DFPattern {
+ public:
+ TVM_DLL CallPattern(DFPattern op, Array<DFPattern> args, bool varg_default_wildcard = false);
+ TVM_DEFINE_OBJECT_REF_METHODS(CallPattern, DFPattern, CallPatternNode);
+};
+
+/*!
+ * \brief A pattern to match an array of PrimExpr.
+ * \sa PrimArrPattern
+ * \note This is often used to match shapes specified as arguments to a function.
+ */
+class PrimArrPatternNode : public DFPatternNode {
+ public:
+ Array<PrimExpr> fields; /*!< The array to match */
+ void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("fields", &fields); }
+ static constexpr const char* _type_key = "relax.dpl.PrimArrPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(PrimArrPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to a PrimArrPattern.
+ * \sa PrimArrPatternNode
+ */
+class PrimArrPattern : public DFPattern {
+ public:
+ TVM_DLL PrimArrPattern(Array<PrimExpr> arr);
+ TVM_DEFINE_OBJECT_REF_METHODS(PrimArrPattern, DFPattern, PrimArrPatternNode);
+};
+
+/*!
+ * \brief A pattern to match a Relax Function
+ * \sa Function
+ * \sa FunctionPattern
+ */
+class FunctionPatternNode : public DFPatternNode {
+ public:
+ tvm::Array<DFPattern> params; /*!< The parameters of the function */
+ /*!
+ * \note Note that in Relax, the function body is a SeqExpr which contains
+ * 1) SeqExprNode::blocks, which is a list of blocks of statements; and 2)
+ * SeqExprNode::body, which is an Expr that can be anything. FunctionPattern
+ * only matches the body of the function (writing patterns to statements is tricky).
+ */
+ DFPattern body; /*!< The body of the function */
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("params", ¶ms);
+ v->Visit("body", &body);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.FunctionPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(FunctionPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to FunctionPatternNode.
+ * \sa FunctionPatternNode
+ */
+class FunctionPattern : public DFPattern {
+ public:
+ /*!
+ * \brief Constructor
+ * \param params The parameters of the function.
+ * \param body The body of the function.
+ */
+ TVM_DLL FunctionPattern(tvm::Array<DFPattern> params, DFPattern body);
+
+ TVM_DEFINE_OBJECT_REF_METHODS(FunctionPattern, DFPattern, FunctionPatternNode);
+};
+
+/*!
+ * \brief Pattern to match a tuple of ordered expressions.
+ * \sa TuplePattern
+ */
+class TuplePatternNode : public DFPatternNode {
+ public:
+ tvm::Array<DFPattern> fields; /*!< The fields of the tuple */
+
+ void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("fields", &fields); }
+
+ static constexpr const char* _type_key = "relax.dpl.TuplePattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(TuplePatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to TuplePatternNode.
+ * \sa TuplePatternNode
+ */
+class TuplePattern : public DFPattern {
+ public:
+ TVM_DLL explicit TuplePattern(tvm::Array<DFPattern> fields);
+ TVM_DEFINE_OBJECT_REF_METHODS(TuplePattern, DFPattern, TuplePatternNode);
+};
+
+/*!
+ * \brief A pattern to match multiple expressions unorderedly.
+ * \sa UnorderedTuplePattern
+ */
+class UnorderedTuplePatternNode : public DFPatternNode {
+ public:
+ tvm::Array<DFPattern> fields; /*!< The fields of the tuple */
+
+ void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("fields", &fields); }
+
+ static constexpr const char* _type_key = "relax.dpl.UnorderedTuplePattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(UnorderedTuplePatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to UnorderedTuplePatternNode.
+ * \sa UnorderedTuplePatternNode
+ */
+class UnorderedTuplePattern : public DFPattern {
+ public:
+ TVM_DLL explicit UnorderedTuplePattern(tvm::Array<DFPattern> fields);
+ TVM_DEFINE_OBJECT_REF_METHODS(UnorderedTuplePattern, DFPattern, UnorderedTuplePatternNode);
+};
+
+/*!
+ * \brief A pattern to match n'th indexing to a tuple.
+ * \sa TupleGetItem
+ * \sa TupleGetItemPattern
+ */
+class TupleGetItemPatternNode : public DFPatternNode {
+ public:
+ DFPattern tuple; /*!< The tuple Expression */
+ int index; /*!< The index of the tuple with -1 meaning arbitrary */
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("tuple", &tuple);
+ v->Visit("index", &index);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.TupleGetItemPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(TupleGetItemPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to TupleGetItemPatternNode.
+ * \sa TupleGetItemPatternNode
+ */
+class TupleGetItemPattern : public DFPattern {
+ public:
+ TVM_DLL TupleGetItemPattern(DFPattern tuple, int index);
+ TVM_DEFINE_OBJECT_REF_METHODS(TupleGetItemPattern, DFPattern, TupleGetItemPatternNode);
+};
+
+/*!
+ * \brief Match a conjunction of other patterns.
+ * \sa AndPattern
+ */
+class AndPatternNode : public DFPatternNode {
+ public:
+ DFPattern left; /*!< The left hand side of the conjunction */
+ DFPattern right; /*!< The right hand side of the conjunction */
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("left", &left);
+ v->Visit("right", &right);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.AndPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(AndPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to AndPatternNode.
+ * \sa AndPatternNode
+ */
+class AndPattern : public DFPattern {
+ public:
+ TVM_DLL AndPattern(DFPattern lhs, DFPattern rhs);
+ TVM_DEFINE_OBJECT_REF_METHODS(AndPattern, DFPattern, AndPatternNode);
+};
+
+/*!
+ * \brief Match a disjunction of other patterns.
+ * \sa OrPattern
+ */
+class OrPatternNode : public DFPatternNode {
+ public:
+ DFPattern left; /*!< The left hand side of the disjunction */
+ DFPattern right; /*!< The right hand side of the disjunction */
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("left", &left);
+ v->Visit("right", &right);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.OrPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(OrPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to OrPatternNode.
+ * \sa OrPatternNode
+ */
+class OrPattern : public DFPattern {
+ public:
+ TVM_DLL OrPattern(DFPattern left, DFPattern right);
+ TVM_DEFINE_OBJECT_REF_METHODS(OrPattern, DFPattern, OrPatternNode);
+};
+
+/*!
+ * \brief Pattern for rejecting a certain pattern.
+ * \sa NotPattern
+ */
+class NotPatternNode : public DFPatternNode {
+ public:
+ DFPattern reject; /*!< The pattern to reject */
+
+ void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("reject", &reject); }
+
+ static constexpr const char* _type_key = "relax.dpl.NotPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(NotPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to NotPatternNode.
+ * \sa NotPatternNode
+ */
+class NotPattern : public DFPattern {
+ public:
+ TVM_DLL NotPattern(DFPattern reject);
+ TVM_DEFINE_OBJECT_REF_METHODS(NotPattern, DFPattern, NotPatternNode);
+};
+
+/*!
+ * \brief Wildcard Pattern is a pattern that can match anything.
+ * \sa WildcardPattern
+ */
+class WildcardPatternNode : public DFPatternNode {
+ public:
+ void VisitAttrs(tvm::AttrVisitor* v) {}
+
+ static constexpr const char* _type_key = "relax.dpl.WildcardPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(WildcardPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to WildcardPatternNode.
+ * \sa WildcardPatternNode
+ */
+class WildcardPattern : public DFPattern {
+ public:
+ TVM_DEFINE_OBJECT_REF_METHODS(WildcardPattern, DFPattern, WildcardPatternNode);
+};
+
+/*!
+ * \brief Pattern for matching a certain type.
+ * \sa TypePattern
+ */
+class TypePatternNode : public DFPatternNode {
+ public:
+ DFPattern pattern; /*!< The pattern to match */
+ Type type; /*!< The type to match */
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("pattern", &pattern);
+ v->Visit("type", &type);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.TypePattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(TypePatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to TypePatternNode.
+ * \sa TypePatternNode
+ */
+class TypePattern : public DFPattern {
+ public:
+ TVM_DLL TypePattern(DFPattern pattern, Type type);
+ TVM_DEFINE_OBJECT_REF_METHODS(TypePattern, DFPattern, TypePatternNode);
+};
+
+/*!
+ * \brief A pattern that asserting a root pattern has a certain shape.
+ * \sa ShapePattern
+ */
+class ShapePatternNode : public DFPatternNode {
+ public:
+ DFPattern pattern; /*!< The root pattern to match */
+ Array<PrimExpr> shape; /*!< The shape to match */
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("pattern", &pattern);
+ v->Visit("shape", &shape);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.ShapePattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(ShapePatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to ShapePatternNode.
+ * \sa ShapePatternNode
+ */
+class ShapePattern : public DFPattern {
+ public:
+ TVM_DLL ShapePattern(DFPattern pattern, Array<PrimExpr> type);
+ TVM_DEFINE_OBJECT_REF_METHODS(ShapePattern, DFPattern, ShapePatternNode);
+};
+
+/*!
+ * \brief A pattern that asserting a root pattern has a certain data type.
+ * \sa DataTypePattern
+ */
+class DataTypePatternNode : public DFPatternNode {
+ public:
+ DFPattern pattern; /*!< The root pattern to match */
+ DataType dtype; /*!< The data type to match */
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("pattern", &pattern);
+ v->Visit("dtype", &dtype);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.DataTypePattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(DataTypePatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to DataTypePatternNode.
+ * \sa DataTypePatternNode
+ */
+class DataTypePattern : public DFPattern {
+ public:
+ TVM_DLL DataTypePattern(DFPattern pattern, DataType dtype);
+ TVM_DEFINE_OBJECT_REF_METHODS(DataTypePattern, DFPattern, DataTypePatternNode);
+};
+
+/*!
+ * \brief A pattern that asserting a root pattern has certain attributes.
+ * \sa AttrPattern
+ */
+class AttrPatternNode : public DFPatternNode {
+ public:
+ DFPattern pattern; /*!< The root pattern to match */
+ DictAttrs attrs; /*!< The attributes (a map/dictionary) to match */
+
+ void VisitAttrs(tvm::AttrVisitor* v) {
+ v->Visit("pattern", &pattern);
+ v->Visit("attrs", &attrs);
+ }
+
+ static constexpr const char* _type_key = "relax.dpl.AttrPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(AttrPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to AttrPatternNode.
+ * \sa AttrPatternNode
+ */
+class AttrPattern : public DFPattern {
+ public:
+ TVM_DLL AttrPattern(DFPattern pattern, DictAttrs attrs);
+ TVM_DEFINE_OBJECT_REF_METHODS(AttrPattern, DFPattern, AttrPatternNode);
+};
+
+/*!
+ * \brief A pattern of external function.
+ * \sa ExternFunc
+ * \sa ExternFuncPattern
+ */
+class ExternFuncPatternNode : public DFPatternNode {
+ public:
+ String global_symbol_; /*!< The global symbol name of the external function */
+
+ /*! \brief The the external function name */
+ const String& global_symbol() const { return global_symbol_; }
+ void VisitAttrs(tvm::AttrVisitor* v) { v->Visit("global_symbol", &global_symbol_); }
+
+ static constexpr const char* _type_key = "relax.dpl.ExternFuncPattern";
+ TVM_DECLARE_FINAL_OBJECT_INFO(ExternFuncPatternNode, DFPatternNode);
+};
+
+/*!
+ * \brief Managed reference to ExternFuncPatternNode.
+ * \sa ExternFuncPatternNode
+ */
+class ExternFuncPattern : public DFPattern {
+ public:
+ TVM_DLL ExternFuncPattern(String global_symbol);
+ TVM_DEFINE_OBJECT_REF_METHODS(ExternFuncPattern, DFPattern, ExternFuncPatternNode);
+};
+
+/*! \brief Syntatic Sugar for creating a VarPattern with a name */
+VarPattern IsVar(const String& name);
+/*! \brief Syntatic Sugar for creating a ConstantPattern */
+ConstantPattern IsConst();
+/*! \brief Syntatic Sugar for creating a WildcardPattern */
+WildcardPattern Wildcard();
+/*! \brief Syntatic Sugar for creating a ExprPattern */
+ExprPattern IsExpr(const Expr& expr);
+/*! \brief Syntatic Sugar for creating a ExprPattern base on an Op */
+ExprPattern IsOp(const String& op_name);
+/*! \brief Syntatic Sugar for call_tir (return a tensor) */
+// Todo(relax-team): Dataflow pattern for StructInfo, and match out_sinfo
+CallPattern IsCallTIR(const String& name, Optional<TuplePattern> args = NullOpt);
+/*! \brief Syntatic Sugar for call_tir (return a tuple of tensor) */
+CallPattern IsCallTIR(const String& name, TuplePattern var_args);
+/*! \brief Syntatic Sugar for creating TuplePattern or UnorderedTuplePattern (unordered=true) */
+DFPattern IsTuple(const Array<DFPattern>& fields, bool unordered = false);
+/*! \brief Syntatic Sugar for creating a TupleGetItemPattern */
+TupleGetItemPattern IsTupleGetItem(const DFPattern tuple, int index = -1);
+
+/*! \brief Implementation of the templated CallPattern syntax sugar */
+template <typename... Args>
+CallPattern DFPattern::operator()(Args&&... args) const {
+ return CallPattern(GetRef<DFPattern>(this->get()),
+ Array<DFPattern>({std::forward<Args>(args)...}));
+}
+
+} // namespace relax
+} // namespace tvm
+#endif // TVM_RELAX_DATAFLOW_PATTERN_H_
diff --git a/include/tvm/relax/dataflow_pattern_functor.h b/include/tvm/relax/dataflow_pattern_functor.h
new file mode 100644
index 0000000000..983881ddc9
--- /dev/null
+++ b/include/tvm/relax/dataflow_pattern_functor.h
@@ -0,0 +1,183 @@
+/*
+ * 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/dataflow_pattern_functor.h
+ * \brief Functors and visitors for dataflow patterns.
+ */
+#ifndef TVM_RELAX_DATAFLOW_PATTERN_FUNCTOR_H_
+#define TVM_RELAX_DATAFLOW_PATTERN_FUNCTOR_H_
+
+#include <tvm/relax/dataflow_pattern.h>
+
+#include <unordered_set>
+#include <utility>
+
+namespace tvm {
+namespace relax {
+
+/*!
+ * \brief A dynamical functor that dispatches on in the first DFPattern argument.
+ *
+ * \tparam FType function signature
+ * This type is only defined for FType with function signature R(const DFPattern&,
+ * Args...)
+ */
+template <typename FType>
+class DFPatternFunctor;
+
+// functions to be overriden.
+#define DFPATTERN_FUNCTOR_DEFAULT \
+ { return VisitDFPatternDefault_(op, std::forward<Args>(args)...); }
+
+#define RELAX_DFPATTERN_FUNCTOR_DISPATCH(OP) \
+ vtable.template set_dispatch<OP>([](const ObjectRef& n, TSelf* self, Args... args) { \
+ return self->VisitDFPattern_(static_cast<const OP*>(n.get()), std::forward<Args>(args)...); \
+ });
+
+template <typename R, typename... Args>
+class DFPatternFunctor<R(const DFPattern& n, Args...)> {
+ private:
+ using TSelf = DFPatternFunctor<R(const DFPattern& n, Args...)>;
+ using FType = tvm::NodeFunctor<R(const ObjectRef& n, TSelf* self, Args...)>;
+
+ public:
+ /*! \brief virtual destructor */
+ virtual ~DFPatternFunctor() {}
+ /*!
+ * \brief Same as call.
+ * \param n The expression node.
+ * \param args Additional arguments.
+ * \return The result of the call
+ */
+ R operator()(const DFPattern& n, Args... args) {
+ return VisitDFPattern(n, std::forward<Args>(args)...);
+ }
+ /*!
+ * \brief The functor call.
+ * \param n The expression node.
+ * \param args Additional arguments.
+ * \return The result of the call
+ */
+ virtual R VisitDFPattern(const DFPattern& n, Args... args) {
+ ICHECK(n.defined());
+ static FType vtable = InitVTable();
+ return vtable(n, this, std::forward<Args>(args)...);
+ }
+ // Functions that can be overriden by subclass
+ virtual R VisitDFPattern_(const OrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const AndPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const NotPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const AttrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const CallPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const ConstantPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const DataTypePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ 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 TupleGetItemPatternNode* op,
+ Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const TuplePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const TypePatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const WildcardPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const VarPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+
+ virtual R VisitDFPattern_(const DataflowVarPatternNode* op,
+ Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const GlobalVarPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const ExternFuncPatternNode* op,
+ Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const PrimArrPatternNode* op, Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+ virtual R VisitDFPattern_(const UnorderedTuplePatternNode* op,
+ Args... args) DFPATTERN_FUNCTOR_DEFAULT;
+
+ virtual R VisitDFPatternDefault_(const Object* op, Args...) {
+ LOG(FATAL) << "Do not have a default for " << op->GetTypeKey();
+ throw;
+ }
+
+ private:
+ // initialize the vtable.
+ static FType InitVTable() {
+ FType vtable;
+ // Set dispatch
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(OrPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(AndPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(NotPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(AttrPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(CallPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(ConstantPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataTypePatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExprPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(FunctionPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(ShapePatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(TupleGetItemPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(TuplePatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(TypePatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(WildcardPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(VarPatternNode);
+
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(DataflowVarPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(GlobalVarPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(ExternFuncPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(PrimArrPatternNode);
+ RELAX_DFPATTERN_FUNCTOR_DISPATCH(UnorderedTuplePatternNode);
+ return vtable;
+ }
+};
+
+/*!
+ * \brief A simple visitor wrapper around DFPatternFunctor.
+ * Recursively visit the content.
+ *
+ * DFPatternVisitor treats the Pattern as dataflow graph,and only visit each Expr node once.
+ */
+class DFPatternVisitor : public DFPatternFunctor<void(const DFPattern&)> {
+ public:
+ void VisitDFPattern(const DFPattern& pattern) override;
+ void VisitDFPattern_(const OrPatternNode* op) override;
+ void VisitDFPattern_(const AndPatternNode* op) override;
+ void VisitDFPattern_(const NotPatternNode* op) override;
+ void VisitDFPattern_(const AttrPatternNode* op) override;
+ void VisitDFPattern_(const CallPatternNode* op) override;
+ void VisitDFPattern_(const ConstantPatternNode* op) override;
+ void VisitDFPattern_(const DataTypePatternNode* op) override;
+ void VisitDFPattern_(const ExprPatternNode* op) override;
+ void VisitDFPattern_(const FunctionPatternNode* op) override;
+ void VisitDFPattern_(const ShapePatternNode* op) override;
+ void VisitDFPattern_(const TupleGetItemPatternNode* op) override;
+ void VisitDFPattern_(const TuplePatternNode* op) override;
+ void VisitDFPattern_(const TypePatternNode* op) override;
+ void VisitDFPattern_(const WildcardPatternNode* op) override;
+ void VisitDFPattern_(const VarPatternNode* op) override;
+
+ void VisitDFPattern_(const DataflowVarPatternNode* op) override;
+ void VisitDFPattern_(const GlobalVarPatternNode* op) override;
+ void VisitDFPattern_(const ExternFuncPatternNode* op) override;
+ void VisitDFPattern_(const PrimArrPatternNode* op) override;
+ void VisitDFPattern_(const UnorderedTuplePatternNode* op) override;
+
+ protected:
+ // set of already-visited nodes
+ std::unordered_set<const Object*> visited_;
+};
+
+} // namespace relax
+} // namespace tvm
+#endif // TVM_RELAX_DATAFLOW_PATTERN_FUNCTOR_H_
diff --git a/python/tvm/relax/analysis/analysis.py b/python/tvm/relax/analysis/analysis.py
index 7107883478..45c5b6f962 100644
--- a/python/tvm/relax/analysis/analysis.py
+++ b/python/tvm/relax/analysis/analysis.py
@@ -21,14 +21,14 @@ This file contains the set of passes for Relax, which exposes an interface for
configuring the passes and scripting them in Python.
"""
-from typing import Dict
+from typing import Dict, List
from enum import IntEnum
from tvm import tir
from tvm import IRModule
from tvm.relax.ty import Type
from tvm.relax.struct_info import StructInfo, FuncStructInfo
-from tvm.relax.expr import Var, Expr, Call
+from tvm.relax.expr import DataflowBlock, Var, Expr, Function, Call
from . import _ffi_api
@@ -210,6 +210,40 @@ def has_reshape_pattern(func: tir.PrimFunc) -> bool:
return _ffi_api.has_reshape_pattern(func) # type: ignore
+def get_var2val(func: Function) -> Dict[Var, Expr]:
+ """
+ Get a mapping from Var to Expr for each variable in the function.
+
+ Parameters
+ ----------
+ func : Function
+ The input function to be analyzed.
+
+ Returns
+ -------
+ Dict[Var, Expr]
+ A mapping from Var to Expr.
+ """
+ return _ffi_api.get_var2val(func) # type: ignore
+
+
+def udchain(dfb: DataflowBlock) -> Dict[Var, List[Var]]:
+ """
+ Analyze the variable use-def chain in a dataflow block.
+
+ Parameters
+ ----------
+ dfb : DataflowBlock
+ The dataflow block to analyze
+
+ Returns
+ -------
+ Dict[Var, List[Var]]
+ A mapping from variable definition to its uses.
+ """
+ return _ffi_api.udchain(dfb) # type: ignore
+
+
def well_formed(mod: IRModule, check_struct_info: bool = True) -> bool:
"""Check if the IRModule is well formed.
diff --git a/python/tvm/relax/dpl/__init__.py b/python/tvm/relax/dpl/__init__.py
new file mode 100644
index 0000000000..e0bbdaff05
--- /dev/null
+++ b/python/tvm/relax/dpl/__init__.py
@@ -0,0 +1,21 @@
+# 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.
+
+"""The Relax Dataflow Pattern Language."""
+
+from .pattern import *
+from .context import *
diff --git a/python/tvm/relax/dpl/_ffi.py b/python/tvm/relax/dpl/_ffi.py
new file mode 100644
index 0000000000..6699e42bee
--- /dev/null
+++ b/python/tvm/relax/dpl/_ffi.py
@@ -0,0 +1,20 @@
+# 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.
+"""DataFlow Pattern Language FFI bindings."""
+import tvm._ffi
+
+tvm._ffi._init_api("relax.dpl", __name__)
diff --git a/python/tvm/relax/dpl/context.py b/python/tvm/relax/dpl/context.py
new file mode 100644
index 0000000000..69a5e70ed0
--- /dev/null
+++ b/python/tvm/relax/dpl/context.py
@@ -0,0 +1,86 @@
+# 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.
+
+"""The Graph Matching Context Manager for Dataflow Pattern Language."""
+
+from typing import Optional, Dict
+
+import tvm
+from ..expr import DataflowBlock, Var
+from .pattern import DFPattern
+from . import _ffi as ffi
+
+
+class PatternContext(tvm.runtime.Object):
+ """A context object for doing graph (topogical) pattern matching."""
+
+ def __init__(self, incremental=False):
+ """
+ Initialize the PatternContext
+
+ Parameters
+ ----------
+ incremental : bool, optional
+ perform incremental matching based on the recent context, by default False
+ """
+ self.__init_handle_by_constructor__(ffi.PatternContext, incremental) # type: ignore
+
+ def __enter__(self):
+ """Enter the context"""
+ ffi.enter_context(self) # type: ignore
+ return self
+
+ def __exit__(self, exc_type, exc_value, traceback):
+ """Exit the context"""
+ ffi.exit_context(self) # type: ignore
+
+ @staticmethod
+ def current() -> "PatternContext":
+ """
+ Get the current context
+
+ Returns
+ -------
+ PatternContext
+ The current context
+ """
+ return ffi.current_context() # type: ignore
+
+ def match_dfb(
+ self,
+ dfb: DataflowBlock,
+ start_hint: Optional[Var] = None,
+ must_include_hint: bool = False,
+ ) -> Dict[DFPattern, Var]:
+ """
+ Match a DataflowBlock via a graph of DFPattern and corresponding constraints
+
+ Parameters
+ ----------
+ dfb : DataflowBlock
+ The DataflowBlock to match
+ start_hint : Optional[Var], optional
+ Indicating the starting expression to match, by default None
+ must_include_hint : bool, optional
+ Whether the start_hint expression must be matched, by default False
+
+ Returns
+ -------
+ Dict[DFPattern, Var]
+ The mapping from DFPattern to matched expression
+ """
+ return ffi.match_dfb(self, dfb, start_hint, must_include_hint) # type: ignore
diff --git a/python/tvm/relax/dpl/pattern.py b/python/tvm/relax/dpl/pattern.py
new file mode 100644
index 0000000000..44faa0c93a
--- /dev/null
+++ b/python/tvm/relax/dpl/pattern.py
@@ -0,0 +1,1095 @@
+# 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.
+
+"""Pattern types in Relax Dataflow Pattern Language"""
+# pylint: disable=no-member
+# pylint: disable=pointless-statement
+
+import typing
+from typing import Dict, List, Optional, Tuple, Union
+
+import tvm
+import tvm._ffi as tvm_ffi
+from tvm.ir.container import Array
+from tvm.ir.expr import PrimExpr
+from tvm.relay.op import get
+
+from ...ir import make_node
+from ...ir.base import Node
+from ...runtime import Object
+from ..expr import Expr, Var
+from . import _ffi as ffi
+
+
+def register_df_node(type_key=None):
+ """
+ Register a Relax node type
+
+ Parameters
+ ----------
+ type_key : str or cls
+ The type key of the node
+ """
+ if not isinstance(type_key, str):
+ return tvm_ffi.register_object("relax.dpl." + type_key.__name__)(type_key)
+ return tvm_ffi.register_object(type_key)
+
+
+class DFPattern(Node):
+ """Base class of all Patterns."""
+
+ def __call__(self, *args, varg_default_wildcard=False) -> "CallPattern":
+ """
+ Syntax sugar for creating a CallPattern with argument patterns
+
+ Returns
+ -------
+ result: CallPattern
+ The resulting CallPattern
+ """
+ return CallPattern(self, args, varg_default_wildcard)
+
+ def __or__(self, other: "DFPattern") -> "OrPattern":
+ """
+ Syntax sugar for creating an OrPattern
+
+ Parameters
+ ----------
+ other: DFPattern
+ Alternative pattern
+
+ Returns
+ -------
+ result: OrPattern
+ The resulting OrPattern
+ """
+ return OrPattern(self, other)
+
+ def __and__(self, other: "DFPattern") -> "AndPattern":
+ """
+ Syntax sugar for creating an AndPattern
+
+ Parameters
+ ----------
+ other: DFPattern
+ Additional pattern to satisfy
+
+ Returns
+ -------
+ result: AndPattern
+ The resulting AndPattern
+ """
+ return AndPattern(self, other)
+
+ def __invert__(self) -> "NotPattern":
+ """
+ Syntax sugar for creating a DFPattern to reject
+
+ Returns
+ -------
+ result: NotPattern
+ The resulting NotPattern
+ """
+ return reject(self)
+
+ def has_attr(self, attrs: Dict[str, Object]) -> "AttrPattern":
+ """
+ Add an attribute constraint to this pattern
+
+ Parameters
+ ----------
+ attrs: Dict[str, Object]
+
+ Returns
+ -------
+ result: AttrPattern
+ The resulting AttrPattern
+ """
+ attrs = make_node("DictAttrs", **attrs)
+ return AttrPattern(self, attrs)
+
+ def has_type(self, ttype: tvm.ir.type.Type) -> "TypePattern":
+ """
+ Add a type constraint to this pattern
+
+ Parameters
+ ----------
+ ttype: tvm.ir.type.Type
+ The type to match
+
+ Returns
+ -------
+ result: TypePattern
+ The resulting TypePattern
+ """
+ return TypePattern(self, ttype)
+
+ def has_dtype(self, dtype: str) -> "DataTypePattern":
+ """
+ Add a type constraint to this pattern
+
+ Parameters
+ ----------
+ dtype: str
+ The dtype to match
+
+ Returns
+ -------
+ result: DataTypePattern
+ The resulting DataTypePattern
+ """
+ return has_dtype(dtype, self)
+
+ def has_shape(self, shape: List[PrimExpr]) -> "ShapePattern":
+ """
+ Add a shape constraint to this pattern
+
+ Parameters
+ ----------
+ shape: List[PrimExpr]
+ Expected shape list
+
+ Returns
+ -------
+ result: ShapePattern
+ The resulting ShapePattern
+
+ Note
+ ----
+ has_shape assumes that the matched relax.Expr only has one
+ output tensor. Use is_tuple for those with multiple outputs.
+ """
+ if not isinstance(shape, (list, tuple, tvm.ir.PrimExpr)):
+ raise ValueError("has_shape takes a list or tuple as input.")
+ return ShapePattern(pattern=self, shape=shape)
+
+ def match(self, expr, var2val: Optional[Dict[Var, Expr]] = None) -> bool:
+ """
+ Match a relax.Expr syntactically
+
+ Parameters
+ ----------
+ expr : tvm.relax.Expr
+ The expression to match
+ var2val : Optional[Dict[tvm.relax.Var, tvm.relax.Expr]]
+ A mapping from relax.Var to relax.Expr for autojump.
+
+ Returns
+ -------
+ result: bool
+ Whether or not the expression matches the pattern
+
+ Note
+ ----
+ Unlike Relay whose function is an expression, functions in Relax consist
+ of blocks of bindings that are not syntactically connected. We use a
+ mapping (i.e., var2val) to mitigate the gap. For example, when matching
+ "relax.add(lv0, lv1)", given var2val, we match lv0's bound expression
+ when the recursive pattern matching goes to check lv0. The var2val mapping
+ can be computed through the tvm.relax.analysis.get_var2val function.
+ """
+ return ffi.match_expr(self, expr, var2val) # type: ignore
+
+ def used_by(self, other: Union["DFPattern", "PatternSeq"], index=-1) -> "PatternSeq":
+ """
+ The current pattern being used by another pattern (sequence)
+
+ Parameters
+ ----------
+ other : Union[DFPattern, DFPattern]
+ The consumer pattern (sequence)
+ index : int, optional
+ The argument index called by the consumer pattern, by default -1
+
+ Returns
+ -------
+ result: PatternSeq
+ A chained pattern sequence
+ """
+ return _used_by(self, other, index)
+
+ def __xor__(self, other: Union["DFPattern", "PatternSeq"]) -> "PatternSeq":
+ """Syntax sugar of DFPattern.used_by"""
+ return self.used_by(other, -1)
+
+ def only_used_by(self, other: Union["DFPattern", "PatternSeq"], index=-1) -> "PatternSeq":
+ """
+ The current pattern being **ONLY** used by another pattern (sequence)
+
+ Parameters
+ ----------
+ other : Union[DFPattern, DFPattern]
+ The consumer pattern (sequence)
+ index : int, optional
+ The argument index called by the consumer pattern, by default -1
+
+ Returns
+ -------
+ result: PatternSeq
+ A chained pattern sequence
+ """
+ return _only_used_by(self, other, index)
+
+ def __rshift__(self, other: Union["DFPattern", "PatternSeq"]) -> "PatternSeq":
+ """Syntax sugar of DFPattern.only_used_by"""
+ return self.only_used_by(other, -1)
+
+ def dup(self) -> "DFPattern":
+ """
+ Duplicate the current pattern (new object under different address)
+
+ Returns
+ -------
+ DFPattern
+ A duplicated pattern
+ """
+ return ffi.dup_pattern(self) # type: ignore
+
+ def fork_to(self, *args) -> None:
+ """Fork the current pattern to multiple pattern branches"""
+ for v in args:
+ self ^ v
+
+
+@register_df_node
+class ExprPattern(DFPattern):
+ """A pattern which matches an expression.
+
+ Parameters
+ ----------
+ expr : tvm.relax.Expr
+ The expression to match.
+ """
+
+ def __init__(self, expr: Expr):
+ self.__init_handle_by_constructor__(ffi.ExprPattern, expr) # type: ignore
+
+
+@register_df_node
+class VarPattern(DFPattern):
+ """A pattern for Var.
+
+ Parameters
+ ----------
+ name_hint: str
+ The name of the variable. Optional, if not provided,
+ the pattern will match any VarNode.
+ """
+
+ def __init__(self, name_hint: str = ""):
+ self.__init_handle_by_constructor__(ffi.VarPattern, name_hint) # type: ignore
+
+
+@register_df_node
+class DataflowVarPattern(DFPattern):
+ """A pattern for DataflowVar.
+
+ Parameters
+ ----------
+ name_hint: str
+ The name of the variable. Optional, if not provided,
+ the pattern will match any VarNode.
+ """
+
+ def __init__(self, name_hint: str = ""):
+ self.__init_handle_by_constructor__(ffi.DataflowVarPattern, name_hint) # type: ignore
+
+
+@register_df_node
+class GlobalVarPattern(DFPattern):
+ """A pattern for GlobalVar.
+
+ Parameters
+ ----------
+ name_hint: str
+ The name of the variable. Optional, if not provided,
+ the pattern will match any GlobalVarNode.
+ """
+
+ def __init__(self, name_hint: str = ""):
+ self.__init_handle_by_constructor__(ffi.GlobalVarPattern, name_hint) # type: ignore
+
+
+@register_df_node
+class ExternFuncPattern(DFPattern):
+ """A external function pattern.
+
+ Parameters
+ ----------
+ global_symbol: str
+ The name of the function. Optional, if not provided,
+ the pattern will match any ExternFuncNode.
+ """
+
+ def __init__(self, global_symbol: str = ""):
+ self.__init_handle_by_constructor__(ffi.ExternFuncPattern, global_symbol) # type: ignore
+
+
+@register_df_node
+class ConstantPattern(DFPattern):
+ """A pattern matching a Relax Constant."""
+
+ def __init__(self):
+ self.__init_handle_by_constructor__(ffi.ConstantPattern) # type: ignore
+
+
+@register_df_node
+class CallPattern(DFPattern):
+ """A pattern matching a function call node.
+
+ Parameters
+ ----------
+ op: tvm.relax.dpl.DFPattern
+ The operation to be called.
+
+ args: List[tvm.relax.dpl.DFPattern]
+ The arguments to the call or None to match any arguments.
+
+ varg_default_wildcard: bool
+ If True, args can be fewer than actual provided arguments.
+
+ Note
+ ----
+ By setting varg_default_wildcard to True, we can only focus on the argument
+ patterns we specified. For example, CallPattern(Op, [A, B]) can match
+ a call of Op(A, B) or Op(A, B, C, ...) that has more arguments. However,
+ the specified argument patterns must be matched (i.e., A and B).
+ """
+
+ def __init__(
+ self,
+ op: "DFPattern",
+ args: Union[List["DFPattern"], typing.Tuple["DFPattern", ...]],
+ varg_default_wildcard: bool = False,
+ ):
+ self.__init_handle_by_constructor__(
+ ffi.CallPattern, op, args, varg_default_wildcard # type: ignore
+ )
+
+
+@register_df_node
+class FunctionPattern(DFPattern):
+ """A pattern matching a function node in Relax.
+
+ Parameters
+ ----------
+ params: List[tvm.relax.dpl.DFPattern]
+ The parameters to the Function or None to match any parameters.
+
+ body: tvm.relax.dpl.DFPattern
+ The body fo the Function
+
+ """
+
+ def __init__(
+ self,
+ params: List["DFPattern"],
+ body: "DFPattern",
+ ):
+ self.__init_handle_by_constructor__(ffi.FunctionPattern, params, body) # type: ignore
+
+
+@register_df_node
+class TuplePattern(DFPattern):
+ """A patern matching a Relax Tuple.
+
+ Parameters
+ ----------
+ fields : Array[tvm.relax.dpl.DFPattern]
+ The fields in the tuple.
+ """
+
+ def __init__(self, fields: Array):
+ self.__init_handle_by_constructor__(ffi.TuplePattern, fields) # type: ignore
+
+ def __getitem__(self, index: Optional[int]) -> "TupleGetItemPattern":
+ if index is not None:
+ # support negative index for being pythonic
+ if index < 0:
+ index += len(self)
+ if index >= len(self):
+ raise IndexError("TuplePattern index out of range")
+ else:
+ index = -1 # -1 means matching any index
+ return TupleGetItemPattern(self, index)
+
+ def __len__(self):
+ return len(self.fields)
+
+
+@register_df_node
+class UnorderedTuplePattern(DFPattern):
+ """A patern matching a Relax Tuple unorderedly.
+
+ Parameters
+ ----------
+ fields : Array[tvm.relax.dpl.DFPattern]
+ The fields in the tuple.
+ """
+
+ def __init__(self, fields: Array):
+ self.__init_handle_by_constructor__(ffi.UnorderedTuplePattern, fields) # type: ignore
+
+ def __len__(self):
+ return len(self.fields)
+
+
+@register_df_node
+class TupleGetItemPattern(DFPattern):
+ """Get index-th item from a TuplePattern.
+
+ Parameters
+ ----------
+ tuple_value: tvm.relax.dpl.DFPattern
+ The input tuple expression.
+
+ index: Optional[int]
+ The index to match; Default (None) to match a TupleGetItem with any index.
+ """
+
+ def __init__(self, tuple_value: "DFPattern", index: Optional[int] = None):
+ match_index = index if index is not None else -1
+ self.__init_handle_by_constructor__(
+ ffi.TupleGetItemPattern, tuple_value, match_index # type: ignore
+ )
+
+
+@register_df_node
+class OrPattern(DFPattern):
+ """Create a Pattern that can match one of two conditions
+
+ Parameters
+ ----------
+ left: tvm.relax.dpl.DFPattern
+ One possible matching pattern.
+ right: tvm.relax.dpl.DFPattern
+ One possible matching pattern.
+ """
+
+ def __init__(self, left: "DFPattern", right: "DFPattern"):
+ self.__init_handle_by_constructor__(ffi.OrPattern, left, right) # type: ignore
+
+
+@register_df_node
+class AndPattern(DFPattern):
+ """Create a Pattern that must match two conditions
+
+ Parameters
+ ----------
+ left: tvm.relax.dpl.DFPattern
+ One must-matching pattern.
+ right: tvm.relax.dpl.DFPattern
+ One must-matching pattern.
+ """
+
+ def __init__(self, left: "DFPattern", right: "DFPattern"):
+ self.__init_handle_by_constructor__(ffi.AndPattern, left, right) # type: ignore
+
+
+@register_df_node
+class NotPattern(DFPattern):
+ """Create a Pattern that matches the negation of a condition.
+
+ Parameters
+ ----------
+ to_reject: tvm.relax.dpl.DFPattern
+ The pattern to deny.
+ """
+
+ def __init__(self, to_reject: "DFPattern"):
+ self.__init_handle_by_constructor__(ffi.NotPattern, to_reject) # type: ignore
+
+
+@register_df_node
+class WildcardPattern(DFPattern):
+ """A pattern which matches anything."""
+
+ def __init__(self):
+ self.__init_handle_by_constructor__(ffi.WildcardPattern) # type: ignore
+
+
+@register_df_node
+class TypePattern(DFPattern):
+ """A pattern that matches another pattern with a certain type annotation.
+
+ Parameters
+ ----------
+ pattern: tvm.relax.dpl.DFPattern
+ The input pattern that needs type annotation.
+
+ ttype: tvm.ir.type.Type
+ The type to match.
+ """
+
+ def __init__(self, pattern: "DFPattern", ttype: tvm.ir.type.Type):
+ self.__init_handle_by_constructor__(ffi.TypePattern, pattern, ttype) # type: ignore
+
+
+@register_df_node
+class DataTypePattern(DFPattern):
+ """A pattern that matches another pattern with certain data type
+
+ Parameters
+ ----------
+ pattern: tvm.relax.dpl.DFPattern
+ The input pattern that needs type annotation.
+
+ dtype: str
+ The dtype to match.
+ """
+
+ def __init__(self, pattern: "DFPattern", dtype: str):
+ self.__init_handle_by_constructor__(ffi.DataTypePattern, pattern, dtype) # type: ignore
+
+
+@register_df_node
+class ShapePattern(DFPattern):
+ """A pattern that matches another pattern with a certain tensor shape
+
+ Parameters
+ ----------
+ pattern: tvm.relax.dpl.DFPattern
+ The input pattern that needs type annotation.
+
+ shape: List[tvm.ir.PrimExpr]
+ The shape to match.
+ """
+
+ def __init__(self, pattern: "DFPattern", shape: List[tvm.ir.PrimExpr]):
+ self.__init_handle_by_constructor__(ffi.ShapePattern, pattern, shape) # type: ignore
+
+
+@register_df_node
+class PrimArrPattern(DFPattern):
+ """
+ A pattern to match an array of PrimExpr
+
+ Parameters
+ ----------
+ shape : List[tvm.ir.PrimExpr]
+ The shape to match.
+ """
+
+ def __init__(self, shape: List[tvm.ir.PrimExpr]):
+ self.__init_handle_by_constructor__(ffi.PrimArrPattern, shape) # type: ignore
+
+ def __getitem__(self, index: int):
+ if index >= len(self):
+ raise IndexError("PrimArrPattern index out of range")
+ return self.fields[index]
+
+ def __len__(self):
+ return len(self.fields)
+
+
+@register_df_node
+class AttrPattern(DFPattern):
+ """Get match an expression with a certain attributes.
+ Currently only supports Op Attributes, not call Attributes.
+
+ Parameters
+ ----------
+ pattern: tvm.relax.dpl.DFPattern
+ The input pattern.
+
+ attrs: tvm.ir.attrs.Attrs
+ The attributes to match.
+ """
+
+ def __init__(self, pattern: "DFPattern", attrs: tvm.ir.attrs.Attrs):
+ self.__init_handle_by_constructor__(ffi.AttrPattern, pattern, attrs) # type: ignore
+
+
+def is_var(name: str = "") -> VarPattern:
+ """
+ Syntatic sugar for creating an optionally named VarPattern.
+
+ Parameters
+ ----------
+ name: str
+ The name of the input pattern to match.
+
+ Returns
+ -------
+ result: tvm.relax.dpl.VarPattern
+ The resulting pattern.
+ """
+ return VarPattern(name)
+
+
+def is_gv(name: str = "") -> GlobalVarPattern:
+ """Syntax sugar for creating an optionally (if name is empty) named GlobalVarPattern."""
+ return GlobalVarPattern(name)
+
+
+def is_dfv(name: str = "") -> DataflowVarPattern:
+ """Syntax sugar for creating an optionally (if name is empty) named DataflowVarPattern."""
+ return DataflowVarPattern(name)
+
+
+def is_const() -> ConstantPattern:
+ """
+ Syntatic sugar for creating a ConstantPattern.
+
+ Parameters
+ ----------
+ name: str
+ The name of the input pattern to match.
+
+ Returns
+ -------
+ result: tvm.relax.dpl.ConstantPattern
+ The resulting pattern.
+ """
+ return ConstantPattern()
+
+
+def is_expr(expr: Expr) -> ExprPattern:
+ """
+ Syntatic sugar for creating an ExprPattern.
+
+ Parameters
+ ----------
+ expr: Expr
+ The Relax expression to match.
+
+ Returns
+ -------
+ result: tvm.relax.dpl.ExprPattern
+ The resulting pattern.
+ """
+ return ExprPattern(expr)
+
+
+def is_op(op_name: str) -> ExprPattern:
+ """
+ Syntatic sugar for creating an operator ExprPattern.
+
+ Parameters
+ ----------
+ op_name: String
+ The name of the tvm.ir.op.Op object
+
+ Returns
+ -------
+ result: tvm.relax.dpl.ExprPattern
+ The resulting ExprPattern
+ """
+ op = get(op_name)
+ return ExprPattern(op)
+
+
+def is_tuple(
+ fields: Union[Array, List, Tuple], unordered=False
+) -> Union[TuplePattern, UnorderedTuplePattern]:
+ """
+ Syntatic sugar for creating an ExprPattern.
+
+ Parameters
+ ----------
+ fields : Array[tvm.relax.dpl.DFPattern]
+ The fields in the tuple.
+
+ Returns
+ -------
+ result: tvm.relax.dpl.DFPattern
+ The resulting pattern.
+ """
+ if not isinstance(fields, (list, tuple, Array)):
+ raise ValueError("fields must be a list, tuple, or Array")
+ if unordered:
+ return UnorderedTuplePattern(fields)
+ return TuplePattern(fields)
+
+
+def is_tuple_get_item(tuple_value: DFPattern, index: Optional[int] = None) -> TupleGetItemPattern:
+ """
+ Syntatic sugar for creating an ExprPattern.
+
+ Parameters
+ ----------
+ tuple_value: tvm.relax.dpl.DFPattern
+ The input tuple expression.
+
+ index: Optional[int]
+ The index to match; Default (None) to match a TupleGetItem with any index.
+
+ Returns
+ -------
+ result: tvm.relax.dpl.TupleGetItemPattern
+ The resulting pattern.
+ """
+ return TupleGetItemPattern(tuple_value, index)
+
+
+def wildcard() -> WildcardPattern:
+ """
+ Syntatic sugar for creating a WildcardPattern.
+
+ Returns
+ -------
+ result: tvm.relax.dpl.WildcardPattern
+ The resulting pattern.
+ """
+ return WildcardPattern()
+
+
+def has_dtype(dtype: str, pattern: DFPattern = None) -> DataTypePattern:
+ """
+ Syntatic sugar for creating a DataTypePattern
+
+ Parameters
+ ----------
+ dtype: str
+ The dtype to match
+
+ pattern: tvm.relax.dpl.DFPattern
+ The pattern that needs type annotation
+
+ Returns
+ -------
+ result: tvm.relax.dpl.DataTypePattern
+ The resulting DataTypePattern
+ """
+ if pattern is None:
+ pattern = wildcard()
+ return DataTypePattern(pattern, dtype)
+
+
+def is_shape(shape: List[tvm.ir.PrimExpr]) -> "PrimArrPattern":
+ """
+ Directly matches a shape which is an array of PrimExpr
+
+ Parameters
+ ----------
+ shape : List[tvm.ir.PrimExpr]
+ The expected shape
+
+ Returns
+ -------
+ PrimArrPattern
+ The resulting PrimArrPattern pattern
+
+ Raises
+ ------
+ ValueError
+ If the argument shape is not a list/tuple/tvm.ir.Array
+
+ Note
+ ----
+ The difference between p.has_shape(s) and is_shape(s) is that: has_shape
+ puts assumptions on the shape of the tensor matched by pattern p. While
+ is_shape directly matches the shape (an array of PrimExpr).
+ """
+ if not isinstance(shape, (list, tuple, tvm.ir.Array)):
+ raise ValueError("is_shape takes a list or tuple as input.")
+ return PrimArrPattern(shape)
+
+
+# Todo(relax-team): Dataflow pattern for StructInfo, and match out_sinfo
+def _is_call_tir(
+ func_pattern: DFPattern,
+ args: Union[List, Tuple, TuplePattern] = None,
+) -> CallPattern:
+ if args is None:
+ args = wildcard()
+ elif isinstance(args, (list, tuple)):
+ args = TuplePattern(args)
+
+ return is_op("relax.call_tir")(func_pattern, args)
+
+
+# Todo(relax-team): Dataflow pattern for StructInfo, and match out_sinfo
+def is_call_tir(
+ func_name: str,
+ args: Union[List, Tuple, TuplePattern] = None,
+) -> CallPattern:
+ """
+ Syntax sugar for creating a CallPattern for call_tir that calls an function through global var.
+
+ Parameters
+ ----------
+ func_name : str
+ Name of the CPS function to call.
+ args : Union[List[DFPattern], Tuple[DFPattern]], optional
+ Arguments in expected call_packed, by default None meaning arbitrary (number of) arguments
+
+ Returns
+ -------
+ CallPattern
+ The resulting CallPattern
+ """
+ func_pattern = GlobalVarPattern(func_name)
+ return _is_call_tir(func_pattern, args)
+
+
+def is_call_tir_extern(
+ func_name: str,
+ args: Union[List, Tuple, TuplePattern] = None,
+) -> CallPattern:
+ """Syntax sugar for creating a CallPattern for call_tir that calls an extern function
+
+ Parameters
+ ----------
+ func_name : str
+ Name of the CPS function to call.
+ args : Union[List[DFPattern], Tuple[DFPattern]], optional
+ Arguments in expected call_packed, by default None meaning arbitrary (number of) arguments
+
+ Returns
+ -------
+ CallPattern
+ The resulting CallPattern
+ """
+ func_pattern = ExternFuncPattern(func_name)
+ return _is_call_tir(func_pattern, args)
+
+
+def is_call_packed(
+ func_name: str, args: Union[List[DFPattern], Tuple[DFPattern]] = None
+) -> CallPattern:
+ """
+ Syntax sugar for creating a CallPattern for call_packed
+
+ Parameters
+ ----------
+ func_name : str
+ Name of the external function to call
+ args : Union[List[DFPattern], Tuple[DFPattern]], optional
+ Arguments in expected call_packed, by default None meaning arbitrary (number of) arguments
+
+ Returns
+ -------
+ CallPattern
+ The resulting CallPattern
+ """
+ if args is None:
+ return ExternFuncPattern(func_name)(varg_default_wildcard=True)
+ return ExternFuncPattern(func_name)(*args)
+
+
+def reject(pattern: DFPattern) -> NotPattern:
+ """
+ Syntax sugar for creating a DFPattern to reject
+
+ Parameters
+ ----------
+ pattern : DFPattern
+ The pattern to deny
+
+ Returns
+ -------
+ result: NotPattern
+ The resulting NotPattern
+ """
+ return NotPattern(pattern)
+
+
+def has_attr(attrs, pattern=None) -> AttrPattern:
+ """
+ Syntatic sugar for creating an AttrPattern
+
+ Parameters
+ ----------
+ attrs: Dict[str, Object]
+ The attributes to match
+
+ pattern: Optional[tvm.relax.dpl.DFPattern]
+ The input pattern.
+
+ Returns
+ -------
+ result: tvm.relax.dpl.DFPattern
+ The resulting AttrPattern
+ """
+ if pattern is None:
+ pattern = wildcard()
+ return pattern.has_attr(attrs)
+
+
+@register_df_node
+class PatternSeq(Node):
+ """A sequence of patterns with consecutive constraints"""
+
+ def __init__(self, patterns: List[DFPattern], only_use=False):
+ """
+ Initializer to PatternSeq
+
+ Parameters
+ ----------
+ patterns : List[DFPattern]
+ A chain of patterns
+ only_use : bool, optional
+ Whether the patterns follows only-used-by relations consecutively, by default False
+ """
+ self.__init_handle_by_constructor__(ffi.PatternSeq, patterns, only_use) # type: ignore
+
+ def used_by(self, other: Union[DFPattern, "PatternSeq"], index=-1) -> "PatternSeq":
+ """
+ Assuming the right-most pattern must be used by the `other` pattern as a producer
+
+ Parameters
+ ----------
+ other : Union[DFPattern, PatternSeq]
+ The consumer pattern (sequence)
+ index : int, optional
+ The argument index called by the consumer pattern, by default -1
+
+ Returns
+ -------
+ PatternSeq
+ A chained pattern sequence
+
+ Note
+ ----
+ If other is PatternSeq, it means the right-most pattern must be used by the left-most
+ pattern of the other sequence.
+ """
+ return _used_by(self, other, index)
+
+ def only_used_by(self, other: Union[DFPattern, "PatternSeq"], index=-1) -> "PatternSeq":
+
+ """
+ Assuming the right-most pattern must be **ONLY** used by the `other` pattern as a producer
+
+ Parameters
+ ----------
+ other : Union[DFPattern, PatternSeq]
+ The consumer pattern (sequence)
+ index : int, optional
+ The argument index called by the consumer pattern, by default -1
+
+ Returns
+ -------
+ PatternSeq
+ A chained pattern sequence
+
+ Note
+ ----
+ If other is PatternSeq, it means the right-most pattern must be **ONLY** used by the
+ left-most pattern of the other sequence.
+ """
+ return _only_used_by(self, other, index)
+
+ def __getitem__(self, index: int) -> DFPattern:
+ """
+ Access the pattern at the given index
+
+ Parameters
+ ----------
+ index : int
+ Index of the accessed pattern
+
+ Returns
+ -------
+ DFPattern
+ The accessed pattern
+ """
+ return self.patterns[index]
+
+ def __xor__(self, other) -> "PatternSeq":
+ """Syntax sugar of PatternSeq.used_by"""
+ return self.used_by(other, -1)
+
+ def __rshift__(self, other) -> "PatternSeq":
+ """Syntax sugar of PatternSeq.only_used_by"""
+ return self.only_used_by(other, -1)
+
+ def dup(self) -> "PatternSeq":
+ """
+ Duplicate the pattern sequence (new object under different address)
+
+ Returns
+ -------
+ PatternSeq
+ A duplicated chain
+ """
+ return ffi.dup_seq(self) # type: ignore
+
+
+### Private functions
+
+
+def _used_by(
+ lhs: Union[DFPattern, PatternSeq],
+ rhs: Union[DFPattern, PatternSeq],
+ index=-1,
+) -> PatternSeq:
+ if isinstance(lhs, DFPattern):
+ lhs = PatternSeq([lhs])
+ if isinstance(rhs, DFPattern):
+ rhs = PatternSeq([rhs])
+ return ffi.used_by(lhs, rhs, index) # type: ignore
+
+
+def _only_used_by(
+ lhs: Union[DFPattern, PatternSeq], rhs: Union[DFPattern, PatternSeq], index=-1
+) -> PatternSeq:
+ if isinstance(lhs, DFPattern):
+ lhs = PatternSeq([lhs])
+ if isinstance(rhs, DFPattern):
+ rhs = PatternSeq([rhs])
+ return ffi.only_used_by(lhs, rhs, index) # type: ignore
+
+
+def _add_bias_activation_pattern(out, with_bias=False, activation=None):
+ if with_bias:
+ bias = wildcard()
+ out = is_op("relax.add")(out, bias)
+
+ if activation:
+ return is_op(activation)(out)
+
+ return out
+
+
+def make_fused_bias_activation_pattern(op_name, with_bias=False, activation=None):
+ """
+ A simple utility to create patterns for an operation fused with bias addition and activation.
+
+ Parameters
+ ----------
+ op_name: str
+ The name of a Relax op, such as "relax.nn.conv2d"
+
+ with_bias: bool
+ Whether or not to include bias addition
+
+ activation: str
+ The name of an activation Relax op, such as "relax.nn.relu"
+
+ Returns
+ -------
+ pattern: DFPattern
+ The resulting pattern describing a fused operation
+ """
+ lhs = wildcard()
+ rhs = wildcard()
+ out = is_op(op_name)(lhs, rhs)
+
+ return _add_bias_activation_pattern(out, with_bias, activation)
+
+
+def make_matmul_pattern(with_bias=False, activation=None, transposed_b=False):
+ lhs = wildcard()
+ if transposed_b:
+ rhs = is_op("relax.permute_dims")(wildcard())
+ else:
+ rhs = wildcard()
+ out = is_op("relax.matmul")(lhs, rhs)
+
+ return _add_bias_activation_pattern(out, with_bias, activation)
diff --git a/src/relax/analysis/udchain.cc b/src/relax/analysis/udchain.cc
new file mode 100644
index 0000000000..f3d9b4686b
--- /dev/null
+++ b/src/relax/analysis/udchain.cc
@@ -0,0 +1,102 @@
+/*
+ * 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 src/relax/analysis/udchain.cc
+ * \brief Implementation of use-def analysis.
+ */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+
+#include <cstddef>
+#include <limits>
+#include <type_traits>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace relax {
+
+class UDChain : public relax::ExprVisitor {
+ public:
+ // nullptr users means it is the output of the function.
+ std::map<const VarNode*, std::set<const VarNode*>> to_users;
+
+ const VarNode* cur_user_;
+
+ void VisitBinding_(const VarBindingNode* binding) override {
+ // init
+ cur_user_ = binding->var.get();
+ this->VisitVarDef(binding->var);
+ this->VisitExpr(binding->value);
+ cur_user_ = nullptr;
+ }
+
+ void VisitExpr_(const VarNode* op) override { to_users[op].insert(cur_user_); }
+ void VisitVarDef(const Var& var) override { to_users[var.get()] = {}; }
+ void VisitExpr_(const FunctionNode* op) override { ExprVisitor::VisitExpr_(op); }
+
+ void VisitExpr_(const DataflowVarNode* op) override {
+ VisitExpr_(static_cast<const VarNode*>(op));
+ }
+};
+
+std::pair<runtime::Map<Var, runtime::Array<Var>>, runtime::Array<Var>> FunctionUseDef(
+ const Function& fn) {
+ UDChain udchain;
+ udchain.VisitExpr_(fn.get());
+
+ Map<Var, Array<Var>> user_map;
+ Array<Var> fn_outs;
+
+ for (const auto& kv : udchain.to_users) {
+ Array<Var> uses{};
+ uses.reserve(kv.second.size());
+ for (const auto& v : kv.second) {
+ if (nullptr == v &&
+ fn_outs.end() == std::find(fn_outs.begin(), fn_outs.end(), GetRef<Var>(kv.first))) {
+ fn_outs.push_back(GetRef<Var>(kv.first));
+ } else {
+ uses.push_back(GetRef<Var>(v));
+ }
+ }
+ user_map.Set(GetRef<Var>(kv.first), std::move(uses));
+ }
+ return std::make_pair(std::move(user_map), std::move(fn_outs));
+}
+
+runtime::Map<Var, Array<Var>> DataflowBlockUseDef(const DataflowBlock& dfb) {
+ UDChain udchain;
+ udchain.VisitBindingBlock_(dfb.get());
+ runtime::Map<Var, Array<Var>> ret;
+ for (const auto& kv : udchain.to_users) {
+ Array<Var> uses{};
+ uses.reserve(kv.second.size());
+ for (const auto& v : kv.second) uses.push_back(GetRef<Var>(v));
+ ret.Set(GetRef<Var>(kv.first), std::move(uses));
+ }
+ return ret;
+}
+
+TVM_REGISTER_GLOBAL("relax.analysis.udchain").set_body_typed(DataflowBlockUseDef);
+
+} // namespace relax
+} // namespace tvm
diff --git a/src/relax/analysis/var2value.cc b/src/relax/analysis/var2value.cc
new file mode 100644
index 0000000000..be50e9bdce
--- /dev/null
+++ b/src/relax/analysis/var2value.cc
@@ -0,0 +1,91 @@
+/*
+ * 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.
+ */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+
+namespace tvm {
+namespace relax {
+class Var2ValAnalysis : public relax::ExprVisitor {
+ public:
+ tvm::runtime::Map<Var, Expr> var2value_;
+ void VisitBinding_(const VarBindingNode* binding) override {
+ var2value_.Set(binding->var, binding->value);
+ // Recursively visit the value to handle local functions.
+ VisitExpr(binding->value);
+ }
+};
+
+tvm::runtime::Map<Var, Expr> AnalyzeVar2Value(const Expr& expr) {
+ Var2ValAnalysis var2val_analysis;
+ var2val_analysis.VisitExpr(expr);
+ return std::move(var2val_analysis.var2value_);
+}
+
+tvm::runtime::Map<Var, Expr> AnalyzeVar2Value(const DataflowBlock& dfb) {
+ Var2ValAnalysis var2val_analysis;
+ var2val_analysis.VisitBindingBlock_(dfb.get());
+ return std::move(var2val_analysis.var2value_);
+}
+
+tvm::runtime::Map<Var, Expr> AnalyzeVar2Value(const IRModule& m) {
+ Var2ValAnalysis var2val_analysis;
+
+ for (const auto& it : m->functions) {
+ // visit relax.Function
+ if (auto* n = it.second.as<FunctionNode>()) {
+ var2val_analysis.VisitExpr(GetRef<Function>(n));
+ }
+ }
+
+ return std::move(var2val_analysis.var2value_);
+}
+
+TVM_REGISTER_GLOBAL(("relax.analysis.get_var2val")).set_body_typed([](const Function& f) {
+ return AnalyzeVar2Value(f);
+});
+
+class Name2BindingAnalysis : public relax::ExprVisitor {
+ public:
+ // runtime::Map is not suitable for doing in-place update.
+ // so we use standard container for internal usage.
+ std::map<String, Array<Binding>> name2bindings_;
+ void VisitBinding_(const VarBindingNode* binding) override {
+ const auto& vname = binding->var->name_hint();
+ name2bindings_[vname].push_back(GetRef<VarBinding>(binding));
+ }
+
+ void VisitBinding_(const MatchCastNode* binding) override {
+ const auto& vname = binding->var->name_hint();
+ name2bindings_[vname].push_back(GetRef<MatchCast>(binding));
+ }
+};
+
+Map<String, Array<Binding>> NameToBinding(const Function& fn) {
+ Name2BindingAnalysis analysis{};
+ analysis.VisitExpr_(fn.get());
+ return Map<String, Array<Binding>>(std::make_move_iterator(analysis.name2bindings_.begin()),
+ std::make_move_iterator(analysis.name2bindings_.end()));
+}
+
+TVM_REGISTER_GLOBAL(("relax.analysis.name_to_binding")).set_body_typed(NameToBinding);
+
+} // namespace relax
+} // namespace tvm
diff --git a/src/relax/ir/dataflow_matcher.cc b/src/relax/ir/dataflow_matcher.cc
new file mode 100644
index 0000000000..92eb452a00
--- /dev/null
+++ b/src/relax/ir/dataflow_matcher.cc
@@ -0,0 +1,768 @@
+/*
+ * 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 src/relax/ir/dataflow_matcher.cc
+ * \brief The dataflow pattern matcher for Relax.
+ */
+
+#include <tvm/relax/analysis.h>
+#include <tvm/relax/dataflow_matcher.h>
+#include <tvm/relax/dataflow_pattern.h>
+#include <tvm/relax/expr.h>
+#include <tvm/relax/expr_functor.h>
+#include <tvm/relax/struct_info.h>
+#include <tvm/tir/op.h>
+
+#include <array>
+#include <cstddef>
+#include <limits>
+#include <stack>
+#include <type_traits>
+#include <unordered_map>
+#include <unordered_set>
+#include <utility>
+#include <vector>
+
+#include "dataflow_matcher_impl.h"
+
+namespace tvm {
+namespace relax {
+
+using tvm::arith::Analyzer;
+
+// Pattern Matcher
+bool DFPatternMatcher::Match(const DFPattern& pattern, const Expr& expr) {
+ memo_.clear();
+ matched_nodes_.clear();
+ return VisitDFPattern(pattern, expr);
+}
+
+static Expr TryGetValOfVar(const Expr& expr, const Map<Var, Expr>& var2val) {
+ if (var2val.empty()) return expr;
+
+ // if not match, try to match value of var if expr is a var.
+ if (const VarNode* var = expr.as<VarNode>()) {
+ auto may = var2val.Get(GetRef<Var>(var));
+ if (may.defined()) return may.value();
+ }
+
+ return expr;
+}
+
+void DFPatternMatcher::ClearMap(size_t watermark) {
+ for (size_t i = watermark; i < matched_nodes_.size(); ++i) {
+ memo_.erase(matched_nodes_[i]);
+ }
+ matched_nodes_.erase(matched_nodes_.begin() + watermark, matched_nodes_.end());
+}
+
+bool DFPatternMatcher::VisitDFPattern(const DFPattern& pattern, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ if (memoize_ && memo_.count(pattern)) {
+ ICHECK_EQ(memo_[pattern].size(), 1);
+ return expr.same_as(memo_[pattern][0]);
+ } else {
+ size_t watermark = matched_nodes_.size();
+ bool out = DFPatternFunctor::VisitDFPattern(pattern, expr);
+ if (out) {
+ memo_[pattern].push_back(expr);
+ matched_nodes_.push_back(pattern);
+ } else {
+ ClearMap(watermark);
+ }
+ return out;
+ }
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const OrPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ return VisitDFPattern(op->left, expr) || VisitDFPattern(op->right, expr);
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const AndPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ return VisitDFPattern(op->left, expr) && VisitDFPattern(op->right, expr);
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const NotPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ return !VisitDFPattern(op->reject, expr);
+}
+
+bool MatchRetValue(const ObjectRef& lhs, const TVMRetValue& rhs) {
+ switch (rhs.type_code()) {
+ case kDLInt:
+ if (auto* val = lhs.as<IntImmNode>()) {
+ return val->value == rhs.operator int64_t();
+ }
+ break;
+ case kDLFloat:
+ if (auto* val = lhs.as<FloatImmNode>()) {
+ return val->value == rhs.operator double();
+ }
+ break;
+ case kTVMStr:
+ if (auto* val = lhs.as<tir::StringImmNode>()) {
+ return val->value == rhs.operator std::string();
+ } else if (auto* val = lhs.as<StringObj>()) {
+ return val->data == rhs.operator std::string();
+ }
+ break;
+ case kTVMDataType:
+ if (auto* val = lhs.as<tir::StringImmNode>()) {
+ return rhs.operator std::string() == val->value;
+ } else if (auto* val = lhs.as<StringObj>()) {
+ return rhs.operator std::string() == val->data;
+ } else {
+ ICHECK(false) << "PatternMatcher: Unsupported TVMDataType " << lhs;
+ }
+ break;
+ case kTVMObjectHandle:
+ if (rhs.IsObjectRef<String>()) {
+ if (auto* val = lhs.as<tir::StringImmNode>()) {
+ return rhs.operator String() == val->value;
+ } else if (auto* val = lhs.as<StringObj>()) {
+ return rhs.operator String() == val->data;
+ }
+ } else {
+ // Compare the objects for structural equality
+ static auto* structural_equal = runtime::Registry::Get("node.StructuralEqual");
+ ICHECK(structural_equal) << "node.StructuralEqual is not registered.";
+ if ((*structural_equal)(lhs, GetRef<ObjectRef>(rhs.ptr<Object>()), false, true)) {
+ return true;
+ }
+ }
+ break;
+ default:
+ ICHECK(false) << "Unsupported type code in Pattern Node " << rhs.type_code();
+ }
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const AttrPatternNode* attr_pattern, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ bool matches = VisitDFPattern(attr_pattern->pattern, expr);
+ if (!matches) return matches;
+ VLOG(1) << "considering AttrPatternNode at:\n" << expr;
+ auto attributes = attr_pattern->attrs.as<DictAttrsNode>()->dict;
+ if (const auto* op_node = expr.as<OpNode>()) {
+ Op op = GetRef<Op>(op_node);
+ for (auto kv : attributes) {
+ auto attr_name = kv.first;
+ auto attr_value = kv.second;
+ if (Op::HasAttrMap(attr_name)) {
+ auto op_map = Op::GetAttrMap<TVMRetValue>(attr_name);
+ if (op_map.count(op)) {
+ matches &= MatchRetValue(attr_value, op_map[op]);
+ } else {
+ matches = false;
+ }
+ } else {
+ matches = false;
+ }
+ }
+ } else if (auto* op = expr.as<CallNode>()) {
+ matches = true;
+ // TODO(mbrookhart): When OpNode Attrs move from TVMRetValue to the Object system, remove this
+ // and replace the whole thing with a Visitor-based approach
+ ReflectionVTable* reflection = ReflectionVTable::Global();
+ auto attrs_node = const_cast<BaseAttrsNode*>(op->attrs.get());
+ // attrs may be undefined on non-op calls so we check first
+ std::vector<std::string> attr_names;
+ if (attrs_node) {
+ attr_names = reflection->ListAttrNames(attrs_node);
+ }
+ for (auto kv : attributes) {
+ std::string attr = kv.first;
+ if (matches && std::find(attr_names.begin(), attr_names.end(), attr) != attr_names.end()) {
+ matches &= MatchRetValue(kv.second, reflection->GetAttr(attrs_node, attr));
+ } else {
+ matches = false;
+ break;
+ }
+ }
+ } else if (auto* op = expr.as<FunctionNode>()) {
+ matches = true;
+ for (auto kv : attributes) {
+ if (matches && op->attrs.defined() && op->attrs->dict.count(kv.first)) {
+ matches &= StructuralEqual()(kv.second, op->attrs->dict[kv.first]);
+ } else {
+ matches = false;
+ break;
+ }
+ }
+ } else {
+ matches = false;
+ }
+ return matches;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ // utilities
+ auto get_op_node = [](const CallPatternNode* op) -> const tvm::OpNode* {
+ if (op) {
+ if (auto* expr_pattern = op->op.as<ExprPatternNode>()) {
+ return expr_pattern->expr.as<OpNode>();
+ }
+ }
+ return nullptr;
+ };
+ auto is_pattern_op = [&get_op_node](const CallPatternNode* op, std::string op_type) {
+ if (const auto* op_node = get_op_node(op)) {
+ if (op_node->name == op_type) {
+ return true;
+ }
+ }
+ return false;
+ };
+ auto is_expr_op = [](const Expr& expr, std::string op_type) {
+ if (const auto* call_node = expr.as<CallNode>()) {
+ if (const auto* op_node = call_node->op.as<OpNode>()) {
+ if (op_node->name == op_type) {
+ return true;
+ }
+ }
+ }
+ return false;
+ };
+
+ // logic
+ auto watermark = matched_nodes_.size();
+ if (const auto* call_node = expr.as<CallNode>()) {
+ auto matches_op = VisitDFPattern(op->op, call_node->op);
+ if (matches_op) {
+ auto watermark2 = matched_nodes_.size();
+
+ auto match_args = [this, &watermark2](const Array<DFPattern>& pattern_args, auto expr_begin,
+ auto expr_end) {
+ bool matches = true;
+ auto pattern_it = pattern_args.begin();
+ auto expr_it = expr_begin;
+ if (pattern_args.defined()) {
+ while (matches && pattern_it != pattern_args.end())
+ matches &= VisitDFPattern(*(pattern_it++), *(expr_it++));
+ }
+ if (!matches) ClearMap(watermark2);
+ return matches;
+ };
+
+ const size_t n_arg_pattern = op->args.size();
+ const size_t n_arg_expr = call_node->args.size();
+ // if allow variable args, #pattern must >= #expr.
+ if (op->varg_default_wildcard && n_arg_expr < n_arg_pattern) return false;
+ // if variable args are not allowed, #pattern must == #expr.
+ if (!op->varg_default_wildcard && n_arg_expr != n_arg_pattern) return false;
+
+ // Standard case
+ if (match_args(op->args, call_node->args.begin(), call_node->args.end())) return true;
+
+ // Commutative Matching
+ if (const OpNode* op_node = get_op_node(op)) {
+ if ((op_node->name == "relax.add") || (op_node->name == "relax.multiply")) {
+ if (match_args(op->args, call_node->args.rbegin(), call_node->args.rend())) {
+ return true;
+ }
+ }
+ }
+ } else {
+ ClearMap(watermark);
+ // associate divide/multiply
+ if (is_pattern_op(op, "relax.divide")) {
+ if (const auto* arg_node = op->args[0].as<CallPatternNode>()) {
+ if (is_pattern_op(arg_node, "relax.multiply") && is_expr_op(expr, "relax.multiply") &&
+ (is_expr_op(call_node->args[0], "relax.divide") ||
+ is_expr_op(call_node->args[1], "relax.divide"))) {
+ bool out = false;
+ for (size_t arg_id = 0; arg_id < 2; ++arg_id) {
+ auto div = CallPattern(op->op, {arg_node->args[arg_id], op->args[1]});
+ auto mul = CallPattern(arg_node->op, {arg_node->args[(arg_id + 1) % 2], div});
+ out = VisitDFPattern(mul, expr);
+ if (out) {
+ return true;
+ } else {
+ ClearMap(watermark);
+ }
+ }
+ return out;
+ }
+ }
+ }
+ if (is_pattern_op(op, "relax.multiply")) {
+ // associate multiply/divide
+ for (size_t arg_id = 0; arg_id < 2; ++arg_id) {
+ if (auto* arg_node = op->args[arg_id].as<CallPatternNode>()) {
+ if (is_pattern_op(arg_node, "relax.divide") && is_expr_op(expr, "relax.divide") &&
+ (is_expr_op(call_node->args[0], "relax.multiply") ||
+ is_expr_op(call_node->args[1], "relax.multiply"))) {
+ auto mul = CallPattern(op->op, {arg_node->args[0], op->args[(arg_id + 1) % 2]});
+ auto div = CallPattern(arg_node->op, {mul, arg_node->args[1]});
+ return VisitDFPattern(div, expr);
+ }
+ }
+ }
+ }
+ }
+ }
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const ExprPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ return StructuralEqual()(op->expr, expr);
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const FunctionPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ bool matches = false;
+ if (const auto* func = expr.as<FunctionNode>()) {
+ matches = true;
+ if (op->params.defined()) {
+ size_t i = 0;
+ if (op->params.size() == func->params.size()) {
+ while (matches && i < op->params.size()) {
+ matches &= VisitDFPattern(op->params[i], func->params[i]);
+ ++i;
+ }
+ } else {
+ matches = false;
+ }
+ }
+ if (matches) {
+ matches &= VisitDFPattern(op->body, func->body);
+ }
+ }
+ return matches;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const TupleGetItemPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ if (const auto* tuple_get_item_node = expr.as<TupleGetItemNode>()) {
+ return (op->index == -1 || op->index == tuple_get_item_node->index) &&
+ VisitDFPattern(op->tuple, tuple_get_item_node->tuple);
+ }
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const TuplePatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ bool matches = false;
+ if (const auto* tuple_node = expr.as<TupleNode>()) {
+ matches = true;
+ if (op->fields.size() == tuple_node->fields.size()) {
+ size_t i = 0;
+ while (matches && i < op->fields.size()) {
+ matches &= VisitDFPattern(op->fields[i], tuple_node->fields[i]);
+ ++i;
+ }
+ } else {
+ matches = false;
+ }
+ }
+ return matches;
+}
+
+bool DFPatternMatcher::TryUnorderedMatch(size_t idx, const tvm::Array<DFPattern> patterns,
+ const tvm::Array<Expr> fields,
+ std::vector<int8_t>& match_cache,
+ std::vector<bool>& matched) {
+ if (idx >= patterns.size()) return true;
+ constexpr int8_t kUnknown = -1;
+ auto this_pattern = patterns[idx];
+ for (size_t i = 0; i < fields.size(); ++i) {
+ if (matched[i]) continue;
+ const size_t table_idx = idx * fields.size() + i;
+ match_cache[table_idx] =
+ kUnknown ? VisitDFPattern(this_pattern, fields[i]) : match_cache[table_idx];
+ if (match_cache[table_idx]) {
+ // continue to match the rest;
+ matched[i] = true;
+ if (TryUnorderedMatch(idx + 1, patterns, fields, match_cache, matched)) return true;
+ matched[i] = false;
+ }
+ }
+
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const UnorderedTuplePatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+
+ if (const auto* tuple_node = expr.as<TupleNode>()) {
+ if (op->fields.size() == tuple_node->fields.size()) {
+ constexpr int8_t kUnknown = -1;
+ ICHECK_LE(op->fields.size(), std::numeric_limits<uint8_t>::max()) << "Too many fields!";
+ // dynamic programming.
+ std::vector<int8_t> match_cache(op->fields.size() * op->fields.size(), kUnknown);
+ std::vector<bool> field_match_bitmap(op->fields.size(), false);
+ return TryUnorderedMatch(0, op->fields, tuple_node->fields, match_cache, field_match_bitmap);
+ }
+ }
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const TypePatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ auto expr_type = expr.as<ExprNode>()->checked_type();
+ return (StructuralEqual()(op->type, expr_type)) && VisitDFPattern(op->pattern, expr);
+}
+
+static bool ShapeEqual(Analyzer* analyzer, const Array<PrimExpr>& lhs, const Array<PrimExpr>& rhs) {
+ if (lhs.size() != rhs.size()) return false;
+ for (size_t i = 0; i < lhs.size(); ++i)
+ if (!tir::is_one(analyzer->Simplify(lhs[i] == rhs[i]))) return false;
+ return true;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const ShapePatternNode* op, const Expr& expr) {
+ // no need to jump, as var.shape == value.shape
+ if (const auto* tinfo = GetStructInfoAs<TensorStructInfoNode>(expr)) {
+ if (const ShapeExprNode* shape_expr = tinfo->shape.as<ShapeExprNode>()) {
+ return ShapeEqual(&analyzer_, op->shape, shape_expr->values) &&
+ VisitDFPattern(op->pattern, expr);
+ }
+ }
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const PrimArrPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ if (const ShapeExprNode* shape_expr = expr.as<ShapeExprNode>())
+ return ShapeEqual(&analyzer_, op->fields, shape_expr->values);
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const DataTypePatternNode* op, const Expr& expr) {
+ // no need to jump, as var.dtype == value.dtype
+ auto expr_type = expr.as<ExprNode>()->checked_type();
+ if (const DynTensorTypeNode* tensor_type = expr_type.as<DynTensorTypeNode>()) {
+ return (StructuralEqual()(op->dtype, tensor_type->dtype)) && VisitDFPattern(op->pattern, expr);
+ }
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const VarPatternNode* op, const Expr& expr) {
+ // We don't jump for var pattern, as there's no need to access its value to judge it.
+ if (const auto* var_node = expr.as<VarNode>()) {
+ // "" means any name.
+ return "" == op->name_hint() || op->name_hint() == var_node->name_hint();
+ }
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const ExternFuncPatternNode* op, const Expr& expr0) {
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ if (const auto* extern_fn = expr.as<ExternFuncNode>()) {
+ return "" == op->global_symbol() || op->global_symbol() == extern_fn->global_symbol;
+ }
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const ConstantPatternNode* op, const Expr& expr0) {
+ // constants can be binded to relax.Var as well.
+ auto expr = TryGetValOfVar(expr0, var2val_);
+ return expr.as<ConstantNode>() != nullptr;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const DataflowVarPatternNode* op, const Expr& expr) {
+ // DataflowVar is inherented from Var, so dispatch it to VarPattern.
+ return expr->IsInstance<DataflowVarNode>() &&
+ VisitDFPattern_(static_cast<const VarPatternNode*>(op), expr);
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const GlobalVarPatternNode* op, const Expr& expr) {
+ // GlobalVarPattern is not inherited from Var, so we need to handle it separately.
+ if (const auto* var_node = expr.as<GlobalVarNode>())
+ return "" == op->name_hint() || op->name_hint() == var_node->name_hint;
+ return false;
+}
+
+bool DFPatternMatcher::VisitDFPattern_(const WildcardPatternNode* op, const Expr& expr) {
+ return true;
+}
+
+Optional<Map<DFPattern, Expr>> ExtractMatchedExpr(DFPattern pattern, Expr expr,
+ Optional<Map<Var, Expr>> bindings_opt) {
+ auto bindings = bindings_opt.value_or({});
+ DFPatternMatcher matcher(bindings);
+
+ if (!matcher.Match(pattern, expr)) {
+ return NullOpt;
+ }
+
+ Map<DFPattern, Expr> matching;
+ for (const auto& [pat, matches] : matcher.GetMemo()) {
+ ICHECK_EQ(matches.size(), 1) << "More than one match for the pattern " << pat;
+ matching.Set(pat, matches[0]);
+ }
+ return matching;
+}
+
+bool MatchExpr(DFPattern pattern, Expr expr, Optional<Map<Var, Expr>> bindings_opt) {
+ return static_cast<bool>(ExtractMatchedExpr(pattern, expr, bindings_opt));
+}
+
+TVM_REGISTER_GLOBAL("relax.dpl.match_expr").set_body_typed(MatchExpr);
+
+struct PNode {
+ const DFPatternNode* ptr;
+ const VarNode* matched = nullptr;
+ std::vector<std::pair<PNode*, const std::vector<PairCons>&>> children;
+ std::vector<std::pair<PNode*, const std::vector<PairCons>&>> parents;
+};
+
+struct RNode {
+ const VarNode* ptr;
+ const DFPatternNode* matched = nullptr;
+ std::vector<RNode*> children;
+ std::vector<RNode*> parents;
+};
+
+/**
+ * \brief This method try to match a real node and a pattern node along with its neighbors.
+ */
+static bool try_match(PNode* p, RNode* r, DFPatternMatcher* m,
+ const std::map<const VarNode*, std::set<const VarNode*>>& def2use,
+ const std::map<const VarNode*, std::vector<const VarNode*>>& use2def) {
+ if (nullptr != p->matched && p->matched == r->ptr) return true; // matched before.
+ if (!m->Match(GetRef<DFPattern>(p->ptr), GetRef<Var>(r->ptr))) return false;
+
+ std::stack<std::pair<PNode*, RNode*>> undo_stack{};
+
+ const auto commit = [&undo_stack](PNode* p, RNode* r) {
+ // match with each other.
+ p->matched = r->ptr;
+ r->matched = p->ptr;
+ undo_stack.emplace(p, r);
+ };
+
+ const auto quit = [&undo_stack] {
+ while (!undo_stack.empty()) {
+ auto& top = undo_stack.top();
+ top.first->matched = nullptr;
+ top.second->matched = nullptr;
+ undo_stack.pop();
+ }
+ return false;
+ };
+
+ commit(p, r);
+
+ // match parent patterns.
+ for (auto& pparent_pairs : p->parents) {
+ PNode* pparent = pparent_pairs.first;
+ const std::vector<PairCons>& constraints = pparent_pairs.second;
+
+ bool any_cons_sat = false;
+ for (auto& rparent : r->parents) {
+ // skip if mismatch.
+ if (rparent->matched && rparent->matched != pparent->ptr) continue;
+
+ const auto& uses = def2use.at(rparent->ptr);
+ // skip if `rparent` is not used by `r`.
+ if (uses.cend() == uses.find(r->ptr)) continue;
+
+ // check edge constraints.
+ bool cons_sat = true;
+ for (const auto& cons : constraints) {
+ if (PairCons::kOnlyUsedBy == cons.type && uses.size() != 1) {
+ cons_sat = false;
+ break;
+ }
+
+ if (-1 != cons.index) {
+ const auto& callees = use2def.at(r->ptr);
+ if (static_cast<size_t>(cons.index) >= callees.size() ||
+ rparent->ptr != callees[cons.index]) {
+ cons_sat = false;
+ break;
+ }
+ }
+ }
+ if (!cons_sat) continue;
+ any_cons_sat = true;
+
+ // try all parent R nodes that are not matched yet.
+ // as long as ppattern can match one node.
+ if (!pparent->matched && try_match(pparent, rparent, m, def2use, use2def)) {
+ commit(pparent, rparent);
+ break;
+ }
+ }
+ if (!pparent->matched || !any_cons_sat) return quit();
+ }
+
+ // forward matching;
+ for (auto& pchild_pairs : p->children) {
+ PNode* pchild = pchild_pairs.first;
+ const std::vector<PairCons>& constraints = pchild_pairs.second;
+ bool any_cons_sat = false;
+ for (auto& rchild : r->children) {
+ if (rchild->matched && rchild->matched != pchild->ptr) continue;
+
+ const auto& uses = def2use.at(r->ptr);
+ if (uses.cend() == uses.find(rchild->ptr)) continue;
+
+ // check edge constraints.
+ bool all_cons_pass = true;
+ for (const auto& cons : constraints) {
+ if (PairCons::kOnlyUsedBy == cons.type && uses.size() != 1) {
+ all_cons_pass = false;
+ break;
+ }
+
+ if (-1 != cons.index) {
+ const auto& callees = use2def.at(rchild->ptr);
+ if (static_cast<size_t>(cons.index) >= callees.size() || r->ptr != callees[cons.index]) {
+ all_cons_pass = false;
+ break;
+ }
+ }
+ }
+ if (!all_cons_pass) continue;
+ any_cons_sat = true;
+
+ if (!pchild->matched && try_match(pchild, rchild, m, def2use, use2def)) {
+ commit(pchild, rchild);
+ break;
+ }
+ }
+ if (!pchild->matched || !any_cons_sat) return quit();
+ }
+
+ return true;
+}
+
+class MatcherUseDefAnalysis : public relax::ExprVisitor {
+ public:
+ std::map<const VarNode*, std::set<const VarNode*>> def2use;
+ // caller -> callee table.
+ std::map<const VarNode*, std::vector<const VarNode*>> caller2callees;
+
+ const VarNode* cur_user_;
+
+ void VisitBinding_(const VarBindingNode* binding) override {
+ // init
+ cur_user_ = binding->var.get();
+ this->VisitVarDef(binding->var);
+ this->VisitExpr(binding->value);
+ cur_user_ = nullptr;
+ }
+
+ void VisitExpr_(const VarNode* op) override {
+ if (nullptr == cur_user_) return;
+
+ def2use[op].insert(cur_user_);
+ caller2callees[cur_user_].push_back(op);
+ }
+
+ void VisitExpr_(const DataflowVarNode* op) override {
+ VisitExpr_(static_cast<const VarNode*>(op));
+ }
+};
+
+Map<DFPattern, Var> MatchGraph(const PatternContext& ctx, const DataflowBlock& dfb,
+ Optional<Var> start_hint, bool must_include_hint) {
+ Map<DFPattern, Var> ret;
+ // TODO(@ganler): Handle non-may external use.
+ ICHECK(ctx->allow_extern_use == PatternContextNode::kMay) << "Only kMay is supported yet.";
+ ICHECK(!must_include_hint || start_hint.defined())
+ << "must_include_hint is only supported with start_hint.";
+
+ const auto var2val = AnalyzeVar2Value(dfb);
+ DFPatternMatcher matcher(var2val);
+
+ // std::map<const VarNode*, std::set<const VarNode*>>
+ MatcherUseDefAnalysis ud_analysis;
+ ud_analysis.VisitBindingBlock_(dfb.get());
+ const auto& def2use = ud_analysis.def2use;
+ const auto& caller2callees = ud_analysis.caller2callees;
+
+ // First construct a graph of PNode and RNode.
+ std::unordered_map<const VarNode*, RNode> var2node;
+ var2node.reserve(dfb->bindings.size());
+
+ for (const auto& du : def2use) {
+ const VarNode* cur_var = du.first;
+ const std::set<const VarNode*>& uses = du.second;
+ RNode& cur_node = var2node[cur_var];
+ cur_node.ptr = cur_var;
+ for (const VarNode* use : uses) {
+ auto& use_node = var2node[use];
+ use_node.ptr = use;
+ cur_node.children.push_back(&use_node);
+ use_node.parents.push_back(&cur_node);
+ }
+ }
+
+ std::unordered_map<const DFPatternNode*, PNode> pattern2node;
+ pattern2node.reserve(ctx->constraints.size());
+
+ for (const auto& def2use_pattern : ctx->constraints) {
+ const DFPatternNode* def_pattern = def2use_pattern.first.get();
+ const std::map<DFPattern, std::vector<PairCons>>& uses = def2use_pattern.second;
+ PNode& def_node = pattern2node[def_pattern];
+ def_node.ptr = def_pattern;
+ def_node.children.reserve(uses.size());
+ for (const auto& use : uses) {
+ const auto& cons = use.second;
+ const DFPatternNode* use_pattern = use.first.get();
+ PNode& use_node = pattern2node[use_pattern];
+ use_node.ptr = use_pattern;
+ use_node.parents.emplace_back(&def_node, std::ref(cons));
+ def_node.children.emplace_back(&use_node, std::ref(cons));
+ }
+ }
+
+ if (start_hint.defined()) {
+ Var v = start_hint.value();
+ auto rnode_ptr = var2node.find(v.get());
+ for (auto& ppair : pattern2node) {
+ if (try_match(&ppair.second, &rnode_ptr->second, &matcher, def2use, caller2callees)) {
+ for (auto ppair : pattern2node)
+ ret.Set(GetRef<DFPattern>(ppair.first), GetRef<Var>(ppair.second.matched));
+ return ret;
+ }
+ }
+
+ if (must_include_hint) return ret;
+ }
+
+ PNode* pnode_start = &pattern2node.begin()->second;
+
+ if (!pnode_start->matched) {
+ for (auto& rpair : var2node) {
+ if (start_hint.defined() && start_hint.value().get() == rpair.first) continue;
+ if (try_match(pnode_start, &rpair.second, &matcher, def2use, caller2callees)) {
+ for (auto ppair : pattern2node)
+ ret.Set(GetRef<DFPattern>(ppair.first), GetRef<Var>(ppair.second.matched));
+
+ return ret;
+ }
+ }
+ }
+
+ return ret;
+}
+
+TVM_REGISTER_GLOBAL("relax.dpl.match_dfb").set_body_typed(MatchGraph);
+
+} // namespace relax
+} // namespace tvm
diff --git a/src/relax/ir/dataflow_matcher_impl.h b/src/relax/ir/dataflow_matcher_impl.h
new file mode 100644
index 0000000000..89f3d114c1
--- /dev/null
+++ b/src/relax/ir/dataflow_matcher_impl.h
@@ -0,0 +1,87 @@
+/*
+ * 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 src/tvm/relax/dataflow_matcher_impl.h
+ * \brief The auxiliary data structure for dataflow matcher.
+ */
+#ifndef TVM_RELAX_IR_DATAFLOW_MATCHER_IMPL_H_
+#define TVM_RELAX_IR_DATAFLOW_MATCHER_IMPL_H_
+
+#include <tvm/arith/analyzer.h>
+#include <tvm/relax/dataflow_matcher.h>
+#include <tvm/relax/dataflow_pattern.h>
+#include <tvm/relax/dataflow_pattern_functor.h>
+
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+namespace tvm {
+namespace relax {
+
+class DFPatternMatcher : public DFPatternFunctor<bool(const DFPattern&, const Expr&)> {
+ public:
+ using var2val_t = runtime::Map<Var, Expr>;
+
+ explicit DFPatternMatcher() {}
+ explicit DFPatternMatcher(var2val_t var2val) : var2val_(std::move(var2val)) {}
+ bool Match(const DFPattern& pattern, const Expr& expr);
+ Map<DFPattern, Array<Expr>> GetMemo() { return Map<DFPattern, Array<Expr>>(memo_); }
+
+ protected:
+ bool VisitDFPattern(const DFPattern& pattern, const Expr& expr) override;
+ bool VisitDFPattern_(const OrPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const AndPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const NotPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const AttrPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const CallPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const ConstantPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const DataTypePatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const ExprPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const FunctionPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const ShapePatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const TupleGetItemPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const TuplePatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const TypePatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const WildcardPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const VarPatternNode* op, const Expr& expr) override;
+
+ bool VisitDFPattern_(const DataflowVarPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const GlobalVarPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const ExternFuncPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const PrimArrPatternNode* op, const Expr& expr) override;
+ bool VisitDFPattern_(const UnorderedTuplePatternNode* op, const Expr& expr) override;
+
+ void ClearMap(size_t watermark);
+ bool TryUnorderedMatch(size_t idx, const tvm::Array<DFPattern> patterns,
+ const tvm::Array<Expr> fields, std::vector<int8_t>& match_cache,
+ std::vector<bool>& matched);
+
+ std::unordered_map<DFPattern, Array<Expr>, ObjectPtrHash, ObjectPtrEqual> memo_;
+ var2val_t var2val_;
+ std::vector<DFPattern> matched_nodes_;
+ arith::Analyzer analyzer_;
+ bool memoize_ = true;
+};
+
+} // namespace relax
+} // namespace tvm
+
+#endif // TVM_RELAX_IR_DATAFLOW_MATCHER_IMPL_H_
diff --git a/src/relax/ir/dataflow_pattern.cc b/src/relax/ir/dataflow_pattern.cc
new file mode 100644
index 0000000000..3768627c20
--- /dev/null
+++ b/src/relax/ir/dataflow_pattern.cc
@@ -0,0 +1,607 @@
+/*
+ * 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 src/relax/ir/dataflow_pattern.cc
+ * \brief The dataflow pattern language for Relax (inherited from Relay).
+ */
+
+#include <tvm/relax/dataflow_pattern.h>
+#include <tvm/relax/dataflow_pattern_functor.h>
+
+#include <memory>
+#include <stack>
+#include <string>
+
+#define RELAX_PATTERN_PRINTER_DEF(NODE_TYPE, REPR_LAMBDA) \
+ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable) \
+ .set_dispatch<NODE_TYPE>([](const ObjectRef& ref, ReprPrinter* p) { \
+ auto* node = static_cast<const NODE_TYPE*>(ref.get()); \
+ REPR_LAMBDA(p, node); \
+ })
+
+namespace tvm {
+namespace relax {
+
+TVM_REGISTER_NODE_TYPE(ExternFuncPatternNode);
+ExternFuncPattern::ExternFuncPattern(String global_symbol) {
+ ObjectPtr<ExternFuncPatternNode> n = make_object<ExternFuncPatternNode>();
+ n->global_symbol_ = std::move(global_symbol);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.ExternFuncPattern").set_body_typed([](String global_symbol) {
+ return ExternFuncPattern(global_symbol);
+});
+RELAX_PATTERN_PRINTER_DEF(ExternFuncPatternNode, [](auto p, auto node) {
+ p->stream << "ExternFuncPattern(" << node->global_symbol() << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(VarPatternNode);
+VarPattern::VarPattern(String name_hint) {
+ ObjectPtr<VarPatternNode> n = make_object<VarPatternNode>();
+ n->name = std::move(name_hint);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.VarPattern").set_body_typed([](String name_hint) {
+ return VarPattern(name_hint);
+});
+RELAX_PATTERN_PRINTER_DEF(VarPatternNode, [](auto p, auto node) {
+ p->stream << "VarPattern(" << node->name_hint() << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(DataflowVarPatternNode);
+TVM_REGISTER_GLOBAL("relax.dpl.DataflowVarPattern").set_body_typed([](String name_hint) {
+ return DataflowVarPattern(name_hint);
+});
+DataflowVarPattern::DataflowVarPattern(String name_hint) {
+ ObjectPtr<DataflowVarPatternNode> n = make_object<DataflowVarPatternNode>();
+ n->name = std::move(name_hint);
+ data_ = std::move(n);
+}
+RELAX_PATTERN_PRINTER_DEF(DataflowVarPatternNode, [](auto p, auto node) {
+ p->stream << "DataflowVarPattern(" << node->name_hint() << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(GlobalVarPatternNode);
+GlobalVarPattern::GlobalVarPattern(String name_hint) {
+ ObjectPtr<GlobalVarPatternNode> n = make_object<GlobalVarPatternNode>();
+ n->name = std::move(name_hint);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.GlobalVarPattern").set_body_typed([](String name_hint) {
+ return GlobalVarPattern(name_hint);
+});
+RELAX_PATTERN_PRINTER_DEF(GlobalVarPatternNode, [](auto p, auto node) {
+ p->stream << "GlobalVarPattern(" << node->name_hint() << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(ExprPatternNode);
+ExprPattern::ExprPattern(Expr expr) {
+ ObjectPtr<ExprPatternNode> n = make_object<ExprPatternNode>();
+ n->expr = std::move(expr);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.ExprPattern").set_body_typed([](Expr e) { return ExprPattern(e); });
+RELAX_PATTERN_PRINTER_DEF(ExprPatternNode, [](auto p, auto node) { p->Print(node->expr); });
+
+TVM_REGISTER_NODE_TYPE(ConstantPatternNode);
+TVM_REGISTER_GLOBAL("relax.dpl.ConstantPattern").set_body_typed([]() {
+ auto c = ConstantPattern(make_object<ConstantPatternNode>());
+ return c;
+});
+RELAX_PATTERN_PRINTER_DEF(ConstantPatternNode,
+ [](auto p, auto node) { p->stream << "ConstantPattern()"; });
+
+TVM_REGISTER_NODE_TYPE(CallPatternNode);
+CallPattern::CallPattern(DFPattern op, Array<DFPattern> args, bool varg_default_wildcard) {
+ ObjectPtr<CallPatternNode> n = make_object<CallPatternNode>();
+ n->op = std::move(op);
+ n->args = std::move(args);
+ n->varg_default_wildcard = varg_default_wildcard;
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.CallPattern")
+ .set_body_typed([](DFPattern op, Array<DFPattern> args, bool varg_default_wildcard) {
+ return CallPattern(op, args, varg_default_wildcard);
+ });
+RELAX_PATTERN_PRINTER_DEF(CallPatternNode, [](auto p, auto node) {
+ p->stream << node->op << "(";
+ for (size_t i = 0; i < node->args.size(); ++i) {
+ if (i != 0) p->stream << ", ";
+ p->stream << node->args[i];
+ }
+ if (node->varg_default_wildcard) {
+ if (node->args.size() != 0) p->stream << ", ";
+ p->stream << "...";
+ }
+ p->stream << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(PrimArrPatternNode);
+PrimArrPattern::PrimArrPattern(Array<PrimExpr> arr) {
+ ObjectPtr<PrimArrPatternNode> n = make_object<PrimArrPatternNode>();
+ n->fields = std::move(arr);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.PrimArrPattern").set_body_typed([](Array<PrimExpr> arr) {
+ return PrimArrPattern(std::move(arr));
+});
+RELAX_PATTERN_PRINTER_DEF(PrimArrPatternNode, [](auto p, auto node) {
+ p->stream << "PrimArrPattern(" << node->fields << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(FunctionPatternNode);
+FunctionPattern::FunctionPattern(Array<DFPattern> params, DFPattern body) {
+ ObjectPtr<FunctionPatternNode> n = make_object<FunctionPatternNode>();
+ n->params = std::move(params);
+ n->body = std::move(body);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.FunctionPattern")
+ .set_body_typed([](Array<DFPattern> params, DFPattern body) {
+ return FunctionPattern(params, body);
+ });
+RELAX_PATTERN_PRINTER_DEF(FunctionPatternNode, [](auto p, auto node) {
+ p->stream << "FunctionPattern(" << node->params << ", " << node->body << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(TuplePatternNode);
+TuplePattern::TuplePattern(tvm::Array<DFPattern> fields) {
+ ObjectPtr<TuplePatternNode> n = make_object<TuplePatternNode>();
+ n->fields = std::move(fields);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.TuplePattern").set_body_typed([](tvm::Array<DFPattern> fields) {
+ return TuplePattern(fields);
+});
+RELAX_PATTERN_PRINTER_DEF(TuplePatternNode, [](auto p, auto node) {
+ p->stream << "TuplePattern(" << node->fields << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(UnorderedTuplePatternNode);
+UnorderedTuplePattern::UnorderedTuplePattern(tvm::Array<DFPattern> fields) {
+ ObjectPtr<UnorderedTuplePatternNode> n = make_object<UnorderedTuplePatternNode>();
+ n->fields = std::move(fields);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.UnorderedTuplePattern")
+ .set_body_typed([](tvm::Array<DFPattern> fields) { return UnorderedTuplePattern(fields); });
+RELAX_PATTERN_PRINTER_DEF(UnorderedTuplePatternNode, [](auto p, auto node) {
+ p->stream << "UnorderedTuplePattern(" << node->fields << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(TupleGetItemPatternNode);
+TupleGetItemPattern::TupleGetItemPattern(DFPattern tuple, int index) {
+ ObjectPtr<TupleGetItemPatternNode> n = make_object<TupleGetItemPatternNode>();
+ n->tuple = std::move(tuple);
+ n->index = index;
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.TupleGetItemPattern").set_body_typed([](DFPattern tuple, int index) {
+ return TupleGetItemPattern(tuple, index);
+});
+RELAX_PATTERN_PRINTER_DEF(TupleGetItemPatternNode, [](auto p, auto node) {
+ p->stream << "TupleGetItemPattern(" << node->tuple << ", " << node->index << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(AndPatternNode);
+AndPattern::AndPattern(DFPattern left, DFPattern right) {
+ ObjectPtr<AndPatternNode> n = make_object<AndPatternNode>();
+ n->left = std::move(left);
+ n->right = std::move(right);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.AndPattern").set_body_typed([](DFPattern left, DFPattern right) {
+ return AndPattern(left, right);
+});
+RELAX_PATTERN_PRINTER_DEF(AndPatternNode, [](auto p, auto node) {
+ p->stream << "AndPattern(" << node->left << " & " << node->right << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(OrPatternNode);
+OrPattern::OrPattern(DFPattern left, DFPattern right) {
+ ObjectPtr<OrPatternNode> n = make_object<OrPatternNode>();
+ n->left = std::move(left);
+ n->right = std::move(right);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.OrPattern").set_body_typed([](DFPattern left, DFPattern right) {
+ return OrPattern(left, right);
+});
+RELAX_PATTERN_PRINTER_DEF(OrPatternNode, [](auto p, auto node) {
+ p->stream << "OrPattern(" << node->left << " | " << node->right << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(NotPatternNode);
+NotPattern::NotPattern(DFPattern reject) {
+ ObjectPtr<NotPatternNode> n = make_object<NotPatternNode>();
+ n->reject = std::move(reject);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.NotPattern").set_body_typed([](DFPattern reject) {
+ return NotPattern(reject);
+});
+RELAX_PATTERN_PRINTER_DEF(NotPatternNode,
+ [](auto p, auto node) { p->stream << "!(" << node->reject << ")"; });
+
+TVM_REGISTER_NODE_TYPE(WildcardPatternNode);
+TVM_REGISTER_GLOBAL("relax.dpl.WildcardPattern").set_body_typed([]() {
+ auto w = WildcardPattern(make_object<WildcardPatternNode>());
+ return w;
+});
+RELAX_PATTERN_PRINTER_DEF(WildcardPatternNode, [](auto p, auto node) { p->stream << "*"; });
+
+TVM_REGISTER_NODE_TYPE(TypePatternNode);
+TypePattern::TypePattern(DFPattern pattern, Type type) {
+ ObjectPtr<TypePatternNode> n = make_object<TypePatternNode>();
+ n->pattern = std::move(pattern);
+ n->type = std::move(type);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.TypePattern").set_body_typed([](DFPattern pattern, Type type) {
+ return TypePattern(pattern, type);
+});
+RELAX_PATTERN_PRINTER_DEF(TypePatternNode, [](auto p, auto node) {
+ p->stream << "TypePattern(" << node->pattern << " has type " << node->type << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(ShapePatternNode);
+ShapePattern::ShapePattern(DFPattern pattern, Array<PrimExpr> shape) {
+ ObjectPtr<ShapePatternNode> n = make_object<ShapePatternNode>();
+ n->pattern = std::move(pattern);
+ n->shape = std::move(shape);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.ShapePattern")
+ .set_body_typed([](DFPattern pattern, Array<PrimExpr> shape) {
+ return ShapePattern(pattern, shape);
+ });
+RELAX_PATTERN_PRINTER_DEF(ShapePatternNode, [](auto p, auto node) {
+ p->stream << "ShapePattern(" << node->pattern << " has shape " << node->shape << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(DataTypePatternNode);
+DataTypePattern::DataTypePattern(DFPattern pattern, DataType dtype) {
+ ObjectPtr<DataTypePatternNode> n = make_object<DataTypePatternNode>();
+ n->pattern = std::move(pattern);
+ n->dtype = std::move(dtype);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.DataTypePattern")
+ .set_body_typed([](DFPattern pattern, DataType dtype) {
+ return DataTypePattern(pattern, dtype);
+ });
+RELAX_PATTERN_PRINTER_DEF(DataTypePatternNode, [](auto p, auto node) {
+ p->stream << "DataTypePattern(" << node->pattern << " has dtype " << node->dtype << ")";
+});
+
+TVM_REGISTER_NODE_TYPE(AttrPatternNode);
+AttrPattern::AttrPattern(DFPattern pattern, DictAttrs attrs) {
+ ObjectPtr<AttrPatternNode> n = make_object<AttrPatternNode>();
+ n->pattern = std::move(pattern);
+ n->attrs = std::move(attrs);
+ data_ = std::move(n);
+}
+TVM_REGISTER_GLOBAL("relax.dpl.AttrPattern").set_body_typed([](DFPattern pattern, DictAttrs attrs) {
+ return AttrPattern(pattern, attrs);
+});
+RELAX_PATTERN_PRINTER_DEF(AttrPatternNode, [](auto p, auto node) {
+ p->stream << "AttrPattern(" << node->pattern << " has attributes " << node->attrs << ")";
+});
+
+class DFPatternDuplicator : public DFPatternFunctor<DFPattern(const DFPattern&)> {
+ public:
+ DFPattern VisitDFPattern(const DFPattern& pattern) override {
+ return DFPatternFunctor::VisitDFPattern(pattern);
+ }
+ DFPattern VisitDFPattern_(const OrPatternNode* op) override {
+ return OrPattern(op->left, op->right);
+ }
+ DFPattern VisitDFPattern_(const AndPatternNode* op) override {
+ return AndPattern(op->left, op->right);
+ }
+ DFPattern VisitDFPattern_(const NotPatternNode* op) override { return NotPattern(op->reject); }
+ DFPattern VisitDFPattern_(const VarPatternNode* op) override { return VarPattern(op->name); }
+ DFPattern VisitDFPattern_(const ConstantPatternNode* op) override {
+ return ConstantPattern(make_object<ConstantPatternNode>());
+ }
+ DFPattern VisitDFPattern_(const WildcardPatternNode* op) override {
+ return WildcardPattern(make_object<WildcardPatternNode>());
+ }
+ DFPattern VisitDFPattern_(const ExprPatternNode* op) override { return ExprPattern(op->expr); }
+ DFPattern VisitDFPattern_(const GlobalVarPatternNode* op) override {
+ return GlobalVarPattern(op->name);
+ }
+ DFPattern VisitDFPattern_(const TuplePatternNode* op) override {
+ return TuplePattern(op->fields);
+ }
+ DFPattern VisitDFPattern_(const UnorderedTuplePatternNode* op) override {
+ return UnorderedTuplePattern(op->fields);
+ }
+ DFPattern VisitDFPattern_(const TupleGetItemPatternNode* op) override {
+ return TupleGetItemPattern(op->tuple, op->index);
+ }
+ DFPattern VisitDFPattern_(const CallPatternNode* op) override {
+ return CallPattern(op->op, op->args);
+ }
+ DFPattern VisitDFPattern_(const DataTypePatternNode* op) override {
+ return DataTypePattern(op->pattern, op->dtype);
+ }
+ DFPattern VisitDFPattern_(const FunctionPatternNode* op) override {
+ return FunctionPattern(op->params, op->body);
+ }
+ DFPattern VisitDFPattern_(const ShapePatternNode* op) override {
+ return ShapePattern(op->pattern, op->shape);
+ }
+ DFPattern VisitDFPattern_(const TypePatternNode* op) override {
+ return TypePattern(op->pattern, op->type);
+ }
+ DFPattern VisitDFPattern_(const DataflowVarPatternNode* op) override {
+ return DataflowVarPattern(op->name);
+ }
+ DFPattern VisitDFPattern_(const ExternFuncPatternNode* op) override {
+ return ExternFuncPattern(op->global_symbol());
+ }
+ DFPattern VisitDFPattern_(const PrimArrPatternNode* op) override {
+ return PrimArrPattern(op->fields);
+ }
+};
+
+// Syntatic Sugar
+CallPattern DFPattern::operator()(const std::vector<DFPattern>& args) const {
+ return CallPattern(*this, Array<DFPattern>(args));
+}
+OrPattern DFPattern::operator|(const DFPattern& other) const { return OrPattern(*this, other); }
+
+AndPattern DFPattern::operator&(const DFPattern& other) const { return AndPattern(*this, other); }
+
+NotPattern DFPattern::operator~() const { return NotPattern(*this); }
+
+AttrPattern DFPattern::HasAttr(const Map<String, ObjectRef>& attrs) const {
+ return AttrPattern(*this, DictAttrs(attrs));
+}
+TypePattern DFPattern::HasType(const Type& type) const { return TypePattern(*this, type); }
+DataTypePattern DFPattern::HasDtype(const DataType& dtype) const {
+ return DataTypePattern(*this, dtype);
+}
+DataTypePattern DFPattern::HasDtype(const std::string& dtype) const {
+ return HasDtype(DataType(runtime::String2DLDataType(dtype)));
+}
+ShapePattern DFPattern::HasShape(const Array<PrimExpr>& shape) const {
+ return ShapePattern(*this, shape);
+}
+
+DFPattern::operator PatternSeq() const { return PatternSeq{{*this}}; }
+
+std::stack<PatternContext>& pattern_ctx_stack() {
+ thread_local std::stack<PatternContext> graph_pattern_managers;
+ return graph_pattern_managers;
+}
+
+PatternContext PatternContext::Current() {
+ ICHECK(!pattern_ctx_stack().empty()) << "No active PatternContext found.";
+ return pattern_ctx_stack().top();
+}
+
+PatternContext::PatternContext(bool incremental) {
+ auto n = make_object<PatternContextNode>();
+ if (incremental) {
+ ICHECK(!pattern_ctx_stack().empty())
+ << "Incremental context needs to be built inside a existing context.";
+ n->allow_extern_use = pattern_ctx_stack().top()->allow_extern_use;
+ n->constraints = pattern_ctx_stack().top()->constraints;
+ }
+
+ data_ = std::move(n);
+}
+
+void PatternContext::EnterWithScope() { pattern_ctx_stack().push(*this); }
+
+void PatternContext::ExitWithScope() {
+ ICHECK(pattern_ctx_stack().top().same_as(*this));
+ pattern_ctx_stack().pop();
+}
+
+static void sync_graph_constraints(const DFPattern& lhs, const DFPattern& rhs, PairCons pcon) {
+ PatternContext::Current().add_constraint(lhs, rhs, pcon);
+}
+
+TVM_REGISTER_NODE_TYPE(PatternSeqNode);
+PatternSeq::PatternSeq(DFPattern init_pattern) {
+ ObjectPtr<PatternSeqNode> n = make_object<PatternSeqNode>();
+ n->patterns = {init_pattern};
+ n->pair_constraints = {};
+ data_ = std::move(n);
+}
+PatternSeq::PatternSeq(tvm::Array<DFPattern> patterns, bool only_used_by) {
+ ICHECK_GE(patterns.size(), 1) << "PatternSeq must have at least one pattern";
+ const auto cons = PairCons(only_used_by ? PairCons::kOnlyUsedBy : PairCons::kUsedBy);
+
+ ObjectPtr<PatternSeqNode> n = make_object<PatternSeqNode>();
+ n->patterns = std::move(patterns);
+ n->pair_constraints = std::vector<PairCons>(n->patterns.size() - 1, cons);
+ data_ = std::move(n);
+}
+
+PatternSeq PatternSeq::UsedBy(PatternSeq other, int index) const {
+ return relax::UsedBy(*this, other, index);
+}
+
+PatternSeq PatternSeq::OnlyUsedBy(PatternSeq other, int index) const {
+ return relax::OnlyUsedBy(*this, other, index);
+}
+
+PatternSeq PatternSeq::dup() const {
+ PatternSeq ret;
+
+ ObjectPtr<PatternSeqNode> n = make_object<PatternSeqNode>();
+ n->patterns = Array<DFPattern>{};
+ n->patterns.reserve(get()->patterns.size());
+ n->pair_constraints = this->get()->pair_constraints;
+
+ for (size_t i = 0; i < get()->patterns.size(); ++i) {
+ n->patterns.push_back(get()->patterns[i].dup());
+ if (i >= 1)
+ sync_graph_constraints(n->patterns[i - 1], n->patterns[i], n->pair_constraints[i - 1]);
+ }
+
+ ret.data_ = std::move(n);
+
+ return ret;
+}
+TVM_REGISTER_GLOBAL("relax.dpl.PatternSeq")
+ .set_body_typed([](Array<DFPattern> patterns, bool only_used_by) {
+ return PatternSeq(std::move(patterns), only_used_by);
+ });
+RELAX_PATTERN_PRINTER_DEF(PatternSeqNode, [](auto p, auto node) {
+ p->stream << "[";
+ for (size_t i = 0; i < node->patterns.size(); ++i) {
+ if (i != 0)
+ p->stream << (PairCons::kOnlyUsedBy == node->pair_constraints[i].type ? " >> " : " ^ ");
+ p->stream << node->patterns[i];
+ }
+ p->stream << "]";
+});
+
+TVM_REGISTER_GLOBAL("relax.dpl.used_by")
+ .set_body_typed([](PatternSeq lhs, PatternSeq rhs, int index) {
+ return lhs.UsedBy(rhs, index);
+ });
+
+TVM_REGISTER_GLOBAL("relax.dpl.only_used_by")
+ .set_body_typed([](PatternSeq lhs, PatternSeq rhs, int index) {
+ return lhs.OnlyUsedBy(rhs, index);
+ });
+
+PatternSeq UsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index) {
+ PatternSeq ret;
+
+ const auto constraint = PairCons{PairCons::kOnlyUsedBy, index};
+
+ sync_graph_constraints(lhs->patterns.back(), rhs->patterns.front(),
+ PairCons{PairCons::kUsedBy, index});
+
+ Array<DFPattern> patterns;
+ patterns.reserve(lhs->patterns.size() + rhs->patterns.size());
+ patterns.insert(patterns.end(), lhs->patterns.begin(), lhs->patterns.end());
+ patterns.insert(patterns.end(), rhs->patterns.begin(), rhs->patterns.end());
+
+ std::vector<PairCons> pair_constraints = lhs->pair_constraints;
+ pair_constraints.reserve(pair_constraints.size() + rhs->pair_constraints.size() + 1);
+ pair_constraints.push_back(constraint);
+ pair_constraints.insert(pair_constraints.end(), rhs->pair_constraints.begin(),
+ rhs->pair_constraints.end());
+
+ ObjectPtr<PatternSeqNode> n = make_object<PatternSeqNode>();
+ n->patterns = std::move(patterns);
+ n->pair_constraints = std::move(pair_constraints);
+ ret.data_ = std::move(n);
+
+ return ret;
+}
+PatternSeq operator^(const PatternSeq& lhs, const PatternSeq& rhs) { return lhs.UsedBy(rhs); }
+
+PatternSeq OnlyUsedBy(const PatternSeq& lhs, const PatternSeq& rhs, int index) {
+ PatternSeq ret;
+
+ const auto constraint = PairCons{PairCons::kOnlyUsedBy, index};
+
+ sync_graph_constraints(lhs->patterns.back(), rhs->patterns.front(), constraint);
+
+ Array<DFPattern> patterns;
+ patterns.reserve(lhs->patterns.size() + rhs->patterns.size());
+ patterns.insert(patterns.end(), lhs->patterns.begin(), lhs->patterns.end());
+ patterns.insert(patterns.end(), rhs->patterns.begin(), rhs->patterns.end());
+
+ std::vector<PairCons> pair_constraints = lhs->pair_constraints;
+ pair_constraints.reserve(pair_constraints.size() + rhs->pair_constraints.size() + 1);
+ pair_constraints.push_back(constraint);
+ pair_constraints.insert(pair_constraints.end(), rhs->pair_constraints.begin(),
+ rhs->pair_constraints.end());
+
+ ObjectPtr<PatternSeqNode> n = make_object<PatternSeqNode>();
+ n->patterns = std::move(patterns);
+ n->pair_constraints = std::move(pair_constraints);
+ ret.data_ = std::move(n);
+
+ return ret;
+}
+PatternSeq operator>>(const PatternSeq& lhs, const PatternSeq& rhs) { return lhs.OnlyUsedBy(rhs); }
+
+VarPattern IsVar(const String& name) { return VarPattern(name); }
+ConstantPattern IsConst() { return ConstantPattern(make_object<ConstantPatternNode>()); }
+WildcardPattern Wildcard() { return WildcardPattern(make_object<WildcardPatternNode>()); }
+ExprPattern IsExpr(const Expr& expr) { return ExprPattern(expr); }
+ExprPattern IsOp(const String& op_name) { return IsExpr(Op::Get(op_name)); }
+CallPattern IsCallTIR(const String& name, Optional<TuplePattern> var_args) {
+ DFPattern arg_pattern;
+ if (!var_args.defined()) {
+ arg_pattern = Wildcard();
+ } else {
+ arg_pattern = var_args.value();
+ }
+
+ return IsOp("relax.call_tir")(GlobalVarPattern(name), arg_pattern);
+}
+
+CallPattern IsCallTIR(const String& name, TuplePattern var_args) {
+ return IsOp("relax.call_tir")(GlobalVarPattern(name), var_args);
+}
+
+DFPattern IsTuple(const Array<DFPattern>& fields, bool unordered) {
+ if (unordered)
+ return UnorderedTuplePattern(fields);
+ else
+ return TuplePattern(fields);
+}
+TupleGetItemPattern IsTupleGetItem(const DFPattern tuple, int index) {
+ return TupleGetItemPattern(tuple, index);
+}
+
+DFPattern DFPattern::dup() const {
+ auto pattern = DFPatternDuplicator().VisitDFPattern(*this);
+ return pattern;
+}
+
+TVM_REGISTER_GLOBAL("relax.dpl.dup_pattern").set_body_typed([](DFPattern pattern) {
+ return pattern.dup();
+});
+
+TVM_REGISTER_GLOBAL("relax.dpl.dup_seq").set_body_typed([](PatternSeq seq) { return seq.dup(); });
+
+TVM_REGISTER_GLOBAL("relax.dpl.PatternContext").set_body_typed([](bool incre) {
+ return PatternContext(incre);
+});
+
+TVM_REGISTER_GLOBAL("relax.dpl.current_context").set_body_typed([] {
+ return PatternContext::Current();
+});
+
+class PatternContext::Internal {
+ public:
+ static void EnterScope(PatternContext pass_ctx) { pass_ctx.EnterWithScope(); }
+ static void ExitScope(PatternContext pass_ctx) { pass_ctx.ExitWithScope(); }
+};
+
+TVM_REGISTER_GLOBAL("relax.dpl.enter_context").set_body_typed(PatternContext::Internal::EnterScope);
+
+TVM_REGISTER_GLOBAL("relax.dpl.exit_context").set_body_typed(PatternContext::Internal::ExitScope);
+
+} // namespace relax
+} // namespace tvm
diff --git a/src/relax/ir/dataflow_pattern_functor.cc b/src/relax/ir/dataflow_pattern_functor.cc
new file mode 100644
index 0000000000..37a98f28be
--- /dev/null
+++ b/src/relax/ir/dataflow_pattern_functor.cc
@@ -0,0 +1,111 @@
+/*
+ * 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 src/tvm/relay/dataflow_matcher.cc
+ * \brief The dataflow pattern matcher for Relay.
+ */
+
+#include <tvm/relax/dataflow_pattern_functor.h>
+
+namespace tvm {
+namespace relax {
+
+// DFPatternVisitor
+
+void DFPatternVisitor::VisitDFPattern(const DFPattern& pattern) {
+ if (this->visited_.count(pattern.get()) == 0) {
+ visited_.insert(pattern.get());
+ DFPatternFunctor::VisitDFPattern(pattern);
+ }
+}
+
+void DFPatternVisitor::VisitDFPattern_(const OrPatternNode* op) {
+ VisitDFPattern(op->left);
+ VisitDFPattern(op->right);
+}
+
+void DFPatternVisitor::VisitDFPattern_(const AndPatternNode* op) {
+ VisitDFPattern(op->left);
+ VisitDFPattern(op->right);
+}
+
+void DFPatternVisitor::VisitDFPattern_(const NotPatternNode* op) { VisitDFPattern(op->reject); }
+
+void DFPatternVisitor::VisitDFPattern_(const AttrPatternNode* op) { VisitDFPattern(op->pattern); }
+
+void DFPatternVisitor::VisitDFPattern_(const CallPatternNode* op) {
+ VisitDFPattern(op->op);
+ if (op->args.defined()) {
+ for (auto arg : op->args) {
+ VisitDFPattern(arg);
+ }
+ }
+}
+
+void DFPatternVisitor::VisitDFPattern_(const DataTypePatternNode* op) {
+ VisitDFPattern(op->pattern);
+}
+
+void DFPatternVisitor::VisitDFPattern_(const ExprPatternNode* op) {}
+
+void DFPatternVisitor::VisitDFPattern_(const FunctionPatternNode* op) {
+ if (op->params.defined()) {
+ for (auto param : op->params) {
+ VisitDFPattern(param);
+ }
+ }
+ VisitDFPattern(op->body);
+}
+
+void DFPatternVisitor::VisitDFPattern_(const ShapePatternNode* op) { VisitDFPattern(op->pattern); }
+
+void DFPatternVisitor::VisitDFPattern_(const TupleGetItemPatternNode* op) {
+ VisitDFPattern(op->tuple);
+}
+
+void DFPatternVisitor::VisitDFPattern_(const TuplePatternNode* op) {
+ if (op->fields.defined()) {
+ for (auto field : op->fields) {
+ VisitDFPattern(field);
+ }
+ }
+}
+
+void DFPatternVisitor::VisitDFPattern_(const UnorderedTuplePatternNode* op) {
+ if (op->fields.defined()) {
+ for (auto field : op->fields) {
+ VisitDFPattern(field);
+ }
+ }
+}
+
+void DFPatternVisitor::VisitDFPattern_(const TypePatternNode* op) { VisitDFPattern(op->pattern); }
+
+// leaf nodes.
+void DFPatternVisitor::VisitDFPattern_(const PrimArrPatternNode* op) {}
+void DFPatternVisitor::VisitDFPattern_(const VarPatternNode* op) {}
+void DFPatternVisitor::VisitDFPattern_(const ConstantPatternNode* op) {}
+void DFPatternVisitor::VisitDFPattern_(const DataflowVarPatternNode* op) {}
+void DFPatternVisitor::VisitDFPattern_(const GlobalVarPatternNode* op) {}
+void DFPatternVisitor::VisitDFPattern_(const ExternFuncPatternNode* op) {}
+void DFPatternVisitor::VisitDFPattern_(const WildcardPatternNode* op) {}
+
+} // namespace relax
+} // namespace tvm
diff --git a/tests/python/relax/test_analysis.py b/tests/python/relax/test_analysis.py
index 5dd83f2da2..43558a52be 100644
--- a/tests/python/relax/test_analysis.py
+++ b/tests/python/relax/test_analysis.py
@@ -21,10 +21,31 @@ import tvm
import tvm.testing
from tvm import tir
from tvm import relax as rx
-from tvm.relax.analysis import has_reshape_pattern
+from tvm.relax.analysis import has_reshape_pattern, udchain
from tvm.script import relax as R, tir as T
+def test_use_def():
+ m = tir.Var("m", "int64")
+ n = tir.Var("n", "int64")
+ x = rx.Var("x", R.Tensor([m, n], "float16"))
+ y = rx.Var("y", R.Tensor([n], "float16"))
+ ib = rx.BlockBuilder()
+ with ib.function("func", [x, y]):
+ with ib.dataflow():
+ lv0 = ib.emit(rx.op.add(x, y))
+ lv1 = ib.emit(rx.op.multiply(lv0, y))
+ gv0 = ib.emit_output(lv1)
+ ib.emit_func_output(gv0)
+ dfb = ib.get()["func"].body.blocks[0]
+ udc = udchain(dfb)
+ assert set(udc[x]) == {lv0}
+ assert set(udc[y]) == {lv0, lv1}
+ assert set(udc[lv0]) == {lv1}
+ assert set(udc[lv1]) == {gv0}
+ assert set(udc[gv0]) == set()
+
+
def test_reshape_pattern_reshape():
@T.prim_func
def reshape(
diff --git a/tests/python/relax/test_dataflow_pattern.py b/tests/python/relax/test_dataflow_pattern.py
new file mode 100644
index 0000000000..ab7a5540ad
--- /dev/null
+++ b/tests/python/relax/test_dataflow_pattern.py
@@ -0,0 +1,867 @@
+# 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.
+
+import pytest
+import tvm.testing
+
+from tvm import relay
+from tvm.relax.dpl import *
+from tvm.relax.analysis import get_var2val
+from tvm import relax as rx, tir
+from tvm.script import relax as R, tir as T
+
+
+@tvm.script.ir_module
+class Module:
+ @T.prim_func
+ def tir_matmul(x: T.handle, y: T.handle, z: T.handle) -> None:
+ T.func_attr({"global_symbol": "tir_matmul"})
+ k = T.int32()
+ A = T.match_buffer(x, (32, 32))
+ B = T.match_buffer(y, (32, 32))
+ C = T.match_buffer(z, (32, 32))
+
+ for (i0, j0, k0) in T.grid(32, 32, 32):
+ with T.block():
+ i, j, k = T.axis.remap("SSR", [i0, j0, k0])
+ with T.init():
+ C[i, j] = 0.0
+ C[i, j] += A[i, k] * B[j, k]
+
+ @T.prim_func
+ def tir_relu(x: T.handle, y: T.handle):
+ T.func_attr({"global_symbol": "tir_relu"})
+ A = T.match_buffer(x, (32, 32))
+ B = T.match_buffer(y, (32, 32))
+ for (i, j) in T.grid(32, 32):
+ with T.block():
+ vi, vj = T.axis.remap("SS", [i, j])
+ B[vi, vj] = T.max(A[vi, vj], 0.0)
+
+ @R.function
+ def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ with R.dataflow():
+ lv0 = R.call_tir(tir_matmul, (x, w), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir(tir_relu, (lv0), R.Tensor((32, 32), dtype="float32"))
+ R.output(lv1)
+ return lv1
+
+
+main_fn = Module["main"]
+bindings = main_fn.body.blocks[0].bindings
+
+## Node-wise Matching
+def test_expr_pattern():
+ ep = is_expr(rx.Var("x"))
+ assert isinstance(ep, ExprPattern)
+ assert isinstance(ep.expr, rx.Var)
+
+
+def test_var_pattern():
+ v = is_var("x")
+ assert isinstance(v, VarPattern)
+ assert v.name == "x"
+ assert v.match(rx.Var("x"))
+ assert is_var().match(rx.Var("x"))
+ assert is_var().match(rx.DataflowVar("x")) # DataflowVar is also a Var
+ assert not v.match(rx.GlobalVar("x"))
+
+
+def test_dataflow_var_pattern():
+ v = is_dfv("x")
+ assert isinstance(v, DataflowVarPattern)
+ assert v.name == "x"
+ assert v.match(rx.DataflowVar("x"))
+ assert not v.match(rx.GlobalVar("x"))
+ assert is_dfv().match(bindings[0].var)
+
+
+def test_global_var_pattern():
+ assert is_gv("x").match(rx.GlobalVar("x"))
+ assert is_gv().match(rx.GlobalVar("x"))
+ assert not is_gv("x").match(rx.GlobalVar("y"))
+ assert not is_gv("x").match(rx.Var("x"))
+
+
+def test_constant_pattern():
+ c = is_const()
+ assert isinstance(c, ConstantPattern)
+ assert c.match(rx.const([[0.1, 1.1, 2.1], [3.1, 4.1, 5.1]]))
+
+
+def test_wildcard_pattern():
+ wc = wildcard()
+ assert isinstance(wc, WildcardPattern)
+ assert wc.match(rx.Var("x"))
+
+
+def test_call_pattern():
+ wc1 = wildcard()
+ wc2 = wildcard()
+ c = is_op("relax.add")(wc1, wc2)
+ assert isinstance(c, CallPattern)
+ assert isinstance(c.args[0], WildcardPattern)
+ assert isinstance(c.args[1], WildcardPattern)
+ assert c.match(rx.op.add(rx.Var("x"), rx.Var("y")))
+
+
+def test_function_pattern():
+ wc1 = wildcard()
+ wc2 = wildcard()
+ f = FunctionPattern([wc1, wc2], is_op("relax.add")(wc1, wc2))
+ assert isinstance(f, FunctionPattern)
+ assert isinstance(f.params[0], WildcardPattern)
+ assert isinstance(f.params[1], WildcardPattern)
+ assert isinstance(f.body, CallPattern)
+ assert isinstance(f.body.args[0], WildcardPattern)
+ assert isinstance(f.body.args[1], WildcardPattern)
+ x = rx.Var("x", R.Tensor("float32"))
+ y = rx.Var("y", R.Tensor("float32"))
+ assert f.match(rx.Function([x, y], rx.op.add(x, y), ret_struct_info=R.Tensor("float32")))
+ assert not f.match(
+ rx.Function([x, y], rx.op.multiply(x, y), ret_struct_info=R.Tensor("float32"))
+ )
+
+
+def test_tuple_pattern():
+ wc1 = wildcard()
+ wc2 = is_dfv()
+ t = is_tuple([wc1, wc2])
+ assert isinstance(t, TuplePattern)
+ assert isinstance(t.fields[0], WildcardPattern)
+ assert isinstance(t.fields[1], DataflowVarPattern)
+ assert t.match(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]))
+ assert not t.match(rx.Tuple([rx.DataflowVar("x"), rx.GlobalVar("y")]))
+ assert not t.match(rx.Tuple([]))
+ assert t[0].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 0))
+ assert t[1].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 1))
+ # Negative index is also allowed
+ assert t[-1].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 1))
+ # None means any index.
+ assert t[None].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 0))
+ assert t[None].match(rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 1))
+ with pytest.raises(IndexError):
+ t[2] # index cannot be greater than or equal to the tuple size.
+
+
+def test_unordered_tuple_pattern():
+ t = is_tuple([is_const(), is_dfv()], unordered=True)
+ assert isinstance(t, UnorderedTuplePattern)
+ assert isinstance(t.fields[0], ConstantPattern)
+ assert isinstance(t.fields[1], DataflowVarPattern)
+ assert t.match(rx.Tuple([rx.const([]), rx.DataflowVar("x")]))
+ assert t.match(rx.Tuple([rx.DataflowVar("x"), rx.const([])]))
+ assert not t.match(rx.Tuple([rx.DataflowVar("x"), rx.DataflowVar("y")]))
+ assert not t.match(rx.Tuple([]))
+
+
+def test_tuple_get_item_pattern():
+ assert is_tuple_get_item(is_tuple([is_gv("x"), is_dfv("y")]), 0).match(
+ rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 0)
+ )
+ assert is_tuple_get_item(is_tuple([is_gv("x"), is_dfv("y")]), 0).match(
+ rx.TupleGetItem(rx.Tuple([rx.GlobalVar("x"), rx.DataflowVar("y")]), 0)
+ )
+
+
+def test_or_pattern():
+ dfv_or_gv = is_dfv("x") | is_gv("x")
+ assert isinstance(dfv_or_gv, OrPattern)
+ assert dfv_or_gv.match(rx.DataflowVar("x"))
+ assert dfv_or_gv.match(rx.GlobalVar("x"))
+ assert not dfv_or_gv.match(rx.Var("x"))
+ assert not dfv_or_gv.match(rx.DataflowVar("y"))
+ assert not dfv_or_gv.match(rx.GlobalVar("y"))
+
+
+def test_and_pattern():
+ # float[2, 3, 3]
+ f32_233 = wildcard().has_shape((2, 3, 3)) & has_dtype("float32")
+ assert isinstance(f32_233, AndPattern)
+ assert f32_233.match(rx.Var("x", R.Tensor((2, 3, 3), "float32")))
+ assert not f32_233.match(rx.Var("x", R.Tensor((3, 3, 3), "float32")))
+ assert not f32_233.match(rx.Var("x", R.Tensor("float32", ndim=3)))
+
+
+def test_not_pattern():
+ no_shape233 = ~wildcard().has_shape((2, 3, 3))
+ assert isinstance(no_shape233, NotPattern)
+ assert no_shape233.match(rx.Var("x", R.Tensor((3, 3, 3), "float32")))
+ assert not no_shape233.match(rx.Var("x", R.Tensor((2, 3, 3), "float32")))
+
+
+def test_type_pattern():
+ assert wildcard().has_type(rx.DynTensorType(2, "float32")).match(bindings[0].var)
+
+
+def test_dtype_pattern():
+ dtype = "float16"
+ pattern = has_dtype(dtype)
+ assert isinstance(pattern, DataTypePattern)
+ assert pattern.dtype == dtype
+ assert has_dtype("float32").match(bindings[0].var)
+
+
+def test_shape_pattern():
+ shape = [32, 32]
+ pattern = wildcard().has_shape(shape)
+ assert isinstance(pattern, ShapePattern)
+ tvm.ir.structural_equal(pattern.shape, shape)
+ assert pattern.match(bindings[0].var)
+ assert wildcard().has_shape([32, 32]).match(bindings[0].var)
+ n, m = tir.Var("n", dtype="int64"), tir.Var("m", dtype="int64")
+ symsh_var = rx.Var("x", R.Tensor([n, m, n + m], "float32"))
+ assert wildcard().has_shape([n, m, n + m]).match(symsh_var)
+ assert wildcard().has_shape([n, m, m + n]).match(symsh_var) # + is commutative.
+ assert not wildcard().has_shape([1, 2, 3]).match(symsh_var)
+ assert not wildcard().has_shape([m, n, n + m]).match(symsh_var)
+
+
+def test_prim_arr_pattern():
+ """
+ The difference between is_shape and has_shape is that:
+ 1) is_shape directly matches a shape (e.g., as an argument);
+ 2) has_shape matches a tensor and puts assumptions on the tensor's shape.
+ """
+ pattern = is_shape([32, 32])
+ assert pattern[0] == 32
+ assert pattern[1] == 32
+ assert isinstance(pattern, PrimArrPattern)
+ assert pattern.match(rx.get_shape_of(bindings[0].var))
+ n, m = tir.Var("n", dtype="int64"), tir.Var("m", dtype="int64")
+ symbolic_shape = rx.ShapeExpr([n, m, n + m])
+ assert is_shape([n, m, n + m]).match(symbolic_shape)
+ assert not is_shape([n, m, n * m]).match(symbolic_shape)
+
+
+def test_extern_fn_pattern():
+ pattern = ExternFuncPattern("test.blockbuilder.nop")
+ assert pattern.match(rx.ExternFunc("test.blockbuilder.nop"))
+
+
+def test_op_attr():
+ x = rx.Var("x", R.Tensor("float32"))
+ y = rx.Var("y", R.Tensor("float32"))
+ conv2d = relay.nn.conv2d(x, y, kernel_size=(3, 3))
+ xp = is_var("x")
+ yp = is_var("y")
+ # TODO(@yuchen): reenable the assert after figuring out why it fails
+ # assert is_op("nn.conv2d")(xp, yp).has_attr({"kernel_size": [3, 3]}).match(conv2d)
+ assert not is_op("nn.conv2d")(xp, yp).has_attr({"kernel_size": [4, 3]}).match(conv2d)
+ assert not is_op("nn.conv2d")(xp, yp).has_attr({"kernel_size_": [3, 3]}).match(conv2d)
+
+
+def test_match_call_attr():
+ x = rx.Var("x", R.Tensor("float32"))
+ y = rx.Var("y", R.Tensor("float32"))
+ fn = rx.Function([x, y], rx.op.add(x, y), ret_struct_info=R.Tensor("float32"))
+ annotated_fn = fn.with_attr({"Codegen": "test-codegen", "global_symbol": "test-symbol"})
+ xp = is_var("x")
+ yp = is_var("y")
+ root_pattern = FunctionPattern([xp, yp], is_op("relax.add")(xp, yp))
+ assert root_pattern.has_attr({"Codegen": "test-codegen", "global_symbol": "test-symbol"}).match(
+ annotated_fn
+ )
+
+ assert root_pattern.has_attr({"Codegen": "test-codegen"}).match(annotated_fn)
+ assert not root_pattern.has_attr({"ping": "pong"}).match(annotated_fn)
+ assert root_pattern.has_attr({}).match(annotated_fn)
+
+
+def test_is_call_tir():
+ lv1_val = bindings[1].value
+ var2val = get_var2val(Module["main"])
+ assert is_call_tir("tir_relu").match(lv1_val)
+ assert is_call_tir("tir_relu", [is_call_tir("tir_matmul")]).match(lv1_val, var2val=var2val)
+ assert not is_call_tir("tir_relu", [is_call_tir("tir_relu")]).match(lv1_val, var2val=var2val)
+
+
+@R.function
+def simple_call_packed(
+ x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")
+) -> R.Tensor:
+ gv0 = R.call_packed("test.vm.mul", x, w, sinfo_args=(R.Tensor(ndim=2, dtype="float32")))
+ return gv0
+
+
+def test_varg_default_wildcard():
+ expr = simple_call_packed.body.blocks[0].bindings[0].value
+ yes_pattern_explicit = ExternFuncPattern("test.vm.mul")(wildcard(), wildcard())
+ yes_pattern_implicit = ExternFuncPattern("test.vm.mul")(varg_default_wildcard=True)
+ no_pattern = ExternFuncPattern("test.vm.mul")(wildcard())
+
+ assert yes_pattern_explicit.match(expr)
+ assert yes_pattern_implicit.match(expr)
+ assert not no_pattern.match(expr)
+
+
+def test_simple_call_packed():
+ expr = simple_call_packed.body.blocks[0].bindings[0].value
+ assert is_call_packed("test.vm.mul").match(expr)
+ assert is_call_packed("test.vm.mul", [is_var("x"), is_var("w")]).match(expr)
+
+
+## Graph-wise Matching
+def test_simple_used_by():
+ with PatternContext() as ctx:
+ n0 = is_var("x") # x is a free var (fn arg)
+ n1 = wildcard()
+ n0 ^ n1
+ dfb = main_fn.body.blocks[0]
+ matched = ctx.match_dfb(dfb)
+ assert matched
+ assert matched[n0] == main_fn.params[0]
+ assert matched[n1] == dfb.bindings[0].var
+
+
+def test_simple_call_tir_edge():
+ with PatternContext() as ctx:
+ n0 = is_call_tir("tir_matmul")
+ n1 = is_call_tir("tir_relu")
+ n0.used_by(n1)
+ dfb = main_fn.body.blocks[0]
+ matched = ctx.match_dfb(dfb)
+ assert matched
+ assert matched[n0] == dfb.bindings[0].var
+ assert matched[n1] == dfb.bindings[1].var
+
+
+def test_simple_oub():
+ with PatternContext() as ctx:
+ n0 = is_call_tir("tir_matmul")
+ n1 = is_call_tir("tir_relu")
+ n0 >> n1
+ dfb = main_fn.body.blocks[0]
+ matched = ctx.match_dfb(dfb)
+ assert matched
+ assert matched[n0] == dfb.bindings[0].var
+ assert matched[n1] == dfb.bindings[1].var
+
+
+def test_counter_syntax_match():
+ with PatternContext() as ctx:
+ n0 = is_call_tir_extern("tir_matmul")
+ n1 = is_call_tir_extern("tir_impossible")
+ n0 >> n1
+ dfb = main_fn.body.blocks[0]
+ assert not ctx.match_dfb(dfb)
+
+ with PatternContext() as ctx:
+ n0 = is_call_tir_extern("tir_matmul")
+ n1 = is_call_tir_extern("tir_impossible")
+ n0 ^ n1
+ dfb = main_fn.body.blocks[0]
+ assert not ctx.match_dfb(dfb)
+
+
+@tvm.script.ir_module
+class Diamond:
+ @R.function
+ def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ with R.dataflow():
+ # matmul
+ # / \
+ # relu sigmoid
+ # \ /
+ # add
+ lv0 = R.call_tir("tir_matmul", (x, w), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir("tir_relu", (lv0,), R.Tensor((32, 32), dtype="float32"))
+ lv2 = R.call_tir("tir_sigmoid", (lv0), R.Tensor((32, 32), dtype="float32"))
+ lv3 = R.call_tir("tir_add", (lv1, lv2), R.Tensor((32, 32), dtype="float32"))
+ R.output(lv3)
+ return lv3
+
+
+def test_diamond():
+ with PatternContext() as ctx:
+ n0 = is_call_tir_extern("tir_matmul")
+ n1 = is_call_tir_extern("tir_relu")
+ n2 = is_call_tir_extern("tir_sigmoid")
+ n3 = is_call_tir_extern("tir_add")
+
+ n0 ^ n1
+ n0 ^ n2
+ n1 >> n3
+ n2 >> n3
+
+ dfb = Diamond["main"].body.blocks[0]
+ assert ctx.match_dfb(dfb)
+
+ # simplify it with fork_to
+ with PatternContext() as ctx:
+ n1 = is_call_tir_extern("tir_relu")
+ n2 = is_call_tir_extern("tir_sigmoid")
+ n3 = is_call_tir_extern("tir_add")
+
+ is_call_tir_extern("tir_matmul").fork_to(n1, n2)
+ n1 >> n3
+ n2 >> n3
+
+ dfb = Diamond["main"].body.blocks[0]
+ assert ctx.match_dfb(dfb)
+
+
+def test_diamond_counter_oub():
+ with PatternContext() as ctx:
+ n0 = is_call_tir_extern("tir_matmul")
+ n1 = is_call_tir_extern("tir_relu")
+ n2 = is_call_tir_extern("tir_sigmoid")
+ n3 = is_call_tir_extern("tir_add")
+
+ n0 >> n1
+ n0 >> n2
+ n1 >> n3
+ n2 >> n3
+
+ dfb = Diamond["main"].body.blocks[0]
+ assert not ctx.match_dfb(dfb)
+
+
+@tvm.script.ir_module
+class SmallDiamond:
+ @R.function
+ def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ with R.dataflow():
+ # relu
+ # / \
+ # \ /
+ # add
+ lv0 = R.call_tir("my_relu", (x,), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir("my_add", (lv0, lv0), R.Tensor((32, 32), dtype="float32"))
+ R.output(lv1)
+ return lv1
+
+
+@tvm.script.ir_module
+class SmallParallel:
+ @R.function
+ def main(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ with R.dataflow():
+ # relu relu
+ # \ /
+ # add
+ lv0 = R.call_tir("my_relu", (x,), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir("my_relu", (x,), R.Tensor((32, 32), dtype="float32"))
+ lv2 = R.call_tir("my_add", (lv0, lv1), R.Tensor((32, 32), dtype="float32"))
+ R.output(lv2)
+ return lv2
+
+
+def test_distiguish_diamond_and_parallel():
+ # relay pattern lang cannot distinguish the two cases above.
+ diamond = SmallDiamond["main"].body.blocks[0]
+ parallel = SmallParallel["main"].body.blocks[0]
+
+ with PatternContext() as ctx:
+ # describe a diamond pattern
+ fork = is_call_tir_extern("my_relu")
+ join = is_call_tir_extern("my_add")
+ fork.only_used_by(join, index=0)
+ fork.only_used_by(join, index=1)
+
+ assert ctx.match_dfb(diamond)
+ assert not ctx.match_dfb(parallel)
+
+ with PatternContext() as ctx:
+ # describe a parallel pattern
+ join = is_call_tir_extern("my_add")
+ # Due to one-one mathcing:
+ # is_call_tir_extern("my_relu") creates the 1st relu
+ is_call_tir_extern("my_relu") >> join
+ # is_call_tir_extern("my_relu")
+ # creates the another different relu (obj address is different)
+ is_call_tir_extern("my_relu") >> join
+
+ assert ctx.match_dfb(parallel)
+ assert not ctx.match_dfb(diamond)
+
+
+@tvm.script.ir_module
+class CBRx2:
+ @R.function
+ def main(
+ x: R.Tensor((32, 32), "float32"),
+ w0: R.Tensor((1, 1), "float32"),
+ bias0: R.Tensor((32, 32), "float32"),
+ w1: R.Tensor((1, 1), "float32"),
+ bias1: R.Tensor((32, 32), "float32"),
+ ) -> R.Tensor:
+ # R.TensorRT's CBR Optimization Pattern
+ # input
+ # / \
+ # cbr0 cbr1
+ # \ /
+ # concat
+ with R.dataflow():
+ lv0 = R.call_tir("conv1x1", (x, w0), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir("bias_add", (lv0, bias0), R.Tensor((32, 32), dtype="float32"))
+ lv2 = R.call_tir("my_relu", (lv1), R.Tensor((32, 32), dtype="float32"))
+ lv3 = R.call_tir("conv1x1", (x, w1), R.Tensor((32, 32), dtype="float32"))
+ lv4 = R.call_tir("bias_add", (lv3, bias1), R.Tensor((32, 32), dtype="float32"))
+ lv5 = R.call_tir("my_relu", (lv4), R.Tensor((32, 32), dtype="float32"))
+ lv6 = R.call_tir("concat", (lv2, lv5), R.Tensor((32, 64), dtype="float32"))
+ R.output(lv6)
+ return lv6
+
+
+def test_single_cbr():
+ with PatternContext() as ctx:
+ (
+ is_call_tir_extern("conv1x1")
+ >> is_call_tir_extern("bias_add")
+ >> is_call_tir_extern("my_relu")
+ )
+ dfb = CBRx2["main"].body.blocks[0]
+ matched = ctx.match_dfb(dfb)
+ assert matched
+
+ with PatternContext() as ctx:
+ chain = (
+ is_call_tir_extern("conv1x1")
+ >> is_call_tir_extern("bias_add")
+ >> is_call_tir_extern("my_relu")
+ )
+ dfb = CBRx2["main"].body.blocks[0]
+ # we want to specifically match the first CBR (lv0)
+ matched = ctx.match_dfb(dfb, start_hint=dfb.bindings[0].var)
+ assert matched
+ assert matched[chain[0]] == dfb.bindings[0].var
+ # we want to specifically match the second CBR (lv3)
+ matched = ctx.match_dfb(dfb, start_hint=dfb.bindings[3].var)
+ assert matched
+ assert matched[chain[0]] == dfb.bindings[3].var
+
+
+def test_counter_single_crb():
+ with PatternContext() as ctx:
+ (
+ is_call_tir_extern("conv1x1")
+ >> is_call_tir_extern("my_relu")
+ >> is_call_tir_extern("bias_add")
+ )
+ dfb = CBRx2["main"].body.blocks[0]
+ assert not ctx.match_dfb(dfb)
+ # Quickly fails unpromising matches by assumiung `start_hint` must be matched by a pattern.
+ # This is usually faster than the full match:
+ # Full match: let one pattern to match -> all Var: complexity ~ #Var
+ # must_include_hint: let `start_hint` to match -> all patterns: complexity ~ #patterns
+ # Usually #patterns is much smaller than #Var, so this is faster.
+ assert not ctx.match_dfb(dfb, start_hint=dfb.bindings[0].var, must_include_hint=True)
+
+
+def test_nested_context():
+ dfb = CBRx2["main"].body.blocks[0]
+ with PatternContext() as ctx0:
+ (
+ is_call_tir_extern("conv1x1")
+ >> is_call_tir_extern("bias_add")
+ >> is_call_tir_extern("my_relu")
+ )
+ with PatternContext() as ctx1:
+ is_call_tir_extern("conv1x1") >> is_call_tir_extern("my_relu") # pattern to miss
+ with PatternContext() as ctx2:
+ is_call_tir_extern("bias_add") >> is_call_tir_extern("my_relu")
+ assert ctx2.match_dfb(dfb)
+ assert PatternContext.current() == ctx2
+ assert not ctx1.match_dfb(dfb)
+ assert PatternContext.current() == ctx1
+ assert ctx0.match_dfb(dfb)
+ assert PatternContext.current() == ctx0
+
+
+def test_two_cbr():
+ with PatternContext() as ctx:
+ cbr0 = (
+ is_call_tir_extern("conv1x1")
+ >> is_call_tir_extern("bias_add")
+ >> is_call_tir_extern("my_relu")
+ )
+ cbr1 = cbr0.dup()
+
+ assert cbr0.patterns[0] != cbr1.patterns[0]
+ assert cbr0.patterns[1] != cbr1.patterns[1]
+ assert cbr0.patterns[2] != cbr1.patterns[2]
+
+ is_var("x").fork_to(cbr0, cbr1)
+ dfb = CBRx2["main"].body.blocks[0]
+ assert ctx.match_dfb(dfb)
+
+ with PatternContext() as ctx:
+ # Deny the pattern
+ cbr0 = (
+ is_call_tir_extern("conv1x1")
+ >> is_call_tir_extern("bias_add")
+ >> is_call_tir_extern("my_relu")
+ )
+ cbr1 = cbr0.dup()
+
+ # input has no fork at y.
+ is_var("y").fork_to(cbr0, cbr1)
+ dfb = CBRx2["main"].body.blocks[0]
+ assert not ctx.match_dfb(dfb)
+
+
+def test_two_matmul():
+ # Same as Figure 2(a) in TASO paper.
+ @tvm.script.ir_module
+ class MatMul2:
+ @R.function
+ def main(
+ a: R.Tensor((32, 16), "float32"),
+ b: R.Tensor((16, 48), "float32"),
+ c: R.Tensor((48, 32), "float32"),
+ ) -> R.Tensor:
+ with R.dataflow():
+ lv0 = R.call_tir("matmul", (a, b), R.Tensor((32, 48), dtype="float32"))
+ lv1 = R.call_tir("matmul", (lv0, c), R.Tensor((32, 32), dtype="float32"))
+ R.output(lv1)
+ return lv1
+
+ with PatternContext() as ctx:
+ is_call_tir_extern("matmul") >> is_call_tir_extern("matmul")
+ dfb = MatMul2["main"].body.blocks[0]
+ assert ctx.match_dfb(dfb)
+
+ with PatternContext() as ctx:
+ is_call_tir_extern("matmul").has_shape([32, 48]) >> is_call_tir_extern("matmul").has_shape(
+ [32, 32]
+ )
+ dfb = MatMul2["main"].body.blocks[0]
+ assert ctx.match_dfb(dfb)
+
+ with PatternContext() as ctx:
+ is_call_tir_extern("matmul") >> is_call_tir_extern("matmul") >> is_call_tir_extern("matmul")
+ dfb = MatMul2["main"].body.blocks[0]
+ # Three MatMul cannot match
+ assert not ctx.match_dfb(dfb)
+
+
+def test_concat_mm_split():
+ # Same as Figure 2(b) in TASO paper.
+ @tvm.script.ir_module
+ class CMS:
+ @R.function
+ def main(
+ a: R.Tensor((32, 32), "float32"),
+ b: R.Tensor((16, 32), "float32"),
+ c: R.Tensor((16, 32), "float32"),
+ ) -> R.Tensor:
+ with R.dataflow():
+ lv0 = R.call_tir("my_concat", (b, c), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir("my_matmul", (a, lv0), R.Tensor((32, 32), dtype="float32"))
+ lv2 = R.call_tir(
+ "my_split",
+ (lv1,),
+ [R.Tensor((16, 32), dtype="float32"), R.Tensor((16, 32), dtype="float32")],
+ )
+ lv3 = R.TupleGetItem(lv2, 0)
+ lv4 = R.TupleGetItem(lv2, 1)
+ lv5 = R.add(lv3, lv4)
+ R.output(lv5)
+ return lv5
+
+ with PatternContext() as ctx:
+ (
+ is_call_tir_extern("my_concat")
+ >> is_call_tir_extern("my_matmul")
+ >> is_call_tir_extern("my_split")
+ )
+ dfb = CMS["main"].body.blocks[0]
+ assert ctx.match_dfb(dfb)
+
+ with PatternContext() as ctx:
+ split = is_call_tir_extern("my_split")
+ lv3 = TupleGetItemPattern(split, 0).has_shape([16, 32])
+ lv4 = TupleGetItemPattern(split, 1).has_shape([16, 32])
+ split.fork_to(lv3, lv4)
+ add = is_op("relax.add")(lv3, lv4)
+ # TODO(@ganler): simplify this through implicit graph pattern.
+ lv3 >> add
+ lv4 >> add
+
+ dfb = CMS["main"].body.blocks[0]
+ assert ctx.match_dfb(dfb)
+
+
+def test_self_attention():
+ # The example comes from.
+ # https://developer.nvidia.com/blog/nlu-with-tensorrt-bert/
+ @tvm.script.ir_module
+ class SelfAttention:
+ @R.function
+ def main(
+ x: R.Tensor(("b", "s", "n", "h"), "float32"),
+ wq: R.Tensor(("h", "h"), "float32"),
+ wk: R.Tensor(("h", "h"), "float32"),
+ wv: R.Tensor(("h", "h"), "float32"),
+ ) -> R.Tensor:
+ b, s, n, h = T.int64(), T.int64(), T.int64(), T.int64()
+ with R.dataflow():
+ fcq = R.call_tir("my_fc", (x, wq), R.Tensor((b, s, n, h), dtype="float32"))
+ tpq = R.call_tir("my_transpose", (fcq,), R.Tensor((b, s, h, n), dtype="float32"))
+
+ fck = R.call_tir("my_fc", (x, wk), R.Tensor((b, s, n, h), dtype="float32"))
+ tpk = R.call_tir("my_transpose", (fck,), R.Tensor((b, s, h, n), dtype="float32"))
+
+ mul = R.multiply(tpq, tpk)
+ scale = R.multiply(mul, R.const(1.1, "float32"))
+ softmax = R.call_tir("softmax", (scale,), R.Tensor((b, s, n, h), dtype="float32"))
+
+ fcv = R.call_tir("my_fc", (x, wv), R.Tensor((b, s, n, h), dtype="float32"))
+ tpv = R.call_tir("my_transpose", (fcv,), R.Tensor((b, s, h, n), dtype="float32"))
+
+ out = R.multiply(softmax, tpv)
+ R.output(out)
+
+ return out
+
+ with PatternContext() as ctx:
+ fc_trans_q = is_call_tir_extern("my_fc") >> is_call_tir_extern("my_transpose")
+ fc_trans_k = fc_trans_q.dup()
+ fc_trans_v = fc_trans_q.dup()
+
+ is_var("x").fork_to(fc_trans_q, fc_trans_k, fc_trans_v)
+ dfb = SelfAttention["main"].body.blocks[0]
+ assert ctx.match_dfb(dfb)
+
+
+def test_nested_diamond():
+ @tvm.script.ir_module
+ class DiamondInDiamond:
+ @R.function
+ def main(x: R.Tensor((32, 32), "float32"), w: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ with R.dataflow():
+ # matmul0 matmul1
+ # / \ / \
+ # sigmoid2 add4 sigmoid3
+ # \ / \ /
+ # add5 add6
+ # \ /
+ # add7
+ lv0 = R.call_tir("tir_matmul", (x, w), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir("tir_matmul", (x, w), R.Tensor((32, 32), dtype="float32"))
+ lv2 = R.call_tir("tir_sigmoid", (lv0), R.Tensor((32, 32), dtype="float32"))
+ lv3 = R.call_tir("tir_sigmoid", (lv1), R.Tensor((32, 32), dtype="float32"))
+ lv4 = R.call_tir("tir_add", (lv0, lv1), R.Tensor((32, 32), dtype="float32"))
+ lv5 = R.call_tir("tir_add", (lv2, lv4), R.Tensor((32, 32), dtype="float32"))
+ lv6 = R.call_tir("tir_add", (lv3, lv4), R.Tensor((32, 32), dtype="float32"))
+ lv7 = R.call_tir("tir_add", (lv5, lv6), R.Tensor((32, 32), dtype="float32"))
+ R.output(lv7)
+ return lv7
+
+ # match matmul0 diamond
+ with PatternContext() as ctx:
+ sigmoid2 = is_call_tir_extern("tir_sigmoid")
+ add4 = is_call_tir_extern("tir_add")
+ is_call_tir_extern("tir_matmul").fork_to(sigmoid2, add4)
+ add5 = is_call_tir_extern("tir_add")
+ sigmoid2 >> add5
+ add4 ^ add5
+ assert ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0])
+
+ # counter case: mis-match matmul0 diamond
+ with PatternContext() as ctx:
+ sigmoid2 = is_call_tir_extern("tir_sigmoid")
+ add4 = is_call_tir_extern("tir_add")
+ is_call_tir_extern("tir_matmul").fork_to(sigmoid2, add4)
+ add5 = is_call_tir_extern("tir_add")
+ sigmoid2 >> add5
+ add4 >> add5 # not only-used-by relation
+ assert not ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0])
+
+ # match matmul1 diamond
+ with PatternContext() as ctx:
+ sigmoid3 = is_call_tir_extern("tir_sigmoid")
+ add4 = is_call_tir_extern("tir_add")
+ is_call_tir_extern("tir_matmul").fork_to(sigmoid3, add4)
+ add6 = is_call_tir_extern("tir_add")
+ sigmoid3 >> add6
+ add4 ^ add6
+ assert ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0])
+
+ # match add-4-5-6-7
+ with PatternContext() as ctx:
+ add5, add6, add7 = (
+ is_call_tir_extern("tir_add"),
+ is_call_tir_extern("tir_add"),
+ is_call_tir_extern("tir_add"),
+ )
+ is_call_tir_extern("tir_add").fork_to(add5, add6) # add4
+ add5 >> add7
+ add6 >> add7
+ assert ctx.match_dfb(DiamondInDiamond["main"].body.blocks[0])
+
+
+def test_incremental_solving():
+ @R.function
+ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ with R.dataflow():
+ # relu -> sigmoid -> neg
+ lv0 = R.call_tir("tir_relu", (x), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir("tir_sigmoid", (lv0), R.Tensor((32, 32), dtype="float32"))
+ lv2 = R.call_tir("tir_neg", (lv1), R.Tensor((32, 32), dtype="float32"))
+ R.output(lv2)
+ return lv2
+
+ relu = is_call_tir_extern("tir_relu")
+ sigmoid = is_call_tir_extern("tir_sigmoid")
+ neg = is_call_tir_extern("tir_neg")
+
+ with PatternContext() as ctx0:
+ relu >> sigmoid
+ with PatternContext(incremental=True) as ctx1:
+ # because we are doing incremental solving
+ # relu >> sigmoid is still a constraint in this context.
+ # that said the total constraint is:
+ # relu >> sigmoid >> neg
+ sigmoid >> neg
+ assert ctx1.match_dfb(simple_chain.body.blocks[0])
+
+ # match relue -> sigmoid
+ assert ctx0.match_dfb(simple_chain.body.blocks[0])
+
+
+def test_incremental_solving_counter():
+ @R.function
+ def simple_chain(x: R.Tensor((32, 32), "float32")) -> R.Tensor:
+ with R.dataflow():
+ # sigmoid -> neg
+ lv0 = R.call_tir("tir_sigmoid", (x), R.Tensor((32, 32), dtype="float32"))
+ lv1 = R.call_tir("tir_neg", (lv0), R.Tensor((32, 32), dtype="float32"))
+ R.output(lv1)
+ return lv1
+
+ relu = is_call_tir_extern("tir_relu")
+ sigmoid = is_call_tir_extern("tir_sigmoid")
+ neg = is_call_tir_extern("tir_neg")
+
+ with PatternContext() as ctx0:
+ relu >> sigmoid # cannot match
+
+ with PatternContext(incremental=False) as ctx1:
+ # total constraint: sigmoid >> neg
+ sigmoid >> neg
+ assert ctx1.match_dfb(simple_chain.body.blocks[0])
+
+ with PatternContext(incremental=True) as ctx1:
+ # total constraint: relu >> sigmoid >> neg
+ sigmoid >> neg
+ assert not ctx1.match_dfb(simple_chain.body.blocks[0])
+
+
+if __name__ == "__main__":
+ tvm.testing.main()