[mlir-tblgen] Add DagNode StaticMatcher.

[Chia-hungDuan]

Summary

Some patterns may share the common DAG structures. Generate a static
function to do the match logic to reduce the binary size.

[jpienaar]

Could you add the impact of this change to the description too?

mlir/tools/mlir-tblgen/RewriterGen.cpp

248

Could you add comments for these?

[jpienaar]

mlir/tools/mlir-tblgen/RewriterGen.cpp

70

Comment?

334

isSrcPattern= ? (for consistency & to enable warnings on mismatch)

397

You can avoid indent(), <<, unindent, by doing

os.scope() << "return ::mlir::failure();\n";

1595

Could you document this one too?

[jpienaar]

Ping?

[Chia-hungDuan]

Address Review Comment

[Chia-hungDuan]

Add comments

[Chia-hungDuan]

Add a data as reference,

The pattern in tensorflow lit(tensorflow/compiler/mlir/lite/transforms/optimize_patterns.td), the generated file size is reduced from 1160004 byte to 932889 byte, about 19% reduction and there's no difference in compilation time. Will do more detailed profiling.

[jpienaar]

Nice, any way to test this locally?

mlir/tools/mlir-tblgen/RewriterGen.cpp

70

param name?

246

Why not pass the RecordOperatorMap by reference too? Is it allowed to null?

287

If memory serves this form is problematic with some C++14 compilers (works in C++17)

291

Not sure, I follow. This description makes it sound like a global topological ordering of DagNodes. But for different patterns the DagNodes are independent, so don't know what it means in that case.

299

Could you add how this is used?

345

/*depth=*/ ?

375

Funnily I consider that we don't support specifying constraint such for op results too as a gap that folks need to work around, so I'd rather say that names are specified both to bind and to provide constraints. So collecting is required in general as names have 2 uses.

383

Same

385

This comment feels out of place, the above one says collecting all bound symbols and loop below does that, not sure how this relates to that.

395

formatv doesn't seem to be doing anything here

405

