[mlir] RewriterGen NativeCodeCall matcher with ConstantOp matcher

[rsuderman]

Summary

Added an underlying matcher for generic constant ops. This
included a rewriter of RewriterGen to make variable use more
clear.

[rriddle]

mlir/include/mlir/IR/Matchers.h

274 ↗

(On Diff #297332)

Is this necessary? Seems like this is just success(m_Constant(&attribute).match(op)) which is something that could be written inline in DRR and not here.

mlir/include/mlir/IR/OpBase.td

2352

I'd prefer that this has something like "Matcher" in the name so that it is obvious this is just for DRR, given that there is a ConstantLike trait.

[rsuderman]

Updated for riverriddle comments.

[jpienaar]

I think you may need to rebase this on head, seems like you were working of an old version here.

mlir/lib/TableGen/Pattern.cpp

220

Would we need to here similarly consider the type as above?

534–535

New error helper functions were added to TableGen, so one can now say PrintFatalError(&def, ...)

564

nested?

569

Maybe phab is confusing me here, but this looks like same then and else bodies.

573

You might see the other change allowing this, where these become constraints to check for equality.

mlir/test/lib/Dialect/Test/TestOps.td

791

You can skip this if empty.

mlir/tools/mlir-tblgen/RewriterGen.cpp

256

os is a raw_indented_ostream here, so you should just need to use a DelimitedScope here [did you check how this looks? I would expect this to have more indenting than expected]

[rsuderman]

Updated for jpienaar@ comments

[rsuderman]

mlir/lib/TableGen/Pattern.cpp

220

We would if it were possible :/. NativeCodeCall doesn't specify types so I need to assume the result (I went with attribute). Long term we need a better solution but I couldn't think of one.

569

The else body statement is the same with the exception of what type of binding is being done. Note bindAttr vs bindValue.

573

I think I implemented as you hoped by adding a verification on the constraint as well. Should be good enough anyway, I am not certain value matching is going to work well for the pattern matcher with native code calls.

mlir/tools/mlir-tblgen/RewriterGen.cpp

256

Indenting is good until the match checks (over indented). I tweaked a little more to decrease the indenting however emitMatchCheck has weird indenting mechanics.

[rsuderman]

Rest of jpienaar@ comments

[jpienaar]

mlir/lib/TableGen/Pattern.cpp

220

Yes I think we need to add support for this ODS side (I don't think it would be too hard). Could you add a TODO here to use a more exact type if available?

533

Could we do an early exist instead with positive logic?

564

Could you test with ASAN too? (you might actually need to add some failing test cases too to verify here and below, a TD while where these are invoked from the run should enable covering these).

566

Lets be consistent with all these we now add and use the new helper methods.

585

Could flipping logic + using continues etc be used to reduce this nesting?

mlir/tools/mlir-tblgen/RewriterGen.cpp

242

You should not need to be doing these manual indent calculations here. These are structured, so you should simply be able to have the scope handle that for you. I see it is still below in the other one too at head used in one call, I think it could be an oversight there too, but I'm not sure why it should not be needed.

448

These should not be needed. As mentioned in previous comments, this appears like you based this on a old revision.

[rsuderman]

Additional jpienaar changes.

[rsuderman]

mlir/lib/TableGen/Pattern.cpp

564

Invoked with ASAN. Added tests look good (excluding already failing tests).

564

ASAN marks things as good. Could you give me a better feeling on what a TD failure case looks like? Most of the tablegen failures are a bit bewildering.

mlir/tools/mlir-tblgen/RewriterGen.cpp

242

Yeah, I missed some of the merge issues when updating. Should be good now.

[jpienaar]

mlir/lib/TableGen/Pattern.cpp

564

Sure, adding a test/tests where we hit these fatal error cases. Given it fatals, one would have to combine something like // RUN: not ... 2>&1 | FileCheck (look at llvm/test/TableGen/defvar.td test for a failure testing case approach that could work well here too)

[rsuderman]

Added tablegen rewriter error tests.

[jpienaar]

Looks good thanks!

mlir/lib/TableGen/Pattern.cpp

534–535

I think you might have missed this one (or uploaded other version)

[rsuderman]

Handled last jpienaar comments

mlir/include/mlir/IR/OpBase.td

 // If used as a DAG node, i.e., `(NativeCodeCall<"..."> <arg0>, ..., <argN>)`,
 // then positional placeholders are also supported; placeholder `$N` in the
 // wrapped C++ expression will be replaced by `<argN>`.
 
 class NativeCodeCall<string expr> {
   string expression = expr;
 }
 
+def ConstantLikeMatcher : NativeCodeCall<"success(matchPattern($0->getResult(0), m_Constant(&$1)))">;
+
 //===----------------------------------------------------------------------===//
 // Rewrite directives
 //===----------------------------------------------------------------------===//
 
 // Directive used in result pattern to specify the location of the generated
 // op. This directive must be used as the last argument to the op creation
 // DAG construct. The arguments to location must be previously captured symbol.
 def location;
 
 // Directive used in result pattern to indicate that no new op are generated,
 // so to replace the matched DAG with an existing SSA value.
 def replaceWithValue;
 
 #endif // OP_BASE

mlir/include/mlir/TableGen/Pattern.h

private:
     // Creates a SymbolInfo instance. `index` is only used for `Attr` and
     // `Operand` so should be negative for `Result` and `Value` kind.
     SymbolInfo(const Operator *op, Kind kind, Optional<int> index);
 
     // Static methods for creating SymbolInfo.
     static SymbolInfo getAttr(const Operator *op, int index) {
       return SymbolInfo(op, Kind::Attr, index);
     }
+    static SymbolInfo getAttr() {
+      return SymbolInfo(nullptr, Kind::Attr, llvm::None);
+    }
     static SymbolInfo getOperand(const Operator *op, int index) {
       return SymbolInfo(op, Kind::Operand, index);
     }
     static SymbolInfo getResult(const Operator *op) {
       return SymbolInfo(op, Kind::Result, llvm::None);
     }
     static SymbolInfo getValue() {
       return SymbolInfo(nullptr, Kind::Value, llvm::None);
public:
   // Binds the given `symbol` to the results the given `op`. Returns false if
   // `symbol` is already bound.
   bool bindOpResult(StringRef symbol, const Operator &op);
 
   // Registers the given `symbol` as bound to a value. Returns false if `symbol`
   // is already bound.
   bool bindValue(StringRef symbol);
 
+  // Registers the given `symbol` as bound to an attr. Returns false if `symbol`
+  // is already bound.
+  bool bindAttr(StringRef symbol);
+
   // Returns true if the given `symbol` is bound.
   bool contains(StringRef symbol) const;
 
   // Returns an iterator to the information of the given symbol named as `key`.
   const_iterator find(StringRef key) const;
 
   // Returns the number of static values of the given `symbol` corresponds to.
   // A static value is an operand/result declared in ODS. Normally a symbol only

mlir/lib/TableGen/Pattern.cpp

int SymbolInfoMap::SymbolInfo::getStaticValueCount() const {
   }
   llvm_unreachable("unknown kind");
 }
 
 std::string SymbolInfoMap::SymbolInfo::getVarDecl(StringRef name) const {
   LLVM_DEBUG(llvm::dbgs() << "getVarDecl for '" << name << "': ");
   switch (kind) {
   case Kind::Attr: {
+    if (op) {
       auto type =
           op->getArg(*argIndex).get<NamedAttribute *>()->attr.getStorageType();
       return std::string(formatv("{0} {1};\n", type, name));
     }
+    return std::string(formatv("Attribute {0};\n", name));
+  }
   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", name));
   }
   case Kind::Value: {
     return std::string(formatv("::llvm::ArrayRef<::mlir::Value> {0};\n", name));
bool SymbolInfoMap::bindOpResult(StringRef symbol, const Operator &op) {
   StringRef name = getValuePackName(symbol);
   return symbolInfoMap.insert({name, SymbolInfo::getResult(&op)}).second;
 }
 
 bool SymbolInfoMap::bindValue(StringRef symbol) {
   return symbolInfoMap.insert({symbol, SymbolInfo::getValue()}).second;
 }
 
+bool SymbolInfoMap::bindAttr(StringRef symbol) {
+  return symbolInfoMap.insert({symbol, SymbolInfo::getAttr()}).second;
+}
+
 bool SymbolInfoMap::contains(StringRef symbol) const {
   return find(symbol) != symbolInfoMap.end();
 }
 
 SymbolInfoMap::const_iterator SymbolInfoMap::find(StringRef key) const {
   StringRef name = getValuePackName(key);
   return symbolInfoMap.find(name);
 }
result.emplace_back(
         llvm::SrcMgr.getLineAndColumn(loc, buf).first);
   }
   return result;
 }
 
 void Pattern::collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
                                   bool isSrcPattern) {
   auto treeName = tree.getSymbol();
-  if (!tree.isOperation()) {
+  auto numTreeArgs = tree.getNumArgs();
+
+  if (!tree.isOperation() && !tree.isNativeCodeCall()) {
+    if (!treeName.empty()) {
+      PrintFatalError(
+          &def,
+          formatv("binding symbol '{0}' to non-operation/native code call "
+                  "unsupported right now",
+                  treeName));
+    }
+    return;
+  }
+
+  if (tree.isNativeCodeCall()) {
     if (!treeName.empty()) {
       PrintFatalError(
           def.getLoc(),
-          formatv("binding symbol '{0}' to non-operation unsupported right now",
-                  treeName));
+          formatv(
+              "binding symbol '{0}' to native code call unsupported right now",
     }
+
+    for (int i = 0; i != numTreeArgs; ++i) {
+      if (auto treeArg = tree.getArgAsNestedDag(i)) {
+        // This DAG node argument is a DAG node itself. Go inside recursively.
+        collectBoundSymbols(treeArg, infoMap, isSrcPattern);
+      } else if (isSrcPattern) {
+        // We can only bind symbols to arguments in source pattern. Those
+        // symbols are referenced in result patterns.
+        auto treeArgName = tree.getArgName(i);
+
+        // `$_` is a special symbol meaning ignore the current argument.
+        if (!treeArgName.empty() && treeArgName != "_") {
+          if (tree.isNestedDagArg(i)) {
+            auto err = formatv("cannot bind '{0}' for nested native call arg",
+                               treeArgName);
+            PrintFatalError(def.getLoc(), err);
+          }
+          DagLeaf leaf = tree.getArgAsLeaf(i);
+          if (leaf.isAttrMatcher() || leaf.isEnumAttrCase() ||
+              leaf.isConstantAttr()) {
+            if (!infoMap.bindAttr(treeArgName)) {
+              auto err =
+                  formatv("symbol '{0}' bound more than once", treeArgName);
+              PrintFatalError(def.getLoc(), err);
+            }
+          } else {
+            if (!infoMap.bindValue(treeArgName)) {
+              auto err =
+                  formatv("symbol '{0}' bound more than once", treeArgName);
+              PrintFatalError(def.getLoc(), err);
+            }
+          }
+        }
+      }
+    }
+
     return;
   }
 
   auto &op = getDialectOp(tree);
   auto numOpArgs = op.getNumArgs();
-  auto numTreeArgs = tree.getNumArgs();
 
   // The pattern might have the last argument specifying the location.
   bool hasLocDirective = false;
   if (numTreeArgs != 0) {
     if (auto lastArg = tree.getArgAsNestedDag(numTreeArgs - 1))
       hasLocDirective = lastArg.isLocationDirective();
   }
 

mlir/test/lib/Dialect/Test/TestDialect.cpp

void TestOpWithRegionPattern::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<TestRemoveOpWithInnerOps>(context);
 }
 
 OpFoldResult TestOpWithRegionFold::fold(ArrayRef<Attribute> operands) {
   return operand();
 }
 
+OpFoldResult TestOpConstant::fold(ArrayRef<Attribute> operands) {
+  return getValue();
+}
+
 LogicalResult TestOpWithVariadicResultsAndFolder::fold(
     ArrayRef<Attribute> operands, SmallVectorImpl<OpFoldResult> &results) {
   for (Value input : this->operands()) {
     results.push_back(input);
   }
   return success();
 }
 

mlir/test/lib/Dialect/Test/TestOps.td

def : Pat<(OpAllAttrConstraint1
           (OpAllAttrConstraint2 $attr)>;
 
 // Op for testing RewritePattern removing op with inner ops.
 def TestOpWithRegionPattern : TEST_Op<"op_with_region_pattern"> {
   let regions = (region SizedRegion<1>:$region);
   let hasCanonicalizer = 1;
 }
 
+def TestOpConstant : TEST_Op<"constant", [ConstantLike, NoSideEffect]> {
+  let arguments = (ins AnyAttr:$value);
+  let results = (outs AnyType);
+  let extraClassDeclaration = [{
+    Attribute getValue() { return getAttr("value"); }
+  }];
+
+  let hasFolder = 1;
+}
+
+def OpN : TEST_Op<"op_n">, Arguments<(ins AnyInteger, AnyInteger)>, Results<(outs AnyInteger)>;
+def OpO : TEST_Op<"op_o">, Arguments<(ins AnyInteger, AnyAttr:$value)>, Results<(outs AnyInteger)>;
+
+def : Pat<(OpN $input1, (ConstantLikeMatcher I32Attr:$input2)),
+          (OpO $input1, $input2)>;
+
 // Op for testing trivial removal via folding of op with inner ops and no uses.
 def TestOpWithRegionFoldNoSideEffect : TEST_Op<
     "op_with_region_fold_no_side_effect", [NoSideEffect]> {
   let regions = (region SizedRegion<1>:$region);
 }
 
 // Op for testing folding of outer op with inner ops.
 def TestOpWithRegionFold : TEST_Op<"op_with_region_fold"> {

mlir/test/lib/Dialect/Test/TestPatterns.cpp

 //===- TestPatterns.cpp - Test dialect pattern driver ---------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "TestDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/StandardOps/Transforms/FuncConversions.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/FoldUtils.h"
 
 using namespace mlir;
 
 // Native function for testing NativeCodeCall

mlir/test/mlir-tblgen/pattern.mlir

func @verifyUnitAttr() -> (i32, i32) {
   // Since the original op doesn't have the unit attribute, the new op
   // only has the constant-constructed unit attribute attr1.
   // CHECK: "test.match_op_attribute4"() {attr1} : () -> i32
   %1 = "test.match_op_attribute3"() : () -> i32
   return %0, %1 : i32, i32
 }
 
 //===----------------------------------------------------------------------===//
+// Test Constant Matching
+//===----------------------------------------------------------------------===//
+
+// CHECK-LABEL: testConstOp
+func @testConstOp() -> (i32) {
+  // CHECK-NEXT: [[C0:%.+]] = constant 1
+  %0 = "test.constant"() {value = 1 : i32} : () -> i32
+
+  // CHECK-NEXT: return [[C0]]
+  return %0 : i32
+}
+
+// CHECK-LABEL: testConstOpUsed
+func @testConstOpUsed() -> (i32) {
+  // CHECK-NEXT: [[C0:%.+]] = constant 1
+  %0 = "test.constant"() {value = 1 : i32} : () -> i32
+
+  // CHECK-NEXT: [[V0:%.+]] = "test.op_o"([[C0]])
+  %1 = "test.op_o"(%0) {value = 1 : i32} : (i32) -> i32
+
+  // CHECK-NEXT: return [[V0]]
+  return %1 : i32
+}
+
+// CHECK-LABEL: testConstOpReplaced
+func @testConstOpReplaced() -> (i32) {
+  // CHECK-NEXT: [[C0:%.+]] = constant 1
+  %0 = "test.constant"() {value = 1 : i32} : () -> i32
+  %1 = "test.constant"() {value = 2 : i32} : () -> i32
+
+  // CHECK: [[V0:%.+]] = "test.op_o"([[C0]]) {value = 2 : i32}
+  %2 = "test.op_n"(%0, %1) : (i32, i32) -> i32
+
+  // CHECK: [[V0]]
+  return %2 : i32
+}
+// CHECK-LABEL: testConstOpMatchFailure
+func @testConstOpMatchFailure() -> (i64) {
+  // CHECK-DAG: [[C0:%.+]] = constant 1
+  %0 = "test.constant"() {value = 1 : i64} : () -> i64
+
+  // CHECK-DAG: [[C1:%.+]] = constant 2
+  %1 = "test.constant"() {value = 2 : i64} : () -> i64
+
+  // CHECK: [[V0:%.+]] = "test.op_n"([[C0]], [[C1]])
+  %2 = "test.op_n"(%0, %1) : (i64, i64) -> i64
+
+  // CHECK: [[V0]]
+  return %2 : i64
+}
+
+//===----------------------------------------------------------------------===//
 // Test Enum Attributes
 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: verifyStrEnumAttr
 func @verifyStrEnumAttr() -> i32 {
   // CHECK: "test.str_enum_attr"() {attr = "B"}
   %0 = "test.str_enum_attr"() {attr = "A"} : () -> i32
   return %0 : i32

mlir/tools/mlir-tblgen/RewriterGen.cpp

 public:
   PatternEmitter(Record *pat, RecordOperatorMap *mapper, raw_ostream &os);
 
   // Emits the mlir::RewritePattern struct named `rewriteName`.
   void emit(StringRef rewriteName);
 
 private:
   // Emits the code for matching ops.
-  void emitMatchLogic(DagNode tree);
+  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);
+
+  // 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`.
-  void emitOpMatch(DagNode tree, int depth);
+  // `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
   // DAG `tree` as an operand.
-  void emitOperandMatch(DagNode tree, int argIndex, int depth);
+  void emitOperandMatch(DagNode tree, StringRef opName, int argIndex,
+                        int depth);
 
   // Emits C++ statements for matching the `argIndex`-th argument of the given
   // DAG `tree` as an attribute.
-  void emitAttributeMatch(DagNode tree, int argIndex, int depth);
+  void emitAttributeMatch(DagNode tree, StringRef opName, int argIndex,
+                          int depth);
 
   // Emits C++ for checking a match with a corresponding match failure
   // diagnostic.
-  void emitMatchCheck(int depth, const FmtObjectBase &matchFmt,
+  void emitMatchCheck(StringRef opName, const FmtObjectBase &matchFmt,
                       const llvm::formatv_object_base &failureFmt);
 
   //===--------------------------------------------------------------------===//
   // Rewrite utilities
   //===--------------------------------------------------------------------===//
 
   // The entry point for handling a result pattern rooted at `resultTree`. This
   // method dispatches to concrete handlers according to `resultTree`'s kind and
   // returns a symbol representing the whole value pack. Callers are expected to
   // further resolve the symbol according to the specific use case.
   //
   // `depth` is the nesting level of `resultTree`; 0 means top-level result
   // pattern. For top-level result pattern, `resultIndex` indicates which result
   // of the matched root op this pattern is intended to replace, which can be
   // used to deduce the result type of the op generated from this result
   // pattern.
   std::string handleResultPattern(DagNode resultTree, int resultIndex,
                                   int depth);
 
   // Emits the C++ statement to replace the matched DAG with a value built via
   // calling native C++ code.
-  std::string handleReplaceWithNativeCodeCall(DagNode resultTree);
+  std::string handleReplaceWithNativeCodeCall(DagNode resultTree, int depth);
 
   // Returns the symbol of the old value serving as the replacement.
   StringRef handleReplaceWithValue(DagNode tree);
 
   // Returns the location value to use.
   std::pair<bool, std::string> getLocation(DagNode tree);
 
   // Returns the location value to use.
private:
 
   // Emits a local variable for each value and attribute to be used for creating
   // an op.
   void createSeparateLocalVarsForOpArgs(DagNode node,
                                         ChildNodeIndexNameMap &childNodeNames);
 
   // Emits the concrete arguments used to call an op's builder.
   void supplyValuesForOpArgs(DagNode node,
-                             const ChildNodeIndexNameMap &childNodeNames);
+                             const ChildNodeIndexNameMap &childNodeNames,
+                             int depth);
 
   // Emits the local variables for holding all values as a whole and all named
   // attributes as a whole to be used for creating an op.
   void createAggregateLocalVarsForOpArgs(
-      DagNode node, const ChildNodeIndexNameMap &childNodeNames);
+      DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth);
 
   // Returns the C++ expression to construct a constant attribute of the given
   // `value` for the given attribute kind `attr`.
   std::string handleConstantAttr(Attribute attr, StringRef value);
 
   // Returns the C++ expression to build an argument from the given DAG `leaf`.
   // `patArgName` is used to bound the argument to the source pattern.
   std::string handleOpArgument(DagLeaf leaf, StringRef patArgName);
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));
 }
 
 // Helper function to match patterns.
-void PatternEmitter::emitOpMatch(DagNode tree, int depth) {
+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.");
+}
+
+// Helper function to match patterns.
+void PatternEmitter::emitNativeCodeMatch(DagNode tree, StringRef opName,
+                                         int depth) {
+  int indent = 4 + 2 * depth;
+  LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall matcher pattern: ");
+  LLVM_DEBUG(tree.print(llvm::dbgs()));
+  LLVM_DEBUG(llvm::dbgs() << '\n');
+
+  // TODO(suderman): iterate through arguments, determine their types, output
+  // names.
+  SmallVector<std::string, 8> capture(8);
+  if (tree.getNumArgs() > 8) {
+    PrintFatalError(loc,
+                    "unsupported NativeCodeCall matcher argument numbers: " +
+                        Twine(tree.getNumArgs()));
+  }
+
+  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);
+    if (DagNode argTree = tree.getArgAsNestedDag(i)) {
+      os.indent(indent) << "Value " << argName << ";\n";
+    } else {
+      auto leaf = tree.getArgAsLeaf(i);
+      if (leaf.isAttrMatcher() || leaf.isConstantAttr()) {
+        os.indent(indent) << "Attribute " << argName << ";\n";
+      } else if (leaf.isOperandMatcher()) {
+        os.indent(indent) << "Operation " << argName << ";\n";
+      }
+    }
+
+    capture[i] = std::move(argName);
+  }
+
+  bool hasLocationDirective;
+  std::string locToUse;
+  std::tie(hasLocationDirective, locToUse) = getLocation(tree);
+
+  auto fmt = tree.getNativeCodeTemplate();
+  auto nativeCodeCall = std::string(tgfmt(
+      fmt, &fmtCtx.addSubst("_loc", locToUse), opName, capture[0], capture[1],
+      capture[2], capture[3], capture[4], capture[5], capture[6], capture[7]));
+
+  os.indent(indent) << "if (failed(" << nativeCodeCall
+                    << ")) return failure();\n";
+
+  for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
+    auto name = tree.getArgName(i);
+    if (!name.empty() && name != "_") {
+      os.indent(indent) << formatv("{0} = {1};\n", name, capture[i]);
+    }
+  }
+
+  for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
+    std::string argName = capture[i];
+
+    // Handle nested DAG construct first
+    if (DagNode argTree = tree.getArgAsNestedDag(i)) {
+      PrintFatalError(
+          loc, formatv("Matching nested tree in NativeCodecall not support for "
+                       "{0} as arg {1}",
+                       argName, i));
+    }
+
+    DagLeaf leaf = tree.getArgAsLeaf(i);
+    auto constraint = leaf.getAsConstraint();
+
+    auto self = formatv("{0}", argName);
+    emitMatchCheck(
+        opName,
+        tgfmt(constraint.getConditionTemplate(), &fmtCtx.withSelf(self)),
+        formatv("\"operand {0} of native code call '{1}' failed to satisfy "
+                "constraint: "
+                "'{2}'\"",
+                i, tree.getNativeCodeTemplate(), constraint.getDescription()));
+  }
+
+  LLVM_DEBUG(llvm::dbgs() << "done emitting match for native code call\n");
+}
+
+// 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');
 
   int indent = 4 + 2 * depth;
