Support NameLoc and CallSiteLoc for mlir::Location

The NameLoc can be used to represent a variable, node or method. The
CallSiteLoc has two fields, one represents the concrete location and another
one represents the caller's location. Multiple CallSiteLocs can be chained as
a call stack.

For example, the following call stack
```
AAA
at file1:1
at file2:135
at file3:34
```

can be formed by call0:

```
auto name = NameLoc::get("AAA");
auto file1 = FileLineColLoc::get("file1", 1);
auto file2 = FileLineColLoc::get("file2", 135);
auto file3 = FileLineColLoc::get("file3", 34);
auto call2 = CallSiteLoc::get(file2, file3);
auto call1 = CallSiteLoc::get(file1, call2);
auto call0 = CallSiteLoc::get(name, call1);
```

PiperOrigin-RevId: 226941797

mlir/include/mlir/IR/Location.h

@@ -30,12 +30,15 @@ namespace mlir {
 
 class Attribute;
 class MLIRContext;
+class Identifier;
 
 namespace detail {
 
 class LocationStorage;
 class UnknownLocationStorage;
 class FileLineColLocationStorage;
+class NameLocationStorage;
+class CallSiteLocationStorage;
 class FusedLocationStorage;
 
 } // namespace detail
@@ -50,6 +53,12 @@ public:
     /// This represents a file/line/column location.
     FileLineCol,
 
+    /// This represents an identity name, such as variable and function name.
+    Name,
+
+    /// This represents a location as call site or variable usage site. .
+    CallSite,
+
     // 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.
@@ -120,7 +129,7 @@ public:
   static UnknownLoc get(MLIRContext *context);
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool kindof(Location::Kind kind) { return kind == Kind::Unknown; }
+  static bool kindof(Kind kind) { return kind == Kind::Unknown; }
 };
 
 /// This class is used to represent a uniqued filename in an MLIRContext.  It is
@@ -158,7 +167,45 @@ public:
   unsigned getColumn() const;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool kindof(Location::Kind kind) { return kind == Kind::FileLineCol; }
+  static bool kindof(Kind kind) { return kind == Kind::FileLineCol; }
+};
+
+/// Represents an identity name. It is usually the callee of a CallLocation.
+class NameLoc : public Location {
+public:
+  using ImplType = detail::NameLocationStorage;
+  /* implicit */ NameLoc(Location::ImplType *ptr);
+
+  /// Return a uniqued name location object.
+  static NameLoc get(Identifier name, MLIRContext *context);
+
+  Identifier getName() const;
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool kindof(Kind kind) { return kind == Kind::Name; }
+};
+
+/// 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
+/// CallLocations can be chained to form a call stack.
+class CallSiteLoc : public Location {
+public:
+  using ImplType = detail::CallSiteLocationStorage;
+  /* implicit */ CallSiteLoc(Location::ImplType *ptr);
+
+  /// Return a uniqued call location object.
+  static CallSiteLoc get(Location callee, Location caller,
+                         MLIRContext *context);
+
+  /// The concrete location information this object presents.
+  Location getCallee() const;
+
+  /// The caller's location.
+  Location getCaller() const;
+
+  /// Methods for support type inquiry through isa, cast, and dyn_cast.
+  static bool kindof(Kind kind) { return kind == Kind::CallSite; }
 };
 
 /// Represents a value composed of multiple source constructs, with an optional
@@ -182,9 +229,7 @@ public:
   Attribute getMetadata() const;
 
   /// Methods for support type inquiry through isa, cast, and dyn_cast.
-  static bool kindof(Location::Kind kind) {
-    return kind == Kind::FusedLocation;
-  }
+  static bool kindof(Kind kind) { return kind == Kind::FusedLocation; }
 };
 
 // Make Location hashable.

mlir/lib/IR/AsmPrinter.cpp

@@ -1683,6 +1683,22 @@ void Location::print(raw_ostream &os) const {
        << fileLoc.getColumn();
     break;
   }