Is another way of saying this, that as we consider nodes in topological order it will only be empty when one needs to emit the static match call? ("theoretically" is throwing me off here, makes it feel like there is some corner case where this could fail but I don't think that is the intention)

515

Which test exercised it?

516

Is it guaranteed to be an op here? Also could we avoid this name being hardcoded here? (perhaps making line 485 a lambda to be called from here and there?)

1577

certain Dag?

1600

Ohhh so refStats is used for computing the order rather than ref as in reference counting/memory. I didn't think of that at definition site.

1622

Would this also generate matchers for non-duplicate DagNode?

[Chia-hungDuan]

Address review comments and add a test

mlir/tools/mlir-tblgen/RewriterGen.cpp

246

This is just to align the usage in PatternEmitter

PatternEmitter(Record *pat, RecordOperatorMap *mapper, raw_ostream &os
                          StaticMatcherHelper &helper);

And the use of mapper are expecting pointer parameter.

287

Do you have more details? I see similar form in framework like,

-- BuiltinTypes.h --

class ShapedType :

class ShapedType : public Type {                                                    
public:                                                                             
  using Type::Type;                                                                 
                                                                                    
  // TODO: merge these two special values in a single one used everywhere.          
  // Unfortunately, uses of `-1` have crept deep into the codebase now and are   
  // hard to track.                                                                 
  static constexpr int64_t kDynamicSize = -1;                                                                                                                                                                      
  static constexpr int64_t kDynamicStrideOrOffset =                                 
      std::numeric_limits<int64_t>::min();                                          
   ...
}

291

If two DagNode in two patterns are identical, they will share the same llvm::DagInit. The reason we need the topological sort is, suppose we have two patterns

DAGRootA
   |
   DAGA   
   |
   DAGC

DAGRootB
   |
   DAGA 
   |
   DAGC

DAGA and DAGC are the common parts between them. When we emit matching logic for DAGA, if DAGC has a static matcher generated, then it'll be a function call in DAGA, otherwise it'll inline the matching logic of DAGC. In this case, inlining is not our expectation, to avoid that, we apply the topological order to ensure all the dependents are generated.

We could have different approaches to do that, this way is the one I think it gives the least impact in the current structure, e.g., impact to PatternEmitter.

I can give more details during our chat.

375

Minor updated the description.

385

Here are two cases we need to collect symbols in a DagNode,

1. Emit the static function for a DagNode
2. Emit function call to a static function of a DagNode.

In 1, we will do things like,

pattern.collectBoundSymbols(tree, symbolInfoMap, /*isSrcPattern=*/true);
symbolInfoMap.assignUniqueAlternativeNames();

In 2, the only overlap is

pattern.collectBoundSymbols(tree, symbolInfoMap, /*isSrcPattern=*/true);

The reason it doesn't need assignUniqueAlternativeNames is we get the alternative name from global SymbolInfoMap.

The comment here is just to indicate that we are not missing calling assignUniqueAlternativeNames like other locations.

Do you think we need to remove it?

405

You're right. That makes people think there's a corner case for topological order. What I want to mention here is, if staticMatcherHelper.useStaticMatcher(node) returns true for a dependent DagNode, then we can just generate the function call to match the dependent DagNode. The exception is when we are going to generating the static function for a DagNode itself.

Comment updated. Let me know if it's better

515

https://github.com/llvm/llvm-project/blob/main/mlir/tools/mlir-tblgen/RewriterGen.cpp#L1133

if (numResults != 0) {
    for (int i = 0; i < numResults; ++i)
      os << formatv("for (auto v: castedOp0.getODSResults({0})) {{\n"
                    "  tblgen_types.push_back(v.getType());\n}\n",
                    resultIndex + i);
  }

Do you think we need to fix this hard code usage?

516

Done.

1622

No, only duplicate DagNodes do.

[jpienaar]

Nice, thanks!

mlir/tools/mlir-tblgen/RewriterGen.cpp

241

s/Exam/Examine/ , or (perhaps) Tracks DagNode's referenced multiple times across patterns.

242

Perhaps: Enables generating static matcher functions for DagNode's referenced multiple times rather than inlining them.

246

I think these may be leftover from when we followed a slightly different convention here. I'm pro making these references as we aren't checking for null anywhere

254

the name of the static DAG matcher function corresponding to the node. ?

260

The interface is for adding an individual Record. Record is the raw TableGen type, is it difficult to use something more specific here? These records represent patterns correct? Could we just say pattern

265

Do we need ostream up until this point? E.g., why not have populateStaticMatchers(raw_ostream& os) and remove it from the (now explicit) class constructor?

272

Could you expand comment to capture that this is a counter used/updated continuously to generate unique names?

287

Sure, https://en.cppreference.com/w/cpp/language/static "a constexpr static data member (since C++11)(until C++17) is odr-used, a definition at namespace scope is still required", although I don't think it would hit here as inside the cpp one.

291

I like this example. Could you add this to the function or class document? So basically, by handling the nodes in order of nesting we can avoid accidentally inlining.

385

Yes I think better to remove, thanks

515

Oh, I meant more: do we have a test that exercises this path? E.g., you have this comment describing a case to avoid, do we have a unit test for it

[Chia-hungDuan]

Address review comments and refine test case(only check invariant)

[Chia-hungDuan]

mlir/tools/mlir-tblgen/RewriterGen.cpp

260

Changed to addPattern. Yes, it looks better to say it pattern. I think it may be a legacy naming.

265

Done. it looks better now.

287

Got it! Besides it's inside cpp, according to the definition of odr-used, https://en.cppreference.com/w/cpp/language/definition

Informally, an object is odr-used if its value is read (unless it is a compile time constant) or written, its address is taken, or a reference is bound to it; ...

I think it's safe here

515

We have a test checking that,

https://github.com/llvm/llvm-project/blob/main/mlir/test/mlir-tblgen/rewriter-indexing.td#L40

But before running the test, it'll fire the compilation error because the result pattern will use castedOp0

mlir/include/mlir/TableGen/Pattern.h

 
 #ifndef MLIR_TABLEGEN_PATTERN_H_
 #define MLIR_TABLEGEN_PATTERN_H_
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/TableGen/Argument.h"
 #include "mlir/TableGen/Operator.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/StringMap.h"
 #include "llvm/ADT/StringSet.h"
 
 #include <unordered_map>
 
 namespace llvm {
 class DagInit;
 class Init;
public:
   // function.
   // Precondition: isNativeCodeCall()
   int getNumReturnsOfNativeCode() const;
 
   void print(raw_ostream &os) const;
 
 private:
   friend class SymbolInfoMap;
+  friend llvm::DenseMapInfo<DagNode>;
   const void *getAsOpaquePointer() const { return node; }
 
   const llvm::DagInit *node; // nullptr means null DagNode
 };
 
 // A class for maintaining information for symbols bound in patterns and
 // provides methods for resolving them according to specific use cases.
 //

 // values in a suitable way.
 class SymbolInfoMap {
 public:
   explicit SymbolInfoMap(ArrayRef<llvm::SMLoc> loc) : loc(loc) {}
 
   // Class for information regarding a symbol.
   class SymbolInfo {
   public:
+    // Returns a type string of a variable.
+    std::string getVarTypeStr(StringRef name) const;
+
     // Returns a string for defining a variable named as `name` to store the
     // value bound by this symbol.
     std::string getVarDecl(StringRef name) const;
 
+    // Returns a string for defining an argument which passes the reference of
+    // the variable.
+    std::string getArgDecl(StringRef name) const;
+
     // Returns a variable name for the symbol named as `name`.
     std::string getVarName(StringRef name) const;
 
   private:
     // Allow SymbolInfoMap to access private methods.
     friend class SymbolInfoMap;
 
     // DagNode and DagLeaf are accessed by value which means it can't be used as
public:
 
   // Returns an iterator to the information of the given symbol named as `key`.
   const_iterator find(StringRef key) const;
 
   // Returns an iterator to the information of the given symbol named as `key`,
   // with index `argIndex` for operator `op`.
   const_iterator findBoundSymbol(StringRef key, DagNode node,
                                  const Operator &op, int argIndex) const;
+  const_iterator findBoundSymbol(StringRef key, SymbolInfo symbolInfo) const;
 
   // Returns the bounds of a range that includes all the elements which
   // bind to the `key`.
   std::pair<iterator, iterator> getRangeOfEqualElements(StringRef key);
 
   // Returns number of times symbol named as `key` was used.
   int count(StringRef key) const;
 
public:
   int getBenefit() const;
 
   using IdentifierLine = std::pair<StringRef, unsigned>;
 
   // Returns the file location of the pattern (buffer identifier + line number
   // pair).
   std::vector<IdentifierLine> getLocation() const;
 
-private:
-  // Helper function to verify variabld binding.
-  void verifyBind(bool result, StringRef symbolName);
-
   // Recursively collects all bound symbols inside the DAG tree rooted
   // at `tree` and updates the given `infoMap`.
   void collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
                            bool isSrcPattern);
 
+private:
+  // Helper function to verify variable binding.
+  void verifyBind(bool result, StringRef symbolName);
+
   // The TableGen definition of this pattern.
   const llvm::Record &def;
 
   // All operators.
   // TODO: we need a proper context manager, like MLIRContext, for managing the
   // lifetime of shared entities.
   RecordOperatorMap *recordOpMap;
 };
 
 } // end namespace tblgen
 } // end namespace mlir
 
