[DRR] Explicit Types and Types Builders in Rewrites

[Mogball]

Summary

Adds a rewrite directive OpWithResultTypeBuilders that allows
operations to be created in rewrites that are not able to infer or
deduce a return type. Types can be specified with a string:

(OpWithResultType<AnOp, "IntegerType::get(getContext(), 32)"> $val)

Or can be built using NativeCodeCall:

(OpWithResultTypeBuilder<AnOp, (NativeCodeCall<"$0.getType()"> $arg)>
$val)

[Mogball]

clang-format files

[Mogball]

fixing commit?

[Mogball]

rebase against HEAD

[jpienaar]

Mostly looks good, thanks! One part to handle to avoid changing return types.

mlir/include/mlir/IR/OpBase.td

2603

s/with/without/ ?

2618

Why can't one use bound values in string one too?

2624

Consider just indenting here and above with 2 spaces rather than lining up ( to avoid such long lines. You could also use op X in the pattern and just in a comment say X is an op without type deduction

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

1341

OOC is there any way to have a name independent one too? E.g., SameTypeAsFirstOperandOfOperatorMatched? :)

mlir/tools/mlir-tblgen/RewriterGen.cpp

1188–1195

In DRR we require that the root node's type during replacement remains as is. So we don't allow the type to change in DRR rules (we wouldn't generally know what casts to assert or type mismatches are allowed)

[Mogball]

Thanks for the review!

mlir/include/mlir/IR/OpBase.td

2618

E.g. you can't do

("doSomething($0)" $arg0)

But you can do

(NativeCodeCall<"doSomething($0)"> $arg0)

Oh, I should improve the wording.

You can use $_builder, for example, in strings since they're wrapped in NativeCodeCall anyways, but of course you'll need to specify NativeCodeCall in order to pass arguments.

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

1341

Yeah. The code snippet is pasted into matchAndRewrite so you could have "castedOp0.getOperand(0).getType()", as long as your okay with hardcoding the table-generated root op name,

mlir/tools/mlir-tblgen/RewriterGen.cpp

1188–1195

True. Let me fix this.

[Mogball]

• prevent explicit types in values replacing root op's return values

[Chia-hungDuan]

About the syntax, I'm wondering if we can make it similar to what we have so far in DRR. For example,

(OpWithResultTypeBuilder<AnOp, (NativeCodeCall<"$0.getType()"> $arg)>
$val)

It nests DAG into template arg. While reading a DRR pattern, we may expect the DAG is a matcher/rewriter for a operation/operand/attribute and now it has different semantics.

I'm thinking a new type of DAG node which indicates the matching/rewriting of the return type

Pat<(FooOp ..., (ReturnType<I32>:$ret)),
        (BarOp arg1, arg2, (ReturnType $ret))>

or

Pat<(FooOp ...),
        (BarOp arg1, arg2, (ReturnType (NativeCodeCall<"IntegerType::get(getContext(), 32)">)))>

What do you think?

[Mogball]

In D109472#2993055, @Chia-hungDuan wrote:

About the syntax, I'm wondering if we can make it similar to what we have so far in DRR. For example,

(OpWithResultTypeBuilder<AnOp, (NativeCodeCall<"$0.getType()"> $arg)>
$val)

It nests DAG into template arg. While reading a DRR pattern, we may expect the DAG is a matcher/rewriter for a operation/operand/attribute and now it has different semantics.

I do agree that the syntax isn't perfect, especially since with template arguments we need 4 different classes: OpWithReturnType, OpWithReturnTypes, OpWithReturnTypeBuilder, and OpWithReturnTypeBuilders. But the DAG still describes a rewriter for a type, it's just placed inside a template arg. Also, the class' names are a bit verbose.

I'm thinking a new type of DAG node which indicates the matching/rewriting of the return type

I like the idea of introducing a matcher for return types, because currently (I think) the only way to match return types is with constraints. The current syntax could also be used for that; it's just not implemented in the rewriter gen (e.g. OpWithResultType<AnOp, I32>:$res). Also, I think I can find a way to deduplicate the classes.

Pat<(FooOp ..., (ReturnType<I32>:$ret)),
    (BarOp arg1, arg2, (ReturnType $ret))>

Return values are typically bound at the DAG node: FooOp:$ret and I assume ReturnType $ret actually means to copy the return type, so it should be (ReturnType (NativeCodeCall<"$0.getType()> $ret)). Alternatively, an explicit type could be like (ReturnType "$_builder.getI32Type()"); I would first propose the following change to your syntax:

Pat<(FooOp:$res $arg1, $arg2, (ReturnType<I32>)),
    (BarOp $arg1, $arg2, (ReturnType "$_builder.getI32Type()"))>;

Assuming I reduce 4 classes down to 1 and introduce matchers for return types as well with the same OpWithResultType syntax, the difference boils down to whether you'd prefer the return type be specified as a trailing DAG node or next to the operation inside a template argument. I personally think it's a little bit weird to have a return type specifier after the operands and attributes, because the expectation of the DAG arguments is that they match the arguments = (ins ...) in the op declaration.

Also, having the return type specifier as a DAG argument would be more complex to implement. You'd have to ensure that there's only one return type specifier, that it's at the end of the argument list (or skip over it if would be allowed to be specified anywhere else), and it would have to be associated back to the operation.

What's your opinion?

[Mogball]

• add test that fails when trying to replace root types

[Chia-hungDuan]

In D109472#2993420, @Mogball wrote:

In D109472#2993055, @Chia-hungDuan wrote:

About the syntax, I'm wondering if we can make it similar to what we have so far in DRR. For example,

(OpWithResultTypeBuilder<AnOp, (NativeCodeCall<"$0.getType()"> $arg)>
$val)

It nests DAG into template arg. While reading a DRR pattern, we may expect the DAG is a matcher/rewriter for a operation/operand/attribute and now it has different semantics.

I do agree that the syntax isn't perfect, especially since with template arguments we need 4 different classes: OpWithReturnType, OpWithReturnTypes, OpWithReturnTypeBuilder, and OpWithReturnTypeBuilders. But the DAG still describes a rewriter for a type, it's just placed inside a template arg. Also, the class' names are a bit verbose.

I'm thinking a new type of DAG node which indicates the matching/rewriting of the return type

I like the idea of introducing a matcher for return types, because currently (I think) the only way to match return types is with constraints. The current syntax could also be used for that; it's just not implemented in the rewriter gen (e.g. OpWithResultType<AnOp, I32>:$res). Also, I think I can find a way to deduplicate the classes.

Pat<(FooOp ..., (ReturnType<I32>:$ret)),
    (BarOp arg1, arg2, (ReturnType $ret))>

Return values are typically bound at the DAG node: FooOp:$ret and I assume ReturnType $ret actually means to copy the return type, so it should be (ReturnType (NativeCodeCall<"$0.getType()> $ret)). Alternatively, an explicit type could be like (ReturnType "$_builder.getI32Type()"); I would first propose the following change to your syntax:

Pat<(FooOp:$res $arg1, $arg2, (ReturnType<I32>)),
    (BarOp $arg1, $arg2, (ReturnType "$_builder.getI32Type()"))>;

I'm assuming if we want to take the return type matched from FooOp, we can use (ReturnType $res). If so, yes, I think it's better.
I'm thinking to keep the meaning of Dag arg, i.e., it'll be binding the operand, specifying the type/attr constraint or another DagNode. An alternative way will be (ReturnType<"$_builder.getI32Type()">), so it aligns the shape of NativeCodeCall.

One question here, I'm not sure the complexity of supporting both ReturnType and ReturnType<...>. I barely remembered that tablegen parser may report an error but I'm not sure.

Assuming I reduce 4 classes down to 1 and introduce matchers for return types as well with the same OpWithResultType syntax, the difference boils down to whether you'd prefer the return type be specified as a trailing DAG node or next to the operation inside a template argument. I personally think it's a little bit weird to have a return type specifier after the operands and attributes, because the expectation of the DAG arguments is that they match the arguments = (ins ...) in the op declaration.

Also, having the return type specifier as a DAG argument would be more complex to implement. You'd have to ensure that there's only one return type specifier, that it's at the end of the argument list (or skip over it if would be allowed to be specified anywhere else), and it would have $_builder.getI32Type()to be associated back to the operation.

I prefer to keep the the DAG operator has the same meaning. It can be pure operation or some utility operations like NativeCodeCall. When we want to add other constraints to the operation, we usually do it in the constraints list (separate from DAG operator).

Yes, so far we have the agreement that a DAG has the form (Op arg0, arg1, ...). I think we could extend it now. When we create a new operation, we needs to specify the type (an exception is the interface inferReturnType). So I think it's ok to extend the same idea that adds an arg position for return type.

I remember that checking the position of ReturnType DAG is not too complicated. We have location which is supposed to be the last DAG arg if presents. So I think it can be verified in similar approach. For the associating the ReturnType with the operation, either in matcher or rewriter, we will go over the argument of the DagNode, so I think it may not be hard to bind the specified return type with an operation. Maybe I missed some invisible complexity :)

What's your opinion?

[mehdi_amini]

mlir/include/mlir/IR/OpBase.td

2641

Slightly off-topic for this revision, but I'm surprised we bundle all of this in OpBase.td, which I'd think would focus on ODS and not DRR: should we split all the DRR stuff in its own header?

[Mogball]

mlir/include/mlir/IR/OpBase.td

2641

+1 I am in favour of this.

[Mogball]

• rename OpWithResultTypeBuilders to ReturnTypeFuncs

[Mogball]

• change syntax to trailing directive

[Chia-hungDuan]

Thanks for this update!

Can you also add a test case which has both location and returnType?

mlir/tools/mlir-tblgen/RewriterGen.cpp

130–144

Do you think if we put this in DagNode class will be better?

1064

Can we lift this out of conditional branch?

1065

This is guarded above, maybe we can omit this check?

1072–1075

How about checking this first so that we can remove one layer of nesting? I.e.,
if (!tree.isNestedDagArg())

break;

// Do everything else

1153

Can we have a variable(maybe tblgen_types?) which stores the return type and we just pass the type to rewriter.create?

[Mogball]

mlir/tools/mlir-tblgen/RewriterGen.cpp

130–144

This I feel belongs in RewriterGen.cpp. It's part of the details of DRR codegen, whereas DagNode and the other classes under mlir/TableGen are purely wrappers for llvm::Init.

One of the functions that break this rule is Pattern::collectBoundSymbols, and it consequently needs to duplicate logic found in RewriterGen.cpp into Pattern.cpp. Moving TrailingDirectives into Pattern.h would prevent this duplication, but that's only because the logic inside collectBoundSymbols is not where it belongs in the first place (IMO).

1064

true

1065

I guess you can, yeah

1153

Yeah, that's what's being done below.

[Mogball]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1072–1075

Yeah sure. (the codebase is not 100% consistent with this though, especially in the messier files like this one XD)

[Mogball]

• minor fixes and adding test case for returnType and location together

[Chia-hungDuan]

mlir/tools/mlir-tblgen/RewriterGen.cpp

130–144

Agree. Thanks!

1153

Sorry, I didn't state this completely. I'm thinking if we have a variable always saves the return type, then do we still need this check !tail.returnType here and down below?

[Mogball]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1153

Do you mean always passing tblgen_types in each of these cases?

The first two cases rely on 1) auto-generated builder that uses OpTraits to infer return types and 2) a user-defined builder that does not require return types to be passed.

Only in the cases of 1) explicitly defined return types and 2) replacing the root op (so we can just copy the return types) will the fully generic builder TypeRange, ValueRange, AttrDict be called.

