[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

1120

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

1120

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?

1013–1016

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

break;

// Do everything else

1027

Can we lift this out of conditional branch?

1028

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

1083–1084

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

1027

true

1028

I guess you can, yeah

1083–1084

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

[Mogball]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1013–1016

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!

1083–1084

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

1083–1084

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

1083–1084

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)

1101–1102

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

1136

Will this introduce a redundant indent?

[jpienaar]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1083–1084

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

1101–1102

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.

1136

yessir let me fix this sir

[Mogball]

• fixing indentation

[Chia-hungDuan]

mlir/tools/mlir-tblgen/RewriterGen.cpp

1101–1102

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

1101–1102

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

1101–1102

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

1101–1102

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

900–902

How about adding one comment for this case?

mlir/include/mlir/IR/OpBase.td

class NativeCodeCall<string expr, int returns = 1> {
   string expression = expr;
   int numReturns = returns;
 }
 
 class NativeCodeCallVoid<string expr> : NativeCodeCall<expr, 0>;
 
 def ConstantLikeMatcher : NativeCodeCall<"success(matchPattern($_self->getResult(0), m_Constant(&$0)))">;
 
+// Create an op in the replacement pattern with explicitly specified return
+// types or type builders. This allows ops to be created with relying on type
+// inference with OpTraits or an op builder with type deduction.
+//
+// ```
+// def : Pat<(SourceOp $val),
+//           (NewOp (OpWithResultType<OpWithoutTypeDeduction,
+//                                    "$_builder.getI32Type()"> $val)
+//           )>;
+// ```
+//
+// Specify multiple explicit return types with `OpWithResultTypes`.
+//
+// ## Type Builders with NativeCodeCall
+//
+// By passing native code calls instead of a string, you can manipulate and/or
+// create types using bound values in the pattern.
+//
+// ```
+// def SameTypeAs : NativeCodeCall<"$0.getType()">;
+//
+// def : Pat<(SourceOp $val),
+//           (NewOp (OpWithResultTypeBuilder<OpWithoutTypeDeduction,
+//                                           (SameTypeAs $val)> $val)
+//           )>;
+// ```
+
+class OpWithResultTypeBuilders<Op op, list<dag> types> {
+  Op operation = op;
+  list<dag> typeBuilders = types;
+}
+
+class OpWithResultTypeBuilder<Op op, dag type> :
+    OpWithResultTypeBuilders<op, [type]>;
+
+class OpWithResultTypes<Op op, list<string> types> :
+    OpWithResultTypeBuilders<op, !foreach(type, types, (NativeCodeCall<type>))>;
+
+class OpWithResultType<Op op, string type> :
+    OpWithResultTypes<op, [type]>;
+
 //===----------------------------------------------------------------------===//
 // 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

public:
   operator bool() const { return node != nullptr; }
 
   // Returns the symbol bound to this DAG node.
   StringRef getSymbol() const;
 
   // Returns the operator wrapper object corresponding to the dialect op matched
   // by this DAG. The operator wrapper will be queried from the given `mapper`
   // and created in it if not existing.
+  //
+  // If this DAG node is an OpWithResultTypeBuilders, return the wrapped
+  // operator object.
   Operator &getDialectOp(RecordOperatorMap *mapper) const;
 
   // Returns the number of operations recursively involved in the DAG tree
   // rooted from this node.
   int getNumOps() const;
 
   // Returns the number of immediate arguments to this DAG node.
   int getNumArgs() const;
public:
   bool isReplaceWithValue() const;
 
   // Returns whether this DAG represents the location of an op creation.
   bool isLocationDirective() const;
 
   // Returns true if this DAG node is wrapping native code call.
   bool isNativeCodeCall() const;
 
-  // Returns true if this DAG node is an operation.
+  // Returns true if this DAG node is an operation with a result type builder.
+  bool isOpWithResultTypeBuilders() const;
+
+  // Returns true if this DAG node is an operation or an operation wrapped in a
+  // OpWithResultTypeBuilders object.
   bool isOperation() const;
 
   // Returns the native code call template inside this DAG node.
   // Precondition: isNativeCodeCall()
   StringRef getNativeCodeTemplate() const;
 
   // Returns the number of values will be returned by the native helper
   // function.
   // Precondition: isNativeCodeCall()
   int getNumReturnsOfNativeCode() const;
 