-  os.indent(indent) << formatv(
-      "auto castedOp{0} = ::llvm::dyn_cast_or_null<{1}>(op{0}); "
-      "(void)castedOp{0};\n",
-      depth, op.getQualCppClassName());
+  std::string castedName = formatv("castedOp{0}", depth);
+  os.indent(indent) << formatv("auto {0} = ::llvm::dyn_cast_or_null<{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) {
     // Skip if there is no defining operation (e.g., arguments to function).
-    os << formatv("if (!castedOp{0})\n  return failure();\n", depth);
+    os.indent(indent) << formatv("if (!{0}) return failure();\n", castedName);
   }
   if (tree.getNumArgs() != op.getNumArgs()) {
     PrintFatalError(loc, formatv("op '{0}' argument number mismatch: {1} in "
                                  "pattern vs. {2} in definition",
                                  op.getOperationName(), tree.getNumArgs(),
                                  op.getNumArgs()));
   }
 
   // If the operand's name is set, set to that variable.
   auto name = tree.getSymbol();
   if (!name.empty())
-    os << formatv("{0} = castedOp{1};\n", name, depth);
+    os.indent(indent) << formatv("{0} = {1};\n", name, castedName);
 
   for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
     auto opArg = op.getArg(i);
+    std::string argName = formatv("op{0}", depth + 1);
 
     // Handle nested DAG construct first
     if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) {
         if (operand->isVariableLength()) {
           auto error = formatv("use nested DAG construct to match op {0}'s "
                                "variadic operand #{1} unsupported now",
                                op.getOperationName(), i);
           PrintFatalError(loc, error);
         }
       }
       os << "{\n";
 
