[mlir] Printing oilist element

[shraiysh]

Summary

This patch attempts to deduce when the oilist element must be printed
based on the optional arguments to it. This especially helps creating
an operation accurately because with the current implementation, the
inferred unit attributes must be manually added to print the clauses
appropriately.

[shraiysh]

Formatting.

[rriddle]

Missing test cases?

[shraiysh]

In D121579#3379996, @rriddle wrote:

Missing test cases?

Where should I add the testcases for this? This is basically an alternative implementation of printing (tests to check printing are already in the repo - mlir/test/IR/traits.mlir - they act as tests for this patch too). The builders will be benefiting from this implementation as they no longer will have to add the inferred attribute. I could not find a test file for testing builder functions or to check the generated C/C++ code for operations, but please let me know if there is one. (AFAICT mlir/test/mlir-tblgen/attr-or-type-format.td checks C++ code generated for attributes and types only).

[shraiysh]

Added handling for unit attributes. Now no extra attributes are added because of the oilist element.

[kiranchandramohan]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1197

Nit: Curlies are probably required here. https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

1897

Nit: Curlies are probably not required here and in the other single line for statements.

2975

Nit: Braces not required.

[shraiysh]

Addressed comments.

[Mogball]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

1197

spell out the auto

1984

drop trivial braces (we should probably have a formatter for this..)

1988

spell out the auto

1995

add braces

2002

add braces

[shraiysh]

Addressed comments, rebase with main.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

oilist( `reduction` `(`
               custom<ReductionVarList>(
                 $reduction_vars, type($reduction_vars), $reductions
               ) `)`
           | `allocate` `(`
               custom<AllocateAndAllocator>(
                 $allocate_vars, type($allocate_vars),
                 $allocators_vars, type($allocators_vars)
               ) `)`
-          | `nowait`
+          | `nowait` $nowait
     ) $region attr-dict
   }];
 
   let hasVerifier = 1;
   let hasRegionVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
let arguments = (ins Optional<I1>:$if_expr,
                        UnitAttr:$nowait);
 
   let regions = (region AnyRegion:$region);
 
   let assemblyFormat = [{
     oilist( `if` `(` $if_expr `)`
           | `device` `(` $device `:` type($device) `)`
           | `thread_limit` `(` $thread_limit `:` type($thread_limit) `)`
-          | `nowait`
+          | `nowait` $nowait
     ) $region attr-dict
   }];
 }
 
 
 //===----------------------------------------------------------------------===//
 // 2.16 master Construct
 //===----------------------------------------------------------------------===//

mlir/test/IR/traits.mlir

