[mlir] Add `OpAsmPrinter::printOptionalLocationSpecifier`

This is the corresponding method to
`OpAsmParser::parseOptionalLocationSpecifier` that prints a location
`loc(...)` based on the op printing flags. Together, these two functions
allow propagating user-level location info outside of their usual spots.

Reviewed By: rriddle

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

mlir/include/mlir/IR/OpImplementation.h

@@ -327,6 +327,9 @@ public:
   using AsmPrinter::AsmPrinter;
   ~OpAsmPrinter() override;
 
+  /// Print a loc(...) specifier if printing debug info is enabled.
+  virtual void printOptionalLocationSpecifier(Location loc) = 0;
+
   /// Print a newline and indent the printer to the start of the current
   /// operation.
   virtual void printNewline() = 0;

mlir/lib/IR/AsmPrinter.cpp

@@ -531,6 +531,11 @@ private:
     return success();
   }
 
+  /// Consider the given location to be printed for an alias.
+  void printOptionalLocationSpecifier(Location loc) override {
+    printAttribute(loc);
+  }
+
   /// Print the given set of attributes with names not included within
   /// 'elidedAttrs'.
   void printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
@@ -2682,6 +2687,12 @@ public:
   // OpAsmPrinter methods
   //===--------------------------------------------------------------------===//
 
+  /// Print a loc(...) specifier if printing debug info is enabled. Locations
+  /// may be deferred with an alias.
+  void printOptionalLocationSpecifier(Location loc) override {
+    printTrailingLocation(loc);
+  }
+
   /// Print a newline and indent the printer to the start of the current
   /// operation.
   void printNewline() override {

mlir/test/IR/locations.mlir

@@ -82,3 +82,10 @@ func.func @location_name_child_is_name() {
 
 // CHECK-ALIAS: #[[LOC]] = loc("out_of_line_location")
 #loc = loc("out_of_line_location")
+
+// CHECK-LABEL: @optional_location_specifier
+// CHECK: test.attr_with_loc("foo" loc("foo_loc"))
+func.func @optional_location_specifier() {
+  test.attr_with_loc("foo" loc("foo_loc"))
+  return
+}

mlir/test/lib/Dialect/Test/TestDialect.cpp

@@ -1030,6 +1030,26 @@ void PolyForOp::getAsmBlockArgumentNames(Region &region,
   }
 }
 
+//===----------------------------------------------------------------------===//
+// TestAttrWithLoc - parse/printOptionalLocationSpecifier
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseOptionalLoc(OpAsmParser &p, Attribute &loc) {
+  Optional<Location> result;
+  SMLoc sourceLoc = p.getCurrentLocation();
+  if (p.parseOptionalLocationSpecifier(result))
+    return failure();
+  if (result)
+    loc = *result;
+  else
+    loc = p.getEncodedSourceLoc(sourceLoc);
+  return success();
+}
+
+static void printOptionalLoc(OpAsmPrinter &p, Operation *op, Attribute loc) {
+  p.printOptionalLocationSpecifier(loc.cast<LocationAttr>());
+}
+
 //===----------------------------------------------------------------------===//
 // Test removing op with inner ops.
 //===----------------------------------------------------------------------===//

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

@@ -1911,6 +1911,12 @@ def PolyForOp : TEST_Op<"polyfor", [OpAsmOpInterface]> {
   let hasCustomAssemblyFormat = 1;
 }
 
+def TestAttrWithLoc : TEST_Op<"attr_with_loc"> {
+  let summary = "op's attribute has a location";
+  let arguments = (ins AnyAttr:$loc, AnyAttr:$value);
+  let assemblyFormat = "`(` $value `` custom<OptionalLoc>($loc) `)` attr-dict";
+}
+
 //===----------------------------------------------------------------------===//
 // Test OpAsmInterface.