-      os.indent() << formatv(
-          "auto *op{0} = "
-          "(*castedOp{1}.getODSOperands({2}).begin()).getDefiningOp();\n",
-          depth + 1, depth, i);
-      emitOpMatch(argTree, depth + 1);
-      os << formatv("tblgen_ops[{0}] = op{1};\n", ++opCounter, depth + 1);
-      os.unindent() << "}\n";
+      os.indent(indent + 2)
+          << formatv("auto *{0} = "
+                     "(*{1}.getODSOperands({2}).begin()).getDefiningOp();\n",
+                     argName, castedName, i);
+      emitMatch(argTree, argName, depth + 1);
+      os.indent(indent + 2)
+          << formatv("tblgen_ops[{0}] = {1};\n", ++opCounter, argName);
+      os.indent(indent) << "}\n";
       continue;
     }
 
     // Next handle DAG leaf: operand or attribute
     if (opArg.is<NamedTypeConstraint *>()) {
-      emitOperandMatch(tree, i, depth);
+      emitOperandMatch(tree, castedName, i, depth);
     } else if (opArg.is<NamedAttribute *>()) {
-      emitAttributeMatch(tree, i, depth);
+      emitAttributeMatch(tree, opName, i, depth);
     } else {
       PrintFatalError(loc, "unhandled case when matching op");
     }
   }
   LLVM_DEBUG(llvm::dbgs() << "done emitting match for op '"
                           << op.getOperationName() << "' at depth " << depth
                           << '\n');
 }
 
