[MLIR] [ODS] Allowing attr-dict in custom directive

[jdd]

Summary

Enhance tblgen's declarative assembly format to allow attr-dict in
custom directives.

[rriddle]

mlir/docs/OpDefinitions.md

805–806

This section needs an update.

836–837

This section needs an update.

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1512

Can we instead change this such that:

• attr-dict would pass MutableDictionaryAttr as the parameter
• we always pass in the operation to the custom directive printer, similarly to the OpAsmPrinter parameter?

?

[jdd]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1512

we always pass in the operation to the custom directive printer, similarly to the OpAsmPrinter parameter?

You mean regardless of the custom directive parameters or just in the attr-dict case? The former is a breaking change. I think you mean the former, which is what I would prefer as well. How much code would it break?

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1512

Regardless, i.e. I think we should do both of those things.

How much code would it break?

Not much given the relative recency of custom directives. That being said, I prefer strongly to not use this as a metric for doing something. I prefer that we do the right thing regardless of whether we have to refactor existing code or not.

[jdd]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1512

I agree. How would we notify users? Do we even have to for breakage this small?

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1512

For really small things like this I usually don't notify users. If you do need to notify in the future, I usually just drop a message in discourse letting people know something big is going to happen and give them a timeline to switch(assuming it's a staged change).

[jdd]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1512

attr-dict would pass MutableDictionaryAttr as the parameter

You mean to the parser, right? Presumably the printer wouldn't be modifying anything.

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1512

I meant the printer here. The "mutable" name is a little weird, but it prevents the need to construct a DictionaryAttr if there aren't any attributes. Though an operation without attributes that implements a custom directive is probably pretty rare, so using DictionaryAttr for the printer is probably fine.

Context: https://github.com/llvm/llvm-project/blob/69efcd03bdb9ee7328dd4e75513ae214c1e23519/mlir/lib/IR/Attributes.cpp#L1359

[jdd]

• Adding test, making suggested changes, fixing breaks

[jdd]

Not sure I'm interpreting all your comments correctly but it's certainly closer.

[rriddle]

Please add a comment to the description detailing the fact that the print function of custom directives pass the operation as well.

mlir/docs/OpDefinitions.md

842

You shouldn't pass references to operations here, they are value-typed(i.e. just smart pointers).

864

This wouldn't be a NamedAttrList, I would expect MutableDictionaryAttr here. NamedAttrList is expensive to construct, whereas MutableDictionaryAttr is free.

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1513

This should be getMutableAttrDict

[jdd]

mlir/docs/OpDefinitions.md

864

Should it be MutableDictionaryAttr or a reference to it? getMutableAttrDict() returns a reference to the internal attribute dictionary and I don't want to give the impression that one should be mutating the object in the printer. Should it be a const reference?

[jdd]

• Changes from feedback #2

[rriddle]

mlir/docs/OpDefinitions.md

864

const MutableDictionaryAttr & SGTM

[jdd]

• const MutableDictionaryAttr&

[mehdi_amini]

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

1647

There is no test actually exercising this op?

[jdd]

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

1647

It would appear not. At least not as part of this patch. I think I would have thought that if the generated file compiled, that'd be good enough.

mlir/docs/OpDefinitions.md

 ```
 
 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.
+(i.e. Attribute, Operand, Successor, etc.), type directives, and `attr-dict`.
+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
+The arguments to the `parse<UserDirective>` method are 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> &`
 *   Region Variables
     -   Single: `Region &`
     -   Variadic: `SmallVectorImpl<std::unique_ptr<Region>> &`
 *   Successor Variables
     -   Single: `Block *&`
     -   Variadic: `SmallVectorImpl<Block *> &`
 *   Type Directives
     -   Single: `Type &`
     -   Optional: `Type &`
     -   Variadic: `SmallVectorImpl<Type> &`
 *   TypeRef Directives
     -   Single: `Type`
     -   Optional: `Type`
     -   Variadic: `const SmallVectorImpl<Type> &`
+*   `attr-dict` Directive: `NamedAttrList &`
 
 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
