[mlir] Add `OpAsmPrinter::printOptionalLocationSpecifier`

[Mogball]

Summary

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.

[rriddle]

mlir/include/mlir/IR/OpImplementation.h

333

I would drop the allowAlias. Locations printed this way should never be deferred, and should always allow aliases.

mlir/lib/IR/AsmPrinter.cpp

567

This should use printAttribute on the provided location, otherwise this will break if there is a deferred alias.

[Mogball]

review comments

[Mogball]

forgot = 0

[rriddle]

mlir/include/mlir/IR/OpImplementation.h

330–331

I would drop this. The part about aliases isn't completely relevant to users here.

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

1039

Can you move the encoding part to later and just grab the SMLoc here? Encoding involves uniquing the attribute (+ some other stuff), which we should only do if we need to (i.e. in the else branch here). It's good to keep this clean in case people look at it as an example usage.

[Mogball]

review comments

mlir/include/mlir/IR/OpImplementation.h

 
 /// This is a pure-virtual base class that exposes the asmprinter hooks
 /// necessary to implement a custom print() method.
 class OpAsmPrinter : public AsmPrinter {
 public:
   using AsmPrinter::AsmPrinter;
   ~OpAsmPrinter() override;
 
+  /// Print a loc(...) specifier if printing debug info is enabled. Locations
+  /// may be deferred with an alias.
+  virtual void printOptionalLocationSpecifier(Location loc,
+                                              bool allowAlias = true) = 0;
+
   /// Print a newline and indent the printer to the start of the current
   /// operation.
   virtual void printNewline() = 0;
 
   /// Print a block argument in the usual format of:
   ///   %ssaName : type {attr1=42} loc("here")
   /// where location printing is controlled by the standard internal option.
   /// You may pass omitType=true to not print a type, and pass an empty

mlir/lib/IR/AsmPrinter.cpp

private:
   raw_ostream &getStream() const override { return os; }
 
   /// The following are hooks of `OpAsmPrinter` that are not necessary for
   /// determining potential aliases.
   void printFloat(const APFloat &value) override {}
   void printAffineMapOfSSAIds(AffineMapAttr, ValueRange) override {}
   void printAffineExprOfSSAIds(AffineExpr, ValueRange, ValueRange) override {}
   void printNewline() override {}
+  void printOptionalLocationSpecifier(Location loc, bool allowAlias) override {}
   void printOperand(Value) override {}
   void printOperand(Value, raw_ostream &os) override {
     // Users expect the output string to have at least the prefixed % to signal
     // a value name. To maintain this invariant, emit a name even if it is
     // guaranteed to go unused.
     os << "%";
   }
   void printKeywordOrString(StringRef) override {}
public:
   /// Print the ID of the given operation.
   void printOperationID(Operation *op,
                         raw_ostream *streamOverride = nullptr) const;
 
   //===--------------------------------------------------------------------===//
   // OpAsmPrinter methods
   //===--------------------------------------------------------------------===//
 
+  /// Print a loc(...) specifier if printing debug info is enabled. Locations
+  /// may be deferred with an alias.
+  void printOptionalLocationSpecifier(Location loc, bool allowAlias) override {
+    printTrailingLocation(loc, allowAlias);
+  }
+
   /// Print a newline and indent the printer to the start of the current
   /// operation.
   void printNewline() override {
     os << newLine;
     os.indent(currentIndent);
   }
 
   /// Print a block argument in the usual format of:

mlir/test/IR/locations.mlir

 // CHECK-LABEL: func @location_name_child_is_name
 func.func @location_name_child_is_name() {
   // CHECK: "foo"("foo")
   return loc("foo"("foo"))
 }
 
 // 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

void PolyForOp::getAsmBlockArgumentNames(Region &region,
   auto e = std::min(arrayAttr.size(), args.size());
   for (unsigned i = 0; i < e; ++i) {
     if (auto strAttr = arrayAttr[i].dyn_cast<StringAttr>())
       setNameFn(args[i], strAttr.getValue());
   }
 }
 
 //===----------------------------------------------------------------------===//
+// TestAttrWithLoc - parse/printOptionalLocationSpecifier
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseOptionalLoc(OpAsmParser &p, Attribute &loc) {
+  Optional<Location> result;
+  Location sourceLoc = p.getEncodedSourceLoc(p.getCurrentLocation());
+  if (p.parseOptionalLocationSpecifier(result))
+    return failure();
+  if (result)
+    loc = *result;
+  else
+    loc = sourceLoc;
+  return success();
+}
+
+static void printOptionalLoc(OpAsmPrinter &p, Operation *op, Attribute loc) {
+  p.printOptionalLocationSpecifier(loc.cast<LocationAttr>());
+}
+
+//===----------------------------------------------------------------------===//
 // Test removing op with inner ops.
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct TestRemoveOpWithInnerOps
     : public OpRewritePattern<TestOpWithRegionPattern> {
   using OpRewritePattern<TestOpWithRegionPattern>::OpRewritePattern;
 

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

def PolyForOp : TEST_Op<"polyfor", [OpAsmOpInterface]> {
   let extraClassDeclaration = [{
     void getAsmBlockArgumentNames(mlir::Region &region,
                                   mlir::OpAsmSetValueNameFn setNameFn);
   }];
   let regions = (region SizedRegion<1>:$region);
   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.
 
 def AsmInterfaceOp : TEST_Op<"asm_interface_op"> {
   let results = (outs AnyType:$first, Variadic<AnyType>:$middle_results,
                       AnyType);
 }