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

[rriddle]

Summary

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.

Depends On D90512

mlir/lib/IR/AsmPrinter.cpp

 class DummyAliasOperationPrinter : private OpAsmPrinter {
 public:
   explicit DummyAliasOperationPrinter(const OpPrintingFlags &flags,
                                       AliasInitializer &initializer)
       : printerFlags(flags), initializer(initializer) {}
 
   /// 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()) {
       // Check to see if this is a known operation.  If so, use the registered
       // custom printer hook.
       if (auto *opInfo = op->getAbstractOperation()) {
         opInfo->printAssembly(op, *this);
         return;
enum class AttrTypeElision {
     Must
   };
 
   /// Print the given attribute.
   void printAttribute(Attribute attr,
                       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
   printAffineExpr(AffineExpr expr,
                   function_ref<void(unsigned, bool)> printValueName = nullptr);
   void printAffineConstraint(AffineExpr expr, bool isEq);
   void printIntegerSet(IntegerSet set);
 

 } // end anonymous namespace
 
 void ModulePrinter::printTrailingLocation(Location loc) {
   // Check to see if we are printing debug information.
   if (!printerFlags.shouldPrintDebugInfo())
     return;
 
   os << " ";
-  printLocation(loc);
+  printLocation(loc, /*allowAlias=*/true);
 }
 
 void ModulePrinter::printLocationInternal(LocationAttr loc, bool pretty) {
   TypeSwitch<LocationAttr>(loc)
       .Case<OpaqueLoc>([&](OpaqueLoc loc) {
         printLocationInternal(loc.getFallbackLocation(), pretty);
       })
       .Case<UnknownLoc>([&](UnknownLoc loc) {
static void printFloatValue(const APFloat &apValue, raw_ostream &os) {
   // in the literal.
   SmallVector<char, 16> str;
   APInt apInt = apValue.bitcastToAPInt();
   apInt.toString(str, /*Radix=*/16, /*Signed=*/false,
                  /*formatAsCLiteral=*/true);
   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);
+
+  if (!allowAlias || !state || failed(state->getAliasState().getAlias(loc, os)))
     printLocationInternal(loc);
-    os << ')';
-  }
-}
 
 /// Returns true if the given dialect symbol data is simple enough to print in
 /// the pretty form, i.e. without the enclosing "".
 static bool isDialectSymbolSimpleEnoughForPrettyForm(StringRef symName) {
   // The name must start with an identifier.
   if (symName.empty() || !isalpha(symName.front()))
     return false;
 

mlir/lib/IR/MLIRContext.cpp

LogicalResult getAlias(Attribute attr, raw_ostream &os) const override {
     if (attr.isa<AffineMapAttr>()) {
       os << "map";
       return success();
     }
     if (attr.isa<IntegerSetAttr>()) {
       os << "set";
       return success();
     }
+    if (attr.isa<LocationAttr>()) {
+      os << "loc";
+      return success();
+    }
     return failure();
   }
 };
 
 /// A builtin dialect to define types/etc that are necessary for the validity of
 /// the IR.
 struct BuiltinDialect : public Dialect {
   BuiltinDialect(MLIRContext *context)

mlir/lib/Parser/AttributeParser.cpp

if (getToken().is(Token::floatliteral))
       return parseFloatAttr(type, /*isNegative=*/true);
 
     return (emitError("expected constant integer or floating point value"),
             nullptr);
   }
 
   // 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.
   case Token::kw_opaque:
     return parseOpaqueElementsAttr(type);
 
   // Parse a sparse elements attribute.
   case Token::kw_sparse:

mlir/lib/Parser/LocationParser.cpp

 //===- LocationParser.cpp - MLIR Location Parser  -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Parser.h"
 
 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 |
 ///                   name-location |
 ///                   callsite-location |
 ///                   fused-location |
 ///                   unknown-location
 /// filelinecol-location ::= string-literal ':' integer-literal
ParseResult Parser::parseLocationInstance(LocationAttr &loc) {
   if (getToken().getSpelling() == "unknown") {
     consumeToken(Token::bare_identifier);
     loc = UnknownLoc::get(getContext());
     return success();
   }
 
   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

public:
 
   /// Parse a sparse elements attribute.
   Attribute parseSparseElementsAttr(Type attrType);
 
   //===--------------------------------------------------------------------===//
   // Location Parsing
   //===--------------------------------------------------------------------===//
 
-  /// Parse an inline location.
-  ParseResult parseLocation(LocationAttr &loc);
-
   /// Parse a raw location instance.
   ParseResult parseLocationInstance(LocationAttr &loc);
 
   /// Parse a callsite location instance.
   ParseResult parseCallSiteLocation(LocationAttr &loc);
 
   /// Parse a fused location instance.
   ParseResult parseFusedLocation(LocationAttr &loc);
 
   /// Parse a name or FileLineCol location instance.
   ParseResult parseNameOrFileLineColLocation(LocationAttr &loc);
 
   /// 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
   //===--------------------------------------------------------------------===//
 
   /// Parse a reference to either an affine map, or an integer set.
   ParseResult parseAffineMapOrIntegerSetReference(AffineMap &map,
                                                   IntegerSet &set);

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

-// 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)
   spv.globalVariable @var0 bind(0, 1) : !spv.ptr<f32, Input>
   spv.func @arithmetic(%arg0 : vector<4xf32>, %arg1 : vector<4xf32>) "None" {
     // CHECK: loc({{".*debug.mlir"}}:8:10)
     %0 = spv.FAdd %arg0, %arg1 : vector<4xf32>
     // CHECK: loc({{".*debug.mlir"}}:10:10)

mlir/test/IR/invalid-locations.mlir

 // RUN: mlir-opt %s -split-input-file -verify-diagnostics
 
 // -----
 
 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}}
 }
 
 // -----
 
 func @location_invalid_instance() {
 ^bb:
   return loc() // expected-error {{expected location instance}}
 }

 }
 
 // -----
 
 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

-// 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)>
 
 // CHECK-LABEL: func @inline_notation
 func @inline_notation() -> i32 {
   // CHECK: -> i32 loc("foo")
   %1 = "foo"() : () -> i32 loc("foo")
 
   // CHECK: constant 4 : index loc(callsite("foo" at "mysource.cc":10:8))
   %2 = constant 4 : index loc(callsite("foo" at "mysource.cc":10:8))
 
   // CHECK: } loc(fused["foo", "mysource.cc":10:8])
   affine.for %i0 = 0 to 8 {
   } loc(fused["foo", "mysource.cc":10:8])
 
   // CHECK: } loc(fused<"myPass">["foo", "foo2"])
   affine.if #set0(%2) {
   } loc(fused<"myPass">["foo", "foo2"])
 
   // 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

-// 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 {
 }
 
 // CHECK: module {
 // CHECK-NEXT: }
 module {

mlir/test/IR/opaque_locations.mlir

-// 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)>
 
 // CHECK: MyLocation: 0: 'foo' op
 // CHECK: nullptr: 'foo' op
 // CHECK: MyLocation: 0: 'foo' op
 // CHECK: MyLocation: 1: 'std.constant' op

mlir/test/IR/parser.mlir

func @scoped_names() {
   }, {
     // CHECK: "foo.unknown_op"
     %scoped_name = "foo.unknown_op"() : () -> none
     "foo.terminator"() : () -> ()
   }) : () -> ()
   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
   "foo.unknown_op"() {foo = #foo.attr : i32} : () -> ()
 }
 
 // CHECK-LABEL: @f16_special_values
 func @f16_special_values() {

mlir/test/IR/wrapping_op.mlir

 // 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"
 func @wrapping_op(%arg0 : i32, %arg1 : f32) -> (i3, i2, i1) {
 // CHECK: %0:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3)
 // CHECK-GENERIC: "test.wrapping_region"() ( {
 // CHECK-GENERIC:   %[[NESTED_RES:.*]]:3 = "some.op"(%arg1, %arg0) {test.attr = "attr"} : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
 // CHECK-GENERIC:   "test.return"(%[[NESTED_RES]]#0, %[[NESTED_RES]]#1, %[[NESTED_RES]]#2) : (i1, i2, i3) -> () loc("some_NameLoc")
 // CHECK-GENERIC: }) : () -> (i1, i2, i3) loc("some_NameLoc")
   %res:3 = test.wrapping_region wraps "some.op"(%arg1, %arg0) { test.attr = "attr" } : (f32, i32) -> (i1, i2, i3) loc("some_NameLoc")
   return %res#2, %res#1, %res#0 : i3, i2, i1
 }
 // CHECK-GENERIC: sym_name = "wrapping_op"

mlir/test/Transforms/inlining.mlir

 // 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.
 func @func_with_arg(%c : i32) -> i32 {
   %b = addi %c, %c : i32
   return %b : i32
 }
 

mlir/test/Transforms/location-snapshot.mlir

-// 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]+}})
 // CHECK-NEXT: loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})
 // CHECK-NEXT: } loc("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})
 
 // TAG: func @function(
 // TAG-NEXT: loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
 // TAG-NEXT: loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
 // TAG-NEXT: } loc(fused["original", "tagged"("[[FILE]]":{{[0-9]+}}:{{[0-9]+}})])
 
 func @function() -> i32 {
   %1 = "foo"() : () -> i32 loc("original")
   return %1 : i32 loc("original")
 } loc("original")

mlir/test/Transforms/strip-debuginfo.mlir

-// 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)>
 
 // CHECK-LABEL: func @inline_notation
 func @inline_notation() -> i32 {
   // CHECK: "foo"() : () -> i32 loc(unknown)
   %1 = "foo"() : () -> i32 loc("foo")

mlir/test/mlir-tblgen/pattern.mlir

-// 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 {
   %0 = "test.op_a"(%arg0) {attr = 10 : i32} : (i32) -> i32 loc("a")
   %result = "test.op_a"(%0) {attr = 20 : i32} : (i32) -> i32 loc("b")
 
   // CHECK: "test.op_b"(%arg0) {attr = 10 : i32} : (i32) -> i32 loc("a")
   // CHECK: "test.op_b"(%arg0) {attr = 20 : i32} : (i32) -> i32 loc(fused["b", "a"])