+  case Kind::Name: {
+    auto nameLoc = cast<NameLoc>();
+    os << nameLoc.getName();
+    break;
+  }
+  case Kind::CallSite: {
+    auto callLocation = cast<CallSiteLoc>();
+    auto callee = callLocation.getCallee();
+    auto caller = callLocation.getCaller();
+    callee.print(os);
+    if (caller.isa<CallSiteLoc>()) {
+      os << "\n at ";
+    }
+    caller.print(os);
+    break;
+  }
   case Kind::FusedLocation: {
     auto fusedLoc = cast<FusedLoc>();
     if (auto metadata = fusedLoc.getMetadata())

mlir/lib/IR/Location.cpp

@@ -37,6 +37,22 @@ unsigned FileLineColLoc::getColumn() const {
   return static_cast<ImplType *>(loc)->column;
 }
 
+NameLoc::NameLoc(Location::ImplType *ptr) : Location(ptr) {}
+
+Identifier NameLoc::getName() const {
+  return static_cast<ImplType *>(loc)->name;
+}
+
+CallSiteLoc::CallSiteLoc(Location::ImplType *ptr) : Location(ptr) {}
+
+Location CallSiteLoc::getCallee() const {
+  return static_cast<ImplType *>(loc)->callee;
+}
+
+Location CallSiteLoc::getCaller() const {
+  return static_cast<ImplType *>(loc)->caller;
+}
+
 FusedLoc::FusedLoc(Location::ImplType *ptr) : Location(ptr) {}
 
 ArrayRef<Location> FusedLoc::getLocations() const {

mlir/lib/IR/LocationDetail.h

@@ -22,6 +22,7 @@
 #define MLIR_IR_LOCATIONDETAIL_H_
 
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/Identifier.h"
 #include "mlir/IR/Location.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/Support/TrailingObjects.h"
@@ -47,6 +48,14 @@ struct FileLineColLocationStorage : public LocationStorage {
   const unsigned line, column;
 };
 
+struct NameLocationStorage : public LocationStorage {
+  const Identifier name;
+};
+
+struct CallSiteLocationStorage : public LocationStorage {
+  const Location callee, caller;
+};
+
 struct FusedLocationStorage final
     : public LocationStorage,
       public llvm::TrailingObjects<FusedLocationStorage, Location> {

mlir/lib/IR/MLIRContext.cpp

@@ -265,6 +265,27 @@ struct OpaqueElementsAttrInfo : DenseMapInfo<OpaqueElementsAttributeStorage *> {
   }
 };
 
+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.
@@ -308,6 +329,12 @@ public:
            FileLineColLocationStorage *>
       fileLineColLocs;
 
+  /// NameLocation uniquing.
+  DenseMap<const char *, NameLocationStorage *> nameLocs;
+
+  /// CallLocation uniquing.
+  DenseSet<CallSiteLocationStorage *, CallSiteLocationKeyInfo> callLocs;
+
   /// FusedLoc uniquing.
   using FusedLocations = DenseSet<FusedLocationStorage *, FusedLocKeyInfo>;
   FusedLocations fusedLocs;
@@ -594,6 +621,39 @@ FileLineColLoc FileLineColLoc::get(UniquedFilename filename, unsigned line,
   return entry;
 }
 
+NameLoc NameLoc::get(Identifier name, MLIRContext *context) {
+  auto &impl = context->getImpl();
+  auto &entry = impl.nameLocs[name.data()];
+  if (!entry) {
+    entry = impl.allocator.Allocate<NameLocationStorage>();
+    new (entry) NameLocationStorage{{Location::Kind::Name}, name};
+  }
+
+  return entry;
+}
+
+CallSiteLoc CallSiteLoc::get(Location callee, Location caller,
+                             MLIRContext *context) {
+  auto &impl = context->getImpl();
+
+  // Look to see if the fused location has been created already.
+  auto existing =
+      impl.callLocs.insert_as(nullptr, std::make_pair(callee, caller));
+
+  // If it has been created, return it.
+  if (!existing.second)
+    return *existing.first;
+
+  // On the first use, we allocate them into the bump pointer.
+  auto *result = impl.allocator.Allocate<detail::CallSiteLocationStorage>();
+
+  // Initialize the memory using placement new.
+  new (result) detail::CallSiteLocationStorage{
+      {Location::Kind::CallSite}, callee, caller};
+
+  return *existing.first = result;
+}
+
 Location FusedLoc::get(ArrayRef<Location> locs, MLIRContext *context) {
   return get(locs, Attribute(), context);
 }