func @succeededOilistTrivial() {
   // CHECK: test.oilist_with_keywords_only keyword
   test.oilist_with_keywords_only keyword
   // CHECK: test.oilist_with_keywords_only otherKeyword
   test.oilist_with_keywords_only otherKeyword
   // CHECK: test.oilist_with_keywords_only keyword otherKeyword
   test.oilist_with_keywords_only keyword otherKeyword
   // CHECK: test.oilist_with_keywords_only keyword otherKeyword
   test.oilist_with_keywords_only otherKeyword keyword
+  // CHECK: test.oilist_with_keywords_only thirdKeyword
+  test.oilist_with_keywords_only thirdKeyword
+  // CHECK: test.oilist_with_keywords_only keyword thirdKeyword
+  test.oilist_with_keywords_only keyword thirdKeyword
   return
 }
 
 // -----
 
 // CHECK-LABEL: @succeededOilistSimple
 func @succeededOilistSimple(%arg0 : i32, %arg1 : i32, %arg2 : i32) {
   // CHECK: test.oilist_with_simple_args keyword %{{.*}} : i32

 // -----
 // CHECK-LABEL: succeededOilistCustom
 // CHECK-SAME: (%[[ARG0:.*]]: i32, %[[ARG1:.*]]: i32, %[[ARG2:.*]]: i32)
 func @succeededOilistCustom(%arg0: i32, %arg1: i32, %arg2: i32) {
   // CHECK: test.oilist_custom private(%[[ARG0]], %[[ARG1]] : i32, i32)
   test.oilist_custom private (%arg0, %arg1 : i32, i32)
   // CHECK: test.oilist_custom private(%[[ARG0]], %[[ARG1]] : i32, i32) nowait
   test.oilist_custom private (%arg0, %arg1 : i32, i32) nowait
-  // CHECK: test.oilist_custom private(%arg0, %arg1 : i32, i32) nowait reduction (%arg1)
+  // CHECK: test.oilist_custom private(%arg0, %arg1 : i32, i32) reduction (%arg1) nowait
   test.oilist_custom nowait reduction (%arg1) private (%arg0, %arg1 : i32, i32)
   return
 }
 
 // -----
 
 func @failedHasDominanceScopeOutsideDominanceFreeScope() -> () {
   "test.ssacfg_region"() ({

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

 }
 
 // This is used to test that the fallback for a custom op's parser and printer
 // is the dialect parser and printer hooks.
 def CustomFormatFallbackOp : TEST_Op<"dialect_custom_format_fallback">;
 
 // Ops related to OIList primitive
 def OIListTrivial : TEST_Op<"oilist_with_keywords_only"> {
+  let arguments = (ins UnitAttr:$keyword, UnitAttr:$otherKeyword,
+                       UnitAttr:$diffNameUnitAttrKeyword);
   let assemblyFormat = [{
-    oilist( `keyword`
-          | `otherKeyword`) attr-dict
+    oilist( `keyword` $keyword
+          | `otherKeyword` $otherKeyword
+          | `thirdKeyword` $diffNameUnitAttrKeyword) attr-dict
   }];
 }
 
 def OIListSimple : TEST_Op<"oilist_with_simple_args", [AttrSizedOperandSegments]> {
   let arguments = (ins Optional<AnyType>:$arg0,
                        Optional<AnyType>:$arg1,
                        Optional<AnyType>:$arg2);
   let assemblyFormat = [{

 }
 
 def OIListCustom : TEST_Op<"oilist_custom", [AttrSizedOperandSegments]> {
   let arguments = (ins Variadic<AnyType>:$arg0,
                        Optional<I32>:$optOperand,
                        UnitAttr:$nowait);
   let assemblyFormat = [{
     oilist( `private` `(` $arg0 `:` type($arg0) `)`
-          | `nowait`
           | `reduction` custom<CustomOptionalOperand>($optOperand)
+          | `nowait` $nowait
     ) attr-dict
   }];
 }
 
 def OIListAllowedLiteral : TEST_Op<"oilist_allowed_literal"> {
   let assemblyFormat = [{
     oilist( `foo` | `bar` ) `buzz` attr-dict
   }];

mlir/tools/mlir-tblgen/OpFormatGen.cpp

ArrayRef<std::vector<FormatElement *>> getParsingElements() const {
     return parsingElements;
   }
 
   /// Returns a range to iterate over tuples of parsing and literal elements.
   auto getClauses() const {
     return llvm::zip(getLiteralElements(), getParsingElements());
   }
 
+  /// If the parsing element is a single UnitAttr element, then it returns the
+  /// attribute variable. Otherwise, returns nullptr.
+  AttributeVariable *
+  getUnitAttrParsingElement(ArrayRef<FormatElement *> pelement) {
+    if (pelement.size() == 1) {
+      auto attrElem = dyn_cast<AttributeVariable>(pelement[0]);
+      if (attrElem && attrElem->isUnitAttr())
+        return attrElem;
+    }
+    return nullptr;
+  }
+
 private:
   /// A vector of `LiteralElement` objects. Each element stores the keyword
   /// for one case of oilist element. For example, an oilist element along with
   /// the `literalElements` vector:
   /// ```
   ///  oilist [ `keyword` `=` `(` $arg0 `)` | `otherKeyword` `<` $arg1 `>`]
   ///  literalElements = { `keyword`, `otherKeyword` }
   /// ```

 /// {0}: literal keyword corresponding to a case for oilist
 const char *oilistParserCode = R"(
   if ({0}Clause) {
     return parser.emitError(parser.getNameLoc())
           << "`{0}` clause can appear at most once in the expansion of the "
              "oilist directive";
   }
   {0}Clause = true;
-  result.addAttribute("{0}", UnitAttr::get(parser.getContext()));
 )";
 
 namespace {
 /// The type of length for a given parse argument.
 enum class ArgumentLengthKind {
   /// The argument is a variadic of a variadic, and may contain 0->N range
   /// elements.
   VariadicOfVariadic,
if (anchor && anchor != elements.front() && anchor->isUnitAttr())
       elidedAnchorElement = anchor;
     for (FormatElement *childElement : elements)
       if (childElement != elidedAnchorElement)
         genElementParserStorage(childElement, op, body);
     for (FormatElement *childElement : optional->getElseElements())
       genElementParserStorage(childElement, op, body);
 
   } else if (auto *oilist = dyn_cast<OIListElement>(element)) {
-    for (ArrayRef<FormatElement *> pelement : oilist->getParsingElements())
+    for (ArrayRef<FormatElement *> pelement : oilist->getParsingElements()) {
+      if (!oilist->getUnitAttrParsingElement(pelement))
         for (FormatElement *element : pelement)
           genElementParserStorage(element, op, body);
+    }
 
   } else if (auto *custom = dyn_cast<CustomDirective>(element)) {
     for (FormatElement *paramElement : custom->getArguments())
       genElementParserStorage(paramElement, op, body);
 
   } else if (isa<OperandsDirective>(element)) {
     body << "  ::mlir::SmallVector<::mlir::OpAsmParser::OperandType, 4> "
             "allOperands;\n";
if (auto *optional = dyn_cast<OptionalElement>(element)) {
     // Generate the parsing loop
     body << "  while(true) {\n";
     for (auto clause : oilist->getClauses()) {
       LiteralElement *lelement = std::get<0>(clause);
       ArrayRef<FormatElement *> pelement = std::get<1>(clause);
       body << "if (succeeded(parser.parseOptional";
       genLiteralParser(lelement->getSpelling(), body);
       body << ")) {\n";
-      StringRef attrName = lelement->getSpelling();
-      body << formatv(oilistParserCode, attrName);
-      inferredAttributes.insert(attrName);
+      StringRef lelementName = lelement->getSpelling();
+      body << formatv(oilistParserCode, lelementName);
+      if (auto unitAttrElem = oilist->getUnitAttrParsingElement(pelement)) {
+        body << "  result.addAttribute(\"" << unitAttrElem->getVar()->name
+             << "\", UnitAttr::get(parser.getContext()));\n";
+      } else {
         for (FormatElement *el : pelement)
           genElementParser(el, body, attrTypeCtx);
+      }
       body << "    } else ";
     }
     body << " {\n";
     body << "    break;\n";
     body << "  }\n";
     body << "}\n";
 
     /// Literals.
TypeSwitch<FormatElement *>(anchor)
         genOptionalGroupPrinterAnchor(element->getInputs(), op, body);
       })
       .Case<AttributeVariable>([&](AttributeVariable *attr) {
         body << "  if ((*this)->getAttr(\"" << attr->getVar()->name
              << "\")) {\n";
       });
 }
 
+void collect(FormatElement *element,
+             SmallVectorImpl<VariableElement *> &variables) {
+  TypeSwitch<FormatElement *>(element)
+      .Case([&](VariableElement *var) { variables.emplace_back(var); })
+      .Case([&](CustomDirective *ele) {
+        for (FormatElement *arg : ele->getArguments())
+          collect(arg, variables);
+      })
+      .Case([&](OptionalElement *ele) {
+        for (FormatElement *arg : ele->getThenElements())
+          collect(arg, variables);
+        for (FormatElement *arg : ele->getElseElements())
+          collect(arg, variables);
+      })
+      .Case([&](FunctionalTypeDirective *funcType) {
+        collect(funcType->getInputs(), variables);
+        collect(funcType->getResults(), variables);
+      })
+      .Case([&](OIListElement *oilist) {
+        for (ArrayRef<FormatElement *> arg : oilist->getParsingElements())
+          for (FormatElement *arg_ : arg)
+            collect(arg_, variables);
+      });
+}
+
 void OperationFormat::genElementPrinter(FormatElement *element,
                                         MethodBody &body, Operator &op,
                                         bool &shouldEmitSpace,
                                         bool &lastWasPunctuation) {
   if (LiteralElement *literal = dyn_cast<LiteralElement>(element))
     return genLiteralPrinter(literal->getSpelling(), body, shouldEmitSpace,
                              lastWasPunctuation);
 
void OperationFormat::genElementPrinter(FormatElement *element,
 
   // Emit the OIList
   if (auto *oilist = dyn_cast<OIListElement>(element)) {
     genLiteralPrinter(" ", body, shouldEmitSpace, lastWasPunctuation);
     for (auto clause : oilist->getClauses()) {
       LiteralElement *lelement = std::get<0>(clause);
       ArrayRef<FormatElement *> pelement = std::get<1>(clause);
 
-      body << "  if ((*this)->hasAttrOfType<UnitAttr>(\""
-           << lelement->getSpelling() << "\")) {\n";
+      SmallVector<VariableElement *> vars;
+      for (FormatElement *el : pelement) {
+        collect(el, vars);
+      }
+      body << "  if (false";
+      for (auto *var : vars) {
+        TypeSwitch<FormatElement *>(var)
+            .Case([&](AttributeVariable *attrEle) {
+              body << " || " << op.getGetterName(attrEle->getVar()->name)
+                   << "Attr()";
+            })
+            .Case([&](OperandVariable *ele) {
+              if (ele->getVar()->isVariadic())
+                body << " || " << op.getGetterName(ele->getVar()->name)
+                     << "().size()";
+              else
+                body << " || " << op.getGetterName(ele->getVar()->name) << "()";
+            })
+            .Case([&](ResultVariable *ele) {
+              if (ele->getVar()->isVariadic())
+                body << " || " << op.getGetterName(ele->getVar()->name)
+                     << "().size()";
+              else
+                body << " || " << op.getGetterName(ele->getVar()->name) << "()";
+            })
+            .Case([&](RegionVariable *reg) {
+              body << " || " << op.getGetterName(reg->getVar()->name) << "()";
+            });
+      }
+
+      body << ") {\n";
       genLiteralPrinter(lelement->getSpelling(), body, shouldEmitSpace,
                         lastWasPunctuation);
-      for (FormatElement *element : pelement) {
+      if (oilist->getUnitAttrParsingElement(pelement) == nullptr) {
+        for (FormatElement *element : pelement)
           genElementPrinter(element, body, op, shouldEmitSpace,
                             lastWasPunctuation);
       }
       body << "  }\n";
     }
     return;
   }
 
   // Emit the attribute dictionary.
   if (auto *attrDict = dyn_cast<AttrDictDirective>(element)) {
OpFormatParser::parseOIListDirective(SMLoc loc, Context context) {
   } while (true);
 
   return create<OIListElement>(std::move(literalElements),
                                std::move(parsingElements));
 }
 
 LogicalResult OpFormatParser::verifyOIListParsingElement(FormatElement *element,
                                                          SMLoc loc) {
-  return TypeSwitch<FormatElement *, LogicalResult>(element)
+  SmallVector<VariableElement *> vars;
+  collect(element, vars);
+  for (VariableElement *elem : vars) {
+    LogicalResult res =
+        TypeSwitch<FormatElement *, LogicalResult>(elem)
             // Only optional attributes can be within an oilist parsing group.
             .Case([&](AttributeVariable *attrEle) {
-        if (!attrEle->getVar()->attr.isOptional())
-          return emitError(loc, "only optional attributes can be used to "
-                                "in an oilist parsing group");
+              if (!attrEle->getVar()->attr.isOptional() &&
+                  !attrEle->getVar()->attr.hasDefaultValue())
+                return emitError(loc, "only optional attributes can be used in "
+                                      "an oilist parsing group");
               return success();
             })
-      // Only optional-like(i.e. variadic) operands can be within an oilist
-      // parsing group.
+            // Only optional-like(i.e. variadic) operands can be within an
+            // oilist parsing group.
             .Case([&](OperandVariable *ele) {
               if (!ele->getVar()->isVariableLength())
                 return emitError(loc, "only variable length operands can be "
                                       "used within an oilist parsing group");
               return success();
             })
             // Only optional-like(i.e. variadic) results can be within an oilist
             // parsing group.
             .Case([&](ResultVariable *ele) {
               if (!ele->getVar()->isVariableLength())
                 return emitError(loc, "only variable length results can be "
                                       "used within an oilist parsing group");
               return success();
             })
-      .Case([&](RegionVariable *) {
-        // TODO: When ODS has proper support for marking "optional" regions, add
-        // a check here.
-        return success();
-      })
-      .Case([&](TypeDirective *ele) {
-        return verifyOIListParsingElement(ele->getArg(), loc);
-      })
-      .Case([&](FunctionalTypeDirective *ele) {
-        if (failed(verifyOIListParsingElement(ele->getInputs(), loc)))
-          return failure();
-        return verifyOIListParsingElement(ele->getResults(), loc);
-      })
-      // Literals, whitespace, and custom directives may be used.
-      .Case<LiteralElement, WhitespaceElement, CustomDirective,
-            FunctionalTypeDirective, OptionalElement>(
-          [&](FormatElement *) { return success(); })
+            .Case([&](RegionVariable *) { return success(); })
             .Default([&](FormatElement *) {
-        return emitError(loc, "only literals, types, and variables can be "
+              return emitError(loc,
+                               "only literals, types, and variables can be "
                                "used within an oilist group");
             });
+    if (failed(res))
+      return failure();
+  }
+  return success();
 }
 
 FailureOr<FormatElement *> OpFormatParser::parseTypeDirective(SMLoc loc,
                                                               Context context) {
   if (context == TypeDirectiveContext)
     return emitError(loc, "'type' cannot be used as a child of another `type`");
 
   bool isRefChild = context == RefDirectiveContext;