[MLIR] Make `parseDimensionListRanked` configurable wrt parsing a trailing `x`

[frgossen]

Summary

This is useful for parsing types similar to tensor or memref but w/o an element type, e.g. a tile<1x?x3>.

[rriddle]

Could you add a test somewhere?

mlir/lib/Parser/TypeParser.cpp

528–529

What do eps and epx mean?

539

555–557

559

Can we early return success here?

566–569

[frgossen]

Address comments

[frgossen]

Thanks!

mlir/lib/Parser/TypeParser.cpp

528–529

eps is the empty word epsilon, epx was a typo. Thanks for noticing.

mlir/include/mlir/IR/OpImplementation.h

public:
   /// success.  This is a helper designed to allow parse methods to be simple
   /// and chain through || operators.
   ParseResult addTypesToList(ArrayRef<Type> types,
                              SmallVectorImpl<Type> &result) {
     result.append(types.begin(), types.end());
     return success();
   }
 
-  /// Parse a 'x' separated dimension list. This populates the dimension list,
-  /// using -1 for the `?` dimensions if `allowDynamic` is set and errors out on
-  /// `?` otherwise.
+  /// Parse a dimension list of a tensor or memref type.  This populates the
+  /// dimension list, using -1 for the `?` dimensions if `allowDynamic` is set
+  /// and errors out on `?` otherwise. Parsing the trailing `x` is configurable.
   ///
-  ///   dimension-list ::= (dimension `x`)*
-  ///   dimension ::= `?` | integer
+  ///   dimension-list ::= eps | dimension (`x` dimension)*
+  ///   dimension-list-with-trailing-x ::= (dimension `x`)*
+  ///   dimension ::= `?` | decimal-literal
   ///
   /// When `allowDynamic` is not set, this is used to parse:
   ///
-  ///   static-dimension-list ::= (integer `x`)*
+  ///   static-dimension-list ::= eps | decimal-literal (`x` decimal-literal)*
+  ///   static-dimension-list-with-trailing-x ::= (dimension `x`)*
   virtual ParseResult parseDimensionList(SmallVectorImpl<int64_t> &dimensions,
-                                         bool allowDynamic = true) = 0;
+                                         bool allowDynamic = true,
+                                         bool withTrailingX = true) = 0;
 
   /// Parse an 'x' token in a dimension list, handling the case where the x is
   /// juxtaposed with an element type, as in "xf32", leaving the "f32" as the
   /// next token.
   virtual ParseResult parseXInDimensionList() = 0;
 
 private:
   AsmParser(const AsmParser &) = delete;

mlir/lib/Parser/AsmParserImpl.h

public:
   ParseResult
   parseOptionalColonTypeList(SmallVectorImpl<Type> &result) override {
     if (!parser.consumeIf(Token::colon))
       return success();
     return parser.parseTypeListNoParens(result);
   }
 
   ParseResult parseDimensionList(SmallVectorImpl<int64_t> &dimensions,
-                                 bool allowDynamic) override {
-    return parser.parseDimensionListRanked(dimensions, allowDynamic);
+                                 bool allowDynamic,
+                                 bool withTrailingX) override {
+    return parser.parseDimensionListRanked(dimensions, allowDynamic,
+                                           withTrailingX);
   }
 
   ParseResult parseXInDimensionList() override {
     return parser.parseXInDimensionList();
   }
 
 protected:
   /// The source location of the dialect symbol.

mlir/lib/Parser/Parser.h

public:
   /// Parse a tuple type.
   Type parseTupleType();
 
   /// Parse a vector type.
   VectorType parseVectorType();
   ParseResult parseVectorDimensionList(SmallVectorImpl<int64_t> &dimensions,
                                        unsigned &numScalableDims);
   ParseResult parseDimensionListRanked(SmallVectorImpl<int64_t> &dimensions,
-                                       bool allowDynamic = true);
+                                       bool allowDynamic = true,
+                                       bool withTrailingX = true);
   ParseResult parseIntegerInDimensionList(int64_t &value);
   ParseResult parseXInDimensionList();
 
   /// Parse strided layout specification.
   ParseResult parseStridedLayout(int64_t &offset,
                                  SmallVectorImpl<int64_t> &strides);
 
   // Parse a brace-delimiter list of comma-separated integers with `?` as an

mlir/lib/Parser/TypeParser.cpp

return emitWrongTokenError(
         "missing ']' closing set of scalable dimensions");
   }
 
   return success();
 }
 
 /// Parse a dimension list of a tensor or memref type.  This populates the
 /// dimension list, using -1 for the `?` dimensions if `allowDynamic` is set and
-/// errors out on `?` otherwise.
+/// errors out on `?` otherwise. Parsing the trailing `x` is configurable.
 ///
-///   dimension-list-ranked ::= (dimension `x`)*
+///   dimension-list ::= eps | dimension (`x` dimension)*
+///   dimension-list-with-trailing-x ::= (dimension `x`)*
 ///   dimension ::= `?` | decimal-literal
 ///
 /// When `allowDynamic` is not set, this is used to parse:
 ///
-///   static-dimension-list ::= (decimal-literal `x`)*
+///   static-dimension-list ::= eps | decimal-literal (`x` decimal-literal)*
+///   static-dimension-list-with-trailing-x ::= (dimension `x`)*
 ParseResult
 Parser::parseDimensionListRanked(SmallVectorImpl<int64_t> &dimensions,
-                                 bool allowDynamic) {
-  while (getToken().isAny(Token::integer, Token::question)) {
+                                 bool allowDynamic, bool withTrailingX) {
+  auto parseDim = [&]() -> LogicalResult {
     auto loc = getToken().getLoc();
     if (consumeIf(Token::question)) {
       if (!allowDynamic)
         return emitError(loc, "expected static shape");
       dimensions.push_back(-1);
     } else {
       int64_t value;
-      if (parseIntegerInDimensionList(value))
+      if (failed(parseIntegerInDimensionList(value)))
         return failure();
       dimensions.push_back(value);
     }
-    // Make sure we have an 'x' or something like 'xbf32'.
-    if (parseXInDimensionList())
+    return success();
+  };
+
+  if (withTrailingX) {
+    while (getToken().isAny(Token::integer, Token::question)) {
+      if (failed(parseDim()) || failed(parseXInDimensionList()))
         return failure();
     }
+    return success();
+  }
 
+  if (getToken().isAny(Token::integer, Token::question)) {
+    if (failed(parseDim()))
+      return failure();
+    while (getToken().is(Token::bare_identifier) &&
+           getTokenSpelling()[0] == 'x') {
+      if (failed(parseXInDimensionList()) || failed(parseDim()))
+        return failure();
+    }
+  }
   return success();
 }
 
 ParseResult Parser::parseIntegerInDimensionList(int64_t &value) {
   // Hexadecimal integer literals (starting with `0x`) are not allowed in
   // aggregate type declarations.  Therefore, `0xf32` should be processed as
   // a sequence of separate elements `0`, `x`, `f32`.
   if (getTokenSpelling().size() > 1 && getTokenSpelling()[1] == 'x') {