[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

2354

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

224

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

629–630

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

657

nested?

662

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

666

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

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

817

You can skip this if empty.

mlir/tools/mlir-tblgen/RewriterGen.cpp

261

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

224

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.

662

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

666

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

261

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

224

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?

628

Could we do an early exist instead with positive logic?

657

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).

659

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

678

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

mlir/tools/mlir-tblgen/RewriterGen.cpp

247

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.

450

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

657

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

657

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

247

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

[jpienaar]

mlir/lib/TableGen/Pattern.cpp

657

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

629–630

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;

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 an iterator to the information of the given symbol named as `key`,
   // with index `argIndex` for operator `op`.
public:
 
   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);
 
   // The TableGen definition of this pattern.
   const llvm::Record &def;
 

mlir/lib/TableGen/Pattern.cpp

 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 << "': ");
   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));
     }
+    // TODO(suderman): Use a more exact type when available.
+    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",
                 getVarName(name)));
   }
   case Kind::Value: {
bool SymbolInfoMap::bindOpResult(StringRef symbol, const Operator &op) {
   return symbolInfoMap.count(inserted->first) == 1;
 }
 
 bool SymbolInfoMap::bindValue(StringRef symbol) {
   auto inserted = symbolInfoMap.emplace(symbol.str(), SymbolInfo::getValue());
   return symbolInfoMap.count(inserted->first) == 1;
 }
 
+bool SymbolInfoMap::bindAttr(StringRef symbol) {
+  auto inserted = symbolInfoMap.emplace(symbol, SymbolInfo::getAttr());
+  return symbolInfoMap.count(inserted->first) == 1;
+}
+
 bool SymbolInfoMap::contains(StringRef symbol) const {
   return find(symbol) != symbolInfoMap.end();
 }
 
 SymbolInfoMap::const_iterator SymbolInfoMap::find(StringRef key) const {
   std::string name = getValuePackName(key).str();
 
   return symbolInfoMap.find(name);

 std::vector<AppliedConstraint> Pattern::getConstraints() const {
   auto *listInit = def.getValueAsListInit("constraints");
   std::vector<AppliedConstraint> ret;
   ret.reserve(listInit->size());
 
   for (auto it : *listInit) {
     auto *dagInit = dyn_cast<llvm::DagInit>(it);
     if (!dagInit)
-      PrintFatalError(def.getLoc(), "all elements in Pattern multi-entity "
-                                    "constraints should be DAG nodes");
+      PrintFatalError(&def, "all elements in Pattern multi-entity "
 
     std::vector<std::string> entities;
     entities.reserve(dagInit->arg_size());
     for (auto *argName : dagInit->getArgNames()) {
       if (!argName) {
         PrintFatalError(
-            def.getLoc(),
+            &def,
             "operands to additional constraints can only be symbol references");
       }
       entities.push_back(std::string(argName->getValue()));
     }
 
     ret.emplace_back(cast<llvm::DefInit>(dagInit->getOperator())->getDef(),
                      dagInit->getNameStr(), std::move(entities));
   }
   return ret;
 }
 
 int Pattern::getBenefit() const {
   // The initial benefit value is a heuristic with number of ops in the source
   // pattern.
   int initBenefit = getSourcePattern().getNumOps();
   llvm::DagInit *delta = def.getValueAsDag("benefitDelta");
   if (delta->getNumArgs() != 1 || !isa<llvm::IntInit>(delta->getArg(0))) {
-    PrintFatalError(def.getLoc(),
+    PrintFatalError(&def,
                     "The 'addBenefit' takes and only takes one integer value");
   }
   return initBenefit + dyn_cast<llvm::IntInit>(delta->getArg(0))->getValue();
 }
 
 std::vector<Pattern::IdentifierLine> Pattern::getLocation() const {
   std::vector<std::pair<StringRef, unsigned>> result;
   result.reserve(def.getLoc().size());
   for (auto loc : def.getLoc()) {
     unsigned buf = llvm::SrcMgr.FindBufferContainingLoc(loc);
     assert(buf && "invalid source location");
     result.emplace_back(
         llvm::SrcMgr.getBufferInfo(buf).Buffer->getBufferIdentifier(),
         llvm::SrcMgr.getLineAndColumn(loc, buf).first);
   }
   return result;
 }
 
+void Pattern::verifyBind(bool result, StringRef symbolName) {
+  if (!result) {
+    auto err = formatv("symbol '{0}' bound more than once", symbolName);
+    PrintFatalError(&def, err);
+  }
+}
+
 void Pattern::collectBoundSymbols(DagNode tree, SymbolInfoMap &infoMap,
                                   bool isSrcPattern) {
   auto treeName = tree.getSymbol();
-  if (!tree.isOperation()) {
+  auto numTreeArgs = tree.getNumArgs();
+
+  if (tree.isNativeCodeCall()) {
     if (!treeName.empty()) {
       PrintFatalError(
-          def.getLoc(),
-          formatv("binding symbol '{0}' to non-operation unsupported right now",
-                  treeName));
+          &def,
+          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);
+        continue;
+      }
+
+      if (!isSrcPattern)
+        continue;
+
+      // 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, err);
+        }
+
+        DagLeaf leaf = tree.getArgAsLeaf(i);
+        auto constraint = leaf.getAsConstraint();
+        bool isAttr = leaf.isAttrMatcher() || leaf.isEnumAttrCase() ||
+                      leaf.isConstantAttr() ||
+                      constraint.getKind() == Constraint::Kind::CK_Attr;
+
+        if (isAttr) {
+          verifyBind(infoMap.bindAttr(treeArgName), treeArgName);
+          continue;
+        }
+
+        verifyBind(infoMap.bindValue(treeArgName), treeArgName);
+      }
+    }
+
     return;
   }
 
+  if (tree.isOperation()) {
     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();
     }
 
     if (numOpArgs != numTreeArgs - hasLocDirective) {
       auto err = formatv("op '{0}' argument number mismatch: "
                          "{1} in pattern vs. {2} in definition",
                          op.getOperationName(), numTreeArgs, numOpArgs);
-    PrintFatalError(def.getLoc(), err);
+      PrintFatalError(&def, err);
     }
 
     // The name attached to the DAG node's operator is for representing the
     // results generated from this op. It should be remembered as bound results.
     if (!treeName.empty()) {
       LLVM_DEBUG(llvm::dbgs()
                  << "found symbol bound to op result: " << treeName << '\n');
-    if (!infoMap.bindOpResult(treeName, op))
-      PrintFatalError(def.getLoc(),
-                      formatv("symbol '{0}' bound more than once", treeName));
+      verifyBind(infoMap.bindOpResult(treeName, op), treeName);
     }
 
     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) {
+        continue;
+      }
+
+      if (isSrcPattern) {
         // We can only bind symbols to op 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 != "_") {
           LLVM_DEBUG(llvm::dbgs() << "found symbol bound to op argument: "
                                   << treeArgName << '\n');
-        if (!infoMap.bindOpArgument(treeArgName, op, i)) {
-          auto err = formatv("symbol '{0}' bound more than once", treeArgName);
-          PrintFatalError(def.getLoc(), err);
+          verifyBind(infoMap.bindOpArgument(treeArgName, op, i), treeArgName);
         }
       }
     }
