[mlir][ods] Add support for custom directive in attr/type formats

[Mogball]

Summary

This patch adds support for custom directives in attribute and type formats. Custom directives dispatch calls to user-defined parser and printer functions.

For example, the assembly format "custom<Foo>($foo, ref($bar))" expects a function with the signature

LogicalResult parseFoo(AsmParser &parser, FailureOr<FooT> &foo, BarT bar);
void printFoo(AsmPrinter &printer, FooT foo, BarT bar);

[rriddle]

Thanks Jeff! Great work as always.

mlir/docs/Tutorials/DefiningAttributesAndTypes.md

669

Can you add an example for say, a StringRef? It would be nice to have more complex/non-primitive test cases. StringRef as before is a good one, and so would something that is "expensive" to copy/non-trivially copyable/etc.

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

803–807

993–1000

1045

unfinished sentence: Collect

mlir/tools/mlir-tblgen/FormatGen.h

341

? To avoid saying reference twice.

[Mogball]

Review comments. And added a test case with StringRefParameter

[rriddle]

mlir/docs/Tutorials/DefiningAttributesAndTypes.md

688

Missing a & here?

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

227

nit: Drop trivial braces here.

239

nit: Spell out auto here.

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

1118

Missing test coverage for all of these new error messages?

[Mogball]

Review comments

[rriddle]

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

1068

Missing coverage for this?

[Mogball]

add missing test coverage

mlir/docs/Tutorials/DefiningAttributesAndTypes.md

 
 Attribute and type assembly formats have the following directives:
 
 *   `params`: capture all parameters of an attribute or type.
 *   `qualified`: mark a parameter to be printed with its leading dialect and
     mnemonic.
 *   `struct`: generate a "struct-like" parser and printer for a list of
     key-value pairs.
