[mlir][ods] Allow custom directives in optional groups

[Mogball]

Summary

Attributes and types only (so far). Since struct and params are
allowed, it makes sense to allow custom directives as long as their
arguments contain at least one bound argument.

[rriddle]

Can you update the docs with the semantics around custom inside of an optional group?

Also, can you add tests for the error cases?

[Mogball]

add error tests and update docs

[rriddle]

mlir/docs/AttributesAndTypes.md

663

664

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

1009–1010

Please don't use if initializers for nullable things. In this case it also doesn't provide any real benefits.

1027–1029

Is this possible? (I think I've seen it work before)

[Mogball]

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

1027–1029

I believe it should be. Let me try

[Mogball]

comments

mlir/docs/AttributesAndTypes.md

 field of the parameter is used, but if one is not specified, the equality
 operator is used.
 
 When using `OptionalParameter`, the default value is set to the C++
 default-constructed value for the C++ storage type. For example, `Optional<int>`
 will be set to `llvm::None` and `Attribute` will be set to `nullptr`. The
 presence of these parameters is tested by comparing them to their "null" values.
 
-Only optional parameters or directives that only capture optional parameters can
-be used in optional groups. An optional group is a set of elements optionally
-printed based on the presence of an anchor. The group in which the anchor is
-placed is printed if it is present, otherwise the other one is printed. Suppose
-parameter `a` is an `IntegerAttr`.
+An optional group is a set of elements optionally printed based on the presence
+of an anchor. Only optional parameters or directives that only capture optional
+parameters can be used in optional groups. The group in which the anchor is
+placed is printed if it is present, otherwise the other one is printed. If a
+directive that captures more than one optional parameter is used as the anchor,
+the optional group is printed if any of the captured parameters is present. For
+example, a `custom` directive may only be used as an optional group anchor if it
+captures at least one optional parameter.
+
+Suppose parameter `a` is an `IntegerAttr`.
 
 ```
 ( `(` $a^ `)` ) : (`x`)?
 ```
 
 In the above assembly format, if `a` is present (non-null), then it will be
 printed as `(5 : i32)`. If it is not present, it will be `x`. Directives that
 are used inside optional groups are allowed only if all captured parameters are

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

def TestSimpleEnum : I32EnumAttr<"SimpleEnum", "", [
   let cppNamespace = "::test";
 }
 def TestSimpleEnumAttr : EnumAttr<Test_Dialect, TestSimpleEnum, "simple_enum"> {
   let assemblyFormat = "`` $value";
 }
 def TestArrayOfEnums : ArrayOfAttr<Test_Dialect, "ArrayOfEnums",
     "array_of_enums", "SimpleEnumAttr">;
 
+// Test custom directive as optional group anchor.
+def TestCustomAnchor : Test_Attr<"TestCustomAnchor"> {
+  let parameters = (ins "int":$a, OptionalParameter<"mlir::Optional<int>">:$b);
+  let mnemonic = "custom_anchor";
+  let assemblyFormat = "`<` $a (`>`) : (`,` ` ` custom<TrueFalse>($b)^ `>`)?";
+}
+
 #endif // TEST_ATTRDEFS

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

 
 ArrayRef<uint64_t> TestExtern1DI64ElementsAttr::getElements() const {
   if (auto *blob = getHandle().getBlob())
     return blob->getDataAs<uint64_t>();
   return llvm::None;
 }
 
 //===----------------------------------------------------------------------===//
+// TestCustomAnchorAttr
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseTrueFalse(AsmParser &p,
+                                  FailureOr<Optional<int>> &result) {
+  bool b;
+  if (p.parseInteger(b))
+    return failure();
+  result = Optional<int>(b);
+  return success();
+}
+
+static void printTrueFalse(AsmPrinter &p, Optional<int> result) {
+  p << (*result ? "true" : "false");
+}
+
+//===----------------------------------------------------------------------===//
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//
 
 #include "TestAttrInterfaces.cpp.inc"
 
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.cpp.inc"
 

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

def InvalidTypeQ : InvalidType<"InvalidTypeQ", "invalid_q"> {
   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))";
 }
+
+def InvalidTypeS : InvalidType<"InvalidTypeS", "invalid_s"> {
+  let parameters = (ins OptionalParameter<"int">:$a, "int":$b);
+  // CHECK: `custom` is only allowed in an optional group if all captured parameters are optional
+  let assemblyFormat = "(`(` custom<Foo>($a, $b)^ `)`)?";
+}
+
+def InvalidTypeT : InvalidType<"InvalidTypeT", "invalid_t"> {
+  let parameters = (ins OptionalParameter<"int">:$a);
+  // CHECK: `custom` directive with no bound parameters cannot be used as optional group anchor
+  let assemblyFormat = "$a (`(` custom<Foo>(ref($a))^ `)`)?";
+}

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