+  // Return the explicit type builders for the op in this DAG node. The type
+  // builders are native code calls that may have arguments.
+  // Precondition: isOpWithResultTypeBuilders()
+  llvm::SmallVector<DagNode, 1> getOpResultTypeBuilders() const;
+
   void print(raw_ostream &os) const;
 
 private:
   friend class SymbolInfoMap;
   const void *getAsOpaquePointer() const { return node; }
 
   const llvm::DagInit *node; // nullptr means null DagNode
 };

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::isOpWithResultTypeBuilders() const {
+  if (auto *defInit = dyn_cast_or_null<llvm::DefInit>(node->getOperator()))
+    return defInit->getDef()->isSubClassOf("OpWithResultTypeBuilders");
+  return false;
+}
+
 bool DagNode::isOperation() const {
-  return !isNativeCodeCall() && !isReplaceWithValue() && !isLocationDirective();
+  // Treat OpWithResultTypeBuilders as an Operator so that it inherits all the
+  // logic of using Operators inside patterns.
+  return isOpWithResultTypeBuilders() ||
+         (!isNativeCodeCall() && !isReplaceWithValue() &&
+          !isLocationDirective());
 }
 
 llvm::StringRef DagNode::getNativeCodeTemplate() const {
   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
   return cast<llvm::DefInit>(node->getOperator())
       ->getDef()
       ->getValueAsString("expression");
 }
 
 int DagNode::getNumReturnsOfNativeCode() const {
   assert(isNativeCodeCall() && "the DAG leaf must be NativeCodeCall");
   return cast<llvm::DefInit>(node->getOperator())
       ->getDef()
       ->getValueAsInt("numReturns");
 }
 
 llvm::StringRef DagNode::getSymbol() const { return node->getNameStr(); }
 
 Operator &DagNode::getDialectOp(RecordOperatorMap *mapper) const {
   llvm::Record *opDef = cast<llvm::DefInit>(node->getOperator())->getDef();
+
+  // Retrieve the wrapped operator object from an OpWithResultTypeBuilders
+  if (isOpWithResultTypeBuilders())
+    opDef = opDef->getValueAsDef("operation");
+
   auto it = mapper->find(opDef);
   if (it != mapper->end())
     return *it->second;
   return *mapper->try_emplace(opDef, std::make_unique<Operator>(opDef))
               .first->second;
 }
 
