[mlir] Add a DialectAsmParser::getChecked method

[rriddle]

Summary

This function simplifies calling the getChecked methods on Attributes and Types from within the parser, and removes any need to use getEncodedSourceLocation for these methods (by using an SMLoc instead). This is much more efficient than using an mlir::Location, as the encoding process to produce an mlir::Location is inefficient and undesirable for parsing (locations used during parsing should not persist afterwards unless otherwise necessary).

[jpienaar]

mlir/include/mlir/IR/DialectImplementation.h

175

SMLoc is a very specific to textual IR format parsing (debug/test format) part of the world. Does it make as much sense in IR? Is there an alternative where Location is cheap to create from SMLoc for just these parsing cases where we know the context outlives the potential use? (Alternatively this as a helper method in utils that just makes the getChecked inside smaller to write at call site?). SMLoc just feels a bit weird here.

[rriddle]

mlir/include/mlir/IR/DialectImplementation.h

175

I'm not sure what you are getting at. This is the parser for attributes and types, this almost exclusively uses llvm::SMLoc. I don't see any other location format being useful here, this class is only used for textual IR format parsing.

[jpienaar]

Nice improvement, thanks

mlir/include/mlir/IR/DialectImplementation.h

175

Yeah I was in wrong file "in my mind" switching tabs and I just looked at the snippets.

mlir/include/mlir/IR/DialectImplementation.h

ParseResult getCurrentLocation(llvm::SMLoc *loc) {
     *loc = getCurrentLocation();
     return success();
   }
 
   /// Return the location of the original name token.
   virtual llvm::SMLoc getNameLoc() const = 0;
 
   /// Re-encode the given source location as an MLIR location and return it.
+  /// Note: This method should only be used when a `Location` is necessary, as
+  /// the encoding process is not efficient. In other cases a more suitable
+  /// alternative should be used, such as the `getChecked` methods defined
+  /// below.
   virtual Location getEncodedSourceLoc(llvm::SMLoc loc) = 0;
 
   /// Returns the full specification of the symbol being parsed. This allows for
   /// using a separate parser if necessary.
   virtual StringRef getFullSymbolSpec() const = 0;
 
   // These methods emit an error and return failure or success. This allows
   // these to be chained together into a linear sequence of || expressions in
OptionalParseResult parseOptionalInteger(IntT &result) {
 
     // Try to convert to the provided integer type.
     result = IntT(uintResult);
     if (uint64_t(result) != uintResult)
       return emitError(loc, "integer value too large");
     return success();
   }
 
+  /// Invoke the `getChecked` method of the given Attribute or Type class, using
+  /// the provided location to emit errors in the case of failure. Note that
+  /// unlike `OpBuilder::getType`, this method does not implicitly insert a
+  /// context parameter.
+  template <typename T, typename... ParamsT>
+  T getChecked(llvm::SMLoc loc, ParamsT &&...params) {
+    return T::getChecked([&] { return emitError(loc); },
+                         std::forward<ParamsT>(params)...);
+  }
+  /// A variant of `getChecked` that uses the result of `getNameLoc` to emit
+  /// errors.
+  template <typename T, typename... ParamsT>
+  T getChecked(ParamsT &&...params) {
+    return T::getChecked([&] { return emitError(getNameLoc()); },
+                         std::forward<ParamsT>(params)...);
+  }
+
   //===--------------------------------------------------------------------===//
   // Token Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a '->' token.
   virtual ParseResult parseArrow() = 0;
 
   /// Parse a '->' token if present