attributes {
   attr1 = #test.attr_with_format<5 : two = "a_string", four = [4, 5, 6, 7, 8] : 8 : i8, [10 : i16]>,
   // CHECK: #test<attr_ugly begin 5 : index end>
   attr2 = #test<attr_ugly begin 5 : index end>,
   // CHECK: #test.attr_params<42, 24>
   attr3 = #test.attr_params<42, 24>,
   // CHECK: #test.attr_with_type<i32, vector<4xi32>>
   attr4 = #test.attr_with_type<i32, vector<4xi32>>,
   // CHECK: #test.attr_self_type_format<5> : i32
-  attr5 = #test.attr_self_type_format<5> : i32
+  attr5 = #test.attr_self_type_format<5> : i32,
+  // CHECK: #test.custom_anchor<5>
+  attr6 = #test.custom_anchor<5>,
+  // CHECK: #test.custom_anchor<5, true>
+  attr7 = #test.custom_anchor<5, true>
 }
 
 // CHECK-LABEL: @test_roundtrip_default_parsers_struct
 // CHECK: !test.no_parser<255, [1, 2, 3, 4, 5], "foobar", 4>
 // CHECK: !test.struct_capture_all<v0 = 0, v1 = 1, v2 = 2, v3 = 3>
 // CHECK: !test.optional_param<, 6>
 // CHECK: !test.optional_param<5, 6>
 // CHECK: !test.optional_params<"a">

mlir/tools/mlir-tblgen/AttrOrTypeFormatGen.cpp

 
 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), el->isInverted());
   } else if (auto *params = dyn_cast<ParamsDirective>(anchor)) {
     guardOnAny(ctx, os, params->getParams(), el->isInverted());
-  } else {
-    auto *strct = cast<StructDirective>(anchor);
+  } else if (auto *strct = dyn_cast<StructDirective>(anchor)) {
     guardOnAny(ctx, os, strct->getParams(), el->isInverted());
+  } else {
+    auto *custom = cast<CustomDirective>(anchor);
+    guardOnAny(ctx, os,
+               llvm::make_filter_range(
+                   llvm::map_range(custom->getArguments(),
+                                   [](FormatElement *el) {
+                                     return dyn_cast<ParameterElement>(el);
+                                   }),
+                   [](ParameterElement *param) { return !!param; }),
+               el->isInverted());
   }
   // Generate the printer for the contained elements.
   {
     llvm::SaveAndRestore<bool> shouldEmitSpaceFlag(shouldEmitSpace);
     llvm::SaveAndRestore<bool> lastWasPunctuationFlag(lastWasPunctuation);
     for (FormatElement *element : el->getThenElements())
       genElementPrinter(element, ctx, os);
   }
if (auto *param = dyn_cast<ParameterElement>(el)) {
         return emitError(loc, "`params` directive allowed in optional group "
                               "only if all parameters are optional");
       }
     } else if (auto *strct = dyn_cast<StructDirective>(el)) {
       if (llvm::any_of(strct->getParams(), paramNotOptional)) {
         return emitError(loc, "`struct` is only allowed in an optional group "
                               "if all captured parameters are optional");
       }
+    } else if (auto *custom = dyn_cast<CustomDirective>(el)) {
+      for (FormatElement *el : custom->getArguments()) {
+        // If the custom argument is a variable, then it must be optional.
+        if (auto param = dyn_cast<ParameterElement>(el))
+          if (!param->isOptional())
+            return emitError(loc,
+                             "`custom` is only allowed in an optional group if "
+                             "all captured parameters are optional");
+      }
     }
   }
   // The anchor must be a parameter or one of the aforementioned directives.
-  if (anchor &&
-      !isa<ParameterElement, ParamsDirective, StructDirective>(anchor)) {
-    return emitError(loc,
-                     "optional group anchor must be a parameter or directive");
+  if (anchor) {
+    if (!isa<ParameterElement, ParamsDirective, StructDirective,
+             CustomDirective>(anchor)) {
+      return emitError(
+          loc, "optional group anchor must be a parameter or directive");
+    }
+    // If the anchor is a custom directive, make sure at least one of its
+    // arguments is a bound parameter.
+    if (auto custom = dyn_cast<CustomDirective>(anchor)) {
+      auto bound = llvm::find_if(custom->getArguments(), [](FormatElement *el) {
+        return isa<ParameterElement>(el);
+      });
+      if (bound == custom->getArguments().end())
+        return emitError(loc, "`custom` directive with no bound parameters "
+                              "cannot be used as optional group anchor");
+    }
   }
   return success();
 }
 
 FailureOr<DefFormat> DefFormatParser::parse() {
   FailureOr<std::vector<FormatElement *>> elements = FormatParser::parse();
   if (failed(elements))
     return failure();