+namespace llvm {
+template <>
+struct DenseMapInfo<mlir::tblgen::DagNode> {
+  static mlir::tblgen::DagNode getEmptyKey() {
+    return mlir::tblgen::DagNode(
+        llvm::DenseMapInfo<llvm::DagInit *>::getEmptyKey());
+  }
+  static mlir::tblgen::DagNode getTombstoneKey() {
+    return mlir::tblgen::DagNode(
+        llvm::DenseMapInfo<llvm::DagInit *>::getTombstoneKey());
+  }
+  static unsigned getHashValue(mlir::tblgen::DagNode node) {
+    return llvm::hash_value(node.getAsOpaquePointer());
+  }
+  static bool isEqual(mlir::tblgen::DagNode lhs, mlir::tblgen::DagNode rhs) {
+    return lhs.node == rhs.node;
+  }
+};
+} // end namespace llvm
+
 #endif // MLIR_TABLEGEN_PATTERN_H_

mlir/lib/TableGen/Pattern.cpp

int SymbolInfoMap::SymbolInfo::getStaticValueCount() const {
   }
   llvm_unreachable("unknown kind");
 }
 
 std::string SymbolInfoMap::SymbolInfo::getVarName(StringRef name) const {
   return alternativeName.hasValue() ? alternativeName.getValue() : name.str();
 }
 
