[mlir][Asm] Add support for using an alias for trailing operation locations

Locations often get very long and clutter up operations when printed inline with them. This revision adds support for using aliases with trailing operation locations, and makes printing with aliases the default behavior. Aliases in the trailing location take the form `loc(<alias>)`, such as `loc(#loc0)`. As with all aliases, using `mlir-print-local-scope` can be used to disable them and get the inline behavior.

Differential Revision: https://reviews.llvm.org/D90652

Author: River707

mlir/lib/IR/AsmPrinter.cpp

@@ -289,7 +289,9 @@ class DummyAliasOperationPrinter : private OpAsmPrinter {
 
   /// Print the given operation.
   void print(Operation *op) {
-    // TODO: Consider the operation location for an alias.
+    // Visit the operation location.
+    if (printerFlags.shouldPrintDebugInfo())
+      initializer.visit(op->getLoc());
 
     // If requested, always print the generic form.
     if (!printerFlags.shouldPrintGenericOpForm()) {
@@ -1089,7 +1091,10 @@ class ModulePrinter {
                       AttrTypeElision typeElision = AttrTypeElision::Never);
 
   void printType(Type type);
-  void printLocation(LocationAttr loc);
+
+  /// Print the given location to the stream. If `allowAlias` is true, this
+  /// allows for the internal location to use an attribute alias.
+  void printLocation(LocationAttr loc, bool allowAlias = false);
 
   void printAffineMap(AffineMap map);
   void
@@ -1152,7 +1157,7 @@ void ModulePrinter::printTrailingLocation(Location loc) {
     return;
 
   os << " ";
-  printLocation(loc);
+  printLocation(loc, /*allowAlias=*/true);
 }
 
 void ModulePrinter::printLocationInternal(LocationAttr loc, bool pretty) {
@@ -1269,14 +1274,14 @@ static void printFloatValue(const APFloat &apValue, raw_ostream &os) {
   os << str;
 }
 
-void ModulePrinter::printLocation(LocationAttr loc) {
-  if (printerFlags.shouldPrintDebugInfoPrettyForm()) {
-    printLocationInternal(loc, /*pretty=*/true);
-  } else {
-    os << "loc(";
+void ModulePrinter::printLocation(LocationAttr loc, bool allowAlias) {
+  if (printerFlags.shouldPrintDebugInfoPrettyForm())
+    return printLocationInternal(loc, /*pretty=*/true);
+
+  os << "loc(";
+  if (!allowAlias || !state || failed(state->getAliasState().getAlias(loc, os)))
     printLocationInternal(loc);
-    os << ')';
-  }
+  os << ')';
 }
 
 /// Returns true if the given dialect symbol data is simple enough to print in

mlir/lib/IR/MLIRContext.cpp

@@ -99,6 +99,10 @@ struct BuiltinOpAsmDialectInterface : public OpAsmDialectInterface {
       os << "set";
       return success();
     }
+    if (attr.isa<LocationAttr>()) {
+      os << "loc";
+      return success();
+    }
     return failure();
   }
 };

mlir/lib/Parser/AttributeParser.cpp

@@ -121,8 +121,14 @@ Attribute Parser::parseAttribute(Type type) {
 
   // Parse a location attribute.
   case Token::kw_loc: {
-    LocationAttr attr;
-    return failed(parseLocation(attr)) ? Attribute() : attr;
+    consumeToken(Token::kw_loc);
+
+    LocationAttr locAttr;
+    if (parseToken(Token::l_paren, "expected '(' in inline location") ||
+        parseLocationInstance(locAttr) ||
+        parseToken(Token::r_paren, "expected ')' in inline location"))
+      return Attribute();
+    return locAttr;
   }
 
   // Parse an opaque elements attribute.

mlir/lib/Parser/LocationParser.cpp

@@ -11,24 +11,6 @@
 using namespace mlir;
 using namespace mlir::detail;
 
-/// Parse a location.
-///
-///   location           ::= `loc` inline-location
-///   inline-location    ::= '(' location-inst ')'
-///
-ParseResult Parser::parseLocation(LocationAttr &loc) {
-  // Check for 'loc' identifier.
-  if (parseToken(Token::kw_loc, "expected 'loc' keyword"))
-    return emitError();
-
-  // Parse the inline-location.
-  if (parseToken(Token::l_paren, "expected '(' in inline location") ||
-      parseLocationInstance(loc) ||
-      parseToken(Token::r_paren, "expected ')' in inline location"))
-    return failure();
-  return success();
-}
-
 /// Specific location instances.
 ///
 /// location-inst ::= filelinecol-location |
@@ -195,3 +177,35 @@ ParseResult Parser::parseLocationInstance(LocationAttr &loc) {
 
   return emitError("expected location instance");
 }
+
+ParseResult Parser::parseOptionalTrailingLocation(Location &loc) {
+  // If there is a 'loc' we parse a trailing location.
+  if (!consumeIf(Token::kw_loc))
+    return success();
+  if (parseToken(Token::l_paren, "expected '(' in location"))
+    return failure();
+  Token tok = getToken();
+
+  // Check to see if we are parsing a location alias.
+  LocationAttr directLoc;
+  if (tok.is(Token::hash_identifier)) {
+    // TODO: This should be reworked a bit to allow for resolving operation
+    // locations to aliases after the operation has already been parsed(i.e.
+    // allow post parse location fixups).
+    Attribute attr = parseExtendedAttr(Type());
+    if (!attr)
+      return failure();
+    if (!(directLoc = attr.dyn_cast<LocationAttr>()))
+      return emitError(tok.getLoc()) << "expected location, but found " << attr;
+
+    // Otherwise, we parse the location directly.
+  } else if (parseLocationInstance(directLoc)) {
+    return failure();
+  }
+
+  if (parseToken(Token::r_paren, "expected ')' in location"))
+    return failure();
+
+  loc = directLoc;
+  return success();
+}

mlir/lib/Parser/Parser.h

@@ -231,9 +231,6 @@ class Parser {
   // Location Parsing
   //===--------------------------------------------------------------------===//
 
-  /// Parse an inline location.
-  ParseResult parseLocation(LocationAttr &loc);
-
   /// Parse a raw location instance.
   ParseResult parseLocationInstance(LocationAttr &loc);
 
@@ -248,20 +245,9 @@ class Parser {
 
   /// Parse an optional trailing location.
   ///
-  ///   trailing-location     ::= (`loc` `(` location `)`)?
+  ///   trailing-location ::= (`loc` (`(` location `)` | attribute-alias))?
   ///
-  ParseResult parseOptionalTrailingLocation(Location &loc) {
-    // If there is a 'loc' we parse a trailing location.
-    if (!getToken().is(Token::kw_loc))
-      return success();
-
-    // Parse the location.
-    LocationAttr directLoc;
-    if (parseLocation(directLoc))
-      return failure();
-    loc = directLoc;
-    return success();
-  }
+  ParseResult parseOptionalTrailingLocation(Location &loc);
 
   //===--------------------------------------------------------------------===//
   // Affine Parsing

mlir/test/Dialect/SPIRV/Serialization/debug.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-translate -test-spirv-roundtrip-debug -mlir-print-debuginfo %s | FileCheck %s
+// RUN: mlir-translate -test-spirv-roundtrip-debug -mlir-print-debuginfo -mlir-print-local-scope %s | FileCheck %s
 
 spv.module Logical GLSL450 requires #spv.vce<v1.0, [Shader], []> {
   // CHECK: loc({{".*debug.mlir"}}:5:3)

mlir/test/IR/invalid-locations.mlir

@@ -4,14 +4,14 @@
 
 func @location_missing_l_paren() {
 ^bb:
-  return loc) // expected-error {{expected '(' in inline location}}
+  return loc) // expected-error {{expected '(' in location}}
 }
 
 // -----
 
 func @location_missing_r_paren() {
 ^bb:
-  return loc(unknown // expected-error@+1 {{expected ')' in inline location}}
+  return loc(unknown // expected-error@+1 {{expected ')' in location}}
 }
 
 // -----
@@ -98,3 +98,11 @@ func @location_fused_missing_r_square() {
 ^bb:
   return loc(fused[unknown) // expected-error {{expected ']' in fused location}}
 }
+
+// -----
+
+func @location_invalid_alias() {
+  // expected-error@+1 {{expected location, but found #foo.loc}}
+  return loc(#foo.loc)
+}
+

mlir/test/IR/locations.mlir

@@ -1,4 +1,5 @@
-// RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-debuginfo | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-debuginfo -mlir-print-local-scope | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-debuginfo | FileCheck %s --check-prefix=CHECK-ALIAS
 // This test verifies that debug locations are round-trippable.
 
 #set0 = affine_set<(d0) : (1 == 0)>
@@ -22,3 +23,12 @@ func @inline_notation() -> i32 {
   // CHECK: return %0 : i32 loc(unknown)
   return %1 : i32 loc(unknown)
 }
+
+// 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))})
+
+// CHECK-ALIAS: #[[LOC:.*]] = loc("out_of_line_location")
+#loc = loc("out_of_line_location")
+
+// CHECK-ALIAS: "foo.op"() : () -> () loc(#[[LOC]])
+"foo.op"() : () -> () loc(#loc)

mlir/test/IR/module-op.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -mlir-print-debuginfo | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -split-input-file -mlir-print-debuginfo -mlir-print-local-scope  | FileCheck %s
 
 // CHECK: module {
 module {

mlir/test/IR/opaque_locations.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -allow-unregistered-dialect %s -test-opaque-loc -mlir-print-debuginfo | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -test-opaque-loc -mlir-print-debuginfo -mlir-print-local-scope | FileCheck %s
 // This test verifies that debug opaque locations can be printed.
 
 #set0 = affine_set<(d0) : (1 == 0)>

mlir/test/IR/parser.mlir

@@ -1002,9 +1002,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))})
-
 // CHECK-LABEL: func @dialect_attribute_with_type
 func @dialect_attribute_with_type() {
   // CHECK-NEXT: foo = #foo.attr : i32

mlir/test/IR/wrapping_op.mlir

@@ -1,5 +1,5 @@
 // RUN: mlir-opt -allow-unregistered-dialect %s | FileCheck %s
-// RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic -mlir-print-debuginfo %s | FileCheck %s --check-prefix=CHECK-GENERIC
+// RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic -mlir-print-debuginfo -mlir-print-local-scope %s | FileCheck %s --check-prefix=CHECK-GENERIC
 
 // CHECK-LABEL: func @wrapping_op
 // CHECK-GENERIC: "func"

mlir/test/Transforms/inlining.mlir

@@ -1,5 +1,5 @@
 // RUN: mlir-opt %s -inline="disable-simplify" | FileCheck %s
-// RUN: mlir-opt %s -inline="disable-simplify" -mlir-print-debuginfo | FileCheck %s --check-prefix INLINE-LOC
+// RUN: mlir-opt %s -inline="disable-simplify" -mlir-print-debuginfo -mlir-print-local-scope | FileCheck %s --check-prefix INLINE-LOC
 // RUN: mlir-opt %s -inline | FileCheck %s --check-prefix INLINE_SIMPLIFY
 
 // Inline a function that takes an argument.

mlir/test/Transforms/location-snapshot.mlir

@@ -1,5 +1,5 @@
-// RUN: mlir-opt -allow-unregistered-dialect -snapshot-op-locations='filename=%/t' -mlir-print-debuginfo %s | FileCheck %s -DFILE=%/t
-// RUN: mlir-opt -allow-unregistered-dialect -snapshot-op-locations='filename=%/t tag='tagged'' -mlir-print-debuginfo %s | FileCheck %s --check-prefix=TAG -DFILE=%/t
+// RUN: mlir-opt -allow-unregistered-dialect -snapshot-op-locations='filename=%/t' -mlir-print-local-scope -mlir-print-debuginfo %s | FileCheck %s -DFILE=%/t
+// RUN: mlir-opt -allow-unregistered-dialect -snapshot-op-locations='filename=%/t tag='tagged'' -mlir-print-local-scope -mlir-print-debuginfo %s | FileCheck %s --check-prefix=TAG -DFILE=%/t
 
 // CHECK: func @function(
 // CHECK-NEXT: loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})

mlir/test/Transforms/strip-debuginfo.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-debuginfo -strip-debuginfo | FileCheck %s
+// RUN: mlir-opt -allow-unregistered-dialect %s -mlir-print-debuginfo -mlir-print-local-scope -strip-debuginfo | FileCheck %s
 // This test verifies that debug locations are stripped.
 
 #set0 = affine_set<(d0) : (1 == 0)>

mlir/test/mlir-tblgen/pattern.mlir

@@ -1,4 +1,4 @@
-// RUN: mlir-opt -test-patterns -mlir-print-debuginfo %s | FileCheck %s
+// RUN: mlir-opt -test-patterns -mlir-print-debuginfo -mlir-print-local-scope %s | FileCheck %s
 
 // CHECK-LABEL: verifyFusedLocs
 func @verifyFusedLocs(%arg0 : i32) -> i32 {