[mlir][OpFormatGen] Add support for specifiy "custom" directives.

[rriddle]

Summary

This revision adds support for custom directives to the declarative assembly format. This allows for users to use C++ for printing and parsing subsections of an otherwise declaratively specified format. The custom directive is structured as follows:

custom-directive ::= `custom` `<` UserDirective `>` `(` Params `)`

UserDirective is used as a suffix when this directive is used during printing and parsing. When parsing, parseUserDirective will be invoked. When printing, printUserDirective will be invoked. The first parameter to these methods must be a reference to either the OpAsmParser, or OpAsmPrinter. The type of rest of the parameters is dependent on the Params specified in the assembly format.

[lattner]

Just saw this in the weekly newsletter - this is really really really cool, nice work!

[rriddle]

Rebase

[rriddle]

format

[rriddle]

Rebase

[jpienaar]

Thanks!

[mehdi_amini]

Neat!

mlir/docs/OpDefinitions.md

781

?

793

Should this be folded in the previous bullet?

[rriddle]

Address comments

[rriddle]

Rebase

mlir/docs/OpDefinitions.md

*   `attr-dict`
 
     -   Represents the attribute dictionary of the operation.
 
 *   `attr-dict-with-keyword`
 
     -   Represents the attribute dictionary of the operation, but prefixes the
         dictionary with an `attributes` keyword.
 
+*   `custom` < UserDirective > ( Params )
+
+    -   Represents a custom directive implemented by the user in C++.
+    -   See the [Custom Directives](#custom-directives) section below for more
+        details.
+
 *   `functional-type` ( inputs , results )
 
     -   Formats the `inputs` and `results` arguments as a
         [function type](LangRef.md#function-type).
     -   The constraints on `inputs` and `results` are the same as the `input` of
         the `type` directive.
 
 *   `operands`

 
 A variable is an entity that has been registered on the operation itself, i.e.
 an argument(attribute or operand), result, successor, etc. In the `CallOp`
 example above, the variables would be `$callee` and `$args`.
 
 Attribute variables are printed with their respective value type, unless that
 value type is buildable. In those cases, the type of the attribute is elided.
 
+#### Custom Directives
+
+The declarative assembly format specification allows for handling a large
+majority of the common cases when formatting an operation. For the operations
+that require or desire specifying parts of the operation in a form not supported
+by the declarative syntax, custom directives may be specified. A custom
+directive essentially allows for users to use C++ for printing and parsing
+subsections of an otherwise declaratively specified format. Looking at the
+specification of a custom directive above:
+
+```
+custom-directive ::= `custom` `<` UserDirective `>` `(` Params `)`
+```
+
+A custom directive has two main parts: The `UserDirective` and the `Params`. A
+custom directive is transformed into a call to a `print*` and a `parse*` method
+when generating the C++ code for the format. The `UserDirective` is an
+identifier used as a suffix to these two calls, i.e., `custom<MyDirective>(...)`
+would result in calls to `parseMyDirective` and `printMyDirective` wihtin the
+parser and printer respectively. `Params` may be any combination of variables
+(i.e. Attribute, Operand, Successor, etc.) and type directives. The type
+directives must refer to a variable, but that variable need not also be a
+parameter to the custom directive.
+
+The arguments to the `parse<UserDirective>` method is firstly a reference to the
+`OpAsmParser`(`OpAsmParser &`), and secondly a set of output parameters
+corresponding to the parameters specified in the format. The mapping of
+declarative parameter to `parse` method argument is detailed below:
+
+*   Attribute Variables
+    -   Single: `<Attribute-Storage-Type>(e.g. Attribute) &`
+    -   Optional: `<Attribute-Storage-Type>(e.g. Attribute) &`
+*   Operand Variables
+    -   Single: `OpAsmParser::OperandType &`
+    -   Optional: `Optional<OpAsmParser::OperandType> &`
+    -   Variadic: `SmallVectorImpl<OpAsmParser::OperandType> &`
+*   Successor Variables
+    -   Single: `Block *&`
+    -   Variadic: `SmallVectorImpl<Block *> &`
+*   Type Directives
+    -   Single: `Type &`
+    -   Optional: `Type &`
+    -   Variadic: `SmallVectorImpl<Type> &`
+
+When a variable is optional, the value should only be specified if the variable
+is present. Otherwise, the value should remain `None` or null.
+
+The arguments to the `print<UserDirective>` method is firstly a reference to the
+`OpAsmPrinter`(`OpAsmPrinter &`), and secondly a set of output parameters
+corresponding to the parameters specified in the format. The mapping of
+declarative parameter to `print` method argument is detailed below:
+
+*   Attribute Variables
+    -   Single: `<Attribute-Storage-Type>(e.g. Attribute)`
+    -   Optional: `<Attribute-Storage-Type>(e.g. Attribute)`
+*   Operand Variables
+    -   Single: `Value`
+    -   Optional: `Value`
+    -   Variadic: `OperandRange`
+*   Successor Variables
+    -   Single: `Block *`
+    -   Variadic: `SuccessorRange`
+*   Type Directives
+    -   Single: `Type`
+    -   Optional: `Type`
+    -   Variadic: `TypeRange`
+
+When a variable is optional, the provided value may be null.
+
 #### Optional Groups
 
 In certain situations operations may have "optional" information, e.g.
 attributes or an empty set of variadic operands. In these situations a section
 of the assembly format can be marked as `optional` based on the presence of this
 information. An optional group is defined by wrapping a set of elements within
 `()` followed by a `?` and has the following requirements:
 
 *   The first element of the group must either be a literal, attribute, or an
     operand.
     -   This is because the first element must be optionally parsable.
 *   Exactly one argument variable within the group must be marked as the anchor
     of the group.
     -   The anchor is the element whose presence controls whether the group
         should be printed/parsed.
     -   An element is marked as the anchor by adding a trailing `^`.
     -   The first element is *not* required to be the anchor of the group.
-*   Literals, variables, and type directives are the only valid elements within
-    the group.
+*   Literals, variables, custom directives, and type directives are the only
+    valid elements within the group.
     -   Any attribute variable may be used, but only optional attributes can be
         marked as the anchor.
     -   Only variadic or optional operand arguments can be used.
     -   The operands to a type directive must be defined within the optional
         group.
 
 An example of an operation with an optional group is `std.return`, which has a
 variadic number of operands.

mlir/include/mlir/IR/OpImplementation.h

 }
 
 template <typename ValueRangeT>
 inline OpAsmPrinter &operator<<(OpAsmPrinter &p,
                                 const ValueTypeRange<ValueRangeT> &types) {
   llvm::interleaveComma(types, p);
   return p;
 }
+inline OpAsmPrinter &operator<<(OpAsmPrinter &p, const TypeRange &types) {
+  llvm::interleaveComma(types, p);
+  return p;
+}
 inline OpAsmPrinter &operator<<(OpAsmPrinter &p, ArrayRef<Type> types) {
   llvm::interleaveComma(types, p);
   return p;
 }
 
 //===----------------------------------------------------------------------===//
 // OpAsmParser
 //===----------------------------------------------------------------------===//

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

 } // end anonymous namespace
 
 void FoldToCallOp::getCanonicalizationPatterns(
     OwningRewritePatternList &results, MLIRContext *context) {
   results.insert<FoldToCallOpPattern>(context);
 }
 
 //===----------------------------------------------------------------------===//