[Chia-hungDuan]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1153

Sorry, I got the wrong memory for the logic here. Just read the code again,

I think it should be a check here, if the op has at least one of those traits, then it shouldn't be able to assign an incompatible return type(Or we can disallow the use returnType here)

1170

I think in this case, we can invoke the builder with return type, right?

1207

Will this introduce a redundant indent?

[jpienaar]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1153

One can provide a more refined type (but compatible) than a local inference function could produce, so we should not disallow it, and checking here for compatible would be tricky and at best only handle a subset of cases (and verification here could require duplicating logic).

We should definitely disallow changing root's return types, but for that I'd just say let's uniformly flag it.

[Mogball]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1170

In this case, if we don't have explicitly specified return types, there's nothing ODS can do to infer return types (no OpTraits, no attributes, etc.). DRR will generate a call to a builder that ODS does has not automatically generated, so the burden is on the developer to provide a builder that fits that signature.

This is the main reason why explicit return types were added, because certain ops (like casts) need explicit return types. The solution before was to use a NativeCodeCall to directly call the correct builder.

1207

yessir let me fix this sir

[Mogball]

• fixing indentation

[Chia-hungDuan]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1170

I mean, in this conditional branch, if the user doesn't provide the explicit return type, then its their responsibility to handle the builder call. If the user provides the explicit return type, then we can call the builder with return type argument.