-void PatternEmitter::emitOperandMatch(DagNode tree, int argIndex, int depth) {
+void PatternEmitter::emitOperandMatch(DagNode tree, StringRef opName,
+                                      int argIndex, int depth) {
   Operator &op = tree.getDialectOp(opMap);
   auto *operand = op.getArg(argIndex).get<NamedTypeConstraint *>();
   auto matcher = tree.getArgAsLeaf(argIndex);
+  int indent = 4 + 2 * depth;
 
   // If a constraint is specified, we need to generate C++ statements to
   // check the constraint.
   if (!matcher.isUnspecified()) {
     if (!matcher.isOperandMatcher()) {
       PrintFatalError(
           loc, formatv("the {1}-th argument of op '{0}' should be an operand",
                        op.getOperationName(), argIndex + 1));
     }
 
     // Only need to verify if the matcher's type is different from the one
     // of op definition.
     Constraint constraint = matcher.getAsConstraint();
     if (operand->constraint != constraint) {
       if (operand->isVariableLength()) {
         auto error = formatv(
             "further constrain op {0}'s variadic operand #{1} unsupported now",
             op.getOperationName(), argIndex);
         PrintFatalError(loc, error);
       }
-      auto self =
-          formatv("(*castedOp{0}.getODSOperands({1}).begin()).getType()", depth,
-                  argIndex);
+      auto self = formatv("(*{0}.getODSOperands({1}).begin()).getType()",
+                          opName, argIndex);
       emitMatchCheck(
-          depth,
+          opName,
           tgfmt(constraint.getConditionTemplate(), &fmtCtx.withSelf(self)),
           formatv("\"operand {0} of op '{1}' failed to satisfy constraint: "
                   "'{2}'\"",
                   operand - op.operand_begin(), op.getOperationName(),
                   constraint.getDescription()));
     }
   }
 
   // Capture the value
   auto name = tree.getArgName(argIndex);
   // `$_` is a special symbol to ignore op argument matching.
   if (!name.empty() && name != "_") {
     // We need to subtract the number of attributes before this operand to get
     // the index in the operand list.
     auto numPrevAttrs = std::count_if(
         op.arg_begin(), op.arg_begin() + argIndex,
         [](const Argument &arg) { return arg.is<NamedAttribute *>(); });
 
-    os << formatv("{0} = castedOp{1}.getODSOperands({2});\n", name, depth,
-                  argIndex - numPrevAttrs);
+    os.indent(indent) << formatv("{0} = {1}.getODSOperands({2});\n", name,
+                                 opName, argIndex - numPrevAttrs);
   }
 }
 
-void PatternEmitter::emitAttributeMatch(DagNode tree, int argIndex, int depth) {
+void PatternEmitter::emitAttributeMatch(DagNode tree, StringRef opName,
+                                        int argIndex, int depth) {
   Operator &op = tree.getDialectOp(opMap);
   auto *namedAttr = op.getArg(argIndex).get<NamedAttribute *>();
   const auto &attr = namedAttr->attr;
+  int indent = 4 + 2 * depth;
 
-  os << "{\n";
-  os.indent() << formatv(
-      "auto tblgen_attr = op{0}->getAttrOfType<{1}>(\"{2}\"); "
+  os.indent(indent) << "{\n";
+  indent += 2;
+  os.indent(indent) << formatv(
+      "auto tblgen_attr = {0}->getAttrOfType<{1}>(\"{2}\");"
       "(void)tblgen_attr;\n",
-      depth, attr.getStorageType(), namedAttr->name);
+      opName, attr.getStorageType(), namedAttr->name);
 
   // TODO: This should use getter method to avoid duplication.
   if (attr.hasDefaultValue()) {
     os << "if (!tblgen_attr) tblgen_attr = "
        << std::string(tgfmt(attr.getConstBuilderTemplate(), &fmtCtx,
                             attr.getDefaultValue()))
        << ";\n";
   } else if (attr.isOptional()) {
     // For a missing attribute that is optional according to definition, we
     // should just capture a mlir::Attribute() to signal the missing state.
     // That is precisely what getAttr() returns on missing attributes.
   } else {
-    emitMatchCheck(depth, tgfmt("tblgen_attr", &fmtCtx),
+    emitMatchCheck(opName, tgfmt("tblgen_attr", &fmtCtx),
                    formatv("\"expected op '{0}' to have attribute '{1}' "
                            "of type '{2}'\"",
                            op.getOperationName(), namedAttr->name,
                            attr.getStorageType()));
   }
 
   auto matcher = tree.getArgAsLeaf(argIndex);
   if (!matcher.isUnspecified()) {
     if (!matcher.isAttrMatcher()) {
       PrintFatalError(
           loc, formatv("the {1}-th argument of op '{0}' should be an attribute",
                        op.getOperationName(), argIndex + 1));
     }
 
     // If a constraint is specified, we need to generate C++ statements to
     // check the constraint.
     emitMatchCheck(
-        depth,
+        opName,
         tgfmt(matcher.getConditionTemplate(), &fmtCtx.withSelf("tblgen_attr")),
         formatv("\"op '{0}' attribute '{1}' failed to satisfy constraint: "
                 "{2}\"",
                 op.getOperationName(), namedAttr->name,
                 matcher.getAsConstraint().getDescription()));
   }
 
   // Capture the value
   auto name = tree.getArgName(argIndex);
   // `$_` is a special symbol to ignore op argument matching.
   if (!name.empty() && name != "_") {
     os << formatv("{0} = tblgen_attr;\n", name);
   }
 
   os.unindent() << "}\n";
 }
 
 void PatternEmitter::emitMatchCheck(
-    int depth, const FmtObjectBase &matchFmt,
+    StringRef opName, const FmtObjectBase &matchFmt,
     const llvm::formatv_object_base &failureFmt) {
-  os << "if (!(" << matchFmt.str() << "))";
-  os.scope("{\n", "\n}\n").os
-      << "return rewriter.notifyMatchFailure(op" << depth
-      << ", [&](::mlir::Diagnostic &diag) {\n  diag << " << failureFmt.str()
-      << ";\n});";
+  const char *matchStr = R"(
+    if (!({1})) {
+      return rewriter.notifyMatchFailure({0}, [&](::mlir::Diagnostic &diag) {
+        diag << {2};
+      });
+    })";
+  os << llvm::formatv(matchStr, opName, matchFmt.str(), failureFmt.str())
+     << "\n";
 }
 
-void PatternEmitter::emitMatchLogic(DagNode tree) {
+void PatternEmitter::emitMatchLogic(DagNode tree, StringRef opName) {
   LLVM_DEBUG(llvm::dbgs() << "--- start emitting match logic ---\n");
   int depth = 0;
-  emitOpMatch(tree, depth);
+  emitMatch(tree, opName, depth);
 
   for (auto &appliedConstraint : pattern.getConstraints()) {
     auto &constraint = appliedConstraint.constraint;
     auto &entities = appliedConstraint.entities;
 
     auto condition = constraint.getConditionTemplate();
     if (isa<TypeConstraint>(constraint)) {
       auto self = formatv("({0}.getType())",
                           symbolInfoMap.getValueAndRangeUse(entities.front()));
       emitMatchCheck(
-          depth, tgfmt(condition, &fmtCtx.withSelf(self.str())),
+          opName, tgfmt(condition, &fmtCtx.withSelf(self.str())),
           formatv("\"value entity '{0}' failed to satisfy constraint: {1}\"",
                   entities.front(), constraint.getDescription()));
 
     } else if (isa<AttrConstraint>(constraint)) {
       PrintFatalError(
           loc, "cannot use AttrConstraint in Pattern multi-entity constraints");
     } else {
       // TODO: replace formatv arguments with the exact specified
       // args.
       if (entities.size() > 4) {
         PrintFatalError(loc, "only support up to 4-entity constraints now");
       }
       SmallVector<std::string, 4> names;
       int i = 0;
       for (int e = entities.size(); i < e; ++i)
         names.push_back(symbolInfoMap.getValueAndRangeUse(entities[i]));
       std::string self = appliedConstraint.self;
       if (!self.empty())
         self = symbolInfoMap.getValueAndRangeUse(self);
       for (; i < 4; ++i)
         names.push_back("<unused>");
-      emitMatchCheck(depth,
+      emitMatchCheck(opName,
                      tgfmt(condition, &fmtCtx.withSelf(self), names[0],
                            names[1], names[2], names[3]),
                      formatv("\"entities '{0}' failed to satisfy constraint: "
                              "{1}\"",
                              llvm::join(entities, ", "),
                              constraint.getDescription()));
     }
   }
::mlir::LogicalResult matchAndRewrite(::mlir::Operation *op0,
       // reused for fused loc.
       os << formatv("::mlir::Operation *tblgen_ops[{0}];\n\n",
                     pattern.getSourcePattern().getNumOps());
       LLVM_DEBUG(llvm::dbgs()
                  << "done creating local variables for capturing matches\n");
 
       os << "// Match\n";
       os << "tblgen_ops[0] = op0;\n";
-      emitMatchLogic(sourceTree);
+      emitMatchLogic(sourceTree, "op0");
 
       os << "\n// Rewrite\n";
       emitRewriteLogic();
 
       os << "return ::mlir::success();\n";
     }
     os << "};\n";
   }
std::string PatternEmitter::handleResultPattern(DagNode resultTree,
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   if (resultTree.isLocationDirective()) {
     PrintFatalError(loc,
                     "location directive can only be used with op creation");
   }
 
   if (resultTree.isNativeCodeCall()) {
-    auto symbol = handleReplaceWithNativeCodeCall(resultTree);
+    auto symbol = handleReplaceWithNativeCodeCall(resultTree, depth);
     symbolInfoMap.bindValue(symbol);
     return symbol;
   }
 
   if (resultTree.isReplaceWithValue())
     return handleReplaceWithValue(resultTree).str();
 
   // Normal op creation.
if (leaf.isNativeCodeCall()) {
     auto repl = tgfmt(leaf.getNativeCodeTemplate(), &fmtCtx.withSelf(argName));
     LLVM_DEBUG(llvm::dbgs() << "replace " << patArgName << " with '" << repl
                             << "' (via NativeCodeCall)\n");
     return std::string(repl);
   }
   PrintFatalError(loc, "unhandled case when rewriting op");
 }
 
-std::string PatternEmitter::handleReplaceWithNativeCodeCall(DagNode tree) {
+std::string PatternEmitter::handleReplaceWithNativeCodeCall(DagNode tree,
+                                                            int depth) {
   LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   auto fmt = tree.getNativeCodeTemplate();
   // TODO: replace formatv arguments with the exact specified args.
   SmallVector<std::string, 8> attrs(8);
   if (tree.getNumArgs() > 8) {
-    PrintFatalError(loc, "unsupported NativeCodeCall argument numbers: " +
-                             Twine(tree.getNumArgs()));
+    PrintFatalError(loc,
+                    "unsupported NativeCodeCall replace argument numbers: " +
   }
   bool hasLocationDirective;
   std::string locToUse;
   std::tie(hasLocationDirective, locToUse) = getLocation(tree);
 
   for (int i = 0, e = tree.getNumArgs() - hasLocationDirective; i != e; ++i) {
+    if (tree.isNestedDagArg(i)) {
+      attrs[i] = handleResultPattern(tree.getArgAsNestedDag(i), i, depth + 1);
+    } else {
       attrs[i] = handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i));
+    }
     LLVM_DEBUG(llvm::dbgs() << "NativeCodeCall argument #" << i
                             << " replacement: " << attrs[i] << "\n");
   }
   return std::string(tgfmt(fmt, &fmtCtx.addSubst("_loc", locToUse), attrs[0],
                            attrs[1], attrs[2], attrs[3], attrs[4], attrs[5],
                            attrs[6], attrs[7]));
 }
 
bool useFirstAttr =
       resultOp.getTrait("::mlir::OpTrait::FirstAttrDerivedResultType");
 
   if (isSameOperandsAndResultType || useFirstAttr) {
     // We know how to deduce the result type for ops with these traits and we've
     // generated builders taking aggregate parameters. Use those builders to
     // create the ops.
 
     // First prepare local variables for op arguments used in builder call.
-    createAggregateLocalVarsForOpArgs(tree, childNodeNames);
+    createAggregateLocalVarsForOpArgs(tree, childNodeNames, depth);
 
     // Then create the op.
     os.scope("", "\n}\n").os << formatv(
         "{0} = rewriter.create<{1}>({2}, tblgen_values, tblgen_attrs);",
         valuePackName, resultOp.getQualCppClassName(), locToUse);
     return resultValue;
   }
 
   bool usePartialResults = valuePackName != resultValue;
 
   if (usePartialResults || depth > 0 || resultIndex < 0) {
     // For these cases (broadcastable ops, op results used both as auxiliary
     // values and replacement values, ops in nested patterns, auxiliary ops), we
     // still need to supply the result types when building the op. But because
     // we don't generate a builder automatically with ODS for them, it's the
     // developer's responsibility to make sure such a builder (with result type
     // deduction ability) exists. We go through the separate-parameter builder
     // here given that it's easier for developers to write compared to
     // aggregate-parameter builders.
     createSeparateLocalVarsForOpArgs(tree, childNodeNames);
 
     os.scope().os << formatv("{0} = rewriter.create<{1}>({2}", valuePackName,
                              resultOp.getQualCppClassName(), locToUse);
-    supplyValuesForOpArgs(tree, childNodeNames);
+    supplyValuesForOpArgs(tree, childNodeNames, depth);
     os << "\n  );\n}\n";
     return resultValue;
   }
 
   // If depth == 0 and resultIndex >= 0, it means we are replacing the values
   // generated from the source pattern root op. Then we can use the source
   // pattern's value types to determine the value type of the generated op
   // here.
 
   // First prepare local variables for op arguments used in builder call.
-  createAggregateLocalVarsForOpArgs(tree, childNodeNames);
+  createAggregateLocalVarsForOpArgs(tree, childNodeNames, depth);
 
   // Then prepare the result types. We need to specify the types for all
   // results.
   os.indent() << formatv("::mlir::SmallVector<::mlir::Type, 4> tblgen_types; "
                          "(void)tblgen_types;\n");
   int numResults = resultOp.getNumResults();
   if (numResults != 0) {
     for (int i = 0; i < numResults; ++i)
for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
     }
 
     // Update to use the newly created local variable for building the op later.
     childNodeNames[argIndex] = varName;
   }
 }
 
 void PatternEmitter::supplyValuesForOpArgs(
-    DagNode node, const ChildNodeIndexNameMap &childNodeNames) {
+    DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth) {
   Operator &resultOp = node.getDialectOp(opMap);
   for (int argIndex = 0, numOpArgs = resultOp.getNumArgs();
        argIndex != numOpArgs; ++argIndex) {
     // Start each argument on its own line.
     os << ",\n    ";
 
     Argument opArg = resultOp.getArg(argIndex);
     // Handle the case of operand first.
     if (auto *operand = opArg.dyn_cast<NamedTypeConstraint *>()) {
       if (!operand->name.empty())
         os << "/*" << operand->name << "=*/";
       os << childNodeNames.lookup(argIndex);
       continue;
     }
 
     // The argument in the op definition.
     auto opArgName = resultOp.getArgName(argIndex);
     if (auto subTree = node.getArgAsNestedDag(argIndex)) {
       if (!subTree.isNativeCodeCall())
         PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "
                              "for creating attribute");
       os << formatv("/*{0}=*/{1}", opArgName,
-                    handleReplaceWithNativeCodeCall(subTree));
+                    handleReplaceWithNativeCodeCall(subTree, depth));
     } else {
       auto leaf = node.getArgAsLeaf(argIndex);
       // The argument in the result DAG pattern.
       auto patArgName = node.getArgName(argIndex);
       if (leaf.isConstantAttr() || leaf.isEnumAttrCase()) {
         // TODO: Refactor out into map to avoid recomputing these.
         if (!opArg.is<NamedAttribute *>())
           PrintFatalError(loc, Twine("expected attribute ") + Twine(argIndex));
         if (!patArgName.empty())
           os << "/*" << patArgName << "=*/";
       } else {
         os << "/*" << opArgName << "=*/";
       }
       os << handleOpArgument(leaf, patArgName);
     }
   }
 }
 
 void PatternEmitter::createAggregateLocalVarsForOpArgs(
-    DagNode node, const ChildNodeIndexNameMap &childNodeNames) {
+    DagNode node, const ChildNodeIndexNameMap &childNodeNames, int depth) {
   Operator &resultOp = node.getDialectOp(opMap);
 
   auto scope = os.scope();
   os << formatv("::mlir::SmallVector<::mlir::Value, 4> "
                 "tblgen_values; (void)tblgen_values;\n");
   os << formatv("::mlir::SmallVector<::mlir::NamedAttribute, 4> "
                 "tblgen_attrs; (void)tblgen_attrs;\n");
 
   const char *addAttrCmd =
       "if (auto tmpAttr = {1}) {\n"
       "  tblgen_attrs.emplace_back(rewriter.getIdentifier(\"{0}\"), "
       "tmpAttr);\n}\n";
   for (int argIndex = 0, e = resultOp.getNumArgs(); argIndex < e; ++argIndex) {
     if (resultOp.getArg(argIndex).is<NamedAttribute *>()) {
       // The argument in the op definition.
       auto opArgName = resultOp.getArgName(argIndex);
       if (auto subTree = node.getArgAsNestedDag(argIndex)) {
         if (!subTree.isNativeCodeCall())
           PrintFatalError(loc, "only NativeCodeCall allowed in nested dag node "
                                "for creating attribute");
         os << formatv(addAttrCmd, opArgName,
-                      handleReplaceWithNativeCodeCall(subTree));
+                      handleReplaceWithNativeCodeCall(subTree, depth + 1));
       } else {
         auto leaf = node.getArgAsLeaf(argIndex);
         // The argument in the result DAG pattern.
         auto patArgName = node.getArgName(argIndex);
         os << formatv(addAttrCmd, opArgName,
                       handleOpArgument(leaf, patArgName));
       }
       continue;