+The arguments to the `print<UserDirective>` method is firstly a reference to
+the `OpAsmPrinter`(`OpAsmPrinter &`), second the op (e.g. `FooOp op` which
+can be `Operation *op` alternatively), and finally 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`
 *   Region Variables
     -   Single: `Region &`
     -   Variadic: `MutableArrayRef<Region>`
 *   Successor Variables
     -   Single: `Block *`
     -   Variadic: `SuccessorRange`
 *   Type Directives
     -   Single: `Type`
     -   Optional: `Type`
     -   Variadic: `TypeRange`
 *   TypeRef Directives
     -   Single: `Type`
     -   Optional: `Type`
     -   Variadic: `TypeRange`
+*   `attr-dict` Directive: `const MutableDictionaryAttr&`
 
 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

mlir/lib/Dialect/Async/IR/Async.cpp

static ParseResult parseAwaitResultType(OpAsmParser &parser, Type &operandType,
 
   // Add unwrapped async.value type to the returned values types.
   if (auto valueType = operandType.dyn_cast<ValueType>())
     resultType = valueType.getValueType();
 
   return success();
 }
 
-static void printAwaitResultType(OpAsmPrinter &p, Type operandType,
-                                 Type resultType) {
+static void printAwaitResultType(OpAsmPrinter &p, Operation *op,
+                                 Type operandType, Type resultType) {
   p << operandType;
 }
 
 static LogicalResult verify(AwaitOp op) {
   Type argType = op.operand().getType();
 
   // Awaiting on a token does not have any results.
   if (argType.isa<TokenType>() && !op.getResultTypes().empty())

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

if (parser.parseOptionalKeyword("args"))
     return success();
   SmallVector<NamedAttrList, 4> argAttrs;
   bool isVariadic = false;
   return impl::parseFunctionArgumentList(parser, /*allowAttributes=*/false,
                                          /*allowVariadic=*/false, argNames,
                                          argTypes, argAttrs, isVariadic);
 }
 
-static void printLaunchFuncOperands(OpAsmPrinter &printer,
+static void printLaunchFuncOperands(OpAsmPrinter &printer, Operation *,
                                     OperandRange operands, TypeRange types) {
   if (operands.empty())
     return;
   printer << "args(";
   llvm::interleaveComma(llvm::zip(operands, types), printer,
                         [&](const auto &pair) {
                           printer.printOperand(std::get<0>(pair));
                           printer << " : ";
if (succeeded(parser.parseOptionalKeyword("async"))) {
     if (parser.getNumResults() == 0)
       return parser.emitError(loc, "needs to be named when marked 'async'");
     asyncTokenType = parser.getBuilder().getType<AsyncTokenType>();
   }
   return parser.parseOperandList(asyncDependencies,
                                  OpAsmParser::Delimiter::OptionalSquare);
 }
 
-static void printAsyncDependencies(OpAsmPrinter &printer, Type asyncTokenType,
+static void printAsyncDependencies(OpAsmPrinter &printer, Operation *op,
+                                   Type asyncTokenType,
                                    OperandRange asyncDependencies) {
   if (asyncTokenType)
     printer << "async ";
   if (asyncDependencies.empty())
     return;
   printer << "[";
   llvm::interleaveComma(asyncDependencies, printer);
   printer << "]";
 }
 
 #include "mlir/Dialect/GPU/GPUOpInterfaces.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"

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

if (parser.parseAttribute(attr))
     return failure();
   if (succeeded(parser.parseOptionalComma())) {
     if (parser.parseAttribute(optAttr))
       return failure();
   }
   return success();
 }
 
+static ParseResult parseCustomDirectiveAttrDict(OpAsmParser &parser,
+                                                NamedAttrList &attrs) {
+  return parser.parseOptionalAttrDict(attrs);
+}
+
 //===----------------------------------------------------------------------===//
 // Printing
 
-static void printCustomDirectiveOperands(OpAsmPrinter &printer, Value operand,
-                                         Value optOperand,
+static void printCustomDirectiveOperands(OpAsmPrinter &printer, Operation *,
+                                         Value operand, Value optOperand,
                                          OperandRange varOperands) {
   printer << operand;
   if (optOperand)
     printer << ", " << optOperand;
   printer << " -> (" << varOperands << ")";
 }
-static void printCustomDirectiveResults(OpAsmPrinter &printer, Type operandType,
-                                        Type optOperandType,
+static void printCustomDirectiveResults(OpAsmPrinter &printer, Operation *,
+                                        Type operandType, Type optOperandType,
                                         TypeRange varOperandTypes) {
   printer << " : " << operandType;
   if (optOperandType)
     printer << ", " << optOperandType;
   printer << " -> (" << varOperandTypes << ")";
 }
 static void printCustomDirectiveWithTypeRefs(OpAsmPrinter &printer,
-                                             Type operandType,
+                                             Operation *op, Type operandType,
                                              Type optOperandType,
                                              TypeRange varOperandTypes) {
   printer << " type_refs_capture ";
-  printCustomDirectiveResults(printer, operandType, optOperandType,
+  printCustomDirectiveResults(printer, op, 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,
+static void printCustomDirectiveOperandsAndTypes(
+    OpAsmPrinter &printer, Operation *op, Value operand, Value optOperand,
+    OperandRange varOperands, Type operandType, Type optOperandType,
+  printCustomDirectiveOperands(printer, op, operand, optOperand, varOperands);
+  printCustomDirectiveResults(printer, op, operandType, optOperandType,
                               varOperandTypes);
 }
-static void printCustomDirectiveRegions(OpAsmPrinter &printer, Region &region,
+static void printCustomDirectiveRegions(OpAsmPrinter &printer, Operation *,
+                                        Region &region,
                                         MutableArrayRef<Region> varRegions) {
   printer.printRegion(region);
   if (!varRegions.empty()) {
     printer << ", ";
     for (Region &region : varRegions)
       printer.printRegion(region);
   }
 }
-static void printCustomDirectiveSuccessors(OpAsmPrinter &printer,
+static void printCustomDirectiveSuccessors(OpAsmPrinter &printer, Operation *,
                                            Block *successor,
                                            SuccessorRange varSuccessors) {
   printer << successor;
   if (!varSuccessors.empty())
     printer << ", " << varSuccessors.front();
 }
-static void printCustomDirectiveAttributes(OpAsmPrinter &printer,
+static void printCustomDirectiveAttributes(OpAsmPrinter &printer, Operation *,
                                            Attribute attribute,
                                            Attribute optAttribute) {
   printer << attribute;
   if (optAttribute)
     printer << ", " << optAttribute;
 }
 
+static void printCustomDirectiveAttrDict(OpAsmPrinter &printer, Operation *op,
+                                         MutableDictionaryAttr attrs) {
+  printer.printOptionalAttrDict(attrs.getAttrs());
+}
 //===----------------------------------------------------------------------===//
 // 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 FormatCustomDirectiveAttributes
   let assemblyFormat = [{
     custom<CustomDirectiveAttributes>(
       $attr, $optAttr
     )
     attr-dict
   }];
 }
 
+def FormatCustomDirectiveAttrDict
+    : TEST_Op<"format_custom_directive_attrdict"> {
+  let arguments = (ins I64Attr:$attr, OptionalAttr<I64Attr>:$optAttr);
+  let assemblyFormat = [{
+    custom<CustomDirectiveAttrDict>( 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);

mlir/tools/mlir-tblgen/OpFormatGen.cpp

static void genElementParserStorage(Element *element, OpMethodBody &body) {
   }
 }
 
 /// 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 (isa<AttrDictDirective>(&param)) {
+    body << "result.attributes";
   } 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
static void genSpacePrinter(OpMethodBody &body, bool &shouldEmitSpace,
   body << "  p << ' ';\n";
   shouldEmitSpace = false;
   lastWasPunctuation = false;
 }
 
 /// Generate the printer for a custom directive.
 static void genCustomDirectivePrinter(CustomDirective *customDir,
                                       OpMethodBody &body) {
-  body << "  print" << customDir->getName() << "(p";
+  body << "  print" << customDir->getName() << "(p, *this";
   for (Element &param : customDir->getArguments()) {
     body << ", ";
     if (auto *attr = dyn_cast<AttributeVariable>(&param)) {
       body << attr->getVar()->name << "Attr()";
 
+    } else if (isa<AttrDictDirective>(&param)) {
+      // Enforce the const-ness since getMutableAttrDict() returns a reference
+      // into the Operations `attr` member.
+      body << "(const "
+              "MutableDictionaryAttr&)getOperation()->getMutableAttrDict()";
+
     } else if (auto *operand = dyn_cast<OperandVariable>(&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 << "()";

 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<TypeRefDirective, TypeDirective, AttributeVariable, OperandVariable,
-           RegionVariable, SuccessorVariable>(parameters.back().get())) {
+  if (!isa<TypeRefDirective, TypeDirective, AttrDictDirective,
+           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,