Refactor the location classes to be attributes instead of separate IR classes.

This will allow for locations to be used in the same contexts as attributes. Given that attributes are nullable types, the 'Location' class now represents a non-nullable wrapper around a 'LocationAttr'. This preserves the desired semantics we have for non-optional locations.

PiperOrigin-RevId: 254505278

mlir/include/mlir/IR/Attributes.h

@@ -50,7 +50,6 @@ struct TypeAttributeStorage;
 struct DenseElementsAttributeStorage;
 struct OpaqueElementsAttributeStorage;
 struct SparseElementsAttributeStorage;
-
 } // namespace detail
 
 /// Attributes are known-constant values of operations and functions.
@@ -122,8 +121,7 @@ public:
   const void *getAsOpaquePointer() const { return impl; }
   /// Construct an attribute from the opaque pointer representation.
   static Attribute getFromOpaquePointer(const void *ptr) {
-    return Attribute(
+    return Attribute(reinterpret_cast<const ImplType *>(ptr));
-        const_cast<ImplType *>(reinterpret_cast<const ImplType *>(ptr)));
   }
 
   friend ::llvm::hash_code hash_value(Attribute arg);
@@ -158,6 +156,25 @@ enum Kind {
   SparseElements,
   FIRST_ELEMENTS_ATTR = DenseElements,
   LAST_ELEMENTS_ATTR = SparseElements,
+
+  /// Locations.
+  CallSiteLocation,
+  FileLineColLocation,
+  FusedLocation,
+  NameLocation,
+  UnknownLocation,
+
+  // Represents a location as a 'void*' pointer to a front-end's opaque
+  // location information, which must live longer than the MLIR objects that
+  // refer to it.  OpaqueLocation's are never serialized.
+  //
+  // TODO: OpaqueLocation,
+
+  // Represents a value inlined through a function call.
+  // TODO: InlinedLocation,
+
+  FIRST_LOCATION_ATTR = CallSiteLocation,
+  LAST_LOCATION_ATTR = UnknownLocation,
 };
 } // namespace StandardAttributes
 

mlir/include/mlir/IR/Location.h

@@ -23,8 +23,7 @@
 #ifndef MLIR_IR_LOCATION_H
 #define MLIR_IR_LOCATION_H
 
-#include "mlir/Support/LLVM.h"
+#include "mlir/IR/Attributes.h"
-#include "llvm/ADT/DenseMapInfo.h"
 
 namespace mlir {
 
@@ -44,74 +43,54 @@ struct FusedLocationStorage;
 } // namespace detail
 
 /// Location objects represent source locations information in MLIR.