So whether or not user provides the explicit return type, we can always handle the case (usePartialResults || depth > 0 || resultIndex < 0) and we may not need the check !tail.returnType, right?

[Mogball]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1170

Well, not if no such builder was defined. "Responsibility" might be the wrong word; certain ops need to have an explicit return type and cannot define a type deduction builder.

If you defined a cast operation e.g. my_cast_op and you were trying to create it nested in a result pattern, without any OpTraits or attributes for type inference, and without explicitly specified return types, the generated pattern will not compile because no type deduction builder was defined. So you simply always have to provide an explicit return type.

And anyways, even if a type deduction builder is defined, explicitly specifying the return types should always take precedence.

[Mogball]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1170

It's structured like this:

if has explicit return types
    if is a root op - throw an error
    else use the return types
else
    if has attributes/optraits - use those
    else if is a root op - copy root return types
    else try to call a type deduction builder

[Chia-hungDuan]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1170

Ok, I thought there were codes for special handling usePartialResults and just noticed that it's nothing special but calling the builder with different signature. Sorry for the misleading.

And I saw you already have test case for this, thanks!

[Chia-hungDuan]

mlir/tools/mlir-tblgen/RewriterGen.cpp

920–922

How about adding one comment for this case?

mlir/include/mlir/IR/OpBase.td

class NativeCodeCall<string expr, int returns = 1> {
   int numReturns = returns;
 }
 
 class NativeCodeCallVoid<string expr> : NativeCodeCall<expr, 0>;
 
 def ConstantLikeMatcher : NativeCodeCall<"success(matchPattern($_self->getResult(0), m_Constant(&$0)))">;
 
 //===----------------------------------------------------------------------===//
 // 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;
 