+// Test Format* operations
+//===----------------------------------------------------------------------===//
+
+//===----------------------------------------------------------------------===//
+// Parsing
+
+static ParseResult parseCustomDirectiveOperands(
+    OpAsmParser &parser, OpAsmParser::OperandType &operand,
+    Optional<OpAsmParser::OperandType> &optOperand,
+    SmallVectorImpl<OpAsmParser::OperandType> &varOperands) {
+  if (parser.parseOperand(operand))
+    return failure();
+  if (succeeded(parser.parseOptionalComma())) {
+    optOperand.emplace();
+    if (parser.parseOperand(*optOperand))
+      return failure();
+  }
+  if (parser.parseArrow() || parser.parseLParen() ||
+      parser.parseOperandList(varOperands) || parser.parseRParen())
+    return failure();
+  return success();
+}
+static ParseResult
+parseCustomDirectiveResults(OpAsmParser &parser, Type &operandType,
+                            Type &optOperandType,
+                            SmallVectorImpl<Type> &varOperandTypes) {
+  if (parser.parseColon())
+    return failure();
+
+  if (parser.parseType(operandType))
+    return failure();
+  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 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))
+    return failure();
+  return success();
+}
+static ParseResult
+parseCustomDirectiveSuccessors(OpAsmParser &parser, Block *&successor,
+                               SmallVectorImpl<Block *> &varSuccessors) {
+  if (parser.parseSuccessor(successor))
+    return failure();
+  if (failed(parser.parseOptionalComma()))
+    return success();
+  Block *varSuccessor;
+  if (parser.parseSuccessor(varSuccessor))
+    return failure();
+  varSuccessors.append(2, varSuccessor);
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// Printing
+
+static void printCustomDirectiveOperands(OpAsmPrinter &printer, Value operand,
+                                         Value optOperand,
+                                         OperandRange varOperands) {
+  printer << operand;
+  if (optOperand)
+    printer << ", " << optOperand;
+  printer << " -> (" << varOperands << ")";
+}
+static void printCustomDirectiveResults(OpAsmPrinter &printer, Type operandType,
+                                        Type optOperandType,
+                                        TypeRange varOperandTypes) {
+  printer << " : " << operandType;
+  if (optOperandType)
+    printer << ", " << optOperandType;
+  printer << " -> (" << 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);
+}
+static void printCustomDirectiveSuccessors(OpAsmPrinter &printer,
+                                           Block *successor,
+                                           SuccessorRange varSuccessors) {
+  printer << successor;
+  if (!varSuccessors.empty())
+    printer << ", " << varSuccessors.front();
+}
+
+//===----------------------------------------------------------------------===//
 // Test IsolatedRegionOp - parse passthrough region arguments.
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseIsolatedRegionOp(OpAsmParser &parser,
                                          OperationState &result) {
   OpAsmParser::OperandType argInfo;
   Type argType = parser.getBuilder().getIndexType();
 

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

 
 def FormatOptionalUnitAttrNoElide
     : TEST_Op<"format_optional_unit_attribute_no_elide"> {
   let arguments = (ins UnitAttr:$is_optional);
   let assemblyFormat = "($is_optional^)? attr-dict";
 }
 
 //===----------------------------------------------------------------------===//
-// AllTypesMatch type inference
+// Custom Directives
+
+def FormatCustomDirectiveOperands
+    : TEST_Op<"format_custom_directive_operands", [AttrSizedOperandSegments]> {
+  let arguments = (ins I64:$operand, Optional<I64>:$optOperand,
+                       Variadic<I64>:$varOperands);
+  let assemblyFormat = [{
+    custom<CustomDirectiveOperands>(
+      $operand, $optOperand, $varOperands
+    )
+    attr-dict
+  }];
+}
+
+def FormatCustomDirectiveOperandsAndTypes
+    : TEST_Op<"format_custom_directive_operands_and_types",
+              [AttrSizedOperandSegments]> {
+  let arguments = (ins AnyType:$operand, Optional<AnyType>:$optOperand,
+                       Variadic<AnyType>:$varOperands);
+  let assemblyFormat = [{
+    custom<CustomDirectiveOperandsAndTypes>(
+      $operand, $optOperand, $varOperands,
+      type($operand), type($optOperand), type($varOperands)
+    )
+    attr-dict
+  }];
+}
+
+def FormatCustomDirectiveResults
+    : TEST_Op<"format_custom_directive_results", [AttrSizedResultSegments]> {
+  let results = (outs AnyType:$result, Optional<AnyType>:$optResult,
+                      Variadic<AnyType>:$varResults);
+  let assemblyFormat = [{
+    custom<CustomDirectiveResults>(
+      type($result), type($optResult), type($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
+    )
+    attr-dict
+  }];
+}
+
 //===----------------------------------------------------------------------===//
+// AllTypesMatch type inference
 
 def FormatAllTypesMatchVarOp : TEST_Op<"format_all_types_match_var", [
     AllTypesMatch<["value1", "value2", "result"]>
   ]> {
   let arguments = (ins AnyType:$value1, AnyType:$value2);
   let results = (outs AnyType:$result);
   let assemblyFormat = "attr-dict $value1 `,` $value2 `:` type($value1)";
 }
 
 def FormatAllTypesMatchAttrOp : TEST_Op<"format_all_types_match_attr", [
     AllTypesMatch<["value1", "value2", "result"]>
   ]> {
   let arguments = (ins AnyAttr:$value1, AnyType:$value2);
   let results = (outs AnyType:$result);
   let assemblyFormat = "attr-dict $value1 `,` $value2";
 }
 
 //===----------------------------------------------------------------------===//
 // TypesMatchWith type inference
-//===----------------------------------------------------------------------===//
 
 def FormatTypesMatchVarOp : TEST_Op<"format_types_match_var", [
     TypesMatchWith<"result type matches operand", "value", "result", "$_self">
   ]> {
   let arguments = (ins AnyType:$value);
   let results = (outs AnyType:$result);
   let assemblyFormat = "attr-dict $value `:` type($value)";
 }

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

 def DirectiveAttrDictValidA : TestFormat_Op<"attrdict_valid_a", [{
   attr-dict
 }]>;
 def DirectiveAttrDictValidB : TestFormat_Op<"attrdict_valid_b", [{
   attr-dict-with-keyword
 }]>;
 
 //===----------------------------------------------------------------------===//
+// custom
+
+// CHECK: error: expected '<' before custom directive name
+def DirectiveCustomInvalidA : TestFormat_Op<"custom_invalid_a", [{
+  custom(
+}]>;
+// CHECK: error: expected custom directive name identifier
+def DirectiveCustomInvalidB : TestFormat_Op<"custom_invalid_b", [{
+  custom<>
+}]>;
+// CHECK: error: expected '>' after custom directive name
+def DirectiveCustomInvalidC : TestFormat_Op<"custom_invalid_c", [{
+  custom<MyDirective(
+}]>;
+// CHECK: error: expected '(' before custom directive parameters
+def DirectiveCustomInvalidD : TestFormat_Op<"custom_invalid_d", [{
+  custom<MyDirective>)
+}]>;
+// CHECK: error: only variables and types may be used as parameters to a custom directive
+def DirectiveCustomInvalidE : TestFormat_Op<"custom_invalid_e", [{
+  custom<MyDirective>(operands)
+}]>;
+// CHECK: error: expected ')' after custom directive parameters
+def DirectiveCustomInvalidF : TestFormat_Op<"custom_invalid_f", [{
+  custom<MyDirective>($operand<
+}]>, Arguments<(ins I64:$operand)>;
+// CHECK: error: type directives within a custom directive may only refer to variables
+def DirectiveCustomInvalidH : TestFormat_Op<"custom_invalid_h", [{
+  custom<MyDirective>(type(operands))
+}]>;
+
+// CHECK-NOT: error
+def DirectiveCustomValidA : TestFormat_Op<"custom_valid_a", [{
+  custom<MyDirective>($operand) attr-dict
+}]>, Arguments<(ins Optional<I64>:$operand)>;
+def DirectiveCustomValidB : TestFormat_Op<"custom_valid_b", [{
+  custom<MyDirective>($operand, type($operand), type($result)) attr-dict
+}]>, Arguments<(ins I64:$operand)>, Results<(outs I64:$result)>;
+def DirectiveCustomValidC : TestFormat_Op<"custom_valid_c", [{
+  custom<MyDirective>($attr) attr-dict
+}]>, Arguments<(ins I64Attr:$attr)>;
+
+//===----------------------------------------------------------------------===//
 // functional-type
 
 // CHECK: error: 'functional-type' is only valid as a top-level directive
 def DirectiveFunctionalTypeInvalidA : TestFormat_Op<"functype_invalid_a", [{
   functional-type(functional-type)
 }]>;
 // CHECK: error: expected '(' before argument list
 def DirectiveFunctionalTypeInvalidB : TestFormat_Op<"functype_invalid_b", [{

 // CHECK: error: only literals, types, and variables can be used within an optional group
 def OptionalInvalidJ : TestFormat_Op<"optional_invalid_j", [{
   (attr-dict)
 }]>;
 // CHECK: error: expected '?' after optional group
 def OptionalInvalidK : TestFormat_Op<"optional_invalid_k", [{
   ($arg^)
 }]>, Arguments<(ins Variadic<I64>:$arg)>;
+// CHECK: error: only variables can be used to anchor an optional group
+def OptionalInvalidL : TestFormat_Op<"optional_invalid_l", [{
+  (custom<MyDirective>($arg)^)?
+}]>, Arguments<(ins I64:$arg)>;
 
 //===----------------------------------------------------------------------===//
 // Variables
 //===----------------------------------------------------------------------===//
 
 // CHECK: error: expected variable to refer to an argument, result, or successor
 def VariableInvalidA : TestFormat_Op<"variable_invalid_a", [{
   $unknown_arg attr-dict

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

 
 // CHECK: test.format_optional_operand_result_b_op : i64
 test.format_optional_operand_result_b_op( : ) : i64
 
 // CHECK: test.format_optional_operand_result_b_op : i64
 test.format_optional_operand_result_b_op : i64
 
 //===----------------------------------------------------------------------===//
+// Format custom directives
+//===----------------------------------------------------------------------===//
+
+// CHECK: test.format_custom_directive_operands %[[I64]], %[[I64]] -> (%[[I64]])
+test.format_custom_directive_operands %i64, %i64 -> (%i64)
+
+// CHECK: test.format_custom_directive_operands %[[I64]] -> (%[[I64]])
+test.format_custom_directive_operands %i64 -> (%i64)
+
+// CHECK: test.format_custom_directive_operands_and_types %[[I64]], %[[I64]] -> (%[[I64]]) : i64, i64 -> (i64)
+test.format_custom_directive_operands_and_types %i64, %i64 -> (%i64) : i64, i64 -> (i64)
+
+// CHECK: test.format_custom_directive_operands_and_types %[[I64]] -> (%[[I64]]) : i64 -> (i64)
+test.format_custom_directive_operands_and_types %i64 -> (%i64) : i64 -> (i64)
+
+// 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)
+
+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
+
+^bb2:
+  return
+}
+
+//===----------------------------------------------------------------------===//
 // Format trait type inference
 //===----------------------------------------------------------------------===//
 
 // CHECK: test.format_infer_variadic_type_from_non_variadic %[[I64]], %[[I64]] : i64
 test.format_infer_variadic_type_from_non_variadic %i64, %i64 : i64
 
 //===----------------------------------------------------------------------===//
 // AllTypesMatch type inference

mlir/tools/mlir-tblgen/OpFormatGen.cpp

 
 namespace {
 /// This class represents a single format element.
 class Element {
 public:
   enum class Kind {
     /// This element is a directive.
     AttrDictDirective,
+    CustomDirective,
     FunctionalTypeDirective,
     OperandsDirective,
     ResultsDirective,
     SuccessorsDirective,
     TypeDirective,
 
     /// This element is a literal.
     Literal,
public:
   explicit AttrDictDirective(bool withKeyword) : withKeyword(withKeyword) {}
   bool isWithKeyword() const { return withKeyword; }
 
 private:
   /// If the dictionary should be printed with the 'attributes' keyword.
   bool withKeyword;
 };
 
+/// This class represents a custom format directive that is implemented by the
+/// user in C++.
+class CustomDirective : public Element {
+public:
+  CustomDirective(StringRef name,
+                  std::vector<std::unique_ptr<Element>> &&arguments)
+      : Element{Kind::CustomDirective}, name(name),
+        arguments(std::move(arguments)) {}
+
+  static bool classof(const Element *element) {
+    return element->getKind() == Kind::CustomDirective;
+  }
+
+  /// Return the name of this optional element.
+  StringRef getName() const { return name; }
+
+  /// Return the arguments to the custom directive.
+  auto getArguments() const { return llvm::make_pointee_range(arguments); }
+
+private:
+  /// The user provided name of the directive.
+  StringRef name;
+
+  /// The arguments to the custom directive.
+  std::vector<std::unique_ptr<Element>> arguments;
+};
+
 /// This class represents the `functional-type` directive. This directive takes
 /// two arguments and formats them, respectively, as the inputs and results of a
 /// FunctionType.
 class FunctionalTypeDirective
     : public DirectiveElement<Element::Kind::FunctionalTypeDirective> {
 public:
   FunctionalTypeDirective(std::unique_ptr<Element> inputs,
                           std::unique_ptr<Element> results)
static bool canFormatEnumAttr(const NamedAttribute *attr) {
 
   // The attribute must have a valid underlying type and a constant builder.
   return !enumAttr->getUnderlyingType().empty() &&
          !enumAttr->getConstBuilderTemplate().empty();
 }
 
 /// The code snippet used to generate a parser call for an attribute.
 ///
-/// {0}: The storage type of the attribute.
-/// {1}: The name of the attribute.
-/// {2}: The type for the attribute.
+/// {0}: The name of the attribute.
+/// {1}: The type for the attribute.
 const char *const attrParserCode = R"(
-  {0} {1}Attr;
-  if (parser.parseAttribute({1}Attr{2}, "{1}", result.attributes))
+  if (parser.parseAttribute({0}Attr{1}, "{0}", result.attributes))
     return failure();
 )";
 const char *const optionalAttrParserCode = R"(
-  {0} {1}Attr;
   {
     ::mlir::OptionalParseResult parseResult =
-      parser.parseOptionalAttribute({1}Attr{2}, "{1}", result.attributes);
+      parser.parseOptionalAttribute({0}Attr{1}, "{0}", result.attributes);
     if (parseResult.hasValue() && failed(*parseResult))
       return failure();
   }
 )";
 
 /// The code snippet used to generate a parser call for an enum attribute.
 ///
 /// {0}: The name of the attribute.
if (parser.parseAttribute(attrVal, parser.getBuilder().getNoneType(),
                               "{0}", attrStorage))
       return failure();
 
     auto attrOptional = {1}::{2}(attrVal.getValue());
     if (!attrOptional)
       return parser.emitError(loc, "invalid ")
              << "{0} attribute specification: " << attrVal;
 
-    result.addAttribute("{0}", {3});
+    {0}Attr = {3};
+    result.addAttribute("{0}", {0}Attr);
   }
 )";
 const char *const optionalEnumAttrParserCode = R"(
-  Attribute {0}Attr;
   {
     ::mlir::StringAttr attrVal;
     ::mlir::NamedAttrList attrStorage;
     auto loc = parser.getCurrentLocation();
 
     ::mlir::OptionalParseResult parseResult =
       parser.parseOptionalAttribute(attrVal, parser.getBuilder().getNoneType(),
                                     "{0}", attrStorage);
const char *const optionalEnumAttrParserCode = R"(
     }
   }
 )";
 
 /// The code snippet used to generate a parser call for an operand.
 ///
 /// {0}: The name of the operand.
 const char *const variadicOperandParserCode = R"(
+  {0}OperandsLoc = parser.getCurrentLocation();
   if (parser.parseOperandList({0}Operands))
     return failure();
 )";
 const char *const optionalOperandParserCode = R"(
   {
+    {0}OperandsLoc = parser.getCurrentLocation();
     ::mlir::OpAsmParser::OperandType operand;
     ::mlir::OptionalParseResult parseResult =
                                     parser.parseOptionalOperand(operand);
     if (parseResult.hasValue()) {
       if (failed(*parseResult))
         return failure();
       {0}Operands.push_back(operand);
     }
   }
 )";
 const char *const operandParserCode = R"(
+  {0}OperandsLoc = parser.getCurrentLocation();
   if (parser.parseOperand({0}RawOperands[0]))
     return failure();
 )";
 
 /// The code snippet used to generate a parser call for a type list.
 ///
 /// {0}: The name for the type list.
 const char *const variadicTypeParserCode = R"(
const char *const functionalTypeParserCode = R"(
   {0}Types = {0}__{1}_functionType.getInputs();
   {1}Types = {0}__{1}_functionType.getResults();
 )";
 
 /// The code snippet used to generate a parser call for a successor list.
 ///
 /// {0}: The name for the successor list.
 const char *successorListParserCode = R"(
