diff --git a/mlir/test/mlir-tblgen/rewriter-indexing.td b/mlir/test/mlir-tblgen/rewriter-indexing.td
--- a/mlir/test/mlir-tblgen/rewriter-indexing.td
+++ b/mlir/test/mlir-tblgen/rewriter-indexing.td
@@ -47,3 +47,11 @@
 // CHECK: castedOp0.getODSOperands(1).begin()).getDefiningOp()
 def test2 : Pat<(COp $attr1, $op1, $attr2, (AOp $op2)),
                 (BOp $attr1, $op2)>;
+
+
+// Check rewriting with a DAG subtree in the result and remapping a location.
+// CHECK: struct test3 : public ::mlir::RewritePattern {
+// CHECK: rewriter.create<test::BOp>((*a.getODSResults(0).begin()).getLoc()
+def test3 : Pat<(BOp $attr, (AOp:$a $input)),
+                (BOp $attr, (AOp $input), (location $a))>;
+
diff --git a/mlir/tools/mlir-tblgen/RewriterGen.cpp b/mlir/tools/mlir-tblgen/RewriterGen.cpp
--- a/mlir/tools/mlir-tblgen/RewriterGen.cpp
+++ b/mlir/tools/mlir-tblgen/RewriterGen.cpp
@@ -998,7 +998,7 @@
   // First go through all the child nodes who are nested DAG constructs to
   // create ops for them and remember the symbol names for them, so that we can
   // use the results in the current node. This happens in a recursive manner.
-  for (int i = 0, e = resultOp.getNumOperands(); i != e; ++i) {
+  for (int i = 0, e = tree.getNumArgs() - hasLocationDirective; i != e; ++i) {
     if (auto child = tree.getArgAsNestedDag(i))
       childNodeNames[i] = handleResultPattern(child, i, depth + 1);
   }