+/// LocationAttr acts as the anchor for all Location based attributes.
+class LocationAttr : public Attribute {
+public:
+  using Attribute::Attribute;
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool classof(Attribute attr) {
+    return attr.getKind() >= StandardAttributes::FIRST_LOCATION_ATTR &&
+           attr.getKind() <= StandardAttributes::LAST_LOCATION_ATTR;
+  }
+};
+
+/// This class defines the main interface for locations in MLIR and acts as a
+/// non-nullable wrapper around a LocationAttr.
 class Location {
 public:
-  enum class Kind : uint8_t {
+  Location(LocationAttr loc) : impl(loc) {
-    /// This represents an unknown location.
-    UnknownLocation,
-
-    /// This represents a file/line/column location.
-    FileLineColLocation,
-
-    /// This represents an identity name attached to a child location.
-    NameLocation,
-
-    /// This represents a location as a call site.
-    CallSiteLocation,
-
-    // Represents a location as a 'void*' pointer to a front-end's opaque
-    // location information, which must live longer than the MLIR objects that
-    // refer to it.  OpaqueLocation's are never serialized.
-    //
-    // TODO: OpaqueLocation,
-
-    // Represents a value inlined through a function call.
-    // TODO: InlinedLocation,
-
-    // Represents a value composed of multiple source constructs.
-    FusedLocation,
-  };
-
-  using ImplType = detail::LocationStorage;
-
-  /* implicit */ Location(const ImplType *loc)
-      : loc(const_cast<ImplType *>(loc)) {
     assert(loc && "location should never be null.");
   }
 
-  Location() = delete;
+  /// Access the impl location attribute.
-  Location(const Location &other) : loc(other.loc) {}
+  operator LocationAttr() const { return impl; }
-  Location &operator=(Location other) {
+  LocationAttr *operator->() const { return const_cast<LocationAttr *>(&impl); }
-    loc = other.loc;
-    return *this;
-  }
 
-  bool operator==(Location other) const { return loc == other.loc; }
+  /// Type casting utilities on the underlying location.
-  bool operator!=(Location other) const { return !(*this == other); }
+  template <typename U> bool isa() const { return impl.isa<U>(); }
+  template <typename U> U dyn_cast() const { return impl.dyn_cast<U>(); }
+  template <typename U> U cast() const { return impl.cast<U>(); }
 
-  template <typename U> bool isa() const;
+  /// Comparison operators.
-  template <typename U> Optional<U> dyn_cast() const;
+  bool operator==(Location rhs) const { return impl == rhs.impl; }
-  template <typename U> U cast() const;
+  bool operator!=(Location rhs) const { return !(*this == rhs); }
-
-  /// Return the classification for this location.
-  Kind getKind() const;
 
   /// Print the location.
-  void print(raw_ostream &os) const;
+  void print(raw_ostream &os) const { impl.print(os); }
-  void dump() const;
+  void dump() const { impl.dump(); }
 
   friend ::llvm::hash_code hash_value(Location arg);
 
   /// Methods for supporting PointerLikeTypeTraits.
-  const void *getAsOpaquePointer() const {
+  const void *getAsOpaquePointer() const { return impl.getAsOpaquePointer(); }
-    return static_cast<const void *>(loc);
-  }
   static Location getFromOpaquePointer(const void *pointer) {
-    return Location(reinterpret_cast<ImplType *>(const_cast<void *>(pointer)));
+    return LocationAttr(reinterpret_cast<const AttributeStorage *>(pointer));
   }
 
 protected:
-  ImplType *loc;
+  /// The internal backing location attribute.
+  LocationAttr impl;
 };
 
 inline raw_ostream &operator<<(raw_ostream &os, const Location &loc) {
@@ -119,73 +98,15 @@ inline raw_ostream &operator<<(raw_ostream &os, const Location &loc) {
   return os;
 }
 
-/// Represents an unknown location.  This is always a singleton for a given
-/// MLIRContext.
-class UnknownLoc : public Location {
-public:
-  using ImplType = detail::UnknownLocationStorage;
-  using Location::Location;
-
-  static UnknownLoc get(MLIRContext *context);
-
-  /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool kindof(Kind kind) { return kind == Kind::UnknownLocation; }
-};
-
-/// Represents a location derived from a file/line/column location.  The column
-/// and line may be zero to represent unknown column and/or unknown line/column
-/// information.
-class FileLineColLoc : public Location {
-public:
-  using ImplType = detail::FileLineColLocationStorage;
-  using Location::Location;
-
-  /// Return a uniqued FileLineCol location object.
-  static FileLineColLoc get(Identifier filename, unsigned line, unsigned column,
-                            MLIRContext *context);
-  static FileLineColLoc get(StringRef filename, unsigned line, unsigned column,
-                            MLIRContext *context);
-
-  StringRef getFilename() const;
-
-  unsigned getLine() const;
-  unsigned getColumn() const;
-
-  /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool kindof(Kind kind) { return kind == Kind::FileLineColLocation; }
-};
-
-/// Represents an identity name attached to a child location.
-class NameLoc : public Location {
-public:
-  using ImplType = detail::NameLocationStorage;
-  using Location::Location;
-
-  /// Return a uniqued name location object. The child location must not be
-  /// another NameLoc.
-  static NameLoc get(Identifier name, Location child, MLIRContext *context);
-
-  /// Return a uniqued name location object with an unknown child.
-  static NameLoc get(Identifier name, MLIRContext *context);
-
-  /// Return the name identifier.
-  Identifier getName() const;
-
-  /// Return the child location.
-  Location getChildLoc() const;
-
-  /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool kindof(Kind kind) { return kind == Kind::NameLocation; }
-};
-
 /// Represents a location as call site. "callee" is the concrete location
 /// (Unknown/NameLocation/FileLineColLoc) and "caller" points to the caller's
 /// location (another CallLocation or a concrete location). Multiple
 /// CallSiteLocs can be chained to form a call stack.
-class CallSiteLoc : public Location {
+class CallSiteLoc
+    : public Attribute::AttrBase<CallSiteLoc, LocationAttr,
+                                 detail::CallSiteLocationStorage> {
 public:
-  using ImplType = detail::CallSiteLocationStorage;
+  using Base::Base;
-  using Location::Location;
 
   /// Return a uniqued call location object.
   static CallSiteLoc get(Location callee, Location caller,
@@ -204,22 +125,52 @@ public:
   Location getCaller() const;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool kindof(Kind kind) { return kind == Kind::CallSiteLocation; }
+  static bool kindof(unsigned kind) {
+    return kind == StandardAttributes::CallSiteLocation;
+  }
+};
+
+/// Represents a location derived from a file/line/column location.  The column
+/// and line may be zero to represent unknown column and/or unknown line/column
+/// information.
+class FileLineColLoc
+    : public Attribute::AttrBase<FileLineColLoc, LocationAttr,
+                                 detail::FileLineColLocationStorage> {
+public:
+  using Base::Base;
+
+  /// Return a uniqued FileLineCol location object.
+  static FileLineColLoc get(Identifier filename, unsigned line, unsigned column,
+                            MLIRContext *context);
+  static FileLineColLoc get(StringRef filename, unsigned line, unsigned column,
+                            MLIRContext *context);
+
+  StringRef getFilename() const;
+
+  unsigned getLine() const;
+  unsigned getColumn() const;
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool kindof(unsigned kind) {
+    return kind == StandardAttributes::FileLineColLocation;
+  }
 };
 
 /// Represents a value composed of multiple source constructs, with an optional
 /// metadata attribute.
-class FusedLoc : public Location {
+class FusedLoc : public Attribute::AttrBase<FusedLoc, LocationAttr,
+                                            detail::FusedLocationStorage> {
 public:
-  using ImplType = detail::FusedLocationStorage;
+  using Base::Base;
-  using Location::Location;
 
   /// Return a uniqued Fused Location object. The first location in the list
   /// will get precedence during diagnostic emission, with the rest being
   /// displayed as supplementary "fused from here" style notes.
-  static Location get(ArrayRef<Location> locs, MLIRContext *context);
+  static LocationAttr get(ArrayRef<Location> locs, Attribute metadata,
-  static Location get(ArrayRef<Location> locs, Attribute metadata,
+                          MLIRContext *context);
-                      MLIRContext *context);
+  static LocationAttr get(ArrayRef<Location> locs, MLIRContext *context) {
+    return get(locs, Attribute(), context);
+  }
 
   ArrayRef<Location> getLocations() const;
 
@@ -228,23 +179,54 @@ public:
   Attribute getMetadata() const;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool kindof(Kind kind) { return kind == Kind::FusedLocation; }
+  static bool kindof(unsigned kind) {
+    return kind == StandardAttributes::FusedLocation;
+  }
+};
+
+/// Represents an identity name attached to a child location.
+class NameLoc : public Attribute::AttrBase<NameLoc, LocationAttr,
+                                           detail::NameLocationStorage> {
+public:
+  using Base::Base;
+
+  /// Return a uniqued name location object. The child location must not be
+  /// another NameLoc.
+  static NameLoc get(Identifier name, Location child, MLIRContext *context);
+
+  /// Return a uniqued name location object with an unknown child.
+  static NameLoc get(Identifier name, MLIRContext *context);
+
+  /// Return the name identifier.
+  Identifier getName() const;
+
+  /// Return the child location.
+  Location getChildLoc() const;
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool kindof(unsigned kind) {
+    return kind == StandardAttributes::NameLocation;
+  }
+};
+
+/// Represents an unknown location.  This is always a singleton for a given
+/// MLIRContext.
+class UnknownLoc : public Attribute::AttrBase<UnknownLoc, LocationAttr> {
+public:
+  using Base::Base;
+
+  /// Get an instance of the UnknownLoc.
+  static UnknownLoc get(MLIRContext *context);
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool kindof(unsigned kind) {
+    return kind == StandardAttributes::UnknownLocation;
+  }
 };
 
 // Make Location hashable.
 inline ::llvm::hash_code hash_value(Location arg) {
-  return ::llvm::hash_value(arg.loc);
+  return hash_value(arg.impl);
-}
-
-template <typename U> bool Location::isa() const {
-  return U::kindof(getKind());
-}
-template <typename U> Optional<U> Location::dyn_cast() const {
-  return isa<U>() ? U(loc) : Optional<U>();
-}
-template <typename U> U Location::cast() const {
-  assert(isa<U>());
-  return U(loc);
 }
 
 } // end namespace mlir
@@ -255,11 +237,11 @@ namespace llvm {
 template <> struct DenseMapInfo<mlir::Location> {
   static mlir::Location getEmptyKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
-    return mlir::Location(static_cast<mlir::Location::ImplType *>(pointer));
+    return mlir::Location::getFromOpaquePointer(pointer);
   }
   static mlir::Location getTombstoneKey() {
     auto pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
-    return mlir::Location(static_cast<mlir::Location::ImplType *>(pointer));
+    return mlir::Location::getFromOpaquePointer(pointer);
   }
   static unsigned getHashValue(mlir::Location val) {
     return mlir::hash_value(val);
@@ -278,7 +260,10 @@ public:
   static inline mlir::Location getFromVoidPointer(void *P) {
     return mlir::Location::getFromOpaquePointer(P);
   }
-  enum { NumLowBitsAvailable = 3 };
+  enum {
+    NumLowBitsAvailable =
+        PointerLikeTypeTraits<mlir::Attribute>::NumLowBitsAvailable
+  };
 };
 
 } // namespace llvm

mlir/include/mlir/IR/StandardTypes.h

@@ -372,12 +372,7 @@ public:
   /// construction failures.
   static MemRefType get(ArrayRef<int64_t> shape, Type elementType,
                         ArrayRef<AffineMap> affineMapComposition = {},
-                        unsigned memorySpace = 0) {
+                        unsigned memorySpace = 0);
-    auto result = getImpl(shape, elementType, affineMapComposition, memorySpace,
-                          /*location=*/llvm::None);
-    assert(result && "Failed to construct instance of MemRefType.");
-    return result;
-  }
 
   /// Get or create a new MemRefType based on shape, element type, affine
   /// map composition, and memory space declared at the given location.
@@ -387,10 +382,7 @@ public:
   /// the error stream) and returns nullptr.
   static MemRefType getChecked(ArrayRef<int64_t> shape, Type elementType,
                                ArrayRef<AffineMap> affineMapComposition,
-                               unsigned memorySpace, Location location) {
+                               unsigned memorySpace, Location location);
-    return getImpl(shape, elementType, affineMapComposition, memorySpace,
-                   location);
-  }
 
   ArrayRef<int64_t> getShape() const;
 

mlir/include/mlir/IR/StorageUniquerSupport.h

@@ -22,12 +22,14 @@
 #ifndef MLIR_IR_STORAGEUNIQUERSUPPORT_H
 #define MLIR_IR_STORAGEUNIQUERSUPPORT_H
 
-#include "mlir/IR/Location.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/STLExtras.h"
 #include "mlir/Support/StorageUniquer.h"
 
 namespace mlir {
+class Location;
+class MLIRContext;
+
 namespace detail {
 /// Utility class for implementing users of storage classes uniqued by a
 /// StorageUniquer. Clients are not expected to interact with this class
@@ -69,8 +71,8 @@ protected:
   /// the given, potentially unknown, location. If the arguments provided are
   /// invalid then emit errors and return a null object.
   template <typename... Args>
-  static ConcreteT getChecked(Location loc, MLIRContext *ctx, unsigned kind,
+  static ConcreteT getChecked(const Location &loc, MLIRContext *ctx,
-                              Args... args) {
+                              unsigned kind, Args... args) {
     // If the construction invariants fail then we return a null attribute.
     if (failed(ConcreteT::verifyConstructionInvariants(loc, ctx, args...)))
       return ConcreteT();
@@ -79,9 +81,7 @@ protected:
 
   /// Default implementation that just returns success.
   template <typename... Args>
-  static LogicalResult
+  static LogicalResult verifyConstructionInvariants(Args... args) {
-  verifyConstructionInvariants(llvm::Optional<Location> loc, MLIRContext *ctx,
-                               Args... args) {
     return success();
   }
 

mlir/include/mlir/IR/Types.h

@@ -24,6 +24,7 @@
 
 namespace mlir {
 class FloatType;
+class Identifier;
 class IndexType;
 class IntegerType;
 class MLIRContext;

mlir/lib/IR/AsmPrinter.cpp

@@ -344,7 +344,7 @@ public:
 
   void printType(Type type);
   void print(Function *fn);
-  void printLocation(Location loc);
+  void printLocation(LocationAttr loc);
 
   void printAffineMap(AffineMap map);
   void printAffineExpr(AffineExpr expr);
@@ -359,7 +359,7 @@ protected:
                              ArrayRef<StringRef> elidedAttrs = {});
   void printAttributeOptionalType(Attribute attr, bool includeType);
   void printTrailingLocation(Location loc);
-  void printLocationInternal(Location loc, bool pretty = false);
+  void printLocationInternal(LocationAttr loc, bool pretty = false);
   void printDenseElementsAttr(DenseElementsAttr attr);
 
   /// This enum is used to represent the binding stength of the enclosing
@@ -383,22 +383,22 @@ void ModulePrinter::printTrailingLocation(Location loc) {
   printLocation(loc);
 }
 
-void ModulePrinter::printLocationInternal(Location loc, bool pretty) {
+void ModulePrinter::printLocationInternal(LocationAttr loc, bool pretty) {
   switch (loc.getKind()) {
-  case Location::Kind::UnknownLocation:
+  case StandardAttributes::UnknownLocation:
     if (pretty)
       os << "[unknown]";
     else
       os << "unknown";
     break;
-  case Location::Kind::FileLineColLocation: {
+  case StandardAttributes::FileLineColLocation: {
     auto fileLoc = loc.cast<FileLineColLoc>();
     auto mayQuote = pretty ? "" : "\"";
     os << mayQuote << fileLoc.getFilename() << mayQuote << ':'
        << fileLoc.getLine() << ':' << fileLoc.getColumn();
     break;
   }
-  case Location::Kind::NameLocation: {
+  case StandardAttributes::NameLocation: {
     auto nameLoc = loc.cast<NameLoc>();
     os << '\"' << nameLoc.getName() << '\"';
 
@@ -411,7 +411,7 @@ void ModulePrinter::printLocationInternal(Location loc, bool pretty) {
     }
     break;
   }
-  case Location::Kind::CallSiteLocation: {
+  case StandardAttributes::CallSiteLocation: {
     auto callLocation = loc.cast<CallSiteLoc>();
     auto caller = callLocation.getCaller();
     auto callee = callLocation.getCallee();
@@ -436,7 +436,7 @@ void ModulePrinter::printLocationInternal(Location loc, bool pretty) {
       os << ")";
     break;
   }
-  case Location::Kind::FusedLocation: {
+  case StandardAttributes::FusedLocation: {
     auto fusedLoc = loc.cast<FusedLoc>();
     if (!pretty)
       os << "fused";
@@ -495,7 +495,7 @@ static void printFloatValue(const APFloat &apValue, raw_ostream &os) {
   os << str;
 }
 
-void ModulePrinter::printLocation(Location loc) {
+void ModulePrinter::printLocation(LocationAttr loc) {
   if (printPrettyDebugInfo) {
     printLocationInternal(loc, /*pretty=*/true);
   } else {
@@ -712,6 +712,15 @@ void ModulePrinter::printAttributeOptionalType(Attribute attr,
     os << '>';
     break;
   }
+
+  // Location attributes.
+  case StandardAttributes::CallSiteLocation:
+  case StandardAttributes::FileLineColLocation:
+  case StandardAttributes::FusedLocation:
+  case StandardAttributes::NameLocation:
+  case StandardAttributes::UnknownLocation:
+    printLocation(attr.cast<LocationAttr>());
+    break;
   }
 }
 
@@ -1727,10 +1736,3 @@ void Module::print(raw_ostream &os) {
 }
 
 void Module::dump() { print(llvm::errs()); }
-
-void Location::print(raw_ostream &os) const {
-  ModuleState state(nullptr);
-  ModulePrinter(os, state).printLocation(*this);
-}
-
-void Location::dump() const { print(llvm::errs()); }

mlir/lib/IR/Diagnostics.cpp

@@ -313,12 +313,12 @@ struct SourceMgrDiagnosticHandlerImpl {
 
 /// Return a processable FileLineColLoc from the given location.
 static llvm::Optional<FileLineColLoc> getFileLineColLoc(Location loc) {
-  switch (loc.getKind()) {
+  switch (loc->getKind()) {
-  case Location::Kind::NameLocation:
+  case StandardAttributes::NameLocation:
     return getFileLineColLoc(loc.cast<NameLoc>().getChildLoc());
-  case Location::Kind::FileLineColLocation:
+  case StandardAttributes::FileLineColLocation:
     return loc.cast<FileLineColLoc>();
-  case Location::Kind::CallSiteLocation:
+  case StandardAttributes::CallSiteLocation:
     // Process the callee of a callsite location.
     return getFileLineColLoc(loc.cast<CallSiteLoc>().getCallee());
   default:
@@ -394,11 +394,11 @@ void SourceMgrDiagnosticHandler::emitDiagnostic(Diagnostic &diag) {
   // stack as well.
   if (auto callLoc = loc.dyn_cast<CallSiteLoc>()) {
     // Print the call stack while valid, or until the limit is reached.
-    Location callerLoc = callLoc->getCaller();
+    Location callerLoc = callLoc.getCaller();
     for (unsigned curDepth = 0; curDepth < callStackLimit; ++curDepth) {
       emitDiagnostic(callerLoc, "called from", DiagnosticSeverity::Note);
       if ((callLoc = callerLoc.dyn_cast<CallSiteLoc>()))
-        callerLoc = callLoc->getCaller();
+        callerLoc = callLoc.getCaller();
       else
         break;
     }

mlir/lib/IR/Location.cpp

@@ -17,58 +17,21 @@
 
 #include "mlir/IR/Location.h"
 #include "LocationDetail.h"
+#include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
-//===----------------------------------------------------------------------===//
-// Location
-//===----------------------------------------------------------------------===//
-
-Location::Kind Location::getKind() const { return loc->kind; }
-
-//===----------------------------------------------------------------------===//
-// FileLineColLoc
-//===----------------------------------------------------------------------===//
-
-FileLineColLoc FileLineColLoc::get(StringRef filename, unsigned line,
-                                   unsigned column, MLIRContext *context) {
-  return get(Identifier::get(filename.empty() ? "-" : filename, context), line,
-             column, context);
-}
-
-StringRef FileLineColLoc::getFilename() const {
-  return static_cast<ImplType *>(loc)->filename;
-}
-unsigned FileLineColLoc::getLine() const {
-  return static_cast<ImplType *>(loc)->line;
-}
-unsigned FileLineColLoc::getColumn() const {
-  return static_cast<ImplType *>(loc)->column;
-}
-
-//===----------------------------------------------------------------------===//
-// NameLoc
-//===----------------------------------------------------------------------===//
-
-NameLoc NameLoc::get(Identifier name, MLIRContext *context) {
-  return get(name, UnknownLoc::get(context), context);
-}
-
-/// Return the name identifier.
-Identifier NameLoc::getName() const {
-  return static_cast<ImplType *>(loc)->name;
-}
-
-/// Return the child location.
-Location NameLoc::getChildLoc() const {
-  return static_cast<ImplType *>(loc)->child;
-}
-
 //===----------------------------------------------------------------------===//
 // CallSiteLoc
 //===----------------------------------------------------------------------===//
 
+CallSiteLoc CallSiteLoc::get(Location callee, Location caller,
+                             MLIRContext *context) {
+  return Base::get(context, StandardAttributes::CallSiteLocation, callee,
+                   caller);
+}
+
 CallSiteLoc CallSiteLoc::get(Location name, ArrayRef<Location> frames,
                              MLIRContext *context) {
   assert(!frames.empty() && "required at least 1 frames");
@@ -78,26 +41,86 @@ CallSiteLoc CallSiteLoc::get(Location name, ArrayRef<Location> frames,
   return CallSiteLoc::get(name, caller, context);
 }
 
-Location CallSiteLoc::getCallee() const {
+Location CallSiteLoc::getCallee() const { return getImpl()->callee; }
-  return static_cast<ImplType *>(loc)->callee;
+
+Location CallSiteLoc::getCaller() const { return getImpl()->caller; }
+
+//===----------------------------------------------------------------------===//
+// FileLineColLoc
+//===----------------------------------------------------------------------===//
+
+FileLineColLoc FileLineColLoc::get(Identifier filename, unsigned line,
+                                   unsigned column, MLIRContext *context) {
+  return Base::get(context, StandardAttributes::FileLineColLocation, filename,
+                   line, column);
 }
 
-Location CallSiteLoc::getCaller() const {
+FileLineColLoc FileLineColLoc::get(StringRef filename, unsigned line,
-  return static_cast<ImplType *>(loc)->caller;
+                                   unsigned column, MLIRContext *context) {
+  return get(Identifier::get(filename.empty() ? "-" : filename, context), line,
+             column, context);
 }
 
+StringRef FileLineColLoc::getFilename() const { return getImpl()->filename; }
+unsigned FileLineColLoc::getLine() const { return getImpl()->line; }
+unsigned FileLineColLoc::getColumn() const { return getImpl()->column; }
+
 //===----------------------------------------------------------------------===//
 // FusedLoc
 //===----------------------------------------------------------------------===//
 
-Location FusedLoc::get(ArrayRef<Location> locs, MLIRContext *context) {
+LocationAttr FusedLoc::get(ArrayRef<Location> locs, Attribute metadata,
-  return get(locs, Attribute(), context);
+                           MLIRContext *context) {
+  // Unique the set of locations to be fused.
+  llvm::SmallSetVector<Location, 4> decomposedLocs;
+  for (auto loc : locs) {
+    // If the location is a fused location we decompose it if it has no
+    // metadata or the metadata is the same as the top level metadata.
+    if (auto fusedLoc = loc.dyn_cast<FusedLoc>()) {
+      if (fusedLoc.getMetadata() == metadata) {
+        // UnknownLoc's have already been removed from FusedLocs so we can
+        // simply add all of the internal locations.
+        decomposedLocs.insert(fusedLoc.getLocations().begin(),
+                              fusedLoc.getLocations().end());
+        continue;
+      }
+    }
+    // Otherwise, only add known locations to the set.
+    if (!loc.isa<UnknownLoc>())
+      decomposedLocs.insert(loc);
+  }
+  locs = decomposedLocs.getArrayRef();
+
+  // Handle the simple cases of less than two locations.
+  if (locs.empty())
+    return UnknownLoc::get(context);
+  if (locs.size() == 1)
+    return locs.front();
+  return Base::get(context, StandardAttributes::FusedLocation, locs, metadata);
 }
 
 ArrayRef<Location> FusedLoc::getLocations() const {
-  return static_cast<ImplType *>(loc)->getLocations();
+  return getImpl()->getLocations();
 }
 
-Attribute FusedLoc::getMetadata() const {
+Attribute FusedLoc::getMetadata() const { return getImpl()->metadata; }
-  return static_cast<ImplType *>(loc)->metadata;
+
+//===----------------------------------------------------------------------===//
+// NameLoc
+//===----------------------------------------------------------------------===//
+
+NameLoc NameLoc::get(Identifier name, Location child, MLIRContext *context) {
+  assert(!child.isa<NameLoc>() &&
+         "a NameLoc cannot be used as a child of another NameLoc");
+  return Base::get(context, StandardAttributes::NameLocation, name, child);
 }
+
+NameLoc NameLoc::get(Identifier name, MLIRContext *context) {
+  return get(name, UnknownLoc::get(context), context);
+}
+
+/// Return the name identifier.
+Identifier NameLoc::getName() const { return getImpl()->name; }
+
+/// Return the child location.
+Location NameLoc::getChildLoc() const { return getImpl()->child; }

mlir/lib/IR/LocationDetail.h

@@ -15,7 +15,7 @@
 // limitations under the License.
 // =============================================================================
 //
-// This holds implementation details of Location.
+// This holds implementation details of the location attributes.
 //
 //===----------------------------------------------------------------------===//
 #ifndef MLIR_IR_LOCATIONDETAIL_H_
@@ -29,60 +29,81 @@
 
 namespace mlir {
 
-class MLIRContext;
-
 namespace detail {
 
-/// Base storage class appearing in a Location.
+struct CallSiteLocationStorage : public AttributeStorage {
-struct alignas(8) LocationStorage {
+  CallSiteLocationStorage(Location callee, Location caller)
-  LocationStorage(Location::Kind kind) : kind(kind) {}
+      : callee(callee), caller(caller) {}
 
-  /// Classification of the subclass, used for type checking.
+  /// The hash key used for uniquing.
-  Location::Kind kind : 8;
+  using KeyTy = std::pair<Location, Location>;
+  bool operator==(const KeyTy &key) const {
+    return key == KeyTy(callee, caller);
+  }
+
+  /// Construct a new storage instance.
+  static CallSiteLocationStorage *
+  construct(AttributeStorageAllocator &allocator, const KeyTy &key) {
+    return new (allocator.allocate<CallSiteLocationStorage>())
+        CallSiteLocationStorage(key.first, key.second);
+  }
+
+  Location callee, caller;
 };
 
-struct UnknownLocationStorage : public LocationStorage {
+struct FileLineColLocationStorage : public AttributeStorage {
-  UnknownLocationStorage() : LocationStorage(Location::Kind::UnknownLocation) {}
-};
-
-struct FileLineColLocationStorage : public LocationStorage {
   FileLineColLocationStorage(Identifier filename, unsigned line,
                              unsigned column)
-      : LocationStorage(Location::Kind::FileLineColLocation),
+      : filename(filename), line(line), column(column) {}
-        filename(filename), line(line), column(column) {}
+
+  /// The hash key used for uniquing.
+  using KeyTy = std::tuple<Identifier, unsigned, unsigned>;
+  bool operator==(const KeyTy &key) const {
+    return key == KeyTy(filename, line, column);
+  }
+
+  /// Construct a new storage instance.
+  static FileLineColLocationStorage *
+  construct(AttributeStorageAllocator &allocator, const KeyTy &key) {
+    return new (allocator.allocate<FileLineColLocationStorage>())
+        FileLineColLocationStorage(std::get<0>(key), std::get<1>(key),
+                                   std::get<2>(key));
+  }
 
   Identifier filename;
   unsigned line, column;
 };
 
-struct NameLocationStorage : public LocationStorage {
-  NameLocationStorage(Identifier name, Location child)
-      : LocationStorage(Location::Kind::NameLocation), name(name),
-        child(child) {}
-
-  Identifier name;
-  Location child;
-};
-
-struct CallSiteLocationStorage : public LocationStorage {
-  CallSiteLocationStorage(Location callee, Location caller)
-      : LocationStorage(Location::Kind::CallSiteLocation), callee(callee),
-        caller(caller) {}
-
-  Location callee, caller;
-};
-
 struct FusedLocationStorage final
-    : public LocationStorage,
+    : public AttributeStorage,
       public llvm::TrailingObjects<FusedLocationStorage, Location> {
   FusedLocationStorage(unsigned numLocs, Attribute metadata)
-      : LocationStorage(Location::Kind::FusedLocation), numLocs(numLocs),
+      : numLocs(numLocs), metadata(metadata) {}
-        metadata(metadata) {}
 
   ArrayRef<Location> getLocations() const {
     return ArrayRef<Location>(getTrailingObjects<Location>(), numLocs);
   }
 
+  /// The hash key used for uniquing.
+  using KeyTy = std::pair<ArrayRef<Location>, Attribute>;
+  bool operator==(const KeyTy &key) const {
+    return key == KeyTy(getLocations(), metadata);
+  }
+
+  /// Construct a new storage instance.
+  static FusedLocationStorage *construct(AttributeStorageAllocator &allocator,
+                                         const KeyTy &key) {
+    ArrayRef<Location> locs = key.first;
+
+    auto byteSize = totalSizeToAlloc<Location>(locs.size());
+    auto rawMem = allocator.allocate(byteSize, alignof(FusedLocationStorage));
+    auto result = new (rawMem) FusedLocationStorage(locs.size(), key.second);
+
+    std::uninitialized_copy(locs.begin(), locs.end(),
+                            result->getTrailingObjects<Location>());
+    return result;
+  }
+
   // This stuff is used by the TrailingObjects template.
   friend llvm::TrailingObjects<FusedLocationStorage, Location>;
   size_t numTrailingObjects(OverloadToken<Location>) const { return numLocs; }
@@ -94,6 +115,26 @@ struct FusedLocationStorage final
   Attribute metadata;
 };
 
+struct NameLocationStorage : public AttributeStorage {
+  NameLocationStorage(Identifier name, Location child)
+      : name(name), child(child) {}
+
+  /// The hash key used for uniquing.
+  using KeyTy = std::pair<Identifier, Location>;
+  bool operator==(const KeyTy &key) const { return key == KeyTy(name, child); }
+
+  /// Construct a new storage instance.
+  static NameLocationStorage *construct(AttributeStorageAllocator &allocator,
+                                        const KeyTy &key) {
+    return new (allocator.allocate<NameLocationStorage>())
+        NameLocationStorage(key.first, key.second);
+  }
+
+  Identifier name;
+  Location child;
+};
+
 } // end namespace detail
 } // end namespace mlir
+
 #endif // MLIR_IR_LOCATIONDETAIL_H_

mlir/lib/IR/MLIRContext.cpp

@@ -75,61 +75,6 @@ static ValueT safeGetOrCreate(DenseSet<ValueT, DenseInfoT> &container,
   return *existing.first = constructorFn();
 }
 
-/// A utility function to thread-safely get or create a uniqued instance within
-/// the given vector container.
-template <typename ValueT, typename ConstructorFn>
-ValueT safeGetOrCreate(std::vector<ValueT> &container, unsigned position,
-                       llvm::sys::SmartRWMutex<true> &mutex,
-                       ConstructorFn &&constructorFn) {
-  { // Check for an existing instance in read-only mode.
-    llvm::sys::SmartScopedReader<true> lock(mutex);
-    if (container.size() > position && container[position])
-      return container[position];
-  }
-
-  // Aquire a writer-lock so that we can safely create the new instance.
-  llvm::sys::SmartScopedWriter<true> lock(mutex);
-
-  // Check if we need to resize.
-  if (position >= container.size())
-    container.resize(position + 1, nullptr);
-
-  // Check for an existing instance again here, because another writer thread
-  // may have already created one.
-  auto *&result = container[position];
-  if (result)
-    return result;
-
-  return result = constructorFn();
-}
-
-/// A utility function to safely get or create a uniqued instance within the
-/// given map container.
-template <typename ContainerTy, typename KeyT, typename ConstructorFn>
-static typename ContainerTy::mapped_type
-safeGetOrCreate(ContainerTy &container, KeyT &&key,
-                llvm::sys::SmartRWMutex<true> &mutex,
-                ConstructorFn &&constructorFn) {
-  { // Check for an existing instance in read-only mode.
-    llvm::sys::SmartScopedReader<true> instanceLock(mutex);
-    auto it = container.find(key);
-    if (it != container.end())
-      return it->second;
-  }
-
-  // Aquire a writer-lock so that we can safely create the new instance.
-  llvm::sys::SmartScopedWriter<true> instanceLock(mutex);
-
-  // Check for an existing instance again here, because another writer thread
-  // may have already created one.
-  auto *&result = container[key];
-  if (result)
-    return result;
-
-  // Otherwise, construct a new instance of the value.
-  return result = constructorFn();
-}
-
 namespace {
 /// A builtin dialect to define types/etc that are necessary for the validity of
 /// the IR.
@@ -139,6 +84,8 @@ struct BuiltinDialect : public Dialect {
                   DictionaryAttr, FloatAttr, FunctionAttr, IntegerAttr,
                   IntegerSetAttr, OpaqueAttr, OpaqueElementsAttr,
                   SparseElementsAttr, StringAttr, TypeAttr, UnitAttr>();
+    addAttributes<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc, UnknownLoc>();
+
     addTypes<ComplexType, FloatType, FunctionType, IndexType, IntegerType,
              MemRefType, NoneType, OpaqueType, RankedTensorType, TupleType,
              UnrankedTensorType, VectorType>();
@@ -200,49 +147,6 @@ struct IntegerSetKeyInfo : DenseMapInfo<IntegerSet> {
                                   rhs.getConstraints(), rhs.getEqFlags());
   }
 };
-
-struct CallSiteLocationKeyInfo : DenseMapInfo<CallSiteLocationStorage *> {
-  // Call locations are uniqued based on their held concret location
-  // and the caller location.
-  using KeyTy = std::pair<Location, Location>;
-  using DenseMapInfo<CallSiteLocationStorage *>::isEqual;
-
-  static unsigned getHashValue(CallSiteLocationStorage *key) {
-    return getHashValue(KeyTy(key->callee, key->caller));
-  }
-
-  static unsigned getHashValue(KeyTy key) {
-    return hash_combine(key.first, key.second);
-  }
-
-  static bool isEqual(const KeyTy &lhs, const CallSiteLocationStorage *rhs) {
-    if (rhs == getEmptyKey() || rhs == getTombstoneKey())
-      return false;
-    return lhs == std::make_pair(rhs->callee, rhs->caller);
-  }
-};
-
-struct FusedLocKeyInfo : DenseMapInfo<FusedLocationStorage *> {
-  // Fused locations are uniqued based on their held locations and an optional
-  // metadata attribute.
-  using KeyTy = std::pair<ArrayRef<Location>, Attribute>;
-  using DenseMapInfo<FusedLocationStorage *>::isEqual;
-
-  static unsigned getHashValue(FusedLocationStorage *key) {
-    return getHashValue(KeyTy(key->getLocations(), key->metadata));
-  }
-
-  static unsigned getHashValue(KeyTy key) {
-    return hash_combine(hash_combine_range(key.first.begin(), key.first.end()),
-                        key.second);
-  }
-
-  static bool isEqual(const KeyTy &lhs, const FusedLocationStorage *rhs) {
-    if (rhs == getEmptyKey() || rhs == getTombstoneKey())
-      return false;
-    return lhs == std::make_pair(rhs->getLocations(), rhs->metadata);
-  }
-};
 } // end anonymous namespace.
 
 namespace mlir {
@@ -250,32 +154,6 @@ namespace mlir {
 /// This class is completely private to this file, so everything is public.
 class MLIRContextImpl {
 public:
-  //===--------------------------------------------------------------------===//
-  // Location uniquing
-  //===--------------------------------------------------------------------===//
-
-  // Location allocator and mutex for thread safety.
-  llvm::BumpPtrAllocator locationAllocator;
-  llvm::sys::SmartRWMutex<true> locationMutex;
-
-  /// The singleton for UnknownLoc.
-  UnknownLocationStorage theUnknownLoc;
-
-  /// FileLineColLoc uniquing.
-  DenseMap<std::tuple<const char *, unsigned, unsigned>,
-           FileLineColLocationStorage *>
-      fileLineColLocs;
-
-  /// NameLocation uniquing.
-  DenseMap<const char *, NameLocationStorage *> nameLocs;
-
-  /// CallLocation uniquing.
-  DenseSet<CallSiteLocationStorage *, CallSiteLocationKeyInfo> callLocs;
-
-  /// FusedLoc uniquing.
-  using FusedLocations = DenseSet<FusedLocationStorage *, FusedLocKeyInfo>;
-  FusedLocations fusedLocs;
-
   //===--------------------------------------------------------------------===//
   // Identifier uniquing
   //===--------------------------------------------------------------------===//
@@ -350,6 +228,7 @@ public:
   /// Cached Attribute Instances.
   BoolAttr falseAttr, trueAttr;
   UnitAttr unitAttr;
+  UnknownLoc unknownLocAttr;
 
 public:
   MLIRContextImpl() : identifiers(identifierAllocator) {}
@@ -397,6 +276,9 @@ MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) {
   /// Unit Attribute.
   impl->unitAttr =
       AttributeUniquer::get<UnitAttr>(this, StandardAttributes::Unit);
+  /// Unknown Location Attribute.
+  impl->unknownLocAttr = AttributeUniquer::get<UnknownLoc>(
+      this, StandardAttributes::UnknownLocation);
 }
 
 MLIRContext::~MLIRContext() {}
@@ -597,95 +479,6 @@ Identifier Identifier::get(StringRef str, MLIRContext *context) {
   return Identifier(it->getKeyData());
 }
 
-//===----------------------------------------------------------------------===//
-// Location uniquing
-//===----------------------------------------------------------------------===//
-
-UnknownLoc UnknownLoc::get(MLIRContext *context) {
-  return &context->getImpl().theUnknownLoc;
-}
-
-FileLineColLoc FileLineColLoc::get(Identifier filename, unsigned line,
-                                   unsigned column, MLIRContext *context) {
-  auto &impl = context->getImpl();
-
-  // Safely get or create a location instance.
-  auto key = std::make_tuple(filename.data(), line, column);
-  return safeGetOrCreate(impl.fileLineColLocs, key, impl.locationMutex, [&] {
-    return new (impl.locationAllocator.Allocate<FileLineColLocationStorage>())
-        FileLineColLocationStorage(filename, line, column);
-  });
-}
-
-NameLoc NameLoc::get(Identifier name, Location child, MLIRContext *context) {
-  auto &impl = context->getImpl();
-  assert(!child.isa<NameLoc>() &&
-         "a NameLoc cannot be used as a child of another NameLoc");
-
-  // Safely get or create a location instance.
-  return safeGetOrCreate(impl.nameLocs, name.data(), impl.locationMutex, [&] {
-    return new (impl.locationAllocator.Allocate<NameLocationStorage>())
-        NameLocationStorage(name, child);
-  });
-}
-
-CallSiteLoc CallSiteLoc::get(Location callee, Location caller,
-                             MLIRContext *context) {
-  auto &impl = context->getImpl();
-
-  // Safely get or create a location instance.
-  auto key = std::make_pair(callee, caller);
-  return safeGetOrCreate(impl.callLocs, key, impl.locationMutex, [&] {
-    return new (impl.locationAllocator.Allocate<CallSiteLocationStorage>())
-        CallSiteLocationStorage(callee, caller);
-  });
-}
-
-Location FusedLoc::get(ArrayRef<Location> locs, Attribute metadata,
-                       MLIRContext *context) {
-  // Unique the set of locations to be fused.
-  llvm::SmallSetVector<Location, 4> decomposedLocs;
-  for (auto loc : locs) {
-    // If the location is a fused location we decompose it if it has no
-    // metadata or the metadata is the same as the top level metadata.
-    if (auto fusedLoc = loc.dyn_cast<FusedLoc>()) {
-      if (fusedLoc->getMetadata() == metadata) {
-        // UnknownLoc's have already been removed from FusedLocs so we can
-        // simply add all of the internal locations.
-        decomposedLocs.insert(fusedLoc->getLocations().begin(),
-                              fusedLoc->getLocations().end());
-        continue;
-      }
-    }
-    // Otherwise, only add known locations to the set.
-    if (!loc.isa<UnknownLoc>())
-      decomposedLocs.insert(loc);
-  }
-  locs = decomposedLocs.getArrayRef();
-
-  // Handle the simple cases of less than two locations.
-  if (locs.empty())
-    return UnknownLoc::get(context);
-  if (locs.size() == 1)
-    return locs.front();
-
-  auto &impl = context->getImpl();
-
-  // Safely get or create a location instance.
-  auto key = std::make_pair(locs, metadata);
-  return safeGetOrCreate(impl.fusedLocs, key, impl.locationMutex, [&] {
-    auto byteSize =
-        FusedLocationStorage::totalSizeToAlloc<Location>(locs.size());
-    auto rawMem = impl.locationAllocator.Allocate(
-        byteSize, alignof(FusedLocationStorage));
-    auto result = new (rawMem) FusedLocationStorage(locs.size(), metadata);
-
-    std::uninitialized_copy(locs.begin(), locs.end(),
-                            result->getTrailingObjects<Location>());
-    return result;
-  });
-}
-
 //===----------------------------------------------------------------------===//
 // Type uniquing
 //===----------------------------------------------------------------------===//
@@ -799,6 +592,10 @@ UnitAttr UnitAttr::get(MLIRContext *context) {
   return context->getImpl().unitAttr;
 }
 
+UnknownLoc UnknownLoc::get(MLIRContext *context) {
+  return context->getImpl().unknownLocAttr;
+}
+
 //===----------------------------------------------------------------------===//
 // AffineMap uniquing
 //===----------------------------------------------------------------------===//

mlir/lib/IR/StandardTypes.cpp

@@ -280,6 +280,32 @@ LogicalResult UnrankedTensorType::verifyConstructionInvariants(
 // MemRefType
 //===----------------------------------------------------------------------===//
 
+/// Get or create a new MemRefType based on shape, element type, affine
+/// map composition, and memory space.  Assumes the arguments define a
+/// well-formed MemRef type.  Use getChecked to gracefully handle MemRefType
+/// construction failures.
+MemRefType MemRefType::get(ArrayRef<int64_t> shape, Type elementType,
+                           ArrayRef<AffineMap> affineMapComposition,
+                           unsigned memorySpace) {
+  auto result = getImpl(shape, elementType, affineMapComposition, memorySpace,
+                        /*location=*/llvm::None);
+  assert(result && "Failed to construct instance of MemRefType.");
+  return result;
+}
+
+/// Get or create a new MemRefType based on shape, element type, affine
+/// map composition, and memory space declared at the given location.
+/// If the location is unknown, the last argument should be an instance of
+/// UnknownLoc.  If the MemRefType defined by the arguments would be
+/// ill-formed, emits errors (to the handler registered with the context or to
+/// the error stream) and returns nullptr.
+MemRefType MemRefType::getChecked(ArrayRef<int64_t> shape, Type elementType,
+                                  ArrayRef<AffineMap> affineMapComposition,
+                                  unsigned memorySpace, Location location) {
+  return getImpl(shape, elementType, affineMapComposition, memorySpace,
+                 location);
+}
+
 /// Get or create a new MemRefType defined by the arguments.  If the resulting
 /// type would be ill-formed, return nullptr.  If the location is provided,
 /// emit detailed error messages.  To emit errors when the location is unknown,

mlir/lib/Parser/Parser.cpp

@@ -243,10 +243,10 @@ public:
   //===--------------------------------------------------------------------===//
 
   /// Parse an inline location.
-  ParseResult parseLocation(llvm::Optional<Location> *loc);
+  ParseResult parseLocation(LocationAttr &loc);
 
   /// Parse a raw location instance.
-  ParseResult parseLocationInstance(llvm::Optional<Location> *loc);
+  ParseResult parseLocationInstance(LocationAttr &loc);
 
   /// Parse an optional trailing location.
   ///
@@ -259,10 +259,10 @@ public:
       return success();
 
     // Parse the location.
-    llvm::Optional<Location> directLoc;
+    LocationAttr directLoc;
-    if (parseLocation(&directLoc))
+    if (parseLocation(directLoc))
       return failure();
-    owner->setLoc(*directLoc);
+    owner->setLoc(directLoc);
     return success();
   }
 
@@ -1010,6 +1010,12 @@ Attribute Parser::parseAttribute(Type type) {
     return builder.getFunctionAttr(nameStr.drop_front());
   }
 
+  // Parse a location attribute.
+  case Token::kw_loc: {
+    LocationAttr attr;
+    return failed(parseLocation(attr)) ? Attribute() : attr;
+  }
+
   // Parse an opaque elements attribute.
   case Token::kw_opaque:
     return parseOpaqueElementsAttr();
@@ -1461,13 +1467,10 @@ Attribute Parser::parseSparseElementsAttr() {
 ///   location           ::= `loc` inline-location
 ///   inline-location    ::= '(' location-inst ')'
 ///
-ParseResult Parser::parseLocation(llvm::Optional<Location> *loc) {
+ParseResult Parser::parseLocation(LocationAttr &loc) {
-  assert(loc && "loc is expected to be non-null");
-
   // Check for 'loc' identifier.
-  if (getToken().isNot(Token::kw_loc))
+  if (parseToken(Token::kw_loc, "expected 'loc' keyword"))
-    return emitError("expected location keyword");
+    return emitError();
-  consumeToken(Token::kw_loc);
 
   // Parse the inline-location.
   if (parseToken(Token::l_paren, "expected '(' in inline location") ||
@@ -1492,7 +1495,7 @@ ParseResult Parser::parseLocation(llvm::Optional<Location> *loc) {
 ///                    '[' location-inst (location-inst ',')* ']'
 /// unknown-location ::= 'unknown'
 ///
-ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
+ParseResult Parser::parseLocationInstance(LocationAttr &loc) {
   auto *ctx = getContext();
 
   // Handle either name or filelinecol locations.
@@ -1522,8 +1525,7 @@ ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
         return emitError("expected integer column number in FileLineColLoc");
       consumeToken(Token::integer);
 
-      auto file = Identifier::get(str, ctx);
+      loc = FileLineColLoc::get(str, line.getValue(), column.getValue(), ctx);
-      *loc = FileLineColLoc::get(file, line.getValue(), column.getValue(), ctx);
       return success();
     }
 
@@ -1534,22 +1536,22 @@ ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
       auto childSourceLoc = getToken().getLoc();
 
       // Parse the child location.
-      llvm::Optional<Location> childLoc;
+      LocationAttr childLoc;
-      if (parseLocationInstance(&childLoc))
+      if (parseLocationInstance(childLoc))
         return failure();
 
       // The child must not be another NameLoc.
-      if (childLoc->isa<NameLoc>())
+      if (childLoc.isa<NameLoc>())
         return emitError(childSourceLoc,
                          "child of NameLoc cannot be another NameLoc");
-      *loc = NameLoc::get(Identifier::get(str, ctx), *childLoc, ctx);
+      loc = NameLoc::get(Identifier::get(str, ctx), childLoc, ctx);
 
       // Parse the closing ')'.
       if (parseToken(Token::r_paren,
                      "expected ')' after child location of NameLoc"))
         return failure();
     } else {
-      *loc = NameLoc::get(Identifier::get(str, ctx), ctx);
+      loc = NameLoc::get(Identifier::get(str, ctx), ctx);
     }
 
     return success();
@@ -1559,7 +1561,7 @@ ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
   if (getToken().is(Token::bare_identifier) &&
       getToken().getSpelling() == "unknown") {
     consumeToken(Token::bare_identifier);
-    *loc = UnknownLoc::get(ctx);
+    loc = UnknownLoc::get(ctx);
     return success();
   }
 
@@ -1580,30 +1582,22 @@ ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
         return failure();
     }
 
-    // Parse the '['.
-    if (parseToken(Token::l_square, "expected '[' in fused location"))
-      return failure();
-
-    // Parse the internal locations.
     llvm::SmallVector<Location, 4> locations;
-    do {
+    auto parseElt = [&] {
-      llvm::Optional<Location> newLoc;
+      LocationAttr newLoc;
-      if (parseLocationInstance(&newLoc))
+      if (parseLocationInstance(newLoc))
         return failure();
-      locations.push_back(*newLoc);
+      locations.push_back(newLoc);
+      return success();
+    };
 
-      // Parse the ','.
+    if (parseToken(Token::l_square, "expected '[' in fused location") ||
-    } while (consumeIf(Token::comma));
+        parseCommaSeparatedList(parseElt) ||
-
+        parseToken(Token::r_square, "expected ']' in fused location"))
-    // Parse the ']'.
-    if (parseToken(Token::r_square, "expected ']' in fused location"))
       return failure();
 
     // Return the fused location.
-    if (metadata)
+    loc = FusedLoc::get(locations, metadata, getContext());
-      *loc = FusedLoc::get(locations, metadata, getContext());
-    else
-      *loc = FusedLoc::get(locations, ctx);
     return success();
   }
 
@@ -1617,8 +1611,8 @@ ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
       return failure();
 
     // Parse the callee location.
-    llvm::Optional<Location> calleeLoc;
+    LocationAttr calleeLoc;
-    if (parseLocationInstance(&calleeLoc))
+    if (parseLocationInstance(calleeLoc))
       return failure();
 
     // Parse the 'at'.
@@ -1628,8 +1622,8 @@ ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
     consumeToken(Token::bare_identifier);
 
     // Parse the caller location.
-    llvm::Optional<Location> callerLoc;
+    LocationAttr callerLoc;
-    if (parseLocationInstance(&callerLoc))
+    if (parseLocationInstance(callerLoc))
       return failure();
 
     // Parse the ')'.
@@ -1637,7 +1631,7 @@ ParseResult Parser::parseLocationInstance(llvm::Optional<Location> *loc) {
       return failure();
 
     // Return the callsite location.
-    *loc = CallSiteLoc::get(*calleeLoc, *callerLoc, ctx);
+    loc = CallSiteLoc::get(calleeLoc, callerLoc, ctx);
     return success();
   }
 

mlir/test/IR/parser.mlir

@@ -933,3 +933,6 @@ func @scoped_names() {
   }) : () -> ()
   return
 }
+
+// CHECK-LABEL: func @loc_attr(i1 {foo.loc_attr: loc(callsite("foo" at "mysource.cc":10:8))})
+func @loc_attr(i1 {foo.loc_attr: loc(callsite("foo" at "mysource.cc":10:8))})