-  ::llvm::SmallVector<::mlir::Block *, 2> {0}Successors;
   {
     ::mlir::Block *succ;
     auto firstSucc = parser.parseOptionalSuccessor(succ);
     if (firstSucc.hasValue()) {
       if (failed(*firstSucc))
         return failure();
       {0}Successors.emplace_back(succ);
 
       // Parse any trailing successors.
       while (succeeded(parser.parseOptionalComma())) {
         if (parser.parseSuccessor(succ))
           return failure();
         {0}Successors.emplace_back(succ);
       }
     }
   }
 )";
 
 /// The code snippet used to generate a parser call for a successor.
 ///
 /// {0}: The name of the successor.
 const char *successorParserCode = R"(
-  ::mlir::Block *{0}Successor = nullptr;
   if (parser.parseSuccessor({0}Successor))
     return failure();
 )";
 
 namespace {
 /// The type of length for a given parse argument.
 enum class ArgumentLengthKind {
   /// The argument is variadic, and may contain 0->N elements.
body << (StringRef)llvm::StringSwitch<StringRef>(value)
               .Case("[", "LSquare()")
               .Case("]", "RSquare()")
               .Case("?", "Question()");
 }
 
 /// Generate the storage code required for parsing the given element.
 static void genElementParserStorage(Element *element, OpMethodBody &body) {
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
-    for (auto &childElement : optional->getElements())
+    auto elements = optional->getElements();
+
+    // If the anchor is a unit attribute, it won't be parsed directly so elide
+    // it.
+    auto *anchor = dyn_cast<AttributeVariable>(optional->getAnchor());
+    Element *elidedAnchorElement = nullptr;
+    if (anchor && anchor != &*elements.begin() && anchor->isUnitAttr())
+      elidedAnchorElement = anchor;
+    for (auto &childElement : elements)
+      if (&childElement != elidedAnchorElement)
         genElementParserStorage(&childElement, body);
+
+  } else if (auto *custom = dyn_cast<CustomDirective>(element)) {
+    for (auto &paramElement : custom->getArguments())
+      genElementParserStorage(&paramElement, body);
+
+  } else if (isa<OperandsDirective>(element)) {
+    body << "  ::mlir::SmallVector<::mlir::OpAsmParser::OperandType, 4> "
+            "allOperands;\n";
+
+  } else if (isa<SuccessorsDirective>(element)) {
+    body << "  ::llvm::SmallVector<::mlir::Block *, 2> fullSuccessors;\n";
+
+  } else if (auto *attr = dyn_cast<AttributeVariable>(element)) {
+    const NamedAttribute *var = attr->getVar();
+    body << llvm::formatv("  {0} {1}Attr;\n", var->attr.getStorageType(),
+                          var->name);
+
   } else if (auto *operand = dyn_cast<OperandVariable>(element)) {
     StringRef name = operand->getVar()->name;
     if (operand->getVar()->isVariableLength()) {
       body << "  ::mlir::SmallVector<::mlir::OpAsmParser::OperandType, 4> "
            << name << "Operands;\n";
     } else {
       body << "  ::mlir::OpAsmParser::OperandType " << name
            << "RawOperands[1];\n"
            << "  ::llvm::ArrayRef<::mlir::OpAsmParser::OperandType> " << name
            << "Operands(" << name << "RawOperands);";
     }
-    body << llvm::formatv(
-        "  ::llvm::SMLoc {0}OperandsLoc = parser.getCurrentLocation();\n"
+    body << llvm::formatv("  ::llvm::SMLoc {0}OperandsLoc;\n"
                           "  (void){0}OperandsLoc;\n",
                           name);
+  } else if (auto *successor = dyn_cast<SuccessorVariable>(element)) {
+    StringRef name = successor->getVar()->name;
+    if (successor->getVar()->isVariadic()) {
+      body << llvm::formatv("  ::llvm::SmallVector<::mlir::Block *, 2> "
+                            "{0}Successors;\n",
+                            name);
+    } else {
+      body << llvm::formatv("  ::mlir::Block *{0}Successor = nullptr;\n", name);
+    }
+
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
     ArgumentLengthKind lengthKind;
     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<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";
   }
 }
 
+/// Generate the parser for a parameter to a custom directive.
+static void genCustomParameterParser(Element &param, OpMethodBody &body) {
+  body << ", ";
+  if (auto *attr = dyn_cast<AttributeVariable>(&param)) {
+    body << attr->getVar()->name << "Attr";
+
+  } else if (auto *operand = dyn_cast<OperandVariable>(&param)) {
+    StringRef name = operand->getVar()->name;
+    ArgumentLengthKind lengthKind = getArgumentLengthKind(operand->getVar());
+    if (lengthKind == ArgumentLengthKind::Variadic)
+      body << llvm::formatv("{0}Operands", name);
+    else if (lengthKind == ArgumentLengthKind::Optional)
+      body << llvm::formatv("{0}Operand", name);
+    else
+      body << formatv("{0}RawOperands[0]", name);
+
+  } 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<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
+      body << formatv("{0}RawTypes[0]", listName);
+  } else {
+    llvm_unreachable("unknown custom directive parameter");
+  }
+}
+
+/// Generate the parser for a custom directive.
+static void genCustomDirectiveParser(CustomDirective *dir, OpMethodBody &body) {
+  body << "  {\n";
+
+  // Preprocess the directive variables.
+  // * Add a local variable for optional operands and types. This provides a
+  //   better API to the user defined parser methods.
+  // * Set the location of operand variables.
+  for (Element &param : dir->getArguments()) {
+    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<TypeDirective>(&param)) {
+      ArgumentLengthKind lengthKind;
+      StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
+      if (lengthKind == ArgumentLengthKind::Optional)
+        body << llvm::formatv("    ::mlir::Type {0}Type;\n", listName);
+    }
+  }
+
+  body << "    if (parse" << dir->getName() << "(parser";
+  for (Element &param : dir->getArguments())
+    genCustomParameterParser(param, body);
+
+  body << "))\n"
+       << "      return failure();\n";
+
+  // After parsing, add handling for any of the optional constructs.
+  for (Element &param : dir->getArguments()) {
+    if (auto *attr = dyn_cast<AttributeVariable>(&param)) {
+      const NamedAttribute *var = attr->getVar();
+      if (var->attr.isOptional())
+        body << llvm::formatv("    if ({0}Attr)\n  ", var->name);
+
+      body << llvm::formatv(
+          "    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<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);
+      }
+    }
+  }
+
+  body << "  }\n";
+}
+
 /// Generate the parser for a single format element.
 static void genElementParser(Element *element, OpMethodBody &body,
                              FmtContext &attrTypeCtx) {
   /// Optional Group.
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
     auto elements = optional->getElements();
 
     // Generate a special optional parser for the first element to gate the
if (auto *optional = dyn_cast<OptionalElement>(element)) {
     std::string attrTypeStr;
     if (Optional<StringRef> typeBuilder = attr->getTypeBuilder()) {
       llvm::raw_string_ostream os(attrTypeStr);
       os << ", " << tgfmt(*typeBuilder, &attrTypeCtx);
     }
 
     body << formatv(var->attr.isOptional() ? optionalAttrParserCode
                                            : attrParserCode,
-                    var->attr.getStorageType(), var->name, attrTypeStr);
+                    var->name, attrTypeStr);
   } else if (auto *operand = dyn_cast<OperandVariable>(element)) {
     ArgumentLengthKind lengthKind = getArgumentLengthKind(operand->getVar());
     StringRef name = operand->getVar()->name;
     if (lengthKind == ArgumentLengthKind::Variadic)
       body << llvm::formatv(variadicOperandParserCode, name);
     else if (lengthKind == ArgumentLengthKind::Optional)
       body << llvm::formatv(optionalOperandParserCode, name);
     else
       body << formatv(operandParserCode, name);
   } else if (auto *successor = dyn_cast<SuccessorVariable>(element)) {
     bool isVariadic = successor->getVar()->isVariadic();
     body << formatv(isVariadic ? successorListParserCode : successorParserCode,
                     successor->getVar()->name);
 
     /// Directives.
   } else if (auto *attrDict = dyn_cast<AttrDictDirective>(element)) {
     body << "  if (parser.parseOptionalAttrDict"
          << (attrDict->isWithKeyword() ? "WithKeyword" : "")
          << "(result.attributes))\n"
          << "    return failure();\n";
+  } else if (auto *customDir = dyn_cast<CustomDirective>(element)) {
+    genCustomDirectiveParser(customDir, body);
+
   } else if (isa<OperandsDirective>(element)) {
     body << "  ::llvm::SMLoc allOperandLoc = parser.getCurrentLocation();\n"
-         << "  ::mlir::SmallVector<::mlir::OpAsmParser::OperandType, 4> "
-            "allOperands;\n"
          << "  if (parser.parseOperandList(allOperands))\n"
          << "    return failure();\n";
   } else if (isa<SuccessorsDirective>(element)) {
     body << llvm::formatv(successorListParserCode, "full");
   } else if (auto *dir = dyn_cast<TypeDirective>(element)) {
     ArgumentLengthKind lengthKind;
     StringRef listName = getTypeListName(dir->getOperand(), lengthKind);
     if (lengthKind == ArgumentLengthKind::Variadic)
auto interleaveFn = [&](const NamedTypeConstraint &operand) {
       if (operand.isVariableLength())
         body << "static_cast<int32_t>(" << operand.name << "Operands.size())";
       else
         body << "1";
     };
     llvm::interleaveComma(op.getOperands(), body, interleaveFn);
     body << "}));\n";
   }
