[mlir][ODS] Add TypeRef directive in Declarative Assembly Format to allow custom UserDirective parser to receive previously parsed types.

[nicolasvasilache]

Summary

Specifying type_ref allows decoupling the specification of the type($foo) from its use in custom UserDirective.

This allows writing custom directives such as like:

let assemblyFormat = [{
       `(` $inputs `:` type($inputs) `)`
       ...
       custom<Foo>(
         type_ref($inputs)
     }];

in which parseFoo can implement special behavior based on the type of inputs without having to parse it at the same time.

[nicolasvasilache]

Finishing the impl and adding unit tests.

[ftynse]

This looks useful.

mlir/docs/OpDefinitions.md

702

Could you also document what is the C++ type passed to the function that backs this directive? We have a list under "Custom directives" below.

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

312–313

Types should be passed by-value

mlir/test/mlir-tblgen/op-format-spec-2.td

3–4

Which parsing, filecheck? Have you tried --split-input-file and placing each def into a separate split.

(I'm also surprised to see filecheck used for checking error messages, we usually do -verify-diagnostics but maybe it's absent from mlir-tblgen).

[nicolasvasilache]

mlir/test/mlir-tblgen/op-format-spec-2.td

3–4

Hmm based on the naming of some of those patterns I figured there seems to be a reordering that happens.
By naming the patterns "properly" I am able to put them in the right position so that the Filecheck agrees.

This is pretty nasty, `--split-input-file`` + `-verify-diagnostics`` sound significantly better.

Orthogonal, future work.

[nicolasvasilache]

Update

[ftynse]

By naming the patterns "properly" I am able to put them in the right position so that the Filecheck agrees.

Assuming there is no hashtable somewhere inside...

[rriddle]

Thanks Nicolas. I have a few comments on the direction, but OOO right now. Setting as blocking for now of that's alright.

[rriddle]

Seems like me being on this review was pointless given that you already submitted it. I'll just go ahead and remove myself.

[nicolasvasilache]

I am deeply sorry @rriddle .. I was working on https://reviews.llvm.org/D87767 which was dependent on this.

I had removed the dependence on the mechanism here and wrote the parser/printer manually:
https://github.com/llvm/llvm-project/blob/9744606614df4ba85a4d546c94b3b5ef9d3a3a96/mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp#L1448

But I confused myself with the stack and ended up submitting a squashed version.

I'll try to revert the part related to this revision, it was not meant to land without discussing with you (and I have additional questions re declarative assembly to allow interleaved operands / types like FuncOp arguments).

[nicolasvasilache]

Sent out https://reviews.llvm.org/D88698 but I I have a bug in the TestOps.inc.
I'll have to wait till tomorrow to finish the revert.

tools/mlir/test/lib/Dialect/Test/TestOps.cpp.inc:3599:23: error: no member named 'addAttribute' in 'mlir::NamedAttrList'
    result.attributes.addAttribute("attr", attrAttr);    if (optAttrAttr)
    ~~~~~~~~~~~~~~~~~ ^
tools/mlir/test/lib/Dialect/Test/TestOps.cpp.inc:3600:25: error: no member named 'addAttribute' in 'mlir::NamedAttrList'
      result.attributes.addAttribute("optAttr", optAttrAttr);  }
      ~~~~~~~~~~~~~~~~~ ^

If you see offhand what is wrong please comment on the other PR.

[rriddle]

In D87776#2307219, @nicolasvasilache wrote:

Sent out https://reviews.llvm.org/D88698 but I I have a bug in the TestOps.inc.
I'll have to wait till tomorrow to finish the revert.

tools/mlir/test/lib/Dialect/Test/TestOps.cpp.inc:3599:23: error: no member named 'addAttribute' in 'mlir::NamedAttrList'
    result.attributes.addAttribute("attr", attrAttr);    if (optAttrAttr)
    ~~~~~~~~~~~~~~~~~ ^
tools/mlir/test/lib/Dialect/Test/TestOps.cpp.inc:3600:25: error: no member named 'addAttribute' in 'mlir::NamedAttrList'
      result.attributes.addAttribute("optAttr", optAttrAttr);  }
      ~~~~~~~~~~~~~~~~~ ^

If you see offhand what is wrong please comment on the other PR.

Ah, don't worry about it. It's not that big of a deal. I'll just add comments here. Apologies for delaying in review, dealing with other things.

[rriddle]

Looks great Nicolas. Main comment is that it seems like it would be better to just have a general ref instead. Would likely simplify a bit of the implementation here.

mlir/docs/OpDefinitions.md

702

Would it be better instead to have a general ref directive that allows for capturing other things? Like regions, operands, etc.?

mlir/tools/mlir-tblgen/OpFormatGen.cpp

