[mlir] Allow omitting spaces in assemblyFormat with a `` literal.

[csigg]

Summary

I would like to use this for D90589 to switch std.alloc to assemblyFormat.
Hopefully it will be useful in other places as well.

[rriddle]

mlir/tools/mlir-tblgen/OpFormatGen.cpp

310

Please use proper accessors functions instead of direct access.

[csigg]

Change to proper accessor.

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

 }
 
 //===----------------------------------------------------------------------===//
 // Test Op Asm Format
 //===----------------------------------------------------------------------===//
 
 def FormatLiteralOp : TEST_Op<"format_literal_op"> {
   let assemblyFormat = [{
-    `keyword_$.` `->` `:` `,` `=` `<` `>` `(` `)` `[` `]` ` ` `(` ` ` `)` ` ` attr-dict
+    `keyword_$.` `->` `:` `,` `=` `<` `>` `(` `)` `[` `]` `` `(` ` ` `)` attr-dict
   }];
 }
 
 // Test that we elide attributes that are within the syntax.
 def FormatAttrOp : TEST_Op<"format_attr_op"> {
   let arguments = (ins I64Attr:$attr);
   let assemblyFormat = "$attr attr-dict";
 }

mlir/test/mlir-tblgen/op-format-spec.td

 }]>;
 // CHECK: error: unexpected end of file in literal
 // CHECK: error: expected directive, literal, variable, or optional group
 def LiteralInvalidB : TestFormat_Op<"literal_invalid_b", [{
   `
 }]>;
 // CHECK-NOT: error
 def LiteralValid : TestFormat_Op<"literal_valid", [{
-  `_` `:` `,` `=` `<` `>` `(` `)` `[` `]` ` ` `->` `abc$._`
+  `_` `:` `,` `=` `<` `>` `(` `)` `[` `]` ` ` `` `->` `abc$._`
   attr-dict
 }]>;
 
 //===----------------------------------------------------------------------===//
 // Optional Groups
 //===----------------------------------------------------------------------===//
 
 // CHECK: error: optional groups can only be used as top-level elements
 def OptionalInvalidA : TestFormat_Op<"optional_invalid_a", [{
   type(($attr^)?) attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
 // CHECK: error: expected directive, literal, variable, or optional group
 def OptionalInvalidB : TestFormat_Op<"optional_invalid_b", [{
   () attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
 // CHECK: error: optional group specified no anchor element
 def OptionalInvalidC : TestFormat_Op<"optional_invalid_c", [{
   ($attr)? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
-// CHECK: error: first element of an operand group must be an attribute, literal, operand, or region
+// CHECK: error: first parsable element of an operand group must be an attribute, literal, operand, or region
 def OptionalInvalidD : TestFormat_Op<"optional_invalid_d", [{
   (type($operand) $operand^)? attr-dict
 }]>, Arguments<(ins Optional<I64>:$operand)>;
 // CHECK: error: type directive can only refer to variables within the optional group
 def OptionalInvalidE : TestFormat_Op<"optional_invalid_e", [{
   (`,` $attr^ type(operands))? attr-dict
 }]>, Arguments<(ins OptionalAttr<I64Attr>:$attr)>;
 // CHECK: error: only one element can be marked as the anchor of an optional group

 def OptionalInvalidL : TestFormat_Op<"optional_invalid_l", [{
   (custom<MyDirective>($arg)^)?
 }]>, Arguments<(ins I64:$arg)>;
 // CHECK: error: only variables can be used to anchor an optional group
 def OptionalInvalidM : TestFormat_Op<"optional_invalid_m", [{
   (` `^)?
 }]>, Arguments<(ins)>;
 
+// CHECK-NOT: error
+def OptionalValidA : TestFormat_Op<"optional_valid_a", [{
+  (` ` `` $arg^)?
+}]>;
+
 //===----------------------------------------------------------------------===//
 // Variables
 //===----------------------------------------------------------------------===//
 
 // CHECK: error: expected variable to refer to an argument, region, result, or successor
 def VariableInvalidA : TestFormat_Op<"variable_invalid_a", [{
   $unknown_arg attr-dict
 }]>;

mlir/test/mlir-tblgen/op-format.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s | mlir-opt -allow-unregistered-dialect -verify-diagnostics | FileCheck %s
 
 // CHECK: %[[I64:.*]] =
 %i64 = "foo.op"() : () -> (i64)
 // CHECK: %[[MEMREF:.*]] =
 %memref = "foo.op"() : () -> (memref<1xf64>)
 
-// CHECK: test.format_literal_op keyword_$. -> :, = <> () [] ( ) {foo.some_attr}
-test.format_literal_op keyword_$. -> :, = <> () [] ( ) {foo.some_attr}
+// CHECK: test.format_literal_op keyword_$. -> :, = <> () []( ) {foo.some_attr}
+test.format_literal_op keyword_$. -> :, = <> () []( ) {foo.some_attr}
 
 // CHECK: test.format_attr_op 10
 // CHECK-NOT: {attr
 test.format_attr_op 10
 
 // CHECK: test.format_opt_attr_op_a(10)
 // CHECK-NOT: {opt_attr
 test.format_opt_attr_op_a(10)

mlir/tools/mlir-tblgen/OpFormatGen.cpp

enum class Kind {
     ResultsDirective,
     SuccessorsDirective,
     TypeDirective,
     TypeRefDirective,
 
     /// This element is a literal.
     Literal,
 
-    /// This element is printed as a space. It is ignored by the parser.
+    /// This element prints or omits a space. It is ignored by the parser.
     Space,
 
     /// This element is an variable value.
     AttributeVariable,
     OperandVariable,
     RegionVariable,
     ResultVariable,
     SuccessorVariable,
bool LiteralElement::isValidLiteral(StringRef value) {
   });
 }
 
 //===----------------------------------------------------------------------===//
 // SpaceElement
 
 namespace {
 /// This class represents an instance of a space element. It's a literal that
-/// is only printed, but ignored by the parser.
+/// prints or omits printing a space. It is ignored by the parser.
 class SpaceElement : public Element {
 public:
-  SpaceElement() : Element{Kind::Space} {}
+  SpaceElement(bool value) : Element{Kind::Space}, value(value) {}
   static bool classof(const Element *element) {
     return element->getKind() == Kind::Space;
   }
+  /// If true, the elements prints as a space. If false, it omits printing a
+  /// space between two elements.
+  bool value;
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // OptionalElement
 
 namespace {
 /// This class represents a group of elements that are optionally emitted based
 /// upon an optional variable of the operation.
 class OptionalElement : public Element {
 public:
   OptionalElement(std::vector<std::unique_ptr<Element>> &&elements,
-                  unsigned anchor)
-      : Element{Kind::Optional}, elements(std::move(elements)), anchor(anchor) {
-  }
+                  unsigned anchor, unsigned parseStart)
+      : Element{Kind::Optional}, elements(std::move(elements)), anchor(anchor),
+        parseStart(parseStart) {}
   static bool classof(const Element *element) {
     return element->getKind() == Kind::Optional;
   }
 
   /// Return the nested elements of this grouping.
   auto getElements() const { return llvm::make_pointee_range(elements); }
 
   /// Return the anchor of this optional group.
   Element *getAnchor() const { return elements[anchor].get(); }
 
+  /// Return the index of the first element that needs to be parsed.
+  unsigned getParseStart() const { return parseStart; }
+
 private:
   /// The child elements of this optional.
   std::vector<std::unique_ptr<Element>> elements;
   /// The index of the element that acts as the anchor for the optional group.
   unsigned anchor;
+  /// The index of the first element that is parsed (is not a SpaceElement).
+  unsigned parseStart;
 };
 } // end anonymous namespace
 
 //===----------------------------------------------------------------------===//
 // OperationFormat
 //===----------------------------------------------------------------------===//
 
 namespace {
void OperationFormat::genParser(Operator &op, OpClass &opClass) {
 
   body << "  return ::mlir::success();\n";
 }
 
 void OperationFormat::genElementParser(Element *element, OpMethodBody &body,
                                        FmtContext &attrTypeCtx) {
   /// Optional Group.
   if (auto *optional = dyn_cast<OptionalElement>(element)) {
-    auto elements = optional->getElements();
+    auto elements =
+        llvm::drop_begin(optional->getElements(), optional->getParseStart());
 
     // Generate a special optional parser for the first element to gate the
     // parsing of the rest of the elements.
     Element *firstElement = &*elements.begin();
     if (auto *attrVar = dyn_cast<AttributeVariable>(firstElement)) {
       genElementParser(attrVar, body, attrTypeCtx);
       body << "  if (" << attrVar->getVar()->name << "Attr) {\n";
     } else if (auto *literal = dyn_cast<LiteralElement>(firstElement)) {
static void genLiteralPrinter(StringRef value, OpMethodBody &body,
   // Insert a space after certain literals.
   shouldEmitSpace =
       value.size() != 1 || !StringRef("<({[").contains(value.front());
   lastWasPunctuation = !(value.front() == '_' || isalpha(value.front()));
 }
 
 /// Generate the printer for a space. `shouldEmitSpace` and `lastWasPunctuation`
 /// are set to false.
-static void genSpacePrinter(OpMethodBody &body, bool &shouldEmitSpace,
-                            bool &lastWasPunctuation) {
+static void genSpacePrinter(bool value, OpMethodBody &body,
+                            bool &shouldEmitSpace, bool &lastWasPunctuation) {
+  if (value) {
     body << "  p << ' ';\n";
-  shouldEmitSpace = false;
     lastWasPunctuation = false;
   }
+  shouldEmitSpace = false;
+}
 
 /// Generate the printer for a custom directive.
 static void genCustomDirectivePrinter(CustomDirective *customDir,
                                       OpMethodBody &body) {
   body << "  print" << customDir->getName() << "(p, *this";
   for (Element &param : customDir->getArguments()) {
     body << ", ";
     if (auto *attr = dyn_cast<AttributeVariable>(&param)) {

 
 void OperationFormat::genElementPrinter(Element *element, OpMethodBody &body,
                                         Operator &op, bool &shouldEmitSpace,
                                         bool &lastWasPunctuation) {
   if (LiteralElement *literal = dyn_cast<LiteralElement>(element))
     return genLiteralPrinter(literal->getLiteral(), body, shouldEmitSpace,
                              lastWasPunctuation);
 
-  if (isa<SpaceElement>(element))
-    return genSpacePrinter(body, shouldEmitSpace, lastWasPunctuation);
+  if (SpaceElement *space = dyn_cast<SpaceElement>(element))
+    return genSpacePrinter(space->value, body, shouldEmitSpace,
+                           lastWasPunctuation);
 
   // Emit an optional group.
   if (OptionalElement *optional = dyn_cast<OptionalElement>(element)) {
     // Emit the check for the presence of the anchor element.
     Element *anchor = optional->getAnchor();
     if (auto *operand = dyn_cast<OperandVariable>(anchor)) {
       const NamedTypeConstraint *var = operand->getVar();
       if (var->isOptional())

 }
 
 LogicalResult FormatParser::parseLiteral(std::unique_ptr<Element> &element) {
   Token literalTok = curToken;
   consumeToken();
 
   StringRef value = literalTok.getSpelling().drop_front().drop_back();
 
-  // The parsed literal is a space.
-  if (value.size() == 1 && value.front() == ' ') {
-    element = std::make_unique<SpaceElement>();
+  // The parsed literal is a space element (`` or ` `).
+  if (value.empty() || (value.size() == 1 && value.front() == ' ')) {
+    element = std::make_unique<SpaceElement>(!value.empty());
     return ::mlir::success();
   }
 
   // Check that the parsed literal is valid.
   if (!LiteralElement::isValidLiteral(value))
     return emitError(literalTok.getLoc(), "expected valid literal");
 
   element = std::make_unique<LiteralElement>(value);
LogicalResult FormatParser::parseOptional(std::unique_ptr<Element> &element,
   consumeToken();
   if (failed(parseToken(Token::question, "expected '?' after optional group")))
     return ::mlir::failure();
 
   // The optional group is required to have an anchor.
   if (!anchorIdx)
     return emitError(curLoc, "optional group specified no anchor element");
 
-  // The first element of the group must be one that can be parsed/printed in an
+  // The first parsable element of the group must be able to be parsed in an
   // optional fashion.
-  Element *firstElement = &*elements.front();
+  auto parseBegin = llvm::find_if_not(
+      elements, [](auto &element) { return isa<SpaceElement>(element.get()); });
+  Element *firstElement = parseBegin->get();
   if (!isa<AttributeVariable>(firstElement) &&
       !isa<LiteralElement>(firstElement) &&
       !isa<OperandVariable>(firstElement) && !isa<RegionVariable>(firstElement))
-    return emitError(curLoc, "first element of an operand group must be an "
-                             "attribute, literal, operand, or region");
+    return emitError(curLoc,
+                     "first parsable element of an operand group must be "
+                     "an attribute, literal, operand, or region");
 
   // After parsing all of the elements, ensure that all type directives refer
   // only to elements within the group.
   auto checkTypeOperand = [&](Element *typeEle) {
     auto *opVar = dyn_cast<OperandVariable>(typeEle);
     const NamedTypeConstraint *var = opVar ? opVar->getVar() : nullptr;
     if (!seenVariables.count(var))
       return emitError(curLoc, "type directive can only refer to variables "
for (auto &ele : elements) {
     } else if (auto *typeEle = dyn_cast<FunctionalTypeDirective>(ele.get())) {
       if (failed(checkTypeOperand(typeEle->getInputs())) ||
           failed(checkTypeOperand(typeEle->getResults())))
         return ::mlir::failure();
     }
   }
 
   optionalVariables.insert(seenVariables.begin(), seenVariables.end());
-  element = std::make_unique<OptionalElement>(std::move(elements), *anchorIdx);
+  auto parseStart = parseBegin - elements.begin();
+  element = std::make_unique<OptionalElement>(std::move(elements), *anchorIdx,
+                                              parseStart);
   return ::mlir::success();
 }
 
 LogicalResult FormatParser::parseOptionalChildElement(
     std::vector<std::unique_ptr<Element>> &childElements,
     SmallPtrSetImpl<const NamedTypeConstraint *> &seenVariables,
     Optional<unsigned> &anchorIdx) {
   llvm::SMLoc childLoc = curToken.getLoc();