[mlir] Add the ability to define dialect-specific location attrs.

mlir/include/mlir/IR/Attributes.h

@@ -322,12 +322,19 @@ class AttributeInterface
 // Core AttributeTrait
 //===----------------------------------------------------------------------===//
 
-/// This trait is used to determine if an attribute is mutable or not. It is
-/// attached on an attribute if the corresponding ImplType defines a `mutate`
-/// function with proper signature.
 namespace AttributeTrait {
+/// This trait is used to determine if an attribute is mutable or not. It is
+/// attached on an attribute if the corresponding ConcreteType defines a
+/// `mutate` function with proper signature.
 template <typename ConcreteType>
 using IsMutable = detail::StorageUserTrait::IsMutable<ConcreteType>;
+
+/// This trait is used to determine if an attribute is a location or not. It is
+/// attached to an attribute by the user if they intend the attribute to be used
+/// as a location.
+template <typename ConcreteType>
+struct IsLocation : public AttributeTrait::TraitBase<ConcreteType, IsLocation> {
+};
 } // namespace AttributeTrait
 
 } // namespace mlir.

mlir/include/mlir/IR/BuiltinLocationAttributes.td

@@ -18,7 +18,8 @@ include "mlir/IR/BuiltinDialect.td"
 
 // Base class for Builtin dialect location attributes.
 class Builtin_LocationAttr<string name, list<Trait> traits = []>
-    : AttrDef<Builtin_Dialect, name, traits, "::mlir::LocationAttr"> {
+    : AttrDef<Builtin_Dialect, name, traits # [NativeAttrTrait<"IsLocation">],
+              "::mlir::LocationAttr"> {
   let cppClassName = name;
   let mnemonic = ?;
 }

mlir/lib/AsmParser/LocationParser.cpp

@@ -153,6 +153,29 @@ ParseResult Parser::parseNameOrFileLineColLocation(LocationAttr &loc) {
   return success();
 }
 
+ParseResult Parser::parseDialectLocation(LocationAttr &loc) {
+  consumeToken(Token::bare_identifier);
+
+  if (parseToken(Token::less,
+                 "expected `<` to start dialect location attribute"))
+    return failure();
+
+  Attribute locAttr = parseAttribute(Type{});
+  // No attribute parsed, someone else has returned an error already.
+  if (!locAttr)
+    return failure();
+
+  loc = llvm::dyn_cast<LocationAttr>(locAttr);
+  if (!loc)
+    return emitError() << "expected a location attribute";
+
+  if (parseToken(Token::greater,
+                 "expected `>` to end dialect location attribute"))
+    return failure();
+
+  return success();
+}
+
 ParseResult Parser::parseLocationInstance(LocationAttr &loc) {
   // Handle aliases.
   if (getToken().is(Token::hash_identifier)) {
@@ -187,5 +210,8 @@ ParseResult Parser::parseLocationInstance(LocationAttr &loc) {
     return success();
   }
 
+  if (getToken().getSpelling() == "dialect")
+    return parseDialectLocation(loc);
+
   return emitWrongTokenError("expected location instance");
 }

mlir/lib/AsmParser/Parser.h

@@ -303,6 +303,9 @@ class Parser {
   /// Parse a name or FileLineCol location instance.
   ParseResult parseNameOrFileLineColLocation(LocationAttr &loc);
 
+  /// Parse a dialect-specific location.
+  ParseResult parseDialectLocation(LocationAttr &loc);
+
   //===--------------------------------------------------------------------===//
   // Affine Parsing
   //===--------------------------------------------------------------------===//

mlir/lib/IR/AsmPrinter.cpp

@@ -2061,6 +2061,12 @@ void AsmPrinter::Impl::printLocationInternal(LocationAttr loc, bool pretty,
             [&](Location loc) { printLocationInternal(loc, pretty); },
             [&]() { os << ", "; });
         os << ']';
+      })
+      .Default([&](LocationAttr loc) {
+        // Assumes that this is a dialect-specific attribute.
+        os << "dialect<";
+        printAttribute(loc);
+        os << ">";
       });
 }
 

mlir/lib/IR/Location.cpp

@@ -64,8 +64,7 @@ WalkResult LocationAttr::walk(function_ref<WalkResult(Location)> walkFn) {
 
 /// Methods for support type inquiry through isa, cast, and dyn_cast.
 bool LocationAttr::classof(Attribute attr) {
-  return llvm::isa<CallSiteLoc, FileLineColLoc, FusedLoc, NameLoc, OpaqueLoc,
-                   UnknownLoc>(attr);
+  return attr.hasTrait<AttributeTrait::IsLocation>();
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/IR/locations.mlir

@@ -89,3 +89,10 @@ func.func @optional_location_specifier() {
   test.attr_with_loc("foo" loc("foo_loc"))
   return
 }
+
+// CHECK-LABEL: @dialect_location
+// CHECK: test.attr_with_loc("dialectLoc" loc(dialect<#test.custom_location<"foo.mlir" * 32>>))
+func.func @dialect_location() {
+  test.attr_with_loc("dialectLoc" loc(dialect<#test.custom_location<"foo.mlir"*32>>))
+  return
+}

mlir/test/IR/pretty-locations.mlir

@@ -24,6 +24,9 @@ func.func @inline_notation() -> i32 {
   affine.if #set0(%2) {
   } loc(fused<"myPass">["foo", "foo2"])
 
+  // CHECK: "foo.op"() : () -> () dialect<#test.custom_location<"foo.mlir" * 1234>>
+  "foo.op"() : () -> () loc(dialect<#test.custom_location<"foo.mlir" * 1234>>)
+
   // CHECK: return %0 : i32 [unknown]
   return %1 : i32 loc(unknown)
 }

mlir/test/lib/Dialect/Test/TestAttrDefs.td

@@ -377,4 +377,15 @@ def NestedPolynomialAttr2 : Test_Attr<"NestedPolynomialAttr2"> {
 }
 
 
+// Test custom location handling.
+def TestCustomLocationAttr
+    : Test_Attr<"TestCustomLocation", [NativeAttrTrait<"IsLocation">]> {
+  let mnemonic = "custom_location";
+  let parameters = (ins "mlir::StringAttr":$file, "unsigned":$line);
+
+  // Choose a silly separator token so we know it's hitting this code path
+  // and not another.
+  let assemblyFormat = "`<` $file `*` $line `>`";
+}
+
 #endif // TEST_ATTRDEFS