Skip to content

Commit

Permalink
Dataflow Pattern Lang: Core Matching Features (tlc-pack#163)
Browse files Browse the repository at this point in the history
The structure is similar to the Relay's pattern matcher (apache/tvm#5231). The main difference is that those pattern types are adopted to be relax-compatible. Relay pattern types, some less used patterns (IfPattern) and df-topological patterns (DominatorPattern) are ignored (some of them will be brought later).

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`);
  • Loading branch information
ganler authored and junrushao committed Jan 26, 2023
1 parent 708f24b commit 271d121
Show file tree
Hide file tree
Showing 18 changed files with 4,937 additions and 1 deletion.
32 changes: 32 additions & 0 deletions include/tvm/relax/analysis.h
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,38 @@ TVM_DLL tvm::Array<GlobalVar> RecGlobalVars(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 runtime::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 runtime::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 runtime::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 definitoins to a set of uses.
*/
TVM_DLL runtime::Map<Var, Array<Var>> UseDefChain(const DataflowBlock& dfb);

} // namespace relax
} // namespace tvm

Expand Down
76 changes: 76 additions & 0 deletions include/tvm/relax/dataflow_matcher.h
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
/*
* 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 var2val 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>> var2val = 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_
Loading

0 comments on commit 271d121

Please sign in to comment.