+// Directive used in result patterns to specify the location of the generated
+// op. This directive must be used as a trailing argument to op creation or
+// native code calls.
+//
+// Usage:
+// * Create a named location: `(location "myLocation")`
+// * Copy the location of a captured symbol: `(location $arg)`
+// * Create a fused location: `(location "metadata", $arg0, $arg1)`
+
+def location;
+
+// Directive used in result patterns to specify return types for a created op.
+// This allows ops to be created without relying on type inference with
+// `OpTraits` or an op builder with deduction.
+//
+// This directive must be used as a trailing argument to op creation.
+//
+// Specify one return type with a string literal:
+//
+// ```
+// (AnOp $val, (returnType "$_builder.getI32Type()"))
+// ```
+//
+// Pass a captured value to copy its return type:
+//
+// ```
+// (AnOp $val, (returnType $val));
+// ```
+//
+// Pass a native code call inside a DAG to create a new type with arguments.
+//
+// ```
+// (AnOp $val,
+//       (returnType (NativeCodeCall<"$_builder.getTupleType({$0})"> $val)));
+// ```
+//
+// Specify multiple return types with multiple of any of the above.
+
+def returnType;
 
 //===----------------------------------------------------------------------===//
 // Attribute and Type generation
 //===----------------------------------------------------------------------===//
 
 // Class for defining a custom getter.
 //
 // TableGen generates several generic getter methods for each attribute and type

mlir/include/mlir/TableGen/Pattern.h

public:
 
   // Returns true if this DAG construct means to replace with an existing SSA
   // value.
   bool isReplaceWithValue() const;
 
   // Returns whether this DAG represents the location of an op creation.
   bool isLocationDirective() const;
 
+  // Returns whether this DAG is a return type specifier.
+  bool isReturnTypeDirective() const;
+
   // Returns true if this DAG node is wrapping native code call.
   bool isNativeCodeCall() const;
 
   // Returns true if this DAG node is an operation.
   bool isOperation() const;
 
   // Returns the native code call template inside this DAG node.
   // Precondition: isNativeCodeCall()

mlir/lib/TableGen/Pattern.cpp

 }
 
 bool DagLeaf::isConstantAttr() const { return isSubClassOf("ConstantAttr"); }
 
 bool DagLeaf::isEnumAttrCase() const {
   return isSubClassOf("EnumAttrCaseInfo");
 }
 
