[mlir] Add concept of alias blocks

This PR is part of https://discourse.llvm.org/t/rfc-supporting-aliases-in-cyclic-types-and-attributes/73236

It implements the concept of "alias blocks", a block of alias definitions which may alias any other alias definitions within the block, regardless of definition order. This is purely a convenience for immutable attributes and types, but is a requirement for supporting aliasing definitions in cyclic mutable attributes and types.

The implementation works by first parsing an alias-block, which is simply subsequent alias definitions, in a syntax-only mode. This syntax-only mode only checks for syntactic validity of the parsed attribute or type but does not verify any parsed data. This allows us to essentially skip over alias definitions for the purpose of first collecting them and associating every alias definition with its source region.

In a second pass, we can start parsing the attributes and types while at the same time attempting to resolve any unknown alias references with our list of yet-to-be-parsed attributes and types, parsing them on demand if required.

A later PR will hook up this mechanism to the `tryStartCyclicParse` method added in b121c26 to early register cyclic attributes and types, breaking the parsing cycles.

Author: zero9178

mlir/docs/LangRef.md

@@ -183,12 +183,14 @@ starting with a `//` and going until the end of the line.
 
 ```
 // Top level production
-toplevel := (operation | attribute-alias-def | type-alias-def)*
+toplevel := (operation | alias-block-def)*
+alias-block-def := (attribute-alias-def | type-alias-def)*
 ```
 
 The production `toplevel` is the top level production that is parsed by any parsing