+
+  if (!allResultTypes && op.getTrait("OpTrait::AttrSizedResultSegments")) {
+    body << "  result.addAttribute(\"result_segment_sizes\", "
+         << "parser.getBuilder().getI32VectorAttr({";
+    auto interleaveFn = [&](const NamedTypeConstraint &result) {
+      // If the result is variadic emit the parsed size.
+      if (result.isVariableLength())
+        body << "static_cast<int32_t>(" << result.name << "Types.size())";
+      else
+        body << "1";
+    };
+    llvm::interleaveComma(op.getResults(), body, interleaveFn);
+    body << "}));\n";
+  }
 }
 
 //===----------------------------------------------------------------------===//
 // PrinterGen
 
 /// Generate the printer for the 'attr-dict' directive.
 static void genAttrDictPrinter(OperationFormat &fmt, Operator &op,
                                OpMethodBody &body, bool withKeyword) {
for (auto &it : fmt.elements) {
     }
   }
 
   body << "  p.printOptionalAttrDict" << (withKeyword ? "WithKeyword" : "")
        << "(getAttrs(), /*elidedAttrs=*/{";
   // Elide the variadic segment size attributes if necessary.
   if (!fmt.allOperands && op.getTrait("OpTrait::AttrSizedOperandSegments"))
     body << "\"operand_segment_sizes\", ";