-bool DagLeaf::isStringAttr() const {
-  return isa<llvm::StringInit>(def);
-}
+bool DagLeaf::isStringAttr() const { return isa<llvm::StringInit>(def); }
 
 Constraint DagLeaf::getAsConstraint() const {
   assert((isOperandMatcher() || isAttrMatcher()) &&
          "the DAG leaf must be operand or attribute");
   return Constraint(cast<llvm::DefInit>(def)->getDef());
 }
 
 ConstantAttr DagLeaf::getAsConstantAttr() const {

 
 bool DagNode::isNativeCodeCall() const {
   if (auto *defInit = dyn_cast_or_null<llvm::DefInit>(node->getOperator()))
     return defInit->getDef()->isSubClassOf("NativeCodeCall");
   return false;
 }
 
 bool DagNode::isOperation() const {
-  return !isNativeCodeCall() && !isReplaceWithValue() && !isLocationDirective();
+  return !isNativeCodeCall() && !isReplaceWithValue() &&
+         !isLocationDirective() && !isReturnTypeDirective();
 }
 
 llvm::StringRef DagNode::getNativeCodeTemplate() const {
   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
   return cast<llvm::DefInit>(node->getOperator())
       ->getDef()
       ->getValueAsString("expression");
 }
bool DagNode::isReplaceWithValue() const {
   return dagOpDef->getName() == "replaceWithValue";
 }
 
 bool DagNode::isLocationDirective() const {
   auto *dagOpDef = cast<llvm::DefInit>(node->getOperator())->getDef();
   return dagOpDef->getName() == "location";
 }
 
+bool DagNode::isReturnTypeDirective() const {
+  auto *dagOpDef = cast<llvm::DefInit>(node->getOperator())->getDef();
+  return dagOpDef->getName() == "returnType";
+}
+
 void DagNode::print(raw_ostream &os) const {
   if (node)
     node->print(os);
 }
 
 //===----------------------------------------------------------------------===//
 // SymbolInfoMap
 //===----------------------------------------------------------------------===//
if (tree.isNativeCodeCall()) {
 
     return;
   }
 
   if (tree.isOperation()) {
     auto &op = getDialectOp(tree);
     auto numOpArgs = op.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();
+    // The pattern might have trailing directives.
+    int numDirectives = 0;
+    for (int i = numTreeArgs - 1; i >= 0; --i) {
+      if (auto dagArg = tree.getArgAsNestedDag(i)) {
+        if (dagArg.isLocationDirective() || dagArg.isReturnTypeDirective())
+          ++numDirectives;
+        else
+          break;
+      }
     }
 
-    if (numOpArgs != numTreeArgs - hasLocDirective) {
+    if (numOpArgs != numTreeArgs - numDirectives) {
       auto err = formatv("op '{0}' argument number mismatch: "
                          "{1} in pattern vs. {2} in definition",
                          op.getOperationName(), numTreeArgs, numOpArgs);
       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.

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

def : Pat<(TestLocationSrcOp:$res1
             (TestLocationSrcOp:$res3 $input))),
           (TestLocationDstOp
             (TestLocationDstOp
               (TestLocationDstOp $input, (location $res1)),
               (location "named")),
             (location "fused", $res2, $res3))>;
 
 //===----------------------------------------------------------------------===//
+// Test Patterns (Type Builders)
+
+def SourceOp : TEST_Op<"source_op"> {
+  let arguments = (ins AnyInteger:$arg, AnyI32Attr:$tag);
+  let results = (outs AnyInteger);
+}
+
+// An op without return type deduction.
+def OpX : TEST_Op<"op_x"> {
+  let arguments = (ins AnyInteger:$input);
+  let results = (outs AnyInteger);
+}
+
+// Test that ops without built-in type deduction can be created in the
+// replacement DAG with an explicitly specified type.
+def : Pat<(SourceOp $val, ConstantAttr<I32Attr, "11">:$attr),
+          (OpX (OpX $val, (returnType "$_builder.getI32Type()")))>;
+// Test NativeCodeCall type builder can accept arguments.
+def SameTypeAs : NativeCodeCall<"$0.getType()">;
+
+def : Pat<(SourceOp $val, ConstantAttr<I32Attr, "22">:$attr),
+          (OpX (OpX $val, (returnType (SameTypeAs $val))))>;
+
+// Test multiple return types.
+def MakeI64Type : NativeCodeCall<"$_builder.getI64Type()">;
+def MakeI32Type : NativeCodeCall<"$_builder.getI32Type()">;
+
+def OneToTwo : TEST_Op<"one_to_two"> {
+  let arguments = (ins AnyInteger);
+  let results = (outs AnyInteger, AnyInteger);
+}
+
+def TwoToOne : TEST_Op<"two_to_one"> {
+  let arguments = (ins AnyInteger, AnyInteger);
+  let results = (outs AnyInteger);
+}
+
+def : Pat<(SourceOp $val, ConstantAttr<I32Attr, "33">:$attr),
+          (TwoToOne (OpX (OneToTwo:$res__0 $val, (returnType (MakeI64Type), (MakeI32Type))), (returnType (MakeI32Type))),
+                    (OpX $res__1, (returnType (MakeI64Type))))>;
+
+// Test copy value return type.
+def : Pat<(SourceOp $val, ConstantAttr<I32Attr, "44">:$attr),
+          (OpX (OpX $val, (returnType $val)))>;
+
+// Test create multiple return types with different methods.
+def : Pat<(SourceOp $val, ConstantAttr<I32Attr, "55">:$attr),
+          (TwoToOne (OneToTwo:$res__0 $val, (returnType $val, "$_builder.getI64Type()")), $res__1)>;
+
+//===----------------------------------------------------------------------===//
 // Test Legalization
 //===----------------------------------------------------------------------===//
 
 def Test_LegalizerEnum_Success : StrEnumAttrCase<"Success">;
 def Test_LegalizerEnum_Failure : StrEnumAttrCase<"Failure">;
 
 def Test_LegalizerEnum : StrEnumAttr<"Success", "Failure",
   [Test_LegalizerEnum_Success, Test_LegalizerEnum_Failure]>;

mlir/test/mlir-tblgen/pattern.mlir

 //===----------------------------------------------------------------------===//
 
 // CHECK-LABEL: redundantTest
 func @redundantTest(%arg0: i32) -> i32 {
   %0 = "test.op_m"(%arg0) : (i32) -> i32
   // CHECK: "test.op_m"(%arg0) {optional_attr = 314159265 : i32} : (i32) -> i32
   return %0 : i32
 }
+
+//===----------------------------------------------------------------------===//
+// Test that ops without type deduction can be created with type builders.
+//===----------------------------------------------------------------------===//
+
+func @explicitReturnTypeTest(%arg0 : i64) -> i8 {
+  %0 = "test.source_op"(%arg0) {tag = 11 : i32} : (i64) -> i8
+  // CHECK: "test.op_x"(%arg0) : (i64) -> i32
+  // CHECK: "test.op_x"(%0) : (i32) -> i8
+  return %0 : i8
+}
+
+func @returnTypeBuilderTest(%arg0 : i1) -> i8 {
+  %0 = "test.source_op"(%arg0) {tag = 22 : i32} : (i1) -> i8
+  // CHECK: "test.op_x"(%arg0) : (i1) -> i1
+  // CHECK: "test.op_x"(%0) : (i1) -> i8
+  return %0 : i8
+}
+
+func @multipleReturnTypeBuildTest(%arg0 : i1) -> i1 {
+  %0 = "test.source_op"(%arg0) {tag = 33 : i32} : (i1) -> i1
+  // CHECK: "test.one_to_two"(%arg0) : (i1) -> (i64, i32)
+  // CHECK: "test.op_x"(%0#0) : (i64) -> i32
+  // CHECK: "test.op_x"(%0#1) : (i32) -> i64
+  // CHECK: "test.two_to_one"(%1, %2) : (i32, i64) -> i1
+  return %0 : i1
+}
+
+func @copyValueType(%arg0 : i8) -> i32 {
+  %0 = "test.source_op"(%arg0) {tag = 44 : i32} : (i8) -> i32
+  // CHECK: "test.op_x"(%arg0) : (i8) -> i8
+  // CHECK: "test.op_x"(%0) : (i8) -> i32
+  return %0 : i32
+}
+
+func @multipleReturnTypeDifferent(%arg0 : i1) -> i64 {
+  %0 = "test.source_op"(%arg0) {tag = 55 : i32} : (i1) -> i64
+  // CHECK: "test.one_to_two"(%arg0) : (i1) -> (i1, i64)
+  // CHECK: "test.two_to_one"(%0#0, %0#1) : (i1, i64) -> i64
+  return %0 : i64
+}

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
 // RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR3 %s 2>&1 | FileCheck --check-prefix=ERROR3 %s
+// RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR4 %s 2>&1 | FileCheck --check-prefix=ERROR4 %s
+// RUN: not mlir-tblgen -gen-rewriters -I %S/../../include -DERROR5 %s 2>&1 | FileCheck --check-prefix=ERROR5 %s
 
 include "mlir/IR/OpBase.td"
 
 // Check using the dialect name as the namespace
 def A_Dialect : Dialect {
   let name = "a";
 }
 

 #endif
 
 #ifdef ERROR3
 def NativeMatcher : NativeCodeCall<"success(nativeCall($_self, $0, $1))">;
 // ERROR3: [[@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
+
+#ifdef ERROR4
+// Check trying to pass op as DAG node inside ReturnTypeFunc fails.
+// ERROR4: [[@LINE+1]]:1: error: nested DAG in `returnType` must be a native code
+def : Pat<(OpB $val, AnyI32Attr:$attr), (OpA (OpA $val, $val, (returnType (OpA $val, $val))), $val)>;
+#endif
+
+#ifdef ERROR5
+// Check that trying to specify explicit types at the root node fails.
+// ERROR5: [[@LINE+1]]:1: error: Cannot specify explicit return types in an op
+def : Pat<(OpB $val, AnyI32Attr:$attr), (OpA $val, $val, (returnType "someType()"))>;
+#endif

mlir/test/mlir-tblgen/rewriter-indexing.td

 // CHECK: nativeCall(rewriter, odsLoc, (*v1.begin()), (*v2.begin()), (*v3.begin()), (*v4.begin()), (*v5.begin()), (*v6.begin()), (*v7.begin()), (*v8.begin()), (*v9.begin()), (*v10.begin()))
 def test4 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
                 (NativeBuilder $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10)>;
 
 // CHECK: struct test5 : public ::mlir::RewritePattern {
 // CHECK: foo(rewriter, (*v4.begin()), (*v5.begin()), (*v6.begin()), (*v7.begin()), (*v8.begin()), (*v9.begin()), (*v10.begin()))
 def test5 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
                 (NativeCodeCall<[{ foo($_builder, $3...) }]> $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10)>;
+
+// Check Pattern with return type builder.
+def SameTypeAs : NativeCodeCall<"$0.getType()">;
+// CHECK: struct test6 : public ::mlir::RewritePattern {
+// CHECK: tblgen_types.push_back((*v2.begin()).getType())
+// CHECK: tblgen_types.push_back(rewriter.getI32Type())
+// CHECK: nativeVar_1 = doSomething((*v3.begin()))
+// CHECK: tblgen_types.push_back(nativeVar_1)
+def test6 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
+                (AOp (AOp $v1, (returnType $v2, "$_builder.getI32Type()", (NativeCodeCall<"doSomething($0)"> $v3))))>;

mlir/tools/mlir-tblgen/RewriterGen.cpp

private:
 
   // Emits the C++ statement to replace the matched DAG with a value built via
   // calling native C++ code.
   std::string handleReplaceWithNativeCodeCall(DagNode resultTree, int depth);
 
   // Returns the symbol of the old value serving as the replacement.
   StringRef handleReplaceWithValue(DagNode tree);
 
+  // Trailing directives are used at the end of DAG node argument lists to
+  // specify additional behaviour for op matchers and creators, etc.
+  struct TrailingDirectives {
+    // DAG node containing the `location` directive. Null if there is none.
+    DagNode location;
+
+    // DAG node containing the `returnType` directive. Null if there is none.
+    DagNode returnType;
+
+    // Number of found trailing directives.
+    int numDirectives;
+  };
+
+  // Collect any trailing directives.
+  TrailingDirectives getTrailingDirectives(DagNode tree);
+
   // Returns the location value to use.
-  std::pair<bool, std::string> getLocation(DagNode tree);
+  std::string getLocation(TrailingDirectives &tail);
 
   // Returns the location value to use.
   std::string handleLocationDirective(DagNode tree);
 
+  // Emit return type argument.
+  std::string handleReturnTypeArg(DagNode returnType, int i, int depth);
+
   // Emits the C++ statement to build a new op out of the given DAG `tree` and
   // returns the variable name that this op is assigned to. If the root op in
   // DAG `tree` has a specified name, the created op will be assigned to a
   // variable of the given name. Otherwise, a unique name will be used as the
   // result value name.
   std::string handleOpCreation(DagNode tree, int resultIndex, int depth);
 
   using ChildNodeIndexNameMap = DenseMap<unsigned, std::string>;
if (DagNode argTree = tree.getArgAsNestedDag(i)) {
       } else {
         os << "Value " << argName << ";\n";
       }
     }
 
     capture.push_back(std::move(argName));
   }
 
-  bool hasLocationDirective;
-  std::string locToUse;
-  std::tie(hasLocationDirective, locToUse) = getLocation(tree);
+  auto tail = getTrailingDirectives(tree);
+  if (tail.returnType)
+    PrintFatalError(loc, "`NativeCodeCall` cannot have return type specifier");
+  auto locToUse = getLocation(tail);
 
   auto fmt = tree.getNativeCodeTemplate();
   if (fmt.count("$_self") != 1)
     PrintFatalError(loc, "NativeCodeCall must have $_self as argument for "
                          "passing the defining Operation");
 
   auto nativeCodeCall = std::string(tgfmt(
       fmt, &fmtCtx.addSubst("_loc", locToUse).withSelf(opName.str()), capture));
 
   emitMatchCheck(opName, formatv("!failed({0})", nativeCodeCall),
                  formatv("\"{0} return failure\"", nativeCodeCall));
 
-  for (int i = 0, e = tree.getNumArgs(); i != e; ++i) {
+  for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; 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) {
+  for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; 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));
std::string PatternEmitter::handleLocationDirective(DagNode tree) {
   os << "}";
   if (!strAttr.empty()) {
     os << ", rewriter.getStringAttr(\"" << strAttr << "\")";
   }
   os << ")";
   return os.str();
 }
 