-consuming the MLIR syntax. [Operations](#operations),
-[Attribute aliases](#attribute-value-aliases), and [Type aliases](#type-aliases)
+consuming the MLIR syntax. [Operations](#operations) and 
+[Alias Blocks](#alias-block-definitions) consisting of 
+[Attribute aliases](#attribute-value-aliases) and [Type aliases](#type-aliases)
 can be declared on the toplevel.
 
 ### Identifiers and keywords
@@ -880,3 +882,26 @@ version using readAttribute and readType methods.
 There is no restriction on what kind of information a dialect is allowed to
 encode to model its versioning. Currently, versioning is supported only for
 bytecode formats.
+
+## Alias Block Definitions
+
+An alias block is a list of subsequent attribute or type alias definitions that
+are conceptually parsed as one unit.
+This allows any alias definition within the block to reference any other alias 
+definition within the block, regardless if defined lexically later or earlier in
+the block.
+
+```mlir
+// Alias block consisting of #array, !integer_type and #integer_attr.
+#array = [#integer_attr, !integer_type]
+!integer_type = i32
+#integer_attr = 8 : !integer_type
+
+// Illegal. !other_type is not part of this alias block and defined later 
+// in the file.
+!tuple = tuple<i32, !other_type>
+
+func.func @foo() { ... }
+
+!other_type = f32
+```

mlir/lib/AsmParser/AttributeParser.cpp

@@ -145,6 +145,10 @@ Attribute Parser::parseAttribute(Type type) {
         parseLocationInstance(locAttr) ||
         parseToken(Token::r_paren, "expected ')' in inline location"))
       return Attribute();
+
+    if (syntaxOnly())
+      return state.syntaxOnlyAttr;
+
     return locAttr;
   }
 
@@ -430,6 +434,9 @@ Attribute Parser::parseDecOrHexAttr(Type type, bool isNegative) {
     return FloatAttr::get(floatType, *result);
   }
 
+  if (syntaxOnly())
+    return state.syntaxOnlyAttr;
+
   if (!isa<IntegerType, IndexType>(type))
     return emitError(loc, "integer literal not valid for specified type"),
            nullptr;
@@ -1003,7 +1010,9 @@ Attribute Parser::parseDenseElementsAttr(Type attrType) {
   auto type = parseElementsLiteralType(attrType);
   if (!type)
     return nullptr;
-  return literalParser.getAttr(loc, type);
+  if (syntaxOnly())
+    return state.syntaxOnlyAttr;
+  return literalParser.getAttr(loc, cast<ShapedType>(type));
 }
 
 Attribute Parser::parseDenseResourceElementsAttr(Type attrType) {
@@ -1030,6 +1039,9 @@ Attribute Parser::parseDenseResourceElementsAttr(Type attrType) {
       return nullptr;
   }
 
+  if (syntaxOnly())
+    return state.syntaxOnlyAttr;
+
   ShapedType shapedType = dyn_cast<ShapedType>(attrType);
   if (!shapedType) {
     emitError(typeLoc, "`dense_resource` expected a shaped type");
@@ -1044,7 +1056,7 @@ Attribute Parser::parseDenseResourceElementsAttr(Type attrType) {
 ///   elements-literal-type ::= vector-type | ranked-tensor-type
 ///
 /// This method also checks the type has static shape.
-ShapedType Parser::parseElementsLiteralType(Type type) {
+Type Parser::parseElementsLiteralType(Type type) {
   // If the user didn't provide a type, parse the colon type for the literal.
   if (!type) {
     if (parseToken(Token::colon, "expected ':'"))
@@ -1053,6 +1065,9 @@ ShapedType Parser::parseElementsLiteralType(Type type) {
       return nullptr;
   }
 
+  if (syntaxOnly())
+    return state.syntaxOnlyType;
+
   auto sType = dyn_cast<ShapedType>(type);
   if (!sType) {
     emitError("elements literal must be a shaped type");
@@ -1077,17 +1092,23 @@ Attribute Parser::parseSparseElementsAttr(Type attrType) {
   // of the type.
   Type indiceEltType = builder.getIntegerType(64);
   if (consumeIf(Token::greater)) {
-    ShapedType type = parseElementsLiteralType(attrType);
+    Type type = parseElementsLiteralType(attrType);
     if (!type)
       return nullptr;
 
+    if (syntaxOnly())
+      return state.syntaxOnlyAttr;
+
     // Construct the sparse elements attr using zero element indice/value
     // attributes.
+    ShapedType shapedType = cast<ShapedType>(type);
     ShapedType indicesType =
-        RankedTensorType::get({0, type.getRank()}, indiceEltType);
-    ShapedType valuesType = RankedTensorType::get({0}, type.getElementType());
+        RankedTensorType::get({0, shapedType.getRank()}, indiceEltType);
+    ShapedType valuesType =
+        RankedTensorType::get({0}, shapedType.getElementType());
     return getChecked<SparseElementsAttr>(
-        loc, type, DenseElementsAttr::get(indicesType, ArrayRef<Attribute>()),
+        loc, shapedType,
+        DenseElementsAttr::get(indicesType, ArrayRef<Attribute>()),
         DenseElementsAttr::get(valuesType, ArrayRef<Attribute>()));
   }
 
@@ -1114,14 +1135,20 @@ Attribute Parser::parseSparseElementsAttr(Type attrType) {
   if (!type)
     return nullptr;
 
+  if (syntaxOnly())
+    return state.syntaxOnlyAttr;
+
+  ShapedType shapedType = cast<ShapedType>(type);
+
   // If the indices are a splat, i.e. the literal parser parsed an element and
   // not a list, we set the shape explicitly. The indices are represented by a
   // 2-dimensional shape where the second dimension is the rank of the type.
   // Given that the parsed indices is a splat, we know that we only have one
   // indice and thus one for the first dimension.
   ShapedType indicesType;
   if (indiceParser.getShape().empty()) {
-    indicesType = RankedTensorType::get({1, type.getRank()}, indiceEltType);
+    indicesType =
+        RankedTensorType::get({1, shapedType.getRank()}, indiceEltType);
   } else {
     // Otherwise, set the shape to the one parsed by the literal parser.
     indicesType = RankedTensorType::get(indiceParser.getShape(), indiceEltType);
@@ -1131,15 +1158,15 @@ Attribute Parser::parseSparseElementsAttr(Type attrType) {
   // If the values are a splat, set the shape explicitly based on the number of
   // indices. The number of indices is encoded in the first dimension of the
   // indice shape type.
-  auto valuesEltType = type.getElementType();
+  auto valuesEltType = shapedType.getElementType();
   ShapedType valuesType =
       valuesParser.getShape().empty()
           ? RankedTensorType::get({indicesType.getDimSize(0)}, valuesEltType)
           : RankedTensorType::get(valuesParser.getShape(), valuesEltType);
   auto values = valuesParser.getAttr(valuesLoc, valuesType);
 
   // Build the sparse elements attribute by the indices and values.
-  return getChecked<SparseElementsAttr>(loc, type, indices, values);
+  return getChecked<SparseElementsAttr>(loc, shapedType, indices, values);
 }
 
 Attribute Parser::parseStridedLayoutAttr() {
@@ -1260,6 +1287,9 @@ Attribute Parser::parseDistinctAttr(Type type) {
       return {};
   }
 
+  if (syntaxOnly())
+    return state.syntaxOnlyAttr;
+
   // Add the distinct attribute to the parser state, if it has not been parsed
   // before. Otherwise, check if the parsed reference attribute matches the one
   // found in the parser state.

mlir/lib/AsmParser/DialectSymbolParser.cpp

@@ -15,6 +15,7 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
@@ -156,9 +157,11 @@ ParseResult Parser::parseDialectSymbolBody(StringRef &body,
 }
 
 /// Parse an extended dialect symbol.
-template <typename Symbol, typename SymbolAliasMap, typename CreateFn>
+template <typename Symbol, typename SymbolAliasMap, typename ParseAliasFn,
+          typename CreateFn>
 static Symbol parseExtendedSymbol(Parser &p, AsmParserState *asmState,
                                   SymbolAliasMap &aliases,
+                                  ParseAliasFn &parseAliasFn,
                                   CreateFn &&createSymbol) {
   Token tok = p.getToken();
 
@@ -185,12 +188,32 @@ static Symbol parseExtendedSymbol(Parser &p, AsmParserState *asmState,
   // If there is no '<' token following this, and if the typename contains no
   // dot, then we are parsing a symbol alias.
   if (!hasTrailingData && !isPrettyName) {
+
+    // Don't check the validity of alias reference in syntax-only mode.
+    if (p.syntaxOnly()) {
+      if constexpr (std::is_same_v<Symbol, Type>)
+        return p.getState().syntaxOnlyType;
+      else
+        return p.getState().syntaxOnlyAttr;
+    }
+
     // Check for an alias for this type.
     auto aliasIt = aliases.find(identifier);
-    if (aliasIt == aliases.end())
-      return (p.emitWrongTokenError("undefined symbol alias id '" + identifier +
-                                    "'"),
-              nullptr);
+    if (aliasIt == aliases.end()) {
+      FailureOr<Symbol> symbol = failure();
+      // Try the parse alias function if set.
+      if (parseAliasFn)
+        symbol = parseAliasFn(identifier);
+
+      if (failed(symbol)) {
+        p.emitWrongTokenError("undefined symbol alias id '" + identifier + "'");
+        return nullptr;
+      }
+      if (!*symbol)
+        return nullptr;
+
+      aliasIt = aliases.insert({identifier, *symbol}).first;
+    }
     if (asmState) {
       if constexpr (std::is_same_v<Symbol, Type>)
         asmState->addTypeAliasUses(identifier, range);
@@ -241,12 +264,16 @@ Attribute Parser::parseExtendedAttr(Type type) {
   MLIRContext *ctx = getContext();
   Attribute attr = parseExtendedSymbol<Attribute>(
       *this, state.asmState, state.symbols.attributeAliasDefinitions,
+      state.symbols.parseUnknownAttributeAlias,
       [&](StringRef dialectName, StringRef symbolData, SMLoc loc) -> Attribute {
         // Parse an optional trailing colon type.
         Type attrType = type;
         if (consumeIf(Token::colon) && !(attrType = parseType()))
           return Attribute();
 
+        if (syntaxOnly())
+          return state.syntaxOnlyAttr;
+
         // If we found a registered dialect, then ask it to parse the attribute.
         if (Dialect *dialect =
                 builder.getContext()->getOrLoadDialect(dialectName)) {
@@ -288,7 +315,11 @@ Type Parser::parseExtendedType() {
   MLIRContext *ctx = getContext();
   return parseExtendedSymbol<Type>(
       *this, state.asmState, state.symbols.typeAliasDefinitions,
+      state.symbols.parseUnknownTypeAlias,
       [&](StringRef dialectName, StringRef symbolData, SMLoc loc) -> Type {
+        if (syntaxOnly())
+          return state.syntaxOnlyType;
+
         // If we found a registered dialect, then ask it to parse the type.
         if (auto *dialect = ctx->getOrLoadDialect(dialectName)) {
           // Temporarily reset the lexer to let the dialect parse the type.

mlir/lib/AsmParser/LocationParser.cpp

@@ -53,6 +53,9 @@ ParseResult Parser::parseCallSiteLocation(LocationAttr &loc) {
   if (parseToken(Token::r_paren, "expected ')' in callsite location"))
     return failure();
 
+  if (syntaxOnly())
+    return success();
+
   // Return the callsite location.
   loc = CallSiteLoc::get(calleeLoc, callerLoc);
   return success();
@@ -79,6 +82,9 @@ ParseResult Parser::parseFusedLocation(LocationAttr &loc) {
     LocationAttr newLoc;
     if (parseLocationInstance(newLoc))
       return failure();
+    if (syntaxOnly())
+      return success();
+
     locations.push_back(newLoc);
     return success();
   };
@@ -135,12 +141,15 @@ ParseResult Parser::parseNameOrFileLineColLocation(LocationAttr &loc) {
     if (parseLocationInstance(childLoc))
       return failure();
 
-    loc = NameLoc::get(StringAttr::get(ctx, str), childLoc);
-
     // Parse the closing ')'.
     if (parseToken(Token::r_paren,
                    "expected ')' after child location of NameLoc"))
       return failure();
+
+    if (syntaxOnly())
+      return success();
+
+    loc = NameLoc::get(StringAttr::get(ctx, str), childLoc);
   } else {
     loc = NameLoc::get(StringAttr::get(ctx, str));
   }
@@ -154,6 +163,10 @@ ParseResult Parser::parseLocationInstance(LocationAttr &loc) {
     Attribute locAttr = parseExtendedAttr(Type());
     if (!locAttr)
       return failure();
+
+    if (syntaxOnly())
+      return success();
+
     if (!(loc = dyn_cast<LocationAttr>(locAttr)))
       return emitError("expected location attribute, but got") << locAttr;
     return success();