add prop-dict support for custom directive for mlir-tblgen

[fengxie]

According to https://mlir.llvm.org/docs/DefiningDialects/Operations/#custom-directives, custom directive supports attr-dict

attr-dict Directive: NamedAttrList &

But it doesn't support prop-dict which is introduced into MLIR recently. It's useful to have tblgen support prop-dict like attr-dict. This PR enable tblgen to support prop-dict

error: only variables and types may be used as parameters to a custom directive
   ... custom<Print>(prop-dict)

[llvmbot]

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: drazi (fengxie)

Changes

According to https://mlir.llvm.org/docs/DefiningDialects/Operations/#custom-directives, custom directive supports attr-dict which is introduced recently.

> attr-dict Directive: NamedAttrList &

But it doesn't support prop-dict which is introduced into MLIR recently. It's useful to have tblgen support prop-dict like attr-dict. This PR enable tblgen to support prop-dict

error: only variables and types may be used as parameters to a custom directive
   ... custom<Print>(prop-dict)

---

Full diff: https://github.com/llvm/llvm-project/pull/77061.diff

3 Files Affected:

• (modified) mlir/test/mlir-tblgen/op-format-spec.td (+3)
• (modified) mlir/test/mlir-tblgen/op-format.td (+8)
• (modified) mlir/tools/mlir-tblgen/OpFormatGen.cpp (+8-3)

diff --git a/mlir/test/mlir-tblgen/op-format-spec.td b/mlir/test/mlir-tblgen/op-format-spec.td
index 5d786cce5f2801..ad2a055bc78b03 100644
--- a/mlir/test/mlir-tblgen/op-format-spec.td
+++ b/mlir/test/mlir-tblgen/op-format-spec.td
@@ -44,6 +44,9 @@ def DirectiveCustomValidC : TestFormat_Op<[{
 def DirectiveCustomValidD : TestFormat_Op<[{
   (`(` custom<MyDirective>($operand)^ `)`)? attr-dict
 }]>, Arguments<(ins Optional<I64>:$operand)>;
+def DirectiveCustomValidE : TestFormat_Op<[{
+  custom<MyDirective>(prop-dict) attr-dict
+}]>, Arguments<(ins UnitAttr:$flag)>;
 
 //===----------------------------------------------------------------------===//
 // functional-type
diff --git a/mlir/test/mlir-tblgen/op-format.td b/mlir/test/mlir-tblgen/op-format.td
index dd756788098c90..3250589605f348 100644
--- a/mlir/test/mlir-tblgen/op-format.td
+++ b/mlir/test/mlir-tblgen/op-format.td
@@ -42,6 +42,14 @@ def CustomStringLiteralC : TestFormat_Op<[{
   custom<Foo>("$_builder.getStringAttr(\"foo\")") attr-dict
 }]>;
 
+// CHECK-LABEL: CustomStringLiteralD::parse
+// CHECK: parseFoo({{.*}}, result)
+// CHECK-LABEL: CustomStringLiteralD::print
+// CHECK: printFoo({{.*}}, getProperties())
+def CustomStringLiteralD : TestFormat_Op<[{
+  custom<Foo>(prop-dict) attr-dict
+}]>;
+
 //===----------------------------------------------------------------------===//
 // Optional Groups
 //===----------------------------------------------------------------------===//
diff --git a/mlir/tools/mlir-tblgen/OpFormatGen.cpp b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
index 8553920f7713c1..b6405baef28e84 100644
--- a/mlir/tools/mlir-tblgen/OpFormatGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpFormatGen.cpp
@@ -895,6 +895,8 @@ static void genCustomParameterParser(FormatElement *param, MethodBody &body) {
     body << attr->getVar()->name << "Attr";
   } else if (isa<AttrDictDirective>(param)) {
     body << "result.attributes";
+  } else if (isa<PropDictDirective>(param)) {
+    body << "result";
   } else if (auto *operand = dyn_cast<OperandVariable>(param)) {
     StringRef name = operand->getVar()->name;
     ArgumentLengthKind lengthKind = getArgumentLengthKind(operand->getVar());
@@ -1854,6 +1856,9 @@ static void genCustomDirectiveParameterPrinter(FormatElement *element,
   } else if (isa<AttrDictDirective>(element)) {
     body << "getOperation()->getAttrDictionary()";
 
+  } else if (isa<PropDictDirective>(element)) {
+    body << "getProperties()";
+
   } else if (auto *operand = dyn_cast<OperandVariable>(element)) {
     body << op.getGetterName(operand->getVar()->name) << "()";
 
@@ -3138,9 +3143,9 @@ OpFormatParser::parsePropDictDirective(SMLoc loc, Context context) {
 LogicalResult OpFormatParser::verifyCustomDirectiveArguments(
     SMLoc loc, ArrayRef<FormatElement *> arguments) {
   for (FormatElement *argument : arguments) {
-    if (!isa<AttrDictDirective, AttributeVariable, OperandVariable,
-             PropertyVariable, RefDirective, RegionVariable, SuccessorVariable,
-             StringElement, TypeDirective>(argument)) {
+    if (!isa<AttrDictDirective, PropDictDirective, AttributeVariable,
+             OperandVariable, PropertyVariable, RefDirective, RegionVariable,
+             SuccessorVariable, StringElement, TypeDirective>(argument)) {
       // TODO: FormatElement should have location info attached.
       return emitError(loc, "only variables and types may be used as "
                             "parameters to a custom directive");

[joker-eph]

Thanks @fengxie :)

We just need a test: feel free to tweak an op in the test dialect to show-case this.

[joker-eph]

Oh I see you add an ODS test, this is great!

In general I like having TestDialect exercising it because it provides functional end-to-end running example, but it's not a requirement. Let me know and I'll merge this!

[fengxie]

Thanks a lot. Will create end to end example in TestDialect in next step :) Please go ahead to merge.