+SmallVector<DagNode, 1> DagNode::getOpResultTypeBuilders() const {
+  assert(isOpWithResultTypeBuilders() &&
+         "the DAG node must be OpWithResultTypeBuilders");
+  llvm::ListInit *builderDefs = cast<llvm::DefInit>(node->getOperator())
+                                    ->getDef()
+                                    ->getValueAsListInit("typeBuilders");
+  SmallVector<DagNode, 1> builders;
+  for (auto *builderDef : *builderDefs)
+    builders.push_back(DagNode(cast<llvm::DagInit>(builderDef)));
+  return builders;
+}
+
 int DagNode::getNumOps() const {
   int count = isReplaceWithValue() ? 0 : 1;
   for (int i = 0, e = getNumArgs(); i != e; ++i) {
     if (auto child = getArgAsNestedDag(i))
       count += child.getNumOps();
   }
   return count;
 }

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);
+}
+
+def OpWithoutTypeDeduction : TEST_Op<"op_without_type_deduction"> {
+  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),
+          (OpWithoutTypeDeduction
+            (OpWithResultType<OpWithoutTypeDeduction,
+                              "$_builder.getI32Type()"> $val))>;
+
+// Test NativeCodeCall type builder can accept arguments.
+def SameTypeAs : NativeCodeCall<"$0.getType()">;
+
+def : Pat<(SourceOp $val, ConstantAttr<I32Attr, "22">:$attr),
+          (OpWithoutTypeDeduction
+            (OpWithResultTypeBuilder<OpWithoutTypeDeduction,
+                                     (SameTypeAs $val)> $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
+            (OpWithResultTypeBuilder<OpWithoutTypeDeduction, (MakeI32Type)>
+              (OpWithResultTypeBuilders<OneToTwo, [(MakeI64Type), (MakeI32Type)]>:$res__0 $val)),
+            (OpWithResultTypeBuilder<OpWithoutTypeDeduction, (MakeI64Type)> $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_without_type_deduction"(%arg0) : (i64) -> i32
+  // CHECK: "test.op_without_type_deduction"(%0) : (i32) -> i8
+  return %0 : i8
+}
+
+func @returnTypeBuilderTest(%arg0 : i1) -> i8 {
+  %0 = "test.source_op"(%arg0) {tag = 22 : i32} : (i1) -> i8
+  // CHECK: "test.op_without_type_deduction"(%arg0) : (i1) -> i1
+  // CHECK: "test.op_without_type_deduction"(%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_without_type_deduction"(%0#0) : (i64) -> i32
+  // CHECK: "test.op_without_type_deduction"(%0#1) : (i32) -> i64
+  // CHECK: "test.two_to_one"(%1, %2) : (i32, i64) -> i1
+  return %0 : i1
+}

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
 
 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 OpWithResultTypeBuilder fails.
+class OpAWithTypeBuilder<dag builder> : OpWithResultTypeBuilder<OpA, builder>;
+// ERROR4: [[@LINE+1]]:1: error: expected NativeCodeCall as type builder
+def : Pat<(OpB $val, AnyI32Attr:$attr), (OpA (OpAWithTypeBuilder<(OpA $val, $val)> $val, $val), $val)>;
+#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: (*v2.begin()).getType()
+// CHECK: tblgen_types.push_back(nativeVar_1);
+def test6 : Pat<(DOp $v1, $v2, $v3, $v4, $v5, $v6, $v7, $v8, $v9, $v10),
+                (AOp (OpWithResultTypeBuilder<AOp, (SameTypeAs $v2)> $v1))>;

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);
 
   // Returns the location value to use.
   std::pair<bool, std::string> getLocation(DagNode tree);
 
   // Returns the location value to use.
   std::string handleLocationDirective(DagNode tree);
 
   // 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>;
 
   // 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,
if (leaf.isConstantAttr()) {
     auto constAttr = leaf.getAsConstantAttr();
     return handleConstantAttr(constAttr.getAttribute(),
                               constAttr.getConstantValue());
   }
   if (leaf.isEnumAttrCase()) {
     auto enumCase = leaf.getAsEnumAttrCase();
     if (enumCase.isStrCase())
       return handleConstantAttr(enumCase, enumCase.getSymbol());
     // This is an enum case backed by an IntegerAttr. We need to get its value
     // to build the constant.
     std::string val = std::to_string(enumCase.getValue());
     return handleConstantAttr(enumCase, val);
   }
 
   LLVM_DEBUG(llvm::dbgs() << "handle argument '" << patArgName << "'\n");
   auto argName = symbolInfoMap.getValueAndRangeUse(patArgName);
   if (leaf.isUnspecified() || leaf.isOperandMatcher()) {
     LLVM_DEBUG(llvm::dbgs() << "replace " << patArgName << " with '" << argName
                             << "' (via symbol ref)\n");

 std::pair<bool, std::string> PatternEmitter::getLocation(DagNode tree) {
   auto numPatArgs = tree.getNumArgs();
 
   if (numPatArgs != 0) {
     if (auto lastArg = tree.getArgAsNestedDag(numPatArgs - 1))
       if (lastArg.isLocationDirective()) {
         return std::make_pair(true, handleLocationDirective(lastArg));
       }
   }
 
   // If no explicit location is given, use the default, all fused, location.
   return std::make_pair(false, "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 inPattern = numPatArgs - hasLocationDirective;
   if (numOpArgs != inPattern) {
     PrintFatalError(loc,
                     formatv("resultant op '{0}' argument number mismatch: "
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 (!tree.isOpWithResultTypeBuilders() &&
+      (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 (!tree.isOpWithResultTypeBuilders() &&
+      (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.
+  // Otherwise, 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, 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 (tree.isOpWithResultTypeBuilders()) {
+    // Use the explicitly specified types.
+    for (auto &builder : tree.getOpResultTypeBuilders()) {
+      if (!builder.isNativeCodeCall())
+        PrintFatalError(loc, "expected NativeCodeCall as type builder in "
+                             "DAG node OpWithResultTypeBuilders");
+
+      auto typeVar = handleReplaceWithNativeCodeCall(builder, depth + 1);
+      os << formatv("tblgen_types.push_back({0});\n", typeVar);
+    }
+  } 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(