+  if (!fmt.allResultTypes && op.getTrait("OpTrait::AttrSizedResultSegments"))
+    body << "\"result_segment_sizes\", ";
   llvm::interleaveComma(usedAttributes, body, [&](const NamedAttribute *attr) {
     body << "\"" << attr->name << "\"";
   });
   body << "});\n";
 }
 
 /// Generate the printer for a literal value. `shouldEmitSpace` is true if a
 /// space should be emitted before this element. `lastWasPunctuation` is true if
static void genLiteralPrinter(StringRef value, OpMethodBody &body,
   body << " << \"" << value << "\";\n";
 
   // Insert a space after certain literals.
   shouldEmitSpace =
       value.size() != 1 || !StringRef("<({[").contains(value.front());
   lastWasPunctuation = !(value.front() == '_' || isalpha(value.front()));
 }
 
+/// Generate the printer for a literal value. `shouldEmitSpace` is true if a
+/// space should be emitted before this element. `lastWasPunctuation` is true if
+/// the previous element was a punctuation literal.
+static void genCustomDirectivePrinter(CustomDirective *customDir,
+                                      OpMethodBody &body) {
+  body << "  print" << customDir->getName() << "(p";
+  for (Element &param : customDir->getArguments()) {
+    body << ", ";
+    if (auto *attr = dyn_cast<AttributeVariable>(&param)) {
+      body << attr->getVar()->name << "Attr()";
+
+    } else if (auto *operand = dyn_cast<OperandVariable>(&param)) {
+      body << operand->getVar()->name << "()";
+
+    } else if (auto *successor = dyn_cast<SuccessorVariable>(&param)) {
+      body << successor->getVar()->name << "()";
+
+    } 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())
+        body << llvm::formatv("({0}() ? {0}().getType() : Type())", var->name);
+      else
+        body << var->name << "().getType()";
+    } else {
+      llvm_unreachable("unknown custom directive parameter");
+    }
+  }
+
+  body << ");\n";
+}
+
 /// Generate the C++ for an operand to a (*-)type directive.
 static OpMethodBody &genTypeOperandPrinter(Element *arg, OpMethodBody &body) {
   if (isa<OperandsDirective>(arg))
     return body << "getOperation()->getOperandTypes()";
   if (isa<ResultsDirective>(arg))
     return body << "getOperation()->getResultTypes()";
   auto *operand = dyn_cast<OperandVariable>(arg);
   auto *var = operand ? operand->getVar() : cast<ResultVariable>(arg)->getVar();
if (operand->getVar()->isOptional()) {
       body << "  p << " << operand->getVar()->name << "();\n";
     }
   } else if (auto *successor = dyn_cast<SuccessorVariable>(element)) {
     const NamedSuccessor *var = successor->getVar();
     if (var->isVariadic())
       body << "  ::llvm::interleaveComma(" << var->name << "(), p);\n";
     else
       body << "  p << " << var->name << "();\n";
+  } else if (auto *dir = dyn_cast<CustomDirective>(element)) {
+    genCustomDirectivePrinter(dir, body);
   } else if (isa<OperandsDirective>(element)) {
     body << "  p << getOperation()->getOperands();\n";
   } 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<FunctionalTypeDirective>(element)) {
enum Kind {
     error,
 
     // Tokens with no info.
     l_paren,
     r_paren,
     caret,
     comma,
     equal,
+    less,
+    greater,
     question,
 
     // Keywords.
     keyword_start,
     kw_attr_dict,
     kw_attr_dict_w_keyword,
+    kw_custom,
     kw_functional_type,
     kw_operands,
     kw_results,
     kw_successors,
     kw_type,
     keyword_end,
 
     // String valued tokens.
Token FormatLexer::lexToken() {
 
   // Lex punctuation.
   case '^':
     return formToken(Token::caret, tokStart);
   case ',':
     return formToken(Token::comma, tokStart);
   case '=':
     return formToken(Token::equal, tokStart);
+  case '<':
+    return formToken(Token::less, tokStart);
+  case '>':
+    return formToken(Token::greater, tokStart);
   case '?':
     return formToken(Token::question, tokStart);
   case '(':
     return formToken(Token::l_paren, tokStart);
   case ')':
     return formToken(Token::r_paren, tokStart);
 
   // Ignore whitespace characters.
while (isalnum(*curPtr) || *curPtr == '_' || *curPtr == '-')
     ++curPtr;
 
   // Check to see if this identifier is a keyword.
   StringRef str(tokStart, curPtr - tokStart);
   Token::Kind kind =
       llvm::StringSwitch<Token::Kind>(str)
           .Case("attr-dict", Token::kw_attr_dict)
           .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("results", Token::kw_results)
           .Case("successors", Token::kw_successors)
           .Case("type", Token::kw_type)
           .Default(Token::identifier);
   return Token(kind, str);
 }
LogicalResult parseOptionalChildElement(
       std::vector<std::unique_ptr<Element>> &childElements,
       SmallPtrSetImpl<const NamedTypeConstraint *> &seenVariables,
       Optional<unsigned> &anchorIdx);
 
   /// Parse the various different directives.
   LogicalResult parseAttrDictDirective(std::unique_ptr<Element> &element,
                                        llvm::SMLoc loc, bool isTopLevel,
                                        bool withKeyword);
+  LogicalResult parseCustomDirective(std::unique_ptr<Element> &element,
+                                     llvm::SMLoc loc, bool isTopLevel);
+  LogicalResult parseCustomDirectiveParameter(
+      std::vector<std::unique_ptr<Element>> &parameters);
   LogicalResult parseFunctionalTypeDirective(std::unique_ptr<Element> &element,
                                              Token tok, bool isTopLevel);
   LogicalResult parseOperandsDirective(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 FormatParser::parseDirective(std::unique_ptr<Element> &element,
 
   switch (dirTok.getKind()) {
   case Token::kw_attr_dict:
     return parseAttrDictDirective(element, dirTok.getLoc(), isTopLevel,
                                   /*withKeyword=*/false);
   case Token::kw_attr_dict_w_keyword:
     return parseAttrDictDirective(element, dirTok.getLoc(), isTopLevel,
                                   /*withKeyword=*/true);
+  case Token::kw_custom:
+    return parseCustomDirective(element, dirTok.getLoc(), isTopLevel);
   case Token::kw_functional_type:
     return parseFunctionalTypeDirective(element, dirTok, isTopLevel);
   case Token::kw_operands:
     return parseOperandsDirective(element, dirTok.getLoc(), isTopLevel);
   case Token::kw_results:
     return parseResultsDirective(element, dirTok.getLoc(), isTopLevel);
   case Token::kw_successors:
     return parseSuccessorsDirective(element, dirTok.getLoc(), isTopLevel);
return TypeSwitch<Element *, LogicalResult>(childElements.back().get())
       // group.
       .Case<OperandVariable>([&](OperandVariable *ele) {
         if (!ele->getVar()->isVariableLength())
           return emitError(childLoc, "only variable length operands can be "
                                      "used within an optional group");
         seenVariables.insert(ele->getVar());
         return success();
       })
-      // Literals and type directives may be used, but they can't anchor the
-      // group.
-      .Case<LiteralElement, TypeDirective, FunctionalTypeDirective>(
-          [&](Element *) {
-            if (isAnchor)
-              return emitError(childLoc, "only variables can be used to anchor "
-                                         "an optional group");
-            return success();
-          })
+      // Literals, custom directives, and type directives may be used,
+      // but they can't anchor the group.
+      .Case<LiteralElement, CustomDirective, TypeDirective,
+            FunctionalTypeDirective>([&](Element *) {
       .Default([&](Element *) {
         return emitError(childLoc, "only literals, types, and variables can be "
                                    "used within an optional group");
       });
 }
 
 LogicalResult
 FormatParser::parseAttrDictDirective(std::unique_ptr<Element> &element,
                                      llvm::SMLoc loc, bool isTopLevel,
                                      bool withKeyword) {
   if (!isTopLevel)
     return emitError(loc, "'attr-dict' directive can only be used as a "
                           "top-level directive");
   if (hasAttrDict)
     return emitError(loc, "'attr-dict' directive has already been seen");
 
   hasAttrDict = true;
   element = std::make_unique<AttrDictDirective>(withKeyword);
   return success();
 }
 
 LogicalResult
+FormatParser::parseCustomDirective(std::unique_ptr<Element> &element,
+                                   llvm::SMLoc loc, bool isTopLevel) {
+  llvm::SMLoc curLoc = curToken.getLoc();
+
+  // Parse the custom directive name.
+  if (failed(
+          parseToken(Token::less, "expected '<' before custom directive name")))
+    return failure();
+
+  Token nameTok = curToken;
+  if (failed(parseToken(Token::identifier,
+                        "expected custom directive name identifier")) ||
+      failed(parseToken(Token::greater,
+                        "expected '>' after custom directive name")) ||
+      failed(parseToken(Token::l_paren,
+                        "expected '(' before custom directive parameters")))
+    return failure();
+
+  // Parse the child elements for this optional group.=
+  std::vector<std::unique_ptr<Element>> elements;
+  do {
+    if (failed(parseCustomDirectiveParameter(elements)))
+      return failure();
+    if (curToken.getKind() != Token::comma)
+      break;
+    consumeToken();
+  } while (true);
+
+  if (failed(parseToken(Token::r_paren,
+                        "expected ')' after custom directive parameters")))
+    return 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<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 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 failure();
+
+  // Verify that the element can be placed within a custom directive.
+  if (!isa<TypeDirective, AttributeVariable, OperandVariable,
+           SuccessorVariable>(parameters.back().get())) {
+    return emitError(childLoc, "only variables and types may be used as "
+                               "parameters to a custom directive");
+  }
+  return success();
+}
+
+LogicalResult
 FormatParser::parseFunctionalTypeDirective(std::unique_ptr<Element> &element,
                                            Token tok, bool isTopLevel) {
   llvm::SMLoc loc = tok.getLoc();
   if (!isTopLevel)
     return emitError(
         loc, "'functional-type' is only valid as a top-level directive");
 
   // Parse the main operand.