mlir/lib/Dialect/LLVMIR/IR/LLVMTypeSyntax.cpp

 // Parsing.
 //===----------------------------------------------------------------------===//
 
 static ParseResult dispatchParse(DialectAsmParser &parser, Type &type);
 
 /// Parses an LLVM dialect function type.
 ///   llvm-type :: = `func<` llvm-type `(` llvm-type-list `...`? `)>`
 static LLVMFunctionType parseFunctionType(DialectAsmParser &parser) {
-  Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
+  llvm::SMLoc loc = parser.getCurrentLocation();
   Type returnType;
   if (parser.parseLess() || dispatchParse(parser, returnType) ||
       parser.parseLParen())
     return LLVMFunctionType();
 
   // Function type without arguments.
   if (succeeded(parser.parseOptionalRParen())) {
     if (succeeded(parser.parseGreater()))
-      return LLVMFunctionType::getChecked(loc, returnType, llvm::None,
+      return parser.getChecked<LLVMFunctionType>(loc, returnType, llvm::None,
                                                  /*isVarArg=*/false);
     return LLVMFunctionType();
   }
 
   // Parse arguments.
   SmallVector<Type, 8> argTypes;
   do {
     if (succeeded(parser.parseOptionalEllipsis())) {
       if (parser.parseOptionalRParen() || parser.parseOptionalGreater())
         return LLVMFunctionType();
-      return LLVMFunctionType::getChecked(loc, returnType, argTypes,
+      return parser.getChecked<LLVMFunctionType>(loc, returnType, argTypes,
                                                  /*isVarArg=*/true);
     }
 
     Type arg;
     if (dispatchParse(parser, arg))
       return LLVMFunctionType();
     argTypes.push_back(arg);
   } while (succeeded(parser.parseOptionalComma()));
 
   if (parser.parseOptionalRParen() || parser.parseOptionalGreater())
     return LLVMFunctionType();
-  return LLVMFunctionType::getChecked(loc, returnType, argTypes,
+  return parser.getChecked<LLVMFunctionType>(loc, returnType, argTypes,
                                              /*isVarArg=*/false);
 }
 
 /// Parses an LLVM dialect pointer type.
 ///   llvm-type ::= `ptr<` llvm-type (`,` integer)? `>`
 static LLVMPointerType parsePointerType(DialectAsmParser &parser) {
-  Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
+  llvm::SMLoc loc = parser.getCurrentLocation();
   Type elementType;
   if (parser.parseLess() || dispatchParse(parser, elementType))
     return LLVMPointerType();
 
   unsigned addressSpace = 0;
   if (succeeded(parser.parseOptionalComma()) &&
       failed(parser.parseInteger(addressSpace)))
     return LLVMPointerType();
   if (failed(parser.parseGreater()))
     return LLVMPointerType();
-  return LLVMPointerType::getChecked(loc, elementType, addressSpace);
+  return parser.getChecked<LLVMPointerType>(loc, elementType, addressSpace);
 }
 
 /// Parses an LLVM dialect vector type.
 ///   llvm-type ::= `vec<` `? x`? integer `x` llvm-type `>`
 /// Supports both fixed and scalable vectors.
 static Type parseVectorType(DialectAsmParser &parser) {
   SmallVector<int64_t, 2> dims;
   llvm::SMLoc dimPos, typePos;
   Type elementType;
-  Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
+  llvm::SMLoc loc = parser.getCurrentLocation();
   if (parser.parseLess() || parser.getCurrentLocation(&dimPos) ||
       parser.parseDimensionList(dims, /*allowDynamic=*/true) ||
       parser.getCurrentLocation(&typePos) ||
       dispatchParse(parser, elementType) || parser.parseGreater())
     return Type();
 
   // We parsed a generic dimension list, but vectors only support two forms:
   //  - single non-dynamic entry in the list (fixed vector);
   //  - two elements, the first dynamic (indicated by -1) and the second
   //    non-dynamic (scalable vector).
   if (dims.empty() || dims.size() > 2 ||
       ((dims.size() == 2) ^ (dims[0] == -1)) ||
       (dims.size() == 2 && dims[1] == -1)) {
     parser.emitError(dimPos)
         << "expected '? x <integer> x <type>' or '<integer> x <type>'";
     return Type();
   }
 
   bool isScalable = dims.size() == 2;
   if (isScalable)
-    return LLVMScalableVectorType::getChecked(loc, elementType, dims[1]);
+    return parser.getChecked<LLVMScalableVectorType>(loc, elementType, dims[1]);
   if (elementType.isSignlessIntOrFloat()) {
     parser.emitError(typePos)
         << "cannot use !llvm.vec for built-in primitives, use 'vector' instead";
     return Type();
   }
-  return LLVMFixedVectorType::getChecked(loc, elementType, dims[0]);
+  return parser.getChecked<LLVMFixedVectorType>(loc, elementType, dims[0]);
 }
 
 /// Parses an LLVM dialect array type.
 ///   llvm-type ::= `array<` integer `x` llvm-type `>`
 static LLVMArrayType parseArrayType(DialectAsmParser &parser) {
   SmallVector<int64_t, 1> dims;
   llvm::SMLoc sizePos;
   Type elementType;
-  Location loc = parser.getEncodedSourceLoc(parser.getCurrentLocation());
+  llvm::SMLoc loc = parser.getCurrentLocation();
   if (parser.parseLess() || parser.getCurrentLocation(&sizePos) ||
       parser.parseDimensionList(dims, /*allowDynamic=*/false) ||
       dispatchParse(parser, elementType) || parser.parseGreater())
     return LLVMArrayType();
 
   if (dims.size() != 1) {
     parser.emitError(sizePos) << "expected ? x <type>";
     return LLVMArrayType();
   }
 
-  return LLVMArrayType::getChecked(loc, elementType, dims[0]);
+  return parser.getChecked<LLVMArrayType>(loc, elementType, dims[0]);
 }
 
 /// Attempts to set the body of an identified structure type. Reports a parsing
 /// error at `subtypesLoc` in case of failure.
 static LLVMStructType trySetStructBody(LLVMStructType type,
                                        ArrayRef<Type> subtypes, bool isPacked,
                                        DialectAsmParser &parser,
                                        llvm::SMLoc subtypesLoc) {

mlir/lib/Dialect/Quant/IR/TypeParser.cpp

 
 /// Parses an AnyQuantizedType.
 ///
 ///   any ::= `any<` storage-spec (expressed-type-spec)?`>`
 ///   storage-spec ::= storage-type (`<` storage-range `>`)?
 ///   storage-range ::= integer-literal `:` integer-literal
 ///   storage-type ::= (`i` | `u`) integer-literal
 ///   expressed-type-spec ::= `:` `f` integer-literal
-static Type parseAnyType(DialectAsmParser &parser, Location loc) {
+static Type parseAnyType(DialectAsmParser &parser) {
   IntegerType storageType;
   FloatType expressedType;
   unsigned typeFlags = 0;
   int64_t storageTypeMin;
   int64_t storageTypeMax;
 
   // Type specification.
   if (parser.parseLess())
if (parser.parseType(expressedType)) {
       return nullptr;
     }
   }
 
   if (parser.parseGreater()) {
     return nullptr;
   }
 
-  return AnyQuantizedType::getChecked(loc, typeFlags, storageType,
-                                      expressedType, storageTypeMin,
-                                      storageTypeMax);
+  return parser.getChecked<AnyQuantizedType>(
+      typeFlags, storageType, expressedType, storageTypeMin, storageTypeMax);
 }
 
 static ParseResult parseQuantParams(DialectAsmParser &parser, double &scale,
                                     int64_t &zeroPoint) {
   // scale[:zeroPoint]?
   // scale.
   if (parser.parseFloat(scale))
     return failure();

 ///                        axis-spec `,` scale-zero-list `>`
 ///   storage-spec ::= storage-type (`<` storage-range `>`)?
 ///   storage-range ::= integer-literal `:` integer-literal
 ///   storage-type ::= (`i` | `u`) integer-literal
 ///   expressed-type-spec ::= `:` `f` integer-literal
 ///   axis-spec ::= `:` integer-literal
 ///   scale-zero ::= float-literal `:` integer-literal
 ///   scale-zero-list ::= `{` scale-zero (`,` scale-zero)* `}`
-static Type parseUniformType(DialectAsmParser &parser, Location loc) {
+static Type parseUniformType(DialectAsmParser &parser) {
   IntegerType storageType;
   FloatType expressedType;
   unsigned typeFlags = 0;
   int64_t storageTypeMin;
   int64_t storageTypeMax;
   bool isPerAxis = false;
   int32_t quantizedDimension;
   SmallVector<double, 1> scales;
return (parser.emitError(scaleZPLoc,
                              "multiple scales/zeroPoints provided, but "
                              "quantizedDimension wasn't specified"),
             nullptr);
   }
 
   if (isPerAxis) {
     ArrayRef<double> scalesRef(scales.begin(), scales.end());
     ArrayRef<int64_t> zeroPointsRef(zeroPoints.begin(), zeroPoints.end());
-    return UniformQuantizedPerAxisType::getChecked(
-        loc, typeFlags, storageType, expressedType, scalesRef, zeroPointsRef,
+    return parser.getChecked<UniformQuantizedPerAxisType>(
+        typeFlags, storageType, expressedType, scalesRef, zeroPointsRef,
         quantizedDimension, storageTypeMin, storageTypeMax);
   }
 
-  return UniformQuantizedType::getChecked(
-      loc, typeFlags, storageType, expressedType, scales.front(),
-      zeroPoints.front(), storageTypeMin, storageTypeMax);
+  return parser.getChecked<UniformQuantizedType>(
+      typeFlags, storageType, expressedType, scales.front(), zeroPoints.front(),
+      storageTypeMin, storageTypeMax);
 }
 
 /// Parses an CalibratedQuantizedType.
 ///
 ///   calibrated ::= `calibrated<` expressed-spec `>`
 ///   expressed-spec ::= expressed-type `<` calibrated-range `>`
 ///   expressed-type ::= `f` integer-literal
 ///   calibrated-range ::= float-literal `:` float-literal
-static Type parseCalibratedType(DialectAsmParser &parser, Location loc) {
+static Type parseCalibratedType(DialectAsmParser &parser) {
   FloatType expressedType;
   double min;
   double max;
 
   // Type specification.
   if (parser.parseLess())
     return nullptr;
 
   // Expressed type.
   expressedType = parseExpressedTypeAndRange(parser, min, max);
   if (!expressedType) {
     return nullptr;
   }
 
   if (parser.parseGreater()) {
     return nullptr;
   }
 
-  return CalibratedQuantizedType::getChecked(loc, expressedType, min, max);
+  return parser.getChecked<CalibratedQuantizedType>(expressedType, min, max);
 }
 
 /// Parse a type registered to this dialect.
 Type QuantizationDialect::parseType(DialectAsmParser &parser) const {
-  Location loc = parser.getEncodedSourceLoc(parser.getNameLoc());
-
   // All types start with an identifier that we switch on.
   StringRef typeNameSpelling;
   if (failed(parser.parseKeyword(&typeNameSpelling)))
     return nullptr;
 
   if (typeNameSpelling == "uniform")
-    return parseUniformType(parser, loc);
+    return parseUniformType(parser);
   if (typeNameSpelling == "any")
-    return parseAnyType(parser, loc);
+    return parseAnyType(parser);
   if (typeNameSpelling == "calibrated")
-    return parseCalibratedType(parser, loc);
+    return parseCalibratedType(parser);
 
   parser.emitError(parser.getNameLoc(),
                    "unknown quantized type " + typeNameSpelling);
   return nullptr;
 }
 
 static void printStorageType(QuantizedType type, DialectAsmPrinter &out) {
   // storage type

mlir/lib/Parser/DialectSymbolParser.cpp

Attribute attr = parseExtendedSymbol<Attribute>(
               symbolData, state.context, state.symbols, [&](Parser &parser) {
                 CustomDialectAsmParser customParser(symbolData, parser);
                 return dialect->parseAttribute(customParser, attrType);
               });
         }
 
         // Otherwise, form a new opaque attribute.
         return OpaqueAttr::getChecked(
-            getEncodedSourceLocation(loc),
+            [&] { return emitError(loc); },
             Identifier::get(dialectName, state.context), symbolData,
             attrType ? attrType : NoneType::get(state.context));
       });
 
   // Ensure that the attribute has the same type as requested.
   if (attr && type && attr.getType() != type) {
     emitError("attribute type different than expected: expected ")
         << type << ", but got " << attr.getType();
return parseExtendedSymbol<Type>(
               symbolData, state.context, state.symbols, [&](Parser &parser) {
                 CustomDialectAsmParser customParser(symbolData, parser);
                 return dialect->parseType(customParser);
               });
         }
 
         // Otherwise, form a new opaque type.
         return OpaqueType::getChecked(
-            getEncodedSourceLocation(loc),
+            [&] { return emitError(loc); },
             Identifier::get(dialectName, state.context), symbolData);
       });
 }
 
 //===----------------------------------------------------------------------===//
 // mlir::parseAttribute/parseType
 //===----------------------------------------------------------------------===//