[Mlir-commits] [mlir] 3840e96 - [mlir] Add `OpAsmPrinter::printOptionalLocationSpecifier`

[Jeff Niu]

Author: Jeff Niu
Date: 2022-09-29T15:58:10-07:00
New Revision: 3840e960bafc175c830963f8250d90339fac0d16

URL: https://github.com/llvm/llvm-project/commit/3840e960bafc175c830963f8250d90339fac0d16
DIFF: https://github.com/llvm/llvm-project/commit/3840e960bafc175c830963f8250d90339fac0d16.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpImplementation.h
    mlir/lib/IR/AsmPrinter.cpp
    mlir/test/IR/locations.mlir
    mlir/test/lib/Dialect/Test/TestDialect.cpp
    mlir/test/lib/Dialect/Test/TestOps.td

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h
index 42d12efc0491a..fcfd2f57cf919 100644
--- a/mlir/include/mlir/IR/OpImplementation.h
+++ b/mlir/include/mlir/IR/OpImplementation.h
@@ -327,6 +327,9 @@ class OpAsmPrinter : public AsmPrinter {
   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;

diff  --git a/mlir/lib/IR/AsmPrinter.cpp b/mlir/lib/IR/AsmPrinter.cpp
index f3010fd9ee159..aaebef4341889 100644
--- a/mlir/lib/IR/AsmPrinter.cpp
+++ b/mlir/lib/IR/AsmPrinter.cpp
@@ -531,6 +531,11 @@ class DummyAliasOperationPrinter : private OpAsmPrinter {
     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 @@ class OperationPrinter : public AsmPrinter::Impl, private OpAsmPrinter {
   // 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 {

diff  --git a/mlir/test/IR/locations.mlir b/mlir/test/IR/locations.mlir
index 60be67f035fe1..8d7c7e4f13ed4 100644
--- a/mlir/test/IR/locations.mlir
+++ b/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
+}

diff  --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index bf78f80cc8f47..c9d50c44d9dd8 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/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.
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 4df9caa347ed9..d10c2c97709c1 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/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.