[mlir] Update Location to use new casting infra

[nkreeger]

Summary

This allows for using the llvm namespace cast methods instead of the ones on the Location class. The Location class method are kept for now, but we'll want to remove these eventually (with a really long lead time).

Related change: https://reviews.llvm.org/D135870

[rriddle]

mlir/include/mlir/IR/Location.h

187

Do we need the base_of? I don't think anything inherits from Location.

mlir/include/mlir/IR/Location.h

public:
 
   /// Access the impl location attribute.
   operator LocationAttr() const { return impl; }
   LocationAttr *operator->() const { return const_cast<LocationAttr *>(&impl); }
 
   /// Type casting utilities on the underlying location.
   template <typename U>
   bool isa() const {
-    return impl.isa<U>();
+    return llvm::isa<U>(*this);
   }
   template <typename U>
   U dyn_cast() const {
-    return impl.dyn_cast<U>();
+    return llvm::dyn_cast<U>(*this);
   }
   template <typename U>
   U cast() const {
-    return impl.cast<U>();
+    return llvm::cast<U>(*this);
   }
 
   /// Comparison operators.
   bool operator==(Location rhs) const { return impl == rhs.impl; }
   bool operator!=(Location rhs) const { return !(*this == rhs); }
 
   /// Print the location.
   void print(raw_ostream &os) const { impl.print(os); }
public:
   }
   static inline mlir::Location getFromVoidPointer(void *P) {
     return mlir::Location::getFromOpaquePointer(P);
   }
   static constexpr int NumLowBitsAvailable =
       PointerLikeTypeTraits<mlir::Attribute>::NumLowBitsAvailable;
 };
 
+/// The constructors in mlir::Location ensure that the class is a non-nullable
+/// wrapper around mlir::LocationAttr. Override default behavior and always
+/// return true for isPresent().
+template <>
+struct ValueIsPresent<mlir::Location> {
+  using UnwrappedType = mlir::Location;
+  static inline bool isPresent(const mlir::Location &location) { return true; }
+};
+
+/// Add support for llvm style casts. We provide a cast between To and From if
+/// From is mlir::Location or derives from it.
+template <typename To, typename From>
+struct CastInfo<To, From,
+                std::enable_if_t<
+                    std::is_same_v<mlir::Location, std::remove_const_t<From>> ||
+                    std::is_base_of_v<mlir::Location, From>>>
+    : DefaultDoCastIfPossible<To, From, CastInfo<To, From>> {
+
+  static inline bool isPossible(mlir::Location location) {
+    /// Return a constant true instead of a dynamic true when casting to self or
+    /// up the hierarchy. Additionally, all casting info is deferred to the
+    /// wrapped mlir::LocationAttr instance stored in mlir::Location.
+    return std::is_same_v<To, std::remove_const_t<From>> ||
+           isa<To>(static_cast<mlir::LocationAttr>(location));
+  }
+
+  static inline To castFailed() { return To(); }
+
+  static inline To doCast(mlir::Location location) {
+    return To(location->getImpl());
+  }
+};
+
 } // namespace llvm
 
 #endif

mlir/lib/AsmParser/Parser.cpp

if (!forwardRefPlaceholders.empty()) {
     }
     return failure();
   }
 
   // Resolve the locations of any deferred operations.
   auto &attributeAliases = state.symbols.attributeAliasDefinitions;
   auto locID = TypeID::get<DeferredLocInfo *>();
   auto resolveLocation = [&, this](auto &opOrArgument) -> LogicalResult {
-    auto fwdLoc = opOrArgument.getLoc().template dyn_cast<OpaqueLoc>();
+    auto fwdLoc = dyn_cast<OpaqueLoc>(opOrArgument.getLoc());
     if (!fwdLoc || fwdLoc.getUnderlyingTypeID() != locID)
       return success();
     auto locInfo = deferredLocsReferences[fwdLoc.getUnderlyingLocation()];
     Attribute attr = attributeAliases.lookup(locInfo.identifier);
     if (!attr)
       return this->emitError(locInfo.loc)
              << "operation location alias was never defined";
-    auto locAttr = attr.dyn_cast<LocationAttr>();
+    auto locAttr = dyn_cast<LocationAttr>(attr);
     if (!locAttr)
       return this->emitError(locInfo.loc)
              << "expected location, but found '" << attr << "'";
     opOrArgument.setLoc(locAttr);
     return success();
   };
 
   auto walkRes = topLevelOp->walk([&](Operation *op) {
if (identifier.contains('.')) {
     return emitError(tok.getLoc())
            << "expected location, but found dialect attribute: '#" << identifier
            << "'";
   }
 
   // If this alias can be resolved, do it now.
   Attribute attr = state.symbols.attributeAliasDefinitions.lookup(identifier);
   if (attr) {
-    if (!(loc = attr.dyn_cast<LocationAttr>()))
+    if (!(loc = dyn_cast<LocationAttr>(attr)))
       return emitError(tok.getLoc())
              << "expected location, but found '" << attr << "'";
   } else {
     // Otherwise, remember this operation and resolve its location later.
     // In the meantime, use a special OpaqueLoc as a marker.
     loc = OpaqueLoc::get(deferredLocsReferences.size(),
                          TypeID::get<DeferredLocInfo *>(),
                          UnknownLoc::get(getContext()));

mlir/lib/IR/Diagnostics.cpp

loc->walk([&](Location loc) {
     }
     return WalkResult::advance();
   });
   return firstFileLoc;
 }
 
 /// Return a processable CallSiteLoc from the given location.
 static Optional<CallSiteLoc> getCallSiteLoc(Location loc) {
-  if (auto nameLoc = loc.dyn_cast<NameLoc>())
-    return getCallSiteLoc(loc.cast<NameLoc>().getChildLoc());
-  if (auto callLoc = loc.dyn_cast<CallSiteLoc>())
+  if (auto nameLoc = dyn_cast<NameLoc>(loc))
+    return getCallSiteLoc(cast<NameLoc>(loc).getChildLoc());
+  if (auto callLoc = dyn_cast<CallSiteLoc>(loc))
     return callLoc;
-  if (auto fusedLoc = loc.dyn_cast<FusedLoc>()) {
-    for (auto subLoc : loc.cast<FusedLoc>().getLocations()) {
+  if (auto fusedLoc = dyn_cast<FusedLoc>(loc)) {
+    for (auto subLoc : cast<FusedLoc>(loc).getLocations()) {
       if (auto callLoc = getCallSiteLoc(subLoc)) {
         return callLoc;
       }
     }
     return llvm::None;
   }
   return llvm::None;
 }

mlir/lib/IR/Location.cpp

 
 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 (auto fusedLoc = llvm::dyn_cast<FusedLoc>(loc)) {
       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;
       }
     }