+*   `custom`: dispatch a call to user-define parser and printer functions
+*   `ref`: in a custom directive, references a previously bound variable
 
 #### `params` Directive
 
 This directive is used to refer to all parameters of an attribute or type. When
 used as a top-level directive, `params` generates a parser and printer for a
 comma-separated list of the parameters. For example:
 
 ```tablegen

 assembly format of `` `<` struct(params) `>` `` will result in:
 
 ```mlir
 !my_dialect.struct<b = 2, sym_name = "foo", c = 3, a = 1>
 ```
 
 The order in which the parameters are printed is the order in which they are
 declared in the attribute's or type's `parameter` list.
+
+#### `custom` and `ref` directive
+
+The `custom` directive is used to dispatch calls to user-defined printer and
+parser functions. For example, suppose we had the following type:
+
+```tablegen
+let parameters = (ins "int":$foo, "int":$bar);
+let assemblyFormat = "custom<Foo>($foo) custom<Bar>($bar, ref($foo))";
+```
+
+The `custom` directive `custom<Foo>($foo)` will in the parser and printer
+respectively generate calls to:
+
+```c++
+LogicalResult parseFoo(AsmParser &parser, FailureOr<int> &foo);
+void printFoo(AsmPrinter &printer, int foo);
+```
+
+A previously bound variable can be passed as a parameter to a `custom` directive
+by wrapping it in a `ref` directive. In the previous example, `$foo` is bound by
+the first directive. The second directive references it and expects the
+following printer and parser signatures:
+
+```c++
+LogicalResult parseBar(AsmParser &parser, FailureOr<int> &bar, int foo);
+void printBar(AsmPrinter &printer, int bar, int foo);
+```
+
+More complex C++ types can be used with the `custom` directive. The only caveat
+is that the parameter for the parser must use the storage type of the parameter.
+For example, `StringRefParameter` expects the parser and printer signatures as:
+
+```c++
+LogicalResult parseStringParam(AsmParser &parser,
+                               FailureOr<std::string> &value);
+void printStringParam(AsmPrinter &printer, StringRef value);
+```
+
+The custom parser is considered to have failed if it returns failure or if any
+bound parameters have failure values afterwards.

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

def TestTypeDefaultValuedType : Test_Type<"TestTypeDefaultValuedType"> {
   let parameters = (ins
     DefaultValuedParameter<"mlir::IntegerType",
                            "mlir::IntegerType::get($_ctx, 32)">:$type
   );
   let mnemonic = "default_valued_type";
   let assemblyFormat = "`<` (`(` $type^ `)`)? `>`";
 }
 
+def TestTypeCustom : Test_Type<"TestTypeCustom"> {
+  let parameters = (ins "int":$a, OptionalParameter<"mlir::Optional<int>">:$b);
+  let mnemonic = "custom_type";
+  let assemblyFormat = [{ `<` custom<CustomTypeA>($a)
+                              custom<CustomTypeB>(ref($a), $b) `>` }];
+}
+
+def TestTypeCustomString : Test_Type<"TestTypeCustomString"> {
+  let parameters = (ins StringRefParameter<>:$foo);
+  let mnemonic = "custom_type_string";
+  let assemblyFormat = [{ `<` custom<FooString>($foo)
+                              custom<BarString>(ref($foo)) `>` }];
+}
+
 #endif // TEST_TYPEDEFS

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

for (DataLayoutEntryInterface entry : params) {
     StringRef kind = pair.front().cast<StringAttr>().getValue();
     if (kind == expectedKind)
       return pair.back().cast<IntegerAttr>().getValue().getZExtValue();
   }
   return 1;
 }
 
 //===----------------------------------------------------------------------===//
+// TestCustomType
+//===----------------------------------------------------------------------===//
+
+static LogicalResult parseCustomTypeA(AsmParser &parser,
+                                      FailureOr<int> &a_result) {
+  a_result.emplace();
+  return parser.parseInteger(*a_result);
+}
+
+static void printCustomTypeA(AsmPrinter &printer, int a) { printer << a; }
+
+static LogicalResult parseCustomTypeB(AsmParser &parser, int a,
+                                      FailureOr<Optional<int>> &b_result) {
+  if (a < 0)
+    return success();
+  for (int i : llvm::seq(0, a))
+    if (failed(parser.parseInteger(i)))
+      return failure();
+  b_result.emplace(0);
+  return parser.parseInteger(**b_result);
+}
+
+static void printCustomTypeB(AsmPrinter &printer, int a, Optional<int> b) {
+  if (a < 0)
+    return;
+  printer << ' ';
+  for (int i : llvm::seq(0, a))
+    printer << i << ' ';
+  printer << *b;
+}
+
+static LogicalResult parseFooString(AsmParser &parser,
+                                    FailureOr<std::string> &foo) {
+  std::string result;
+  if (parser.parseString(&result))
+    return failure();
+  foo = std::move(result);
+  return success();
+}
+
+static void printFooString(AsmPrinter &printer, StringRef foo) {
+  printer << '"' << foo << '"';
+}
+
+static LogicalResult parseBarString(AsmParser &parser, StringRef foo) {
+  return parser.parseKeyword(foo);
+}
+
+static void printBarString(AsmPrinter &printer, StringRef foo) {
+  printer << ' ' << foo;
+}
+
+//===----------------------------------------------------------------------===//
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_TYPEDEF_CLASSES
 #include "TestTypeDefs.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // TestDialect

mlir/test/mlir-tblgen/attr-or-type-format-invalid.td

def InvalidTypeM : InvalidType<"InvalidTypeM", "invalid_m"> {
   let assemblyFormat = "(`(` $a^ `,` $b `)`)?";
 }
 
 def InvalidTypeN : InvalidType<"InvalidTypeN", "invalid_n"> {
   let parameters = (ins OptionalParameter<"int">:$a);
   // CHECK: optional group anchor must be a parameter or directive
   let assemblyFormat = "(`(` $a `)`^)?";
 }
+
+def InvalidTypeO : InvalidType<"InvalidTypeO", "invalid_o"> {
+  let parameters = (ins "int":$a);
+  // CHECK: `ref` is only allowed inside custom directives
+  let assemblyFormat = "$a ref($a)";
+}
+
+def InvalidTypeP : InvalidType<"InvalidTypeP", "invalid_p"> {
+  let parameters = (ins "int":$a);
+  // CHECK: parameter 'a' must be bound before it is referenced
+  let assemblyFormat = "custom<Foo>(ref($a)) $a";
+}
+
+def InvalidTypeQ : InvalidType<"InvalidTypeQ", "invalid_q"> {
+  let parameters = (ins "int":$a);
+  // CHECK: `params` can only be used at the top-level context or within a `struct` directive
+  let assemblyFormat = "custom<Foo>(params)";
+}
+
+def InvalidTypeR : InvalidType<"InvalidTypeR", "invalid_r"> {
+  let parameters = (ins "int":$a);
+  // CHECK: `struct` can only be used at the top-level context
+  let assemblyFormat = "custom<Foo>(struct(params))";
+}

mlir/test/mlir-tblgen/attr-or-type-format-roundtrip.mlir

 // CHECK: !test.optional_group_struct<(b = 5)>
 // CHECK: !test.optional_group_struct<(a = 10, b = 5)>
 // CHECK: !test.spaces< 5
 // CHECK-NEXT: ()() 6>
 // CHECK: !test.ap_float<5.000000e+00>
 // CHECK: !test.ap_float<>
 // CHECK: !test.default_valued_type<(i64)>
 // CHECK: !test.default_valued_type<>
+// CHECK: !test.custom_type<-5>
+// CHECK: !test.custom_type<2 0 1 5>
+// CHECK: !test.custom_type_string<"foo" foo>
+// CHECK: !test.custom_type_string<"bar" bar>
 
 func private @test_roundtrip_default_parsers_struct(
   !test.no_parser<255, [1, 2, 3, 4, 5], "foobar", 4>
 ) -> (
   !test.struct_capture_all<v3 = 3, v1 = 1, v2 = 2, v0 = 0>,
   !test.optional_param<, 6>,
   !test.optional_param<5, 6>,
   !test.optional_params<"a">,
) -> (
   !test.optional_group_params<(5, 6)>,
   !test.optional_group_struct<x>,
   !test.optional_group_struct<(b = 5)>,
   !test.optional_group_struct<(b = 5, a = 10)>,
   !test.spaces<5 ()() 6>,
   !test.ap_float<5.0>,
   !test.ap_float<>,
   !test.default_valued_type<(i64)>,
-  !test.default_valued_type<>
+  !test.default_valued_type<>,
+  !test.custom_type<-5>,
+  !test.custom_type<2 9 9 5>,
+  !test.custom_type_string<"foo" foo>,
+  !test.custom_type_string<"bar" bar>
 )

mlir/test/mlir-tblgen/attr-or-type-format.td

 // TYPE: void TestKType::print(::mlir::AsmPrinter &odsPrinter) const {
 // TYPE:   if ((getA() && !(getA() == 10)))
 
 def TypeI : TestType<"TestK"> {
   let parameters = (ins DefaultValuedParameter<"int", "10">:$a);
   let mnemonic = "type_k";
   let assemblyFormat = "$a";
 }
+
+// TYPE: ::mlir::Type TestLType::parse
+// TYPE:   auto odsCustomLoc = odsParser.getCurrentLocation()
+// TYPE:   auto odsCustomResult = parseA(odsParser,
+// TYPE-NEXT: _result_a
+// TYPE:   if (::mlir::failed(odsCustomResult)) return {}
+// TYPE:   if (::mlir::failed(_result_a))
+// TYPE-NEXT: odsParser.emitError(odsCustomLoc,
+// TYPE:   auto odsCustomResult = parseB(odsParser,
+// TYPE-NEXT: _result_b
+// TYPE-NEXT: *_result_a
+
+// TYPE: void TestLType::print
+// TYPE:   printA(odsPrinter
+// TYPE-NEXT: getA()
+// TYPE:   printB(odsPrinter
+// TYPE-NEXT: getB()
+// TYPE-NEXT: getA()
+
+def TypeJ : TestType<"TestL"> {
+  let parameters = (ins "int":$a, OptionalParameter<"Attribute">:$b);
+  let mnemonic = "type_j";
+  let assemblyFormat = "custom<A>($a) custom<B>($b, ref($a))";
+}

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

private:
   void genLiteralParser(StringRef value, FmtContext &ctx, MethodBody &os,
                         bool isOptional = false);
   /// Generate the parser code for a variable.
   void genVariableParser(ParameterElement *el, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for a `params` directive.
   void genParamsParser(ParamsDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for a `struct` directive.
   void genStructParser(StructDirective *el, FmtContext &ctx, MethodBody &os);
+  /// Generate the parser code for a `custom` directive.
+  void genCustomParser(CustomDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the parser code for an optional group.
   void genOptionalGroupParser(OptionalElement *el, FmtContext &ctx,
                               MethodBody &os);
 
   /// Generate the printer code for a specific format element.
   void genElementPrinter(FormatElement *el, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for a literal.
   void genLiteralPrinter(StringRef value, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for a variable.
   void genVariablePrinter(ParameterElement *el, FmtContext &ctx, MethodBody &os,
                           bool skipGuard = false);
   /// Generate a printer for comma-separated parameters.
   void genCommaSeparatedPrinter(ArrayRef<ParameterElement *> params,
                                 FmtContext &ctx, MethodBody &os,
                                 function_ref<void(ParameterElement *)> extra);
   /// Generate the printer code for a `params` directive.
   void genParamsPrinter(ParamsDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for a `struct` directive.
   void genStructPrinter(StructDirective *el, FmtContext &ctx, MethodBody &os);
+  /// Generate the printer code for a `custom` directive.
+  void genCustomPrinter(CustomDirective *el, FmtContext &ctx, MethodBody &os);
   /// Generate the printer code for an optional group.
   void genOptionalGroupPrinter(OptionalElement *el, FmtContext &ctx,
                                MethodBody &os);
   /// Generate a printer (or space eraser) for a whitespace element.
   void genWhitespacePrinter(WhitespaceElement *el, FmtContext &ctx,
                             MethodBody &os);
 
   /// The ODS definition of the attribute or type whose format is being used to
void DefFormat::genElementParser(FormatElement *el, FmtContext &ctx,
   if (auto *literal = dyn_cast<LiteralElement>(el))
     return genLiteralParser(literal->getSpelling(), ctx, os);
   if (auto *var = dyn_cast<ParameterElement>(el))
     return genVariableParser(var, ctx, os);
   if (auto *params = dyn_cast<ParamsDirective>(el))
     return genParamsParser(params, ctx, os);
   if (auto *strct = dyn_cast<StructDirective>(el))
     return genStructParser(strct, ctx, os);
+  if (auto *custom = dyn_cast<CustomDirective>(el))
+    return genCustomParser(custom, ctx, os);
   if (auto *optional = dyn_cast<OptionalElement>(el))
     return genOptionalGroupParser(optional, ctx, os);
   if (isa<WhitespaceElement>(el))
     return;
 
   llvm_unreachable("unknown format element");
 }
 
if (numOptional) {
     // N flags, successfully exiting the loop means that all parameters have
     // been seen. `parseOptionalComma` would cause issues with any formats that
     // use "struct(...) `,`" beacuse structs aren't sounded by braces.
     os.getStream().printReindented(strfmt(loopTerminator, el->getNumParams()));
   }
   os.unindent() << "}\n";
 }
 
+void DefFormat::genCustomParser(CustomDirective *el, FmtContext &ctx,
+                                MethodBody &os) {
+  os << "{\n";
+  os.indent();
+
+  // Bound variables are passed directly to the parser as `FailureOr<T> &`.
+  // Referenced variables are passed as `T`. The custom parser fails if it
+  // returns failure or if any of the required parameters failed.
+  os << tgfmt("auto odsCustomLoc = $_parser.getCurrentLocation();\n", &ctx);
+  os << "(void)odsCustomLoc;\n";
+  os << tgfmt("auto odsCustomResult = parse$0($_parser", &ctx, el->getName());
+  os.indent();
+  for (FormatElement *arg : el->getArguments()) {
+    os << ",\n";
+    FormatElement *param;
+    if (auto *ref = dyn_cast<RefDirective>(arg)) {
+      os << "*";
+      param = ref->getArg();
+    } else {
+      param = arg;
+    }
+    os << "_result_" << cast<ParameterElement>(param)->getName();
+  }
+  os.unindent() << ");\n";
+  os << "if (::mlir::failed(odsCustomResult)) return {};\n";
+  for (FormatElement *arg : el->getArguments()) {
+    if (auto *param = dyn_cast<ParameterElement>(arg)) {
+      if (param->isOptional())
+        continue;
+      os << formatv("if (::mlir::failed(_result_{0})) {{\n", param->getName());
+      os.indent() << tgfmt("$_parser.emitError(odsCustomLoc, ", &ctx)
+                  << "\"custom parser failed to parse parameter '"
+                  << param->getName() << "'\");\n";
+      os << "return {};\n";
+      os.unindent() << "}\n";
+    }
+  }
+
+  os.unindent() << "}\n";
+}
+
 void DefFormat::genOptionalGroupParser(OptionalElement *el, FmtContext &ctx,
                                        MethodBody &os) {
   ArrayRef<FormatElement *> elements =
       el->getThenElements().drop_front(el->getParseStart());
 
   FormatElement *first = elements.front();
   const auto guardOn = [&](auto params) {
     os << "if (!(";

 void DefFormat::genElementPrinter(FormatElement *el, FmtContext &ctx,
                                   MethodBody &os) {
   if (auto *literal = dyn_cast<LiteralElement>(el))
     return genLiteralPrinter(literal->getSpelling(), ctx, os);
   if (auto *params = dyn_cast<ParamsDirective>(el))
     return genParamsPrinter(params, ctx, os);
   if (auto *strct = dyn_cast<StructDirective>(el))
     return genStructPrinter(strct, ctx, os);
+  if (auto *custom = dyn_cast<CustomDirective>(el))
+    return genCustomPrinter(custom, ctx, os);
   if (auto *var = dyn_cast<ParameterElement>(el))
     return genVariablePrinter(var, ctx, os);
   if (auto *optional = dyn_cast<OptionalElement>(el))
     return genOptionalGroupPrinter(optional, ctx, os);
   if (auto *whitespace = dyn_cast<WhitespaceElement>(el))
     return genWhitespacePrinter(whitespace, ctx, os);
 
   llvm::PrintFatalError("unsupported format element");

 void DefFormat::genStructPrinter(StructDirective *el, FmtContext &ctx,
                                  MethodBody &os) {
   genCommaSeparatedPrinter(
       llvm::to_vector(el->getParams()), ctx, os, [&](ParameterElement *param) {
         os << tgfmt("$_printer << \"$0 = \";\n", &ctx, param->getName());
       });
 }
 
+void DefFormat::genCustomPrinter(CustomDirective *el, FmtContext &ctx,
+                                 MethodBody &os) {
+  os << tgfmt("print$0($_printer", &ctx, el->getName());
+  os.indent();
+  for (FormatElement *arg : el->getArguments()) {
+    FormatElement *param = arg;
+    if (auto *ref = dyn_cast<RefDirective>(arg))
+      param = ref->getArg();
+    os << ",\n"
+       << getParameterAccessorName(cast<ParameterElement>(param)->getName())
+       << "()";
+  }
+  os.unindent() << ");\n";
+}
+
 void DefFormat::genOptionalGroupPrinter(OptionalElement *el, FmtContext &ctx,
                                         MethodBody &os) {
   FormatElement *anchor = el->getAnchor();
   if (auto *param = dyn_cast<ParameterElement>(anchor)) {
     guardOnAny(ctx, os, llvm::makeArrayRef(param));
   } else if (auto *params = dyn_cast<ParamsDirective>(anchor)) {
     guardOnAny(ctx, os, params->getParams());
   } else {
public:
   FailureOr<DefFormat> parse();
 
 protected:
   /// Verify the parsed elements.
   LogicalResult verify(SMLoc loc, ArrayRef<FormatElement *> elements) override;
   /// Verify the elements of a custom directive.
   LogicalResult
   verifyCustomDirectiveArguments(SMLoc loc,
-                                 ArrayRef<FormatElement *> arguments) override {
-    return emitError(loc, "'custom' not supported (yet)");
-  }
+                                 ArrayRef<FormatElement *> arguments) override;
   /// Verify the elements of an optional group.
   LogicalResult
   verifyOptionalGroupElements(SMLoc loc, ArrayRef<FormatElement *> elements,
                               Optional<unsigned> anchorIndex) override;
 
   /// Parse an attribute or type variable.
   FailureOr<FormatElement *> parseVariableImpl(SMLoc loc, StringRef name,
                                                Context ctx) override;
   /// Parse an attribute or type format directive.
   FailureOr<FormatElement *>
   parseDirectiveImpl(SMLoc loc, FormatToken::Kind kind, Context ctx) override;
 
 private:
   /// Parse a `params` directive.
-  FailureOr<FormatElement *> parseParamsDirective(SMLoc loc);
+  FailureOr<FormatElement *> parseParamsDirective(SMLoc loc, Context ctx);
   /// Parse a `qualified` directive.
   FailureOr<FormatElement *> parseQualifiedDirective(SMLoc loc, Context ctx);
   /// Parse a `struct` directive.
-  FailureOr<FormatElement *> parseStructDirective(SMLoc loc);
+  FailureOr<FormatElement *> parseStructDirective(SMLoc loc, Context ctx);
+  /// Parse a `ref` directive.
+  FailureOr<FormatElement *> parseRefDirective(SMLoc loc, Context ctx);
 
   /// Attribute or type tablegen def.
   const AttrOrTypeDef &def;
 
   /// Seen attribute or type parameters.
   BitVector seenParams;
 };
 } // namespace
for (auto it : llvm::zip(elements.drop_back(), elements.drop_front())) {
     if (literalEl->getSpelling() == "," && structEl->hasOptionalParams()) {
       return emitError(loc, "`struct` directive with optional parameters "
                             "cannot be followed by a comma literal");
     }
   }
   return success();
 }
 
+LogicalResult DefFormatParser::verifyCustomDirectiveArguments(
+    SMLoc loc, ArrayRef<FormatElement *> arguments) {
+  // Arguments are fully verified by the parser context.
+  return success();
+}
+
 LogicalResult
 DefFormatParser::verifyOptionalGroupElements(llvm::SMLoc loc,
                                              ArrayRef<FormatElement *> elements,
                                              Optional<unsigned> anchorIndex) {
   // `params` and `struct` directives are allowed only if all the contained
   // parameters are optional.
   for (FormatElement *el : elements) {
     if (auto *param = dyn_cast<ParameterElement>(el)) {
auto *it = llvm::find_if(
       params, [&](auto &param) { return param.getName() == name; });
 
   // Check that the parameter reference is valid.
   if (it == params.end()) {
     return emitError(loc,
                      def.getName() + " has no parameter named '" + name + "'");
   }
   auto idx = std::distance(params.begin(), it);
+
+  if (ctx != RefDirectiveContext) {
+    // Check that the variable has not already been bound.
     if (seenParams.test(idx))
       return emitError(loc, "duplicate parameter '" + name + "'");
     seenParams.set(idx);
 
+    // Otherwise, to be referenced, a variable must have been bound.
+  } else if (!seenParams.test(idx)) {
+    return emitError(loc, "parameter '" + name +
+                              "' must be bound before it is referenced");
+  }
+
   return create<ParameterElement>(*it);
 }
 
 FailureOr<FormatElement *>
 DefFormatParser::parseDirectiveImpl(SMLoc loc, FormatToken::Kind kind,
                                     Context ctx) {
 
   switch (kind) {
   case FormatToken::kw_qualified:
     return parseQualifiedDirective(loc, ctx);
   case FormatToken::kw_params:
-    return parseParamsDirective(loc);
+    return parseParamsDirective(loc, ctx);
   case FormatToken::kw_struct:
-    if (ctx != TopLevelContext) {
-      return emitError(
-          loc,
-          "`struct` may only be used in the top-level section of the format");
-    }
-    return parseStructDirective(loc);
+    return parseStructDirective(loc, ctx);
+  case FormatToken::kw_ref:
+    return parseRefDirective(loc, ctx);
+  case FormatToken::kw_custom:
+    return parseCustomDirective(loc, ctx);
 
   default:
     return emitError(loc, "unsupported directive kind");
   }
 }
 
 FailureOr<FormatElement *>
 DefFormatParser::parseQualifiedDirective(SMLoc loc, Context ctx) {
   if (failed(parseToken(FormatToken::l_paren,
                         "expected '(' before argument list")))
     return failure();
   FailureOr<FormatElement *> var = parseElement(ctx);
   if (failed(var))
     return var;
   if (!isa<ParameterElement>(*var))
     return emitError(loc, "`qualified` argument list expected a variable");
   cast<ParameterElement>(*var)->setShouldBeQualified();
   if (failed(
           parseToken(FormatToken::r_paren, "expected ')' after argument list")))
     return failure();
   return var;
 }
 
-FailureOr<FormatElement *> DefFormatParser::parseParamsDirective(SMLoc loc) {
-  // Collect all of the attribute's or type's parameters.
+FailureOr<FormatElement *> DefFormatParser::parseParamsDirective(SMLoc loc,
+                                                                 Context ctx) {
+  // It doesn't make sense to allow references to all parameters in a custom
+  // directive because parameters are the only things that can be bound.
+  if (ctx != TopLevelContext && ctx != StructDirectiveContext) {
+    return emitError(loc, "`params` can only be used at the top-level context "
+                          "or within a `struct` directive");
+  }
+
+  // Collect all of the attribute's or type's parameters and ensure that none of
+  // the parameters have already been captured.
   std::vector<ParameterElement *> vars;
-  // Ensure that none of the parameters have already been captured.
   for (const auto &it : llvm::enumerate(def.getParameters())) {
     if (seenParams.test(it.index())) {
       return emitError(loc, "`params` captures duplicate parameter: " +
                                 it.value().getName());
     }
     seenParams.set(it.index());
     vars.push_back(create<ParameterElement>(it.value()));
   }
   return create<ParamsDirective>(std::move(vars));
 }
 
-FailureOr<FormatElement *> DefFormatParser::parseStructDirective(SMLoc loc) {
+FailureOr<FormatElement *> DefFormatParser::parseStructDirective(SMLoc loc,
+                                                                 Context ctx) {
+  if (ctx != TopLevelContext)
+    return emitError(loc, "`struct` can only be used at the top-level context");
+
   if (failed(parseToken(FormatToken::l_paren,
                         "expected '(' before `struct` argument list")))
     return failure();
 
   // Parse variables captured by `struct`.
   std::vector<ParameterElement *> vars;
 
   // Parse first captured parameter or a `params` directive.
FailureOr<FormatElement *> DefFormatParser::parseStructDirective(SMLoc loc,
 
   if (failed(parseToken(FormatToken::r_paren,
                         "expected ')' at the end of an argument list")))
     return failure();
 
   return create<StructDirective>(std::move(vars));
 }
 
+FailureOr<FormatElement *> DefFormatParser::parseRefDirective(SMLoc loc,
+                                                              Context ctx) {
+  if (ctx != CustomDirectiveContext)
+    return emitError(loc, "`ref` is only allowed inside custom directives");
+
+  // Parse the child parameter element.
+  FailureOr<FormatElement *> child;
+  if (failed(parseToken(FormatToken::l_paren, "expected '('")) ||
+      failed(child = parseElement(RefDirectiveContext)) ||
+      failed(parseToken(FormatToken::r_paren, "expeced ')'")))
+    return failure();
+
+  // Only parameter elements are allowed to be parsed under a `ref` directive.
+  return create<RefDirective>(*child);
+}
+
 //===----------------------------------------------------------------------===//
 // Interface
 //===----------------------------------------------------------------------===//
 
 void mlir::tblgen::generateAttrOrTypeFormat(const AttrOrTypeDef &def,
                                             MethodBody &parser,
                                             MethodBody &printer) {
   llvm::SourceMgr mgr;

mlir/tools/mlir-tblgen/FormatGen.h

private:
   /// methods: `parse{name}` and `print{name}` with the given arguments.
   StringRef name;
   /// The arguments with which to call the custom functions. These are either
   /// variables (for which the functions are responsible for populating) or
   /// references to variables.
   std::vector<FormatElement *> arguments;
 };
 
+/// This class represents a reference directive. This directive can be used to
+/// reference but not bind a previously bound variable or format object. Its
+/// current only use is to pass variables as arguments to the custom directive.
+class RefDirective : public DirectiveElementBase<DirectiveElement::Ref> {
+public:
+  /// Create a reference directive with the single referenced child.
+  RefDirective(FormatElement *arg) : arg(arg) {}
+
+  /// Get the reference argument.
+  FormatElement *getArg() const { return arg; }
+
+private:
+  /// The referenced argument.
+  FormatElement *arg;
+};
+
 /// This class represents a group of elements that are optionally emitted based
 /// on an optional variable "anchor" and a group of elements that are emitted
 /// when the anchor element is not present.
 class OptionalElement : public FormatElementBase<FormatElement::Optional> {
 public:
   /// Create an optional group with the given child elements.
   OptionalElement(std::vector<FormatElement *> &&thenElements,
                   std::vector<FormatElement *> &&elseElements,

mlir/tools/mlir-tblgen/OpFormatGen.cpp

public:
   FormatElement *getInputs() const { return inputs; }
   FormatElement *getResults() const { return results; }
 
 private:
   /// The input and result arguments.
   FormatElement *inputs, *results;
 };
 
-/// This class represents the `ref` directive.
-class RefDirective : public DirectiveElementBase<DirectiveElement::Ref> {
-public:
-  RefDirective(FormatElement *arg) : arg(arg) {}
-
-  FormatElement *getArg() const { return arg; }
-
-private:
-  /// The argument that is used to format the directive.
-  FormatElement *arg;
-};
-
 /// This class represents the `type` directive.
 class TypeDirective : public DirectiveElementBase<DirectiveElement::Type> {
 public:
   TypeDirective(FormatElement *arg) : arg(arg) {}
 
   FormatElement *getArg() const { return arg; }
 
   /// Indicate if this type is printed "qualified" (that is it is