+std::string PatternEmitter::handleReturnTypeArg(DagNode returnType, int i,
+                                                int depth) {
+  // Nested NativeCodeCall.
+  if (auto dagNode = returnType.getArgAsNestedDag(i)) {
+    if (!dagNode.isNativeCodeCall())
+      PrintFatalError(loc, "nested DAG in `returnType` must be a native code "
+                           "call");
+    return handleReplaceWithNativeCodeCall(dagNode, depth);
+  }
+  // String literal.
+  auto dagLeaf = returnType.getArgAsLeaf(i);
+  if (dagLeaf.isStringAttr())
+    return tgfmt(dagLeaf.getStringAttr(), &fmtCtx);
+  return tgfmt(
+      "$0.getType()", &fmtCtx,
+      handleOpArgument(returnType.getArgAsLeaf(i), returnType.getArgName(i)));
+}
+
 std::string PatternEmitter::handleOpArgument(DagLeaf leaf,
                                              StringRef patArgName) {
   if (leaf.isStringAttr())
     PrintFatalError(loc, "raw string not supported as argument");
   if (leaf.isConstantAttr()) {
     auto constAttr = leaf.getAsConstantAttr();
     return handleConstantAttr(constAttr.getAttribute(),
                               constAttr.getConstantValue());
std::string PatternEmitter::handleReplaceWithNativeCodeCall(DagNode tree,
   LLVM_DEBUG(llvm::dbgs() << "handle NativeCodeCall pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   auto fmt = tree.getNativeCodeTemplate();
 
   SmallVector<std::string, 16> attrs;
 
-  bool hasLocationDirective;
-  std::string locToUse;
-  std::tie(hasLocationDirective, locToUse) = getLocation(tree);
+  auto tail = getTrailingDirectives(tree);
+  if (tail.returnType)
+    PrintFatalError(loc, "`NativeCodeCall` cannot have return type specifier");
+  auto locToUse = getLocation(tail);
 
-  for (int i = 0, e = tree.getNumArgs() - hasLocationDirective; i != e; ++i) {
+  for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; i != e; ++i) {
     if (tree.isNestedDagArg(i)) {
       attrs.push_back(
           handleResultPattern(tree.getArgAsNestedDag(i), i, depth + 1));
     } else {
       attrs.push_back(
           handleOpArgument(tree.getArgAsLeaf(i), tree.getArgName(i)));
     }
     LLVM_DEBUG(llvm::dbgs() << "NativeCodeCall argument #" << i
int PatternEmitter::getNodeValueCount(DagNode node) {
   }
 
   if (node.isNativeCodeCall())
     return node.getNumReturnsOfNativeCode();
 
   return 1;
 }
 
-std::pair<bool, std::string> PatternEmitter::getLocation(DagNode tree) {
-  auto numPatArgs = tree.getNumArgs();
+PatternEmitter::TrailingDirectives
+PatternEmitter::getTrailingDirectives(DagNode tree) {
+  TrailingDirectives tail = {DagNode(nullptr), DagNode(nullptr), 0};
 
-  if (numPatArgs != 0) {
-    if (auto lastArg = tree.getArgAsNestedDag(numPatArgs - 1))
-      if (lastArg.isLocationDirective()) {
-        return std::make_pair(true, handleLocationDirective(lastArg));
-      }
-  }
+  // Look backwards through the arguments.
+  auto numPatArgs = tree.getNumArgs();
+  for (int i = numPatArgs - 1; i >= 0; --i) {
+    if (auto dagArg = tree.getArgAsNestedDag(i)) {
+      auto isLocation = dagArg.isLocationDirective();
+      auto isReturnType = dagArg.isReturnTypeDirective();
+      // If encountered a non-trailing-directive, stop looking.
+      if (!(isLocation || isReturnType))
+        break;
+      // Save the directive, but error if one of the same type was already
+      // found.
+      if (isLocation) {
+        if (tail.location)
+          PrintFatalError(loc, "`location` directive can only be specified "
+                               "once");
+        tail.location = dagArg;
+        ++tail.numDirectives;
+      } else if (isReturnType) {
+        if (tail.returnType)
+          PrintFatalError(loc, "`returnType` directive can only be specified "
+                               "once");
+        tail.returnType = dagArg;
+        ++tail.numDirectives;
+      }
+    } else {
+      // A leaf is not a directive.
+      break;
+    }
+
+  return tail;
 
+std::string
+PatternEmitter::getLocation(PatternEmitter::TrailingDirectives &tail) {
+  if (tail.location)
+    return handleLocationDirective(tail.location);
+
   // If no explicit location is given, use the default, all fused, location.
-  return std::make_pair(false, "odsLoc");
+  return "odsLoc";
 }
 
 std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,
                                              int depth) {
   LLVM_DEBUG(llvm::dbgs() << "create op for pattern: ");
   LLVM_DEBUG(tree.print(llvm::dbgs()));
   LLVM_DEBUG(llvm::dbgs() << '\n');
 
   Operator &resultOp = tree.getDialectOp(opMap);
   auto numOpArgs = resultOp.getNumArgs();
   auto numPatArgs = tree.getNumArgs();
 
-  bool hasLocationDirective;
-  std::string locToUse;
-  std::tie(hasLocationDirective, locToUse) = getLocation(tree);
+  auto tail = getTrailingDirectives(tree);
+  auto locToUse = getLocation(tail);
 
-  auto inPattern = numPatArgs - hasLocationDirective;
+  auto inPattern = numPatArgs - tail.numDirectives;
   if (numOpArgs != inPattern) {
     PrintFatalError(loc,
                     formatv("resultant op '{0}' argument number mismatch: "
                             "{1} in pattern vs. {2} in definition",
                             resultOp.getOperationName(), inPattern, numOpArgs));
   }
 
   // A map to collect all nested DAG child nodes' names, with operand index as
   // the key. This includes both bound and unbound child nodes.
   ChildNodeIndexNameMap childNodeNames;
 
   // 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 = tree.getNumArgs() - hasLocationDirective; i != e; ++i) {
+  for (int i = 0, e = tree.getNumArgs() - tail.numDirectives; i != e; ++i) {
     if (auto child = tree.getArgAsNestedDag(i))
       childNodeNames[i] = handleResultPattern(child, i, depth + 1);
   }
 
   // The name of the local variable holding this op.
   std::string valuePackName;
   // The symbol for holding the result of this pattern. Note that the result of
   // this pattern is not necessarily the same as the variable created by this
std::string PatternEmitter::handleOpCreation(DagNode tree, int resultIndex,
   // Right now ODS don't have general type inference support. Except a few
   // special cases listed below, DRR needs to supply types for all results
   // when building an op.
   bool isSameOperandsAndResultType =
       resultOp.getTrait("::mlir::OpTrait::SameOperandsAndResultType");
   bool useFirstAttr =
       resultOp.getTrait("::mlir::OpTrait::FirstAttrDerivedResultType");
 
-  if (isSameOperandsAndResultType || useFirstAttr) {
+  if (!tail.returnType && (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, 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) {
+  if (!tail.returnType && (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, 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.
+  // If we are provided explicit return types, use them to build the op.
+  // However, if depth == 0 and resultIndex >= 0, it means we are replacing
+  // the values generated from the source pattern root op. Then we must use the
+  // source pattern's value types to determine the value type of the generated
+  // op here.
+  if (depth == 0 && resultIndex >= 0 && tail.returnType)
+    PrintFatalError(loc, "Cannot specify explicit return types in an op whose "
+                         "return values replace the source pattern's root op");
 
   // First prepare local variables for op arguments used in builder call.
   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 (tail.returnType) {
+    auto numRetTys = tail.returnType.getNumArgs();
+    os.indent();
+    for (int i = 0; i < numRetTys; ++i) {
+      auto varName = handleReturnTypeArg(tail.returnType, i, depth + 1);
+      os << "tblgen_types.push_back(" << varName << ");\n";
+    }
+    os.unindent();
+  } else {
     if (numResults != 0) {
+      // Copy the result types from the source pattern.
       for (int i = 0; i < numResults; ++i)
         os << formatv("for (auto v: castedOp0.getODSResults({0})) {{\n"
                       "  tblgen_types.push_back(v.getType());\n}\n",
                       resultIndex + i);
     }
+  }
   os << formatv("{0} = rewriter.create<{1}>({2}, tblgen_types, "
                 "tblgen_values, tblgen_attrs);\n",
                 valuePackName, resultOp.getQualCppClassName(), locToUse);
   os.unindent() << "}\n";
   return resultValue;
 }
 
 void PatternEmitter::createSeparateLocalVarsForOpArgs(