822

Do we need to generate storage for these here? Or can we do it inside of the custom directive instead?

2822–2825

If we use a general ref instead, that would help remove some of this awkwardness.

[nicolasvasilache]

The whole mechanism has been generalized by @rriddle and this is now obsolete.

mlir/docs/OpDefinitions.md

*   `successors`
     -   Represents all of the successors of an operation.
 
 *   `type` ( input )
 
     -   Represents the type of the given input.
     -   `input` must be either an operand or result [variable](#variables), the
         `operands` directive, or the `results` directive.
 
+*   `type_ref` ( input )
+
+    -   Represents a reference to the type of the given input that must have
+        already been resolved.
+    -   `input` must be either an operand or result [variable](#variables), the
+        `operands` directive, or the `results` directive.
+    -   Used to pass previously parsed types to custom directives.
+
 #### Literals
 
 A literal is either a keyword or punctuation surrounded by \`\`.
 
 The following are the set of valid punctuation:
 
 `:`, `,`, `=`, `<`, `>`, `(`, `)`, `{`, `}`, `[`, `]`, `->`
 

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

if (succeeded(parser.parseOptionalComma())) {
     if (parser.parseType(optOperandType))
       return failure();
   }
   if (parser.parseArrow() || parser.parseLParen() ||
       parser.parseTypeList(varOperandTypes) || parser.parseRParen())
     return failure();
   return success();
 }
+static ParseResult
+parseCustomDirectiveWithTypeRefs(OpAsmParser &parser, const Type &operandType,
+                                 const Type &optOperandType,
+                                 const SmallVectorImpl<Type> &varOperandTypes) {
+  if (parser.parseKeyword("type_refs_capture"))
+    return failure();
+
+  Type operandType2, optOperandType2;
+  SmallVector<Type, 1> varOperandTypes2;
+  if (parseCustomDirectiveResults(parser, operandType2, optOperandType2,
+                                  varOperandTypes2))
+    return failure();
+
+  if (operandType != operandType2 || optOperandType != optOperandType2 ||
+      varOperandTypes != varOperandTypes2)
+    return failure();
+
+  return success();
+}
 static ParseResult parseCustomDirectiveOperandsAndTypes(
     OpAsmParser &parser, OpAsmParser::OperandType &operand,
     Optional<OpAsmParser::OperandType> &optOperand,
     SmallVectorImpl<OpAsmParser::OperandType> &varOperands, Type &operandType,
     Type &optOperandType, SmallVectorImpl<Type> &varOperandTypes) {
   if (parseCustomDirectiveOperands(parser, operand, optOperand, varOperands) ||
       parseCustomDirectiveResults(parser, operandType, optOperandType,
                                   varOperandTypes))

 static void printCustomDirectiveResults(OpAsmPrinter &printer, Type operandType,
                                         Type optOperandType,
                                         TypeRange varOperandTypes) {
   printer << " : " << operandType;
   if (optOperandType)
     printer << ", " << optOperandType;
   printer << " -> (" << varOperandTypes << ")";
 }
+static void printCustomDirectiveWithTypeRefs(OpAsmPrinter &printer,
+                                             Type operandType,
+                                             Type optOperandType,
+                                             TypeRange varOperandTypes) {
+  printer << " type_refs_capture ";
+  printCustomDirectiveResults(printer, operandType, optOperandType,
+                              varOperandTypes);
+}
 static void
 printCustomDirectiveOperandsAndTypes(OpAsmPrinter &printer, Value operand,
                                      Value optOperand, OperandRange varOperands,
                                      Type operandType, Type optOperandType,
                                      TypeRange varOperandTypes) {
   printCustomDirectiveOperands(printer, operand, optOperand, varOperands);
   printCustomDirectiveResults(printer, operandType, optOperandType,
                               varOperandTypes);

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

def FormatCustomDirectiveResults
   let assemblyFormat = [{
     custom<CustomDirectiveResults>(
       type($result), type($optResult), type($varResults)
     )
     attr-dict
   }];
 }
 
+def FormatCustomDirectiveResultsWithTypeRefs
+    : TEST_Op<"format_custom_directive_results_with_type_refs",
+              [AttrSizedResultSegments]> {
+  let results = (outs AnyType:$result, Optional<AnyType>:$optResult,
+                      Variadic<AnyType>:$varResults);
+  let assemblyFormat = [{
+    custom<CustomDirectiveResults>(
+      type($result), type($optResult), type($varResults)
+    )
+    custom<CustomDirectiveWithTypeRefs>(
+      type_ref($result), type_ref($optResult), type_ref($varResults)
+    )
+    attr-dict
+  }];
+}
+
 def FormatCustomDirectiveSuccessors
     : TEST_Op<"format_custom_directive_successors", [Terminator]> {
   let successors = (successor AnySuccessor:$successor,
                               VariadicSuccessor<AnySuccessor>:$successors);
   let assemblyFormat = [{
     custom<CustomDirectiveSuccessors>(
       $successor, $successors
     )

mlir/test/mlir-tblgen/op-format-spec-2.td

+// RUN: mlir-tblgen -gen-op-decls -asmformat-error-is-fatal=false -I %S/../../include %s -o=%t 2>&1 | FileCheck %s
+
+// TODO: merge into ope-format-spec.td. For some reason, this seems to expose
+// weird parsing issues when mixed with other examples.
+// This file contains tests for the specification of the declarative op format.
+
+include "mlir/IR/OpBase.td"
+
+def TestDialect : Dialect {
+  let name = "test";
+}
+class TestFormat_Op<string name, string fmt, list<OpTrait> traits = []>
+    : Op<TestDialect, name, traits> {
+  let assemblyFormat = fmt;
+}
+
+//===----------------------------------------------------------------------===//
+// Directives
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// type_ref
+
+// CHECK: error: 'type_ref' of 'operand' is not bound by a prior 'type' directive
+def TypeRefOperandInvalidC : TestFormat_Op<"type_ref_operand_invalid_c", [{
+  type_ref($operand) type(operands)
+}]>, Arguments<(ins I64:$operand)>;
+// CHECK: error: 'operands' 'type_ref' is not bound by a prior 'type' directive
+def TypeRefOperandInvalidD : TestFormat_Op<"type_ref_operand_invalid_d", [{
+  type_ref(operands) type($operand)
+}]>, Arguments<(ins I64:$operand)>;
+// CHECK: error: 'type_ref' of 'operand' is not bound by a prior 'type' directive
+def TypeRefOperandInvalidE : TestFormat_Op<"type_ref_operand_invalid_e", [{
+  type_ref($operand) type($operand)
+}]>, Arguments<(ins I64:$operand)>;
+// CHECK: error: 'type_ref' of 'result' is not bound by a prior 'type' directive
+def TypeRefOperandInvalidG : TestFormat_Op<"type_ref_operand_invalid_g", [{
+  type_ref($result) type(results)
+}]>, Results<(outs I64:$result)>;
+// CHECK: error: 'results' 'type_ref' is not bound by a prior 'type' directive
+def TypeRefOperandInvalidH : TestFormat_Op<"type_ref_operand_invalid_h", [{
+  type_ref(results) type($result)
+}]>, Results<(outs I64:$result)>;
+// CHECK: error: 'type_ref' of 'result' is not bound by a prior 'type' directive
+def TypeRefOperandInvalidI : TestFormat_Op<"type_ref_operand_invalid_i", [{
+  type_ref($result) type($result)
+}]>, Results<(outs I64:$result)>;
+
+// CHECK-NOT: error
+def TypeRefOperandB : TestFormat_Op<"type_ref_operand_valid_b", [{
+  type_ref(operands) attr-dict
+}]>;
+// CHECK-NOT: error
+def TypeRefOperandD : TestFormat_Op<"type_ref_operand_valid_d", [{
+  type(operands) type_ref($operand) attr-dict
+}]>, Arguments<(ins I64:$operand)>;
+// CHECK-NOT: error
+def TypeRefOperandE : TestFormat_Op<"type_ref_operand_valid_e", [{
+  type($operand) type_ref($operand) attr-dict
+}]>, Arguments<(ins I64:$operand)>;
+// CHECK-NOT: error
+def TypeRefOperandF : TestFormat_Op<"type_ref_operand_valid_f", [{
+  type(results) type_ref(results) attr-dict
+}]>;
+// CHECK-NOT: error
+def TypeRefOperandG : TestFormat_Op<"type_ref_operand_valid_g", [{
+  type($result) type_ref(results) attr-dict
+}]>, Results<(outs I64:$result)>;
+// CHECK-NOT: error
+def TypeRefOperandH : TestFormat_Op<"type_ref_operand_valid_h", [{
+  type(results) type_ref($result) attr-dict
+}]>, Results<(outs I64:$result)>;
+// CHECK-NOT: error
+def TypeRefOperandI : TestFormat_Op<"type_ref_operand_valid_i", [{
+  type($result) type_ref($result) attr-dict
+}]>, Results<(outs I64:$result)>;
+

mlir/test/mlir-tblgen/op-format-spec.td

 def ZCoverageValidG : TestFormat_Op<"variable_valid_g", [{
   operands type($other) attr-dict
 }], [AllTypesMatch<["operand", "other"]>]>,
      Arguments<(ins AnyMemRef:$operand, AnyMemRef:$other)>;
 def ZCoverageValidH : TestFormat_Op<"variable_valid_h", [{
   operands type($result) attr-dict
 }], [AllTypesMatch<["operand", "result"]>]>,
      Arguments<(ins AnyMemRef:$operand)>, Results<(outs AnyMemRef:$result)>;
-

mlir/test/mlir-tblgen/op-format.mlir

 }
 
 // CHECK: test.format_custom_directive_results : i64, i64 -> (i64)
 test.format_custom_directive_results : i64, i64 -> (i64)
 
 // CHECK: test.format_custom_directive_results : i64 -> (i64)
 test.format_custom_directive_results : i64 -> (i64)
 
+// CHECK: test.format_custom_directive_results_with_type_refs : i64, i64 -> (i64) type_refs_capture : i64, i64 -> (i64)
+test.format_custom_directive_results_with_type_refs : i64, i64 -> (i64) type_refs_capture : i64, i64 -> (i64)
+
+// CHECK: test.format_custom_directive_results_with_type_refs : i64 -> (i64) type_refs_capture : i64 -> (i64)
+test.format_custom_directive_results_with_type_refs : i64 -> (i64) type_refs_capture : i64 -> (i64)
+
 func @foo() {
   // CHECK: test.format_custom_directive_successors ^bb1, ^bb2
   test.format_custom_directive_successors ^bb1, ^bb2
 
 ^bb1:
   // CHECK: test.format_custom_directive_successors ^bb2
   test.format_custom_directive_successors ^bb2
 

mlir/tools/mlir-tblgen/OpFormatGen.cpp

enum class Kind {
     AttrDictDirective,
     CustomDirective,
     FunctionalTypeDirective,
     OperandsDirective,
     RegionsDirective,
     ResultsDirective,
     SuccessorsDirective,
     TypeDirective,
+    TypeRefDirective,
 
     /// This element is a literal.
     Literal,
 
     /// This element is an variable value.
     AttributeVariable,
     OperandVariable,
     RegionVariable,

 public:
   TypeDirective(std::unique_ptr<Element> arg) : operand(std::move(arg)) {}
   Element *getOperand() const { return operand.get(); }
 
 private:
   /// The operand that is used to format the directive.
   std::unique_ptr<Element> operand;
 };
-} // end anonymous namespace
+
+/// This class represents the `type_ref` directive.
+class TypeRefDirective
+    : public DirectiveElement<Element::Kind::TypeRefDirective> {
+public:
+  TypeRefDirective(std::unique_ptr<Element> arg) : operand(std::move(arg)) {}
+  Element *getOperand() const { return operand.get(); }
+
+private:
+  /// The operand that is used to format the directive.
+  std::unique_ptr<Element> operand;
+};
+} // namespace
 
 //===----------------------------------------------------------------------===//
 // LiteralElement
 
 namespace {
 /// This class represents an instance of a literal element.
 class LiteralElement : public Element {
 public:
if (auto *optional = dyn_cast<OptionalElement>(element)) {
     StringRef name = getTypeListName(dir->getOperand(), lengthKind);
     if (lengthKind != ArgumentLengthKind::Single)
       body << "  ::mlir::SmallVector<::mlir::Type, 1> " << name << "Types;\n";
     else
       body << llvm::formatv("  ::mlir::Type {0}RawTypes[1];\n", name)
            << llvm::formatv(
                   "  ::llvm::ArrayRef<::mlir::Type> {0}Types({0}RawTypes);\n",
                   name);
+  } else if (auto *dir = dyn_cast<TypeRefDirective>(element)) {
+    ArgumentLengthKind lengthKind;
+    StringRef name = getTypeListName(dir->getOperand(), lengthKind);
+    // Refer to the previously encountered TypeDirective for name.
+    // Take a `const ::mlir::SmallVector<::mlir::Type, 1> &` in the declaration
+    // to properly track the types that will be parsed and pushed later on.
+    if (lengthKind != ArgumentLengthKind::Single)
+      body << "  const ::mlir::SmallVector<::mlir::Type, 1> &" << name
+           << "TypesRef(" << name << "Types);\n";
+    else
+      body << llvm::formatv(
+          "  ::llvm::ArrayRef<::mlir::Type> {0}RawTypesRef({0}RawTypes);\n",
+          name);
   } else if (auto *dir = dyn_cast<FunctionalTypeDirective>(element)) {
     ArgumentLengthKind ignored;
     body << "  ::llvm::ArrayRef<::mlir::Type> "
          << getTypeListName(dir->getInputs(), ignored) << "Types;\n";
     body << "  ::llvm::ArrayRef<::mlir::Type> "
          << getTypeListName(dir->getResults(), ignored) << "Types;\n";
   }
 }
static void genCustomParameterParser(Element &param, OpMethodBody &body) {
 
   } else if (auto *successor = dyn_cast<SuccessorVariable>(&param)) {
     StringRef name = successor->getVar()->name;
     if (successor->getVar()->isVariadic())
       body << llvm::formatv("{0}Successors", name);
     else
       body << llvm::formatv("{0}Successor", name);
 
+  } else if (auto *dir = dyn_cast<TypeRefDirective>(&param)) {
+    ArgumentLengthKind lengthKind;
+    StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
+    if (lengthKind == ArgumentLengthKind::Variadic)
+      body << llvm::formatv("{0}TypesRef", listName);
+    else if (lengthKind == ArgumentLengthKind::Optional)
+      body << llvm::formatv("{0}TypeRef", listName);
+    else
+      body << formatv("{0}RawTypesRef[0]", listName);
   } else if (auto *dir = dyn_cast<TypeDirective>(&param)) {
     ArgumentLengthKind lengthKind;
     StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
     if (lengthKind == ArgumentLengthKind::Variadic)
       body << llvm::formatv("{0}Types", listName);
     else if (lengthKind == ArgumentLengthKind::Optional)
       body << llvm::formatv("{0}Type", listName);
     else
if (auto *operand = dyn_cast<OperandVariable>(&param)) {
       body << "    " << operand->getVar()->name
            << "OperandsLoc = parser.getCurrentLocation();\n";
       if (operand->getVar()->isOptional()) {
         body << llvm::formatv(
             "    llvm::Optional<::mlir::OpAsmParser::OperandType> "
             "{0}Operand;\n",
             operand->getVar()->name);
       }
+    } else if (auto *dir = dyn_cast<TypeRefDirective>(&param)) {
+      // Reference to an optional which may or may not have been set.
+      // Retrieve from vector if not empty.
+      ArgumentLengthKind lengthKind;
+      StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
+      if (lengthKind == ArgumentLengthKind::Optional)
+        body << llvm::formatv(
+            "    ::mlir::Type {0}TypeRef = {0}TypesRef.empty() "
+            "? Type() : {0}TypesRef[0];\n",
+            listName);
     } else if (auto *dir = dyn_cast<TypeDirective>(&param)) {
       ArgumentLengthKind lengthKind;
       StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
       if (lengthKind == ArgumentLengthKind::Optional)
         body << llvm::formatv("    ::mlir::Type {0}Type;\n", listName);
     }
   }
 
if (auto *attr = dyn_cast<AttributeVariable>(&param)) {
           "    result.attributes.addAttribute(\"{0}\", {0}Attr);", var->name);
     } else if (auto *operand = dyn_cast<OperandVariable>(&param)) {
       const NamedTypeConstraint *var = operand->getVar();
       if (!var->isOptional())
         continue;
       body << llvm::formatv("    if ({0}Operand.hasValue())\n"
                             "      {0}Operands.push_back(*{0}Operand);\n",
                             var->name);
+    } else if (auto *dir = dyn_cast<TypeRefDirective>(&param)) {
+      // In the `type_ref` case, do not parse a new Type that needs to be added.
+      // Just do nothing here.
     } else if (auto *dir = dyn_cast<TypeDirective>(&param)) {
       ArgumentLengthKind lengthKind;
       StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
       if (lengthKind == ArgumentLengthKind::Optional) {
         body << llvm::formatv("    if ({0}Type)\n"
                               "      {0}Types.push_back({0}Type);\n",
                               listName);
       }
void OperationFormat::genElementParser(Element *element, OpMethodBody &body,
   } else if (isa<RegionsDirective>(element)) {
     body << llvm::formatv(regionListParserCode, "full");
     if (hasImplicitTermTrait)
       body << llvm::formatv(regionListEnsureTerminatorParserCode, "full");
 
   } else if (isa<SuccessorsDirective>(element)) {
     body << llvm::formatv(successorListParserCode, "full");
 
+  } else if (auto *dir = dyn_cast<TypeRefDirective>(element)) {
+    ArgumentLengthKind lengthKind;
+    StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
+    if (lengthKind == ArgumentLengthKind::Variadic)
+      body << llvm::formatv(variadicTypeParserCode, listName);
+    else if (lengthKind == ArgumentLengthKind::Optional)
+      body << llvm::formatv(optionalTypeParserCode, listName);
+    else
+      body << formatv(typeParserCode, listName);
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
     ArgumentLengthKind lengthKind;
     StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
     if (lengthKind == ArgumentLengthKind::Variadic)
       body << llvm::formatv(variadicTypeParserCode, listName);
     else if (lengthKind == ArgumentLengthKind::Optional)
       body << llvm::formatv(optionalTypeParserCode, listName);
     else
if (auto *attr = dyn_cast<AttributeVariable>(&param)) {
       body << operand->getVar()->name << "()";
 
     } else if (auto *region = dyn_cast<RegionVariable>(&param)) {
       body << region->getVar()->name << "()";
 
     } else if (auto *successor = dyn_cast<SuccessorVariable>(&param)) {
       body << successor->getVar()->name << "()";
 
+    } else if (auto *dir = dyn_cast<TypeRefDirective>(&param)) {
+      auto *typeOperand = dir->getOperand();
+      auto *operand = dyn_cast<OperandVariable>(typeOperand);
+      auto *var = operand ? operand->getVar()
+                          : cast<ResultVariable>(typeOperand)->getVar();
+      if (var->isVariadic())
+        body << var->name << "().getTypes()";
+      else if (var->isOptional())
+        body << llvm::formatv("({0}() ? {0}().getType() : Type())", var->name);
+      else
+        body << var->name << "().getType()";
     } else if (auto *dir = dyn_cast<TypeDirective>(&param)) {
       auto *typeOperand = dir->getOperand();
       auto *operand = dyn_cast<OperandVariable>(typeOperand);
       auto *var = operand ? operand->getVar()
                           : cast<ResultVariable>(typeOperand)->getVar();
       if (var->isVariadic())
         body << var->name << "().getTypes()";
       else if (var->isOptional())
void OperationFormat::genElementPrinter(Element *element, OpMethodBody &body,
   } else if (isa<RegionsDirective>(element)) {
     genVariadicRegionPrinter("getOperation()->getRegions()", body,
                              hasImplicitTermTrait);
   } else if (isa<SuccessorsDirective>(element)) {
     body << "  ::llvm::interleaveComma(getOperation()->getSuccessors(), p);\n";
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
     body << "  p << ";
     genTypeOperandPrinter(dir->getOperand(), body) << ";\n";
+  } else if (auto *dir = dyn_cast<TypeRefDirective>(element)) {
+    body << "  p << ";
+    genTypeOperandPrinter(dir->getOperand(), body) << ";\n";
   } else if (auto *dir = dyn_cast<FunctionalTypeDirective>(element)) {
     body << "  p.printFunctionalType(";
     genTypeOperandPrinter(dir->getInputs(), body) << ", ";
     genTypeOperandPrinter(dir->getResults(), body) << ");\n";
   } else {
     llvm_unreachable("unknown format element");
   }
 }
enum Kind {
     kw_attr_dict_w_keyword,
     kw_custom,
     kw_functional_type,
     kw_operands,
     kw_regions,
     kw_results,
     kw_successors,
     kw_type,
+    kw_type_ref,
     keyword_end,
 
     // String valued tokens.
     identifier,
     literal,
     variable,
   };
   Token(Kind kind, StringRef spelling) : kind(kind), spelling(spelling) {}
Token::Kind kind =
           .Case("attr-dict-with-keyword", Token::kw_attr_dict_w_keyword)
           .Case("custom", Token::kw_custom)
           .Case("functional-type", Token::kw_functional_type)
           .Case("operands", Token::kw_operands)
           .Case("regions", Token::kw_regions)
           .Case("results", Token::kw_results)
           .Case("successors", Token::kw_successors)
           .Case("type", Token::kw_type)
+          .Case("type_ref", Token::kw_type_ref)
           .Default(Token::identifier);
   return Token(kind, str);
 }
 
 //===----------------------------------------------------------------------===//
 // FormatParser
 //===----------------------------------------------------------------------===//
 
LogicalResult parseOperandsDirective(std::unique_ptr<Element> &element,
                                        llvm::SMLoc loc, bool isTopLevel);
   LogicalResult parseRegionsDirective(std::unique_ptr<Element> &element,
                                       llvm::SMLoc loc, bool isTopLevel);
   LogicalResult parseResultsDirective(std::unique_ptr<Element> &element,
                                       llvm::SMLoc loc, bool isTopLevel);
   LogicalResult parseSuccessorsDirective(std::unique_ptr<Element> &element,
                                          llvm::SMLoc loc, bool isTopLevel);
   LogicalResult parseTypeDirective(std::unique_ptr<Element> &element, Token tok,
-                                   bool isTopLevel);
-  LogicalResult parseTypeDirectiveOperand(std::unique_ptr<Element> &element);
+                                   bool isTopLevel, bool isTypeRef = false);
+  LogicalResult parseTypeDirectiveOperand(std::unique_ptr<Element> &element,
+                                          bool isTypeRef = false);
 
   //===--------------------------------------------------------------------===//
   // Lexer Utilities
   //===--------------------------------------------------------------------===//
 
   /// Advance the current lexer onto the next token.
   void consumeToken() {
     assert(curToken.getKind() != Token::eof &&
LogicalResult FormatParser::parseDirective(std::unique_ptr<Element> &element,
   case Token::kw_operands:
     return parseOperandsDirective(element, dirTok.getLoc(), isTopLevel);
   case Token::kw_regions:
     return parseRegionsDirective(element, dirTok.getLoc(), isTopLevel);
   case Token::kw_results:
     return parseResultsDirective(element, dirTok.getLoc(), isTopLevel);
   case Token::kw_successors:
     return parseSuccessorsDirective(element, dirTok.getLoc(), isTopLevel);
+  case Token::kw_type_ref:
+    return parseTypeDirective(element, dirTok, isTopLevel, /*isTypeRef=*/true);
   case Token::kw_type:
     return parseTypeDirective(element, dirTok, isTopLevel);
 
   default:
     llvm_unreachable("unknown directive token");
   }
 }
 
auto checkTypeOperand = [&](Element *typeEle) {
     auto *opVar = dyn_cast<OperandVariable>(typeEle);
     const NamedTypeConstraint *var = opVar ? opVar->getVar() : nullptr;
     if (!seenVariables.count(var))
       return emitError(curLoc, "type directive can only refer to variables "
                                "within the optional group");
     return ::mlir::success();
   };
   for (auto &ele : elements) {
-    if (auto *typeEle = dyn_cast<TypeDirective>(ele.get())) {
+    if (auto *typeEle = dyn_cast<TypeRefDirective>(ele.get())) {
+      if (failed(checkTypeOperand(typeEle->getOperand())))
+        return failure();
+    } else if (auto *typeEle = dyn_cast<TypeDirective>(ele.get())) {
       if (failed(checkTypeOperand(typeEle->getOperand())))
         return ::mlir::failure();
     } else if (auto *typeEle = dyn_cast<FunctionalTypeDirective>(ele.get())) {
       if (failed(checkTypeOperand(typeEle->getInputs())) ||
           failed(checkTypeOperand(typeEle->getResults())))
         return ::mlir::failure();
     }
   }
return TypeSwitch<Element *, LogicalResult>(childElements.back().get())
       .Case<RegionVariable>([&](RegionVariable *) {
         // TODO: When ODS has proper support for marking "optional" regions, add
         // a check here.
         return ::mlir::success();
       })
       // Literals, custom directives, and type directives may be used,
       // but they can't anchor the group.
       .Case<LiteralElement, CustomDirective, FunctionalTypeDirective,
-            OptionalElement, TypeDirective>([&](Element *) {
+            OptionalElement, TypeRefDirective, TypeDirective>([&](Element *) {
         if (isAnchor)
           return emitError(childLoc, "only variables can be used to anchor "
                                      "an optional group");
         return ::mlir::success();
       })
       .Default([&](Element *) {
         return emitError(childLoc, "only literals, types, and variables can be "
                                    "used within an optional group");
FormatParser::parseCustomDirective(std::unique_ptr<Element> &element,
 
   if (failed(parseToken(Token::r_paren,
                         "expected ')' after custom directive parameters")))
     return ::mlir::failure();
 
   // After parsing all of the elements, ensure that all type directives refer
   // only to variables.
   for (auto &ele : elements) {
+    if (auto *typeEle = dyn_cast<TypeRefDirective>(ele.get())) {
+      if (!isa<OperandVariable, ResultVariable>(typeEle->getOperand())) {
+        return emitError(curLoc,
+                         "type_ref directives within a custom directive "
+                         "may only refer to variables");
+      }
+    }
     if (auto *typeEle = dyn_cast<TypeDirective>(ele.get())) {
       if (!isa<OperandVariable, ResultVariable>(typeEle->getOperand())) {
         return emitError(curLoc, "type directives within a custom directive "
                                  "may only refer to variables");
       }
     }
   }
 
   element = std::make_unique<CustomDirective>(nameTok.getSpelling(),
                                               std::move(elements));
   return ::mlir::success();
 }
 
 LogicalResult FormatParser::parseCustomDirectiveParameter(
     std::vector<std::unique_ptr<Element>> &parameters) {
   llvm::SMLoc childLoc = curToken.getLoc();
   parameters.push_back({});
   if (failed(parseElement(parameters.back(), /*isTopLevel=*/true)))
     return ::mlir::failure();
 
   // Verify that the element can be placed within a custom directive.
-  if (!isa<TypeDirective, AttributeVariable, OperandVariable, RegionVariable,
-           SuccessorVariable>(parameters.back().get())) {
+  if (!isa<TypeRefDirective, TypeDirective, AttributeVariable, OperandVariable,
+           RegionVariable, SuccessorVariable>(parameters.back().get())) {
     return emitError(childLoc, "only variables and types may be used as "
                                "parameters to a custom directive");
   }
   return ::mlir::success();
 }
 
 LogicalResult
 FormatParser::parseFunctionalTypeDirective(std::unique_ptr<Element> &element,
if (hasAllSuccessors || !seenSuccessors.empty())
     return emitError(loc, "'successors' directive creates overlap in format");
   hasAllSuccessors = true;
   element = std::make_unique<SuccessorsDirective>();
   return ::mlir::success();
 }
 
 LogicalResult
 FormatParser::parseTypeDirective(std::unique_ptr<Element> &element, Token tok,
-                                 bool isTopLevel) {
+                                 bool isTopLevel, bool isTypeRef) {
   llvm::SMLoc loc = tok.getLoc();
   if (!isTopLevel)
     return emitError(loc, "'type' is only valid as a top-level directive");
 
   std::unique_ptr<Element> operand;
   if (failed(parseToken(Token::l_paren, "expected '(' before argument list")) ||
-      failed(parseTypeDirectiveOperand(operand)) ||
+      failed(parseTypeDirectiveOperand(operand, isTypeRef)) ||
       failed(parseToken(Token::r_paren, "expected ')' after argument list")))
     return ::mlir::failure();
+  if (isTypeRef)
+    element = std::make_unique<TypeRefDirective>(std::move(operand));
+  else
     element = std::make_unique<TypeDirective>(std::move(operand));
   return ::mlir::success();
 }
 
 LogicalResult
-FormatParser::parseTypeDirectiveOperand(std::unique_ptr<Element> &element) {
+FormatParser::parseTypeDirectiveOperand(std::unique_ptr<Element> &element,
+                                        bool isTypeRef) {
   llvm::SMLoc loc = curToken.getLoc();
   if (failed(parseElement(element, /*isTopLevel=*/false)))
     return ::mlir::failure();
   if (isa<LiteralElement>(element.get()))
     return emitError(
         loc, "'type' directive operand expects variable or directive operand");
 
   if (auto *var = dyn_cast<OperandVariable>(element.get())) {
     unsigned opIdx = var->getVar() - op.operand_begin();
-    if (fmt.allOperandTypes || seenOperandTypes.test(opIdx))
+    if (!isTypeRef && (fmt.allOperandTypes || seenOperandTypes.test(opIdx)))
       return emitError(loc, "'type' of '" + var->getVar()->name +
                                 "' is already bound");
+    if (isTypeRef && !(fmt.allOperandTypes || seenOperandTypes.test(opIdx)))
+      return emitError(loc, "'type_ref' of '" + var->getVar()->name +
+                                "' is not bound by a prior 'type' directive");
     seenOperandTypes.set(opIdx);
   } else if (auto *var = dyn_cast<ResultVariable>(element.get())) {
     unsigned resIdx = var->getVar() - op.result_begin();
-    if (fmt.allResultTypes || seenResultTypes.test(resIdx))
+    if (!isTypeRef && (fmt.allResultTypes || seenResultTypes.test(resIdx)))
       return emitError(loc, "'type' of '" + var->getVar()->name +
                                 "' is already bound");
+    if (isTypeRef && !(fmt.allResultTypes || seenResultTypes.test(resIdx)))
+      return emitError(loc, "'type_ref' of '" + var->getVar()->name +
+                                "' is not bound by a prior 'type' directive");
     seenResultTypes.set(resIdx);
   } else if (isa<OperandsDirective>(&*element)) {
-    if (fmt.allOperandTypes || seenOperandTypes.any())
+    if (!isTypeRef && (fmt.allOperandTypes || seenOperandTypes.any()))
       return emitError(loc, "'operands' 'type' is already bound");
+    if (isTypeRef && !(fmt.allOperandTypes || seenOperandTypes.all()))
+      return emitError(
+          loc,
+          "'operands' 'type_ref' is not bound by a prior 'type' directive");
     fmt.allOperandTypes = true;
   } else if (isa<ResultsDirective>(&*element)) {
-    if (fmt.allResultTypes || seenResultTypes.any())
+    if (!isTypeRef && (fmt.allResultTypes || seenResultTypes.any()))
       return emitError(loc, "'results' 'type' is already bound");
+    if (isTypeRef && !(fmt.allResultTypes || seenResultTypes.all()))
+      return emitError(
+          loc, "'results' 'type_ref' is not bound by a prior 'type' directive");
     fmt.allResultTypes = true;
   } else {
     return emitError(loc, "invalid argument to 'type' directive");
   }
   return ::mlir::success();
 }
 
 //===----------------------------------------------------------------------===//