+    return;
+  }
+
+  if (!treeName.empty()) {
+    PrintFatalError(
+        &def, formatv("binding symbol '{0}' to non-operation/native code call "
+                      "unsupported right now",
+                      treeName));
   }
+  return;
 }

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 OpR : TEST_Op<"op_r">, Arguments<(ins AnyInteger, AnyInteger)>, Results<(outs AnyInteger)>;
+def OpS : TEST_Op<"op_s">, Arguments<(ins AnyInteger, AnyAttr:$value)>, Results<(outs AnyInteger)>;
+
+def : Pat<(OpR $input1, (ConstantLikeMatcher I32Attr:$input2)),
+          (OpS:$unused $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_s"([[C0]])
+  %1 = "test.op_s"(%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_s"([[C0]]) {value = 2 : i32}
+  %2 = "test.op_r"(%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_r"([[C0]], [[C1]])
+  %2 = "test.op_r"(%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/test/mlir-tblgen/rewriter-errors.td

+// RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR1 %s 2>&1 | FileCheck --check-prefix=ERROR1 %s
+// RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR2 %s 2>&1 | FileCheck --check-prefix=ERROR2 %s
+
+include "mlir/IR/OpBase.td"
+
+// Check using the dialect name as the namespace
+def A_Dialect : Dialect {
+  let name = "a";
+}
+
+class A_Op<string mnemonic, list<OpTrait> traits = []> :
+    Op<A_Dialect, mnemonic, traits>;
+
+def OpA : A_Op<"op_a">, Arguments<(ins AnyInteger, AnyInteger)>, Results<(outs AnyInteger)>;
+def OpB : A_Op<"op_b">, Arguments<(ins AnyInteger, AnyAttr:$value)>, Results<(outs AnyInteger)>;
+
+#ifdef ERROR1
+def NativeMatcher : NativeCodeCall<"success(nativeCall($0, $1))">;
+// ERROR1: [[@LINE+1]]:1: error: binding symbol 'error' to native code call unsupported right now
+def : Pat<(OpA (NativeMatcher:$error $val), AnyI32Attr:$arg),
+          (OpB $val, $arg)>;
+#endif
+
+#ifdef ERROR2
+def NativeMatcher : NativeCodeCall<"success(nativeCall($0, $1))">;
+// ERROR2: [[@LINE+1]]:1: error: Matching nested tree in NativeCodecall not support for 
+def : Pat<(OpA (NativeMatcher (OpB $val, $unused)), AnyI32Attr:$arg),
+          (OpB $val, $arg)>;
+#endif

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);
 
   // Emits C++ for checking a match with a corresponding match failure
   // diagnostics.
-  void emitMatchCheck(int depth, const std::string &matchStr,
+  void emitMatchCheck(StringRef opName, const std::string &matchStr,
                       const std::string &failureStr);
 
   //===--------------------------------------------------------------------===//
   // 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) {
+  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 << "Value " << argName << ";\n";
+    } else {
+      auto leaf = tree.getArgAsLeaf(i);
+      if (leaf.isAttrMatcher() || leaf.isConstantAttr()) {
+        os << "Attribute " << argName << ";\n";
+      } else if (leaf.isOperandMatcher()) {
+        os << "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 << "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 << 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 << 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 << 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 << 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);
+          "auto *{0} = "
+          "(*{1}.getODSOperands({2}).begin()).getDefiningOp();\n",
+          argName, castedName, i);
+      emitMatch(argTree, argName, depth + 1);
+      os << formatv("tblgen_ops[{0}] = {1};\n", ++opCounter, argName);
       os.unindent() << "}\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);
 
   // 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 *>(); });
 
     auto res = symbolInfoMap.findBoundSymbol(name, op, argIndex);
-    os << formatv("{0} = castedOp{1}.getODSOperands({2});\n",
-                  res->second.getVarName(name), depth, argIndex - numPrevAttrs);
+    os << formatv("{0} = {1}.getODSOperands({2});\n",
+                  res->second.getVarName(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;
 
   os << "{\n";
-  os.indent() << formatv(
-      "auto tblgen_attr = op{0}->getAttrOfType<{1}>(\"{2}\"); "
+  os.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) {
-  emitMatchCheck(depth, matchFmt.str(), failureFmt.str());
+  emitMatchCheck(opName, matchFmt.str(), failureFmt.str());
 }
 
-void PatternEmitter::emitMatchCheck(int depth, const std::string &matchStr,
+void PatternEmitter::emitMatchCheck(StringRef opName,
+                                    const std::string &matchStr,
                                     const std::string &failureStr) {
+
   os << "if (!(" << matchStr << "))";
-  os.scope("{\n", "\n}\n").os
-      << "return rewriter.notifyMatchFailure(op" << depth
-      << ", [&](::mlir::Diagnostic &diag) {\n  diag << " << failureStr
-      << ";\n});";
+  os.scope("{\n", "\n}\n").os << "return rewriter.notifyMatchFailure(" << opName
+                              << ", [&](::mlir::Diagnostic &diag) {\n  diag << "
+                              << failureStr << ";\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()));
     }
   }
 
   // Some of the operands could be bound to the same symbol name, we need
   // to enforce equality constraint on those.
   // TODO: we should be able to emit equality checks early
   // and short circuit unnecessary work if vars are not equal.
   for (auto symbolInfoIt = symbolInfoMap.begin();
        symbolInfoIt != symbolInfoMap.end();) {
     auto range = symbolInfoMap.getRangeOfEqualElements(symbolInfoIt->first);
     auto startRange = range.first;
     auto endRange = range.second;
 
     auto firstOperand = symbolInfoIt->second.getVarName(symbolInfoIt->first);
     for (++startRange; startRange != endRange; ++startRange) {
       auto secondOperand = startRange->second.getVarName(symbolInfoIt->first);
       emitMatchCheck(
-          depth,
+          opName,
           formatv("*{0}.begin() == *{1}.begin()", firstOperand, secondOperand),
           formatv("\"Operands '{0}' and '{1}' must be equal\"", firstOperand,
                   secondOperand));
     }
 
     symbolInfoIt = endRange;
   }
 
::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;