-std::string SymbolInfoMap::SymbolInfo::getVarDecl(StringRef name) const {
-  LLVM_DEBUG(llvm::dbgs() << "getVarDecl for '" << name << "': ");
+std::string SymbolInfoMap::SymbolInfo::getVarTypeStr(StringRef name) const {
+  LLVM_DEBUG(llvm::dbgs() << "getVarTypeStr for '" << name << "': ");
   switch (kind) {
   case Kind::Attr: {
-    if (op) {
-      auto type = op->getArg(getArgIndex())
-                      .get<NamedAttribute *>()
-                      ->attr.getStorageType();
-      return std::string(formatv("{0} {1};\n", type, name));
-    }
+    if (op)
+      return op->getArg(getArgIndex())
+          ->attr.getStorageType()
+          .str();
     // TODO(suderman): Use a more exact type when available.
-    return std::string(formatv("Attribute {0};\n", name));
+    return "Attribute";
   }
   case Kind::Operand: {
     // Use operand range for captured operands (to support potential variadic
     // operands).
-    return std::string(
-        formatv("::mlir::Operation::operand_range {0}(op0->getOperands());\n",
-                getVarName(name)));
+    return "::mlir::Operation::operand_range";
   }
   case Kind::Value: {
-    return std::string(formatv("::mlir::Value {0};\n", name));
+    return "::mlir::Value";
   }
   case Kind::MultipleValues: {
-    // This is for the variable used in the source pattern. Each named value in
-    // source pattern will only be bound to a Value. The others in the result
-    // pattern may be associated with multiple Values as we will use `auto` to
-    // do the type inference.
-    return std::string(formatv(
-        "::mlir::Value {0}_raw; ::mlir::ValueRange {0}({0}_raw);\n", name));
+    return "::mlir::ValueRange";
   }
   case Kind::Result: {
     // Use the op itself for captured results.
-    return std::string(formatv("{0} {1};\n", op->getQualCppClassName(), name));
+    return op->getQualCppClassName();
   }
   }
   llvm_unreachable("unknown kind");
 }
 
+std::string SymbolInfoMap::SymbolInfo::getVarDecl(StringRef name) const {
+  LLVM_DEBUG(llvm::dbgs() << "getVarDecl for '" << name << "': ");
+  std::string varInit = kind == Kind::Operand ? "(op0->getOperands())" : "";
+  return std::string(
+      formatv("{0} {1}{2};\n", getVarTypeStr(name), getVarName(name), varInit));
+}
+
+std::string SymbolInfoMap::SymbolInfo::getArgDecl(StringRef name) const {
+  LLVM_DEBUG(llvm::dbgs() << "getArgDecl for '" << name << "': ");
+  return std::string(
+      formatv("{0} &{1}", getVarTypeStr(name), getVarName(name)));
+}
+
 std::string SymbolInfoMap::SymbolInfo::getValueAndRangeUse(
     StringRef name, int index, const char *fmt, const char *separator) const {
   LLVM_DEBUG(llvm::dbgs() << "getValueAndRangeUse for '" << name << "': ");
   switch (kind) {
   case Kind::Attr: {
     assert(index < 0);
     auto repl = formatv(fmt, name);
     LLVM_DEBUG(llvm::dbgs() << repl << " (Attr)\n");
SymbolInfoMap::const_iterator SymbolInfoMap::find(StringRef key) const {
   std::string name = getValuePackName(key).str();
 
   return symbolInfoMap.find(name);
 }
 
 SymbolInfoMap::const_iterator
 SymbolInfoMap::findBoundSymbol(StringRef key, DagNode node, const Operator &op,
                                int argIndex) const {
+  return findBoundSymbol(key, SymbolInfo::getOperand(node, &op, argIndex));
+}
+
+SymbolInfoMap::const_iterator
+SymbolInfoMap::findBoundSymbol(StringRef key, SymbolInfo symbolInfo) const {
   std::string name = getValuePackName(key).str();
   auto range = symbolInfoMap.equal_range(name);
 
-  const auto symbolInfo = SymbolInfo::getOperand(node, &op, argIndex);
-
   for (auto it = range.first; it != range.second; ++it)
     if (it->second.dagAndConstant == symbolInfo.dagAndConstant)
       return it;
 
   return symbolInfoMap.end();
 }
 
 std::pair<SymbolInfoMap::iterator, SymbolInfoMap::iterator>

mlir/test/mlir-tblgen/rewriter-static-matcher.td

+// RUN: mlir-tblgen -gen-rewriters -I %S/../../include %s | FileCheck %s
+
+include "mlir/IR/OpBase.td"
+
+def Test_Dialect : Dialect {
+  let name = "test";
+}
+class NS_Op<string mnemonic, list<OpTrait> traits> :
+    Op<Test_Dialect, mnemonic, traits>;
+
+def AOp : NS_Op<"a_op", []> {
+  let arguments = (ins
+    AnyInteger:$any_integer
+  );
+
+  let results = (outs AnyInteger);
+}
+
+def BOp : NS_Op<"b_op", []> {
+  let arguments = (ins
+    AnyAttr: $any_attr,
+    AnyInteger
+  );
+
+  let results = (outs AnyInteger);
+}
+
+def COp : NS_Op<"c_op", []> {
+  let arguments = (ins
+    AnyAttr: $any_attr,
+    AnyInteger
+  );
+
+  let results = (outs AnyInteger);
+}
+
+// Test static matcher for duplicate DagNode
+// ---
+
+// CHECK: static ::mlir::LogicalResult static_dag_matcher_0
+
+// CHECK: if(failed(static_dag_matcher_0(rewriter, op1, tblgen_ops
+def : Pat<(AOp (BOp I32Attr:$attr, I32:$int)),
+          (AOp $int)>;
+
+// CHECK: if(failed(static_dag_matcher_0(rewriter, op1, tblgen_ops
+def : Pat<(COp $_, (BOp I32Attr:$attr, I32:$int)),
+          (COp $attr, $int)>;

mlir/tools/mlir-tblgen/RewriterGen.cpp

 #include "mlir/Support/IndentedOstream.h"
 #include "mlir/TableGen/Attribute.h"
 #include "mlir/TableGen/Format.h"
 #include "mlir/TableGen/GenInfo.h"
 #include "mlir/TableGen/Operator.h"
 #include "mlir/TableGen/Pattern.h"
 #include "mlir/TableGen/Predicate.h"
 #include "mlir/TableGen/Type.h"
+#include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/FormatAdapters.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/Signals.h"
 #include "llvm/TableGen/Error.h"

 };
 } // end namespace llvm
 
 //===----------------------------------------------------------------------===//
 // PatternEmitter
 //===----------------------------------------------------------------------===//
 
 namespace {
+
+class StaticMatcherHelper;
+
 class PatternEmitter {
 public:
-  PatternEmitter(Record *pat, RecordOperatorMap *mapper, raw_ostream &os);
+  PatternEmitter(Record *pat, RecordOperatorMap *mapper, raw_ostream &os,
+                 StaticMatcherHelper &helper);
 
   // Emits the mlir::RewritePattern struct named `rewriteName`.
   void emit(StringRef rewriteName);
 
+  // Emits the static function of DAG matcher.
+  void emitStaticMatcher(DagNode tree, std::string funcName);
+
 private:
   // Emits the code for matching ops.
   void emitMatchLogic(DagNode tree, StringRef opName);
 
   // Emits the code for rewriting ops.
   void emitRewriteLogic();
 
   //===--------------------------------------------------------------------===//
   // Match utilities
   //===--------------------------------------------------------------------===//
 
   // Emits C++ statements for matching the DAG structure.
   void emitMatch(DagNode tree, StringRef name, int depth);
 
+  // Emit C++ function call to static DAG matcher.
+  void emitStaticMatchCall(DagNode tree, StringRef name);
+
   // Emits C++ statements for matching using a native code call.
   void emitNativeCodeMatch(DagNode tree, StringRef name, int depth);
 
   // Emits C++ statements for matching the op constrained by the given DAG
   // `tree` returning the op's variable name.
   void emitOpMatch(DagNode tree, StringRef opName, int depth);
 
   // Emits C++ statements for matching the `argIndex`-th argument of the given
private:
   RecordOperatorMap *opMap;
 
   // Handy wrapper for pattern being emitted.
   Pattern pattern;
 
   // Map for all bound symbols' info.
   SymbolInfoMap symbolInfoMap;
 
+  StaticMatcherHelper &staticMatcherHelper;
+
   // The next unused ID for newly created values.
   unsigned nextValueId;
 
   raw_indented_ostream os;
 
   // Format contexts containing placeholder substitutions.
   FmtContext fmtCtx;
+};
+
+// Tracks DagNode's reference multiple times across patterns. Enables generating
+// static matcher functions for DagNode's referenced multiple times rather than
+// inlining them.
+class StaticMatcherHelper {
+public:
+  StaticMatcherHelper(RecordOperatorMap &mapper);
+
+  // Determine if we should inline the match logic or delegate to a static
+  // function.
+  bool useStaticMatcher(DagNode node) {
+    return refStats[node] > kStaticMatcherThreshold;
+  }
+
+  // Get the name of the static DAG matcher function corresponding to the node.
+  std::string getMatcherName(DagNode node) {
+    assert(useStaticMatcher(node));
+    return matcherNames[node];
+  }
+
+  // Collect the `Record`s, i.e., the DRR, so that we can get the information of
+  // the duplicated DAGs.
+  void addPattern(Record *record);
+
+  // Emit all static functions of DAG Matcher.
+  void populateStaticMatchers(raw_ostream &os);
 
-  // Number of op processed.
-  int opCounter = 0;
+private:
+  static constexpr unsigned kStaticMatcherThreshold = 1;
+
+  // Consider two patterns as down below,
+  //   DagNode_Root_A    DagNode_Root_B
+  //       \                 \
+  //     DagNode_C         DagNode_C
+  //         \                 \
+  //       DagNode_D         DagNode_D
+  //
+  // DagNode_Root_A and DagNode_Root_B share the same subtree which consists of
+  // DagNode_C and DagNode_D. Both DagNode_C and DagNode_D are referenced
+  // multiple times so we'll have static matchers for both of them. When we're
+  // emitting the match logic for DagNode_C, we will check if DagNode_D has the
+  // static matcher generated. If so, then we'll generate a call to the
+  // function, inline otherwise. In this case, inlining is not what we want. As
+  // a result, generate the static matcher in topological order to ensure all
+  // the dependent static matchers are generated and we can avoid accidentally
+  // inlining.
+  //
+  // The topological order of all the DagNodes among all patterns.
+  SmallVector<std::pair<DagNode, Record *>> topologicalOrder;
+
+  RecordOperatorMap &opMap;
+
+  // Records of the static function name of each DagNode
+  DenseMap<DagNode, std::string> matcherNames;
+
+  // After collecting all the DagNode in each pattern, `refStats` records the
+  // number of users for each DagNode. We will generate the static matcher for a
+  // DagNode while the number of users exceeds a certain threshold.
+  DenseMap<DagNode, unsigned> refStats;
+
+  // Number of static matcher generated. This is used to generate a unique name
+  // for each DagNode.
+  int staticMatcherCounter = 0;
 };
+
 } // end anonymous namespace
 
 PatternEmitter::PatternEmitter(Record *pat, RecordOperatorMap *mapper,
-                               raw_ostream &os)
+                               raw_ostream &os, StaticMatcherHelper &helper)
     : loc(pat->getLoc()), opMap(mapper), pattern(pat, mapper),
-      symbolInfoMap(pat->getLoc()), nextValueId(0), os(os) {
+      symbolInfoMap(pat->getLoc()), staticMatcherHelper(helper), nextValueId(0),
+      os(os) {
   fmtCtx.withBuilder("rewriter");
 }
 
 std::string PatternEmitter::handleConstantAttr(Attribute attr,
                                                StringRef value) {
   if (!attr.isConstBuildable())
     PrintFatalError(loc, "Attribute " + attr.getAttrDefName() +
                              " does not have the 'constBuilderCall' field");
 
   // TODO: Verify the constants here
   return std::string(tgfmt(attr.getConstBuilderTemplate(), &fmtCtx, value));
 }
 
+void PatternEmitter::emitStaticMatcher(DagNode tree, std::string funcName) {
+  os << formatv(
+      "static ::mlir::LogicalResult {0}(::mlir::PatternRewriter &rewriter, "
+      "::mlir::Operation *op0, ::llvm::SmallVector<::mlir::Operation "
+      "*, 4> &tblgen_ops",
+      funcName);
+
+  // We pass the reference of the variables that need to be captured. Hence we
+  // need to collect all the symbols in the tree first.
+  pattern.collectBoundSymbols(tree, symbolInfoMap, /*isSrcPattern=*/true);
+  symbolInfoMap.assignUniqueAlternativeNames();
+  for (const auto &info : symbolInfoMap)
+    os << formatv(", {0}", info.second.getArgDecl(info.first));
+
+  os << ") {\n";
+  os.indent();
+  os << "(void)tblgen_ops;\n";
+
+  // Note that a static matcher is considered at least one step from the match
+  // entry.
+  emitMatch(tree, "op0", /*depth=*/1);
+
+  os << "return ::mlir::success();\n";
+  os.unindent();
+  os << "}\n\n";
+}
+
 // Helper function to match patterns.
 void PatternEmitter::emitMatch(DagNode tree, StringRef name, int depth) {
   if (tree.isNativeCodeCall()) {
     emitNativeCodeMatch(tree, name, depth);
     return;
   }
 
   if (tree.isOperation()) {
     emitOpMatch(tree, name, depth);
     return;
   }
 
   PrintFatalError(loc, "encountered non-op, non-NativeCodeCall match.");
 }
 
+void PatternEmitter::emitStaticMatchCall(DagNode tree, StringRef opName) {
+  std::string funcName = staticMatcherHelper.getMatcherName(tree);
+  os << formatv("if(failed({0}(rewriter, {1}, tblgen_ops", funcName, opName);
+
+  // TODO(chiahungduan): Add a lookupBoundSymbols() to do the subtree lookup in
+  // one pass.
+
+  // In general, bound symbol should have the unique name in the pattern but
+  // for the operand, binding same symbol to multiple operands imply a
+  // constraint at the same time. In this case, we will rename those operands
+  // with different names. As a result, we need to collect all the symbolInfos
+  // from the DagNode then get the updated name of the local variables from the
+  // global symbolInfoMap.
+
+  // Collect all the bound symbols in the Dag
+  SymbolInfoMap localSymbolMap(loc);
+  pattern.collectBoundSymbols(tree, localSymbolMap, /*isSrcPattern=*/true);
+
+  for (const auto &info : localSymbolMap) {
+    auto name = info.first;
+    auto symboInfo = info.second;
+    auto ret = symbolInfoMap.findBoundSymbol(name, symboInfo);
+    os << formatv(", {0}", ret->second.getVarName(name));
+  }
+
+  os << "))) {\n";
+  os.scope().os << "return ::mlir::failure();\n";
+  os << "}\n";
+}
+
 // Helper function to match patterns.
 void PatternEmitter::emitNativeCodeMatch(DagNode tree, StringRef opName,
                                          int depth) {
   LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall matcher pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
+  // The order of generating static matcher follows the topological order so
+  // that for every dependent DagNode already have their static matcher
+  // generated if needed. The reason we check if `getMatcherName(tree).empty()`
+  // is when we are generating the static matcher for a DagNode itself. In this
+  // case, we need to emit the function body rather than a function call.
+  if (staticMatcherHelper.useStaticMatcher(tree) &&
+      !staticMatcherHelper.getMatcherName(tree).empty()) {
+    emitStaticMatchCall(tree, opName);
+
+    // NativeCodeCall will never be at depth 0 so that we don't need to catch
+    // the root operation as emitOpMatch();
+
+    return;
+  }
+
   // TODO(suderman): iterate through arguments, determine their types, output
   // names.
   SmallVector<std::string, 8> capture;
 
   raw_indented_ostream::DelimitedScope scope(os);
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
     std::string argName = formatv("arg{0}_{1}", depth, i);

 
 // Helper function to match patterns.
 void PatternEmitter::emitOpMatch(DagNode tree, StringRef opName, int depth) {
   Operator &op = tree.getDialectOp(opMap);
   LLVM_DEBUG(llvm::dbgs() << "start emitting match for op '"
                           << op.getOperationName() << "' at depth " << depth
                           << '\n');
 
-  std::string castedName = formatv("castedOp{0}", depth);
+  auto getCastedName = [depth]() -> std::string {
+    return formatv("castedOp{0}", depth);
+  };
+
+  // The order of generating static matcher follows the topological order so
+  // that for every dependent DagNode already have their static matcher
+  // generated if needed. The reason we check if `getMatcherName(tree).empty()`
+  // is when we are generating the static matcher for a DagNode itself. In this
+  // case, we need to emit the function body rather than a function call.
+  if (staticMatcherHelper.useStaticMatcher(tree) &&
+      !staticMatcherHelper.getMatcherName(tree).empty()) {
+    emitStaticMatchCall(tree, opName);
+    // In the codegen of rewriter, we suppose that castedOp0 will capture the
+    // root operation. Manually add it if the root DagNode is a static matcher.
+    if (depth == 0)
+      os << formatv("auto {2} = ::llvm::dyn_cast_or_null<{1}>({0}); "
+                    "(void){2};\n",
+                    opName, op.getQualCppClassName(), getCastedName());
+    return;
+  }
+
+  std::string castedName = getCastedName();
   os << formatv("auto {0} = ::llvm::dyn_cast<{2}>({1}); "
                 "(void){0};\n",
                 castedName, opName, op.getQualCppClassName());
 
   // Skip the operand matching at depth 0 as the pattern rewriter already does.
   if (depth != 0)
     emitMatchCheck(opName, /*matchStr=*/castedName,
                    formatv("\"{0} is not {1} type\"", castedName,
if (DagNode argTree = tree.getArgAsNestedDag(i)) {
           "auto *{0} = "
           "(*{1}.getODSOperands({2}).begin()).getDefiningOp();\n",
           argName, castedName, nextOperand);
       // Null check of operand's definingOp
       emitMatchCheck(castedName, /*matchStr=*/argName,
                      formatv("\"Operand {0} of {1} has null definingOp\"",
                              nextOperand++, castedName));
       emitMatch(argTree, argName, depth + 1);
-      os << formatv("tblgen_ops[{0}] = {1};\n", ++opCounter, argName);
+      os << formatv("tblgen_ops.push_back({0});\n", argName);
       os.unindent() << "}\n";
       continue;
     }
 
     // Next handle DAG leaf: operand or attribute
     if (opArg.is<NamedTypeConstraint *>()) {
       // emitOperandMatch's argument indexing counts attributes.
       emitOperandMatch(tree, castedName, i, nextOperand, depth);
::mlir::LogicalResult matchAndRewrite(::mlir::Operation *op0,
       for (const auto &symbolInfoPair : symbolInfoMap) {
         const auto &symbol = symbolInfoPair.first;
         const auto &info = symbolInfoPair.second;
 
         os << info.getVarDecl(symbol);
       }
       // TODO: capture ops with consistent numbering so that it can be
       // reused for fused loc.
-      os << formatv("::mlir::Operation *tblgen_ops[{0}];\n\n",
-                    pattern.getSourcePattern().getNumOps());
+      os << "::llvm::SmallVector<::mlir::Operation *, 4> tblgen_ops;\n\n";
       LLVM_DEBUG(llvm::dbgs()
                  << "done creating local variables for capturing matches\n");
 
       os << "// Match\n";
-      os << "tblgen_ops[0] = op0;\n";
+      os << "tblgen_ops.push_back(op0);\n";
       emitMatchLogic(sourceTree, "op0");
 
       os << "\n// Rewrite\n";
       emitRewriteLogic();
 
       os << "return ::mlir::success();\n";
     }
     os << "};\n";
if (operand->isVariadic()) {
           os << symbol;
         }
       }
       os << ");\n";
     }
   }
 }
 
+StaticMatcherHelper::StaticMatcherHelper(RecordOperatorMap &mapper)
+    : opMap(mapper) {}
+
+void StaticMatcherHelper::populateStaticMatchers(raw_ostream &os) {
+  // PatternEmitter will use the static matcher if there's one generated. To
+  // ensure that all the dependent static matchers are generated before emitting
+  // the matching logic of the DagNode, we use topological order to achieve it.
+  for (auto &dagInfo : topologicalOrder) {
+    DagNode node = dagInfo.first;
+    if (!useStaticMatcher(node))
+      continue;
+
+    std::string funcName =
+        formatv("static_dag_matcher_{0}", staticMatcherCounter++);
+    assert(matcherNames.find(node) == matcherNames.end());
+    PatternEmitter(dagInfo.second, &opMap, os, *this)
+        .emitStaticMatcher(node, funcName);
+    matcherNames[node] = funcName;
+  }
+}
+
+void StaticMatcherHelper::addPattern(Record *record) {
+  Pattern pat(record, &opMap);
+
+  // While generating the function body of the DAG matcher, it may depends on
+  // other DAG matchers. To ensure the dependent matchers are ready, we compute
+  // the topological order for all the DAGs and emit the DAG matchers in this
+  // order.
+  llvm::unique_function<void(DagNode)> dfs = [&](DagNode node) {
+    ++refStats[node];
+
+    if (refStats[node] != 1)
+      return;
+
+    for (unsigned i = 0, e = node.getNumArgs(); i < e; ++i)
+      if (DagNode sibling = node.getArgAsNestedDag(i))
+        dfs(sibling);
+
+    topologicalOrder.push_back(std::make_pair(node, record));
+  };
+
+  dfs(pat.getSourcePattern());
+}
+
 static void emitRewriters(const RecordKeeper &recordKeeper, raw_ostream &os) {
   emitSourceFileHeader("Rewriters", os);
 
   const auto &patterns = recordKeeper.getAllDerivedDefinitions("Pattern");
-  auto numPatterns = patterns.size();
 
   // We put the map here because it can be shared among multiple patterns.
   RecordOperatorMap recordOpMap;
 
+  // Exam all the patterns and generate static matcher for the duplicated
+  // DagNode.
+  StaticMatcherHelper staticMatcher(recordOpMap);
+  for (Record *p : patterns)
+    staticMatcher.addPattern(p);
+  staticMatcher.populateStaticMatchers(os);
+
   std::vector<std::string> rewriterNames;
-  rewriterNames.reserve(numPatterns);
+  rewriterNames.reserve(patterns.size());
 
   std::string baseRewriterName = "GeneratedConvert";
   int rewriterIndex = 0;
 
   for (Record *p : patterns) {
     std::string name;
     if (p->isAnonymous()) {
       // If no name is provided, ensure unique rewriter names simply by
       // appending unique suffix.
       name = baseRewriterName + llvm::utostr(rewriterIndex++);
     } else {
       name = std::string(p->getName());
     }
     LLVM_DEBUG(llvm::dbgs()
                << "=== start generating pattern '" << name << "' ===\n");
-    PatternEmitter(p, &recordOpMap, os).emit(name);
+    PatternEmitter(p, &recordOpMap, os, staticMatcher).emit(name);
     LLVM_DEBUG(llvm::dbgs()
                << "=== done generating pattern '" << name << "' ===\n");
     rewriterNames.push_back(std::move(name));
   }
 
   // Emit function to add the generated matchers to the pattern list.
   os << "void LLVM_ATTRIBUTE_UNUSED populateWithGenerated("
         "::mlir::RewritePatternSet &patterns) {\n";