Add support for custom op parser/printer hooks to know about result names.

[lattner]

Summary

This allows the custom parser/printer hooks to do interesting things with
the SSA names. This patch:

• Adds a new 'getResultName' method to OpAsmParser that allows a parser implementation to get information about its result names.
• Adds a OpAsmPrinter::printOperand overload that takes an explicit stream.
• Adds a test.string_attr_pretty_name operation that uses these hooks to do fancy things with the result name.

[lattner]

Ok, build 53570 looks like it has two unrelated failures, but the parser.mlir failure does look like a problem. This:

/mnt/disks/ssd0/agent/workspace/amd64_debian_testing_clang8/mlir/test/IR/parser.mlir:1198:16: error: CHECK-NOT: excluded string found in input
// CHECK-NOT: attributes

^

<stdin>:564:36: note: found here
%x = test.string_attr_pretty_name attributes {name = "x"}

^~~~~~~~~~

Failure means that the logic in "print(OpAsmPrinter &p, StringAttrPrettyNameOp op) " is not working. I don't see how it can fail given that the result name is printing correctly as 'x'. The logic in question is this:

if (tmpStream.str().drop_front() == op.name())

elidedAttrs = {"name"};

Does anyone have any ideas? Needless to say, this works fine on my machine (a mac).

[lattner]

Change to raw_string_ostream to see if it fixes the linux builder, also
add some debug info to catch the problem if it doesn't work.

[lattner]

Remove the test logic. This patch is good to review.

[lattner]

Fix a clang tidy warning, fix the memory issue that caused it to fail on some builders not others.

[lattner]

Ok, this patch is good to go, @rriddle can you take a look when you get a chance?

[mehdi_amini]

mlir/lib/Parser/Parser.cpp

4020

I don't get this loop, in the example you git in the comment above: %x, %y:2, %z = foo.op

I expect to loop over { {"x", 1}, {"y", 2}, {"z", 1} }, but your test if (resultNo < std::get<1>(entry)) { will compare against 1, 2, and 1 ; so getResultName(2) and getResultName(3) would never match, what do I miss?

[lattner]

mlir/lib/Parser/Parser.cpp

4020

You're not missing anything, this is a bug! Thank you!

[lattner]

Fix a bug that Mehdi noticed in patch review in getResultName(). Expand
the testcase to handle and show the multiresult case explicitly, testing
the fix. Add a new OpAsmParser::getNumResults() method.

The later allows ops with variadic results to know how many they have, which
seems like a widely useful thing. IIRC, something in TFRT had to have an
explicit result count passed in as an attribute because we didn't have this
before.

[lattner]

Thank you again so much for noticing this bug Mehdi. I've expanded out the test dialect op to be more interesting (allowing testing the multi-result case). I believe this is ready for review again, thanks!

[rriddle]

Sorry for the delay. Looks good, just some nits

mlir/include/mlir/IR/OpImplementation.h

250

nit: sub-element

250

typo: empty string and

264

nit: getResultName

mlir/lib/IR/AsmPrinter.cpp

1994

nit: Remove the semi-colon here.

2206

nit: is it cleaner to just inline as streamOverride ? streamOverride : &os?

mlir/lib/Parser/Parser.cpp

4019

nit: auto -> StringRef

mlir/test/lib/TestDialect/TestDialect.cpp

399

nit: Remove the extra newline here.

412

llvm::any_of(result.attributes, [](NamedAttribute attr) { return attr.first.is("names"); })
?

420

nit: Can we early exit here instead?

422

You could collect a vector of StringRefs and use Builder::getStrArrayAttr

428

nit: Drop trivial braces

[lattner]

Incorporate all the great suggestions from River's review.

[lattner]

Thank you so much for the fantastic review River! I've incorporated all of the suggested changes. I'll wait for Mehdi's ack before landing since he expressed concerns on the discourse thread. Thanks!

[mehdi_amini]

mlir/lib/Parser/Parser.cpp

4033

FYI: clang-tidy flags this and suggests a const auto &entry here.

[lattner]

Patch landed here, thanks!
https://github.com/llvm/llvm-project/commit/596da62d21ede197dd95eca5146d2ddf0497275c

mlir/include/mlir/IR/OpImplementation.h

 class OpAsmPrinter {
 public:
   OpAsmPrinter() {}
   virtual ~OpAsmPrinter();
   virtual raw_ostream &getStream() const = 0;
 
   /// Print implementations for various things an operation contains.
   virtual void printOperand(Value value) = 0;
+  virtual void printOperand(Value value, raw_ostream &os) = 0;
 
   /// Print a comma separated list of operands.
   template <typename ContainerType>
   void printOperands(const ContainerType &container) {
     printOperands(container.begin(), container.end());
   }
 
   /// Print a comma separated list of operands.
public:
   /// Get the location of the next token and store it into the argument.  This
   /// always succeeds.
   virtual llvm::SMLoc getCurrentLocation() = 0;
   ParseResult getCurrentLocation(llvm::SMLoc *loc) {
     *loc = getCurrentLocation();
     return success();
   }
 
+  /// Return the name of the specified result in the specified syntax, as well
+  /// as the subelement in the name.  It returns an empty string an ~0U for
+  /// invalid result numbers.  For example, in this operation:
+  ///
+  ///  %x, %y:2, %z = foo.op
+  ///
+  ///    getResultName(0) == {"x", 0 }
+  ///    getResultName(1) == {"y", 0 }
+  ///    getResultName(2) == {"y", 1 }
+  ///    getResultName(3) == {"z", 0 }
+  ///    getResultName(4) == {"", ~0U }
+  virtual std::pair<StringRef, unsigned>
+  getResultName(unsigned resultNo) const = 0;
+
   /// Return the location of the original name token.
   virtual llvm::SMLoc getNameLoc() const = 0;
 
   // These methods emit an error and return failure or success. This allows
   // these to be chained together into a linear sequence of || expressions in
   // many cases.
 
   /// Parse an operation in its generic form.
   /// The parsed operation is parsed in the current context and inserted in the
   /// provided block and insertion point. The results produced by this operation

mlir/lib/IR/AsmPrinter.cpp

public:
 
   /// Print the given block. If 'printBlockArgs' is false, the arguments of the
   /// block are not printed. If 'printBlockTerminator' is false, the terminator
   /// operation of the block is not printed.
   void print(Block *block, bool printBlockArgs = true,
              bool printBlockTerminator = true);
 
   /// Print the ID of the given value, optionally with its result number.
-  void printValueID(Value value, bool printResultNo = true) const;
+  void printValueID(Value value, bool printResultNo = true,
+                    raw_ostream *streamOverride = nullptr) const;
 
   //===--------------------------------------------------------------------===//
   // OpAsmPrinter methods
   //===--------------------------------------------------------------------===//
 
   /// Return the current stream of the printer.
   raw_ostream &getStream() const override { return os; }
 
   /// Print the given type.
   void printType(Type type) override { ModulePrinter::printType(type); }
 
   /// Print the given attribute.
   void printAttribute(Attribute attr) override {
     ModulePrinter::printAttribute(attr);
   }
 
   /// Print the given attribute without its type. The corresponding parser must
   /// provide a valid type for the attribute.
   void printAttributeWithoutType(Attribute attr) override {
     ModulePrinter::printAttribute(attr, AttrTypeElision::Must);
   }
 
   /// Print the ID for the given value.
   void printOperand(Value value) override { printValueID(value); }
+  void printOperand(Value value, raw_ostream &os) override {
+    printValueID(value, /*printResultNo=*/true, &os);
+  };
 
   /// Print an optional attribute dictionary with a given set of elided values.
   void printOptionalAttrDict(ArrayRef<NamedAttribute> attrs,
                              ArrayRef<StringRef> elidedAttrs = {}) override {
     ModulePrinter::printOptionalAttrDict(attrs, elidedAttrs);
   }
   void printOptionalAttrDictWithKeyword(
       ArrayRef<NamedAttribute> attrs,
auto range = llvm::make_range(
       std::prev(block->getOperations().end(), printBlockTerminator ? 0 : 1));
   for (auto &op : range) {
     print(&op);
     os << newLine;
   }
   currentIndent -= indentWidth;
 }
 
-void OperationPrinter::printValueID(Value value, bool printResultNo) const {
-  state->getSSANameState().printValueID(value, printResultNo, os);
+void OperationPrinter::printValueID(Value value, bool printResultNo,
+                                    raw_ostream *streamOverride) const {
+  if (!streamOverride)
+    streamOverride = &os;
+  state->getSSANameState().printValueID(value, printResultNo, *streamOverride);
 }
 
 void OperationPrinter::printSuccessor(Block *successor) {
   printBlockName(successor);
 }
 
 void OperationPrinter::printSuccessorAndUseList(Block *successor,
                                                 ValueRange succOperands) {

mlir/lib/Parser/Parser.cpp

public:
   /// Parse an operation instance that is in the generic form.
   Operation *parseGenericOperation();
 
   /// Parse an operation instance that is in the generic form and insert it at
   /// the provided insertion point.
   Operation *parseGenericOperation(Block *insertBlock,
                                    Block::iterator insertPt);
 
+  /// This is the structure of a result specifier in the assembly syntax,
+  /// including the name, number of results, and location.
+  typedef std::tuple<StringRef, unsigned, SMLoc> ResultRecord;
+
   /// Parse an operation instance that is in the op-defined custom form.
-  Operation *parseCustomOperation();
+  /// resultInfo specifies information about the "%name =" specifiers.
+  Operation *parseCustomOperation(ArrayRef<ResultRecord> resultInfo);
 
   //===--------------------------------------------------------------------===//
   // Region Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a region into 'region' with the provided entry block arguments.
   /// 'isIsolatedNameScope' indicates if the naming scope of this region is
   /// isolated from those above.

 ///                        successor-list? (`(` region-list `)`)?
 ///                        attribute-dict? `:` function-type
 ///  custom-operation  ::= bare-id custom-operation-format
 ///  op-result-list    ::= op-result (`,` op-result)* `=`
 ///  op-result         ::= ssa-id (`:` integer-literal)
 ///
 ParseResult OperationParser::parseOperation() {
   auto loc = getToken().getLoc();
-  SmallVector<std::tuple<StringRef, unsigned, SMLoc>, 1> resultIDs;
+  SmallVector<ResultRecord, 1> resultIDs;
   size_t numExpectedResults = 0;
   if (getToken().is(Token::percent_identifier)) {
     // Parse the group of result ids.
     auto parseNextResult = [&]() -> ParseResult {
       // Parse the next result id.
       if (!getToken().is(Token::percent_identifier))
         return emitError("expected valid ssa identifier");
 
if (parseCommaSeparatedList(parseNextResult))
       return failure();
 
     if (parseToken(Token::equal, "expected '=' after SSA name"))
       return failure();
   }
 
   Operation *op;
   if (getToken().is(Token::bare_identifier) || getToken().isKeyword())
-    op = parseCustomOperation();
+    op = parseCustomOperation(resultIDs);
   else if (getToken().is(Token::string))
     op = parseGenericOperation();
   else
     return emitError("expected operation name in quotes");
 
   // If parsing of the basic operation failed, then this whole thing fails.
   if (!op)
     return failure();
 
   // If the operation had a name, register it.
   if (!resultIDs.empty()) {
     if (op->getNumResults() == 0)
       return emitError(loc, "cannot name an operation with no results");
     if (numExpectedResults != op->getNumResults())
       return emitError(loc, "operation defines ")
              << op->getNumResults() << " results but was provided "
              << numExpectedResults << " to bind";
 
     // Add definitions for each of the result groups.
     unsigned opResI = 0;
-    for (std::tuple<StringRef, unsigned, SMLoc> &resIt : resultIDs) {
+    for (ResultRecord &resIt : resultIDs) {
       for (unsigned subRes : llvm::seq<unsigned>(0, std::get<1>(resIt))) {
         if (addDefinition({std::get<0>(resIt), subRes, std::get<2>(resIt)},
                           op->getResult(opResI++)))
           return failure();
       }
     }
   }
 
Operation *OperationParser::parseGenericOperation(Block *insertBlock,
   OpBuilder::InsertionGuard restoreInsertionPoint(opBuilder);
   opBuilder.setInsertionPoint(insertBlock, insertPt);
   return parseGenericOperation();
 }
 
 namespace {
 class CustomOpAsmParser : public OpAsmParser {
 public:
-  CustomOpAsmParser(SMLoc nameLoc, const AbstractOperation *opDefinition,
+  CustomOpAsmParser(SMLoc nameLoc,
+                    ArrayRef<OperationParser::ResultRecord> resultIDs,
+                    const AbstractOperation *opDefinition,
                     OperationParser &parser)
-      : nameLoc(nameLoc), opDefinition(opDefinition), parser(parser) {}
+      : nameLoc(nameLoc), resultIDs(resultIDs), opDefinition(opDefinition),
+        parser(parser) {}
 
   /// Parse an instance of the operation described by 'opDefinition' into the
   /// provided operation state.
   ParseResult parseOperation(OperationState &opState) {
     if (opDefinition->parseAssembly(*this, opState))
       return failure();
     return success();
   }
public:
   }
 
   llvm::SMLoc getCurrentLocation() override {
     return parser.getToken().getLoc();
   }
 
   Builder &getBuilder() const override { return parser.builder; }
 
+  /// Return the name of the specified result in the specified syntax, as well
+  /// as the subelement in the name.  For example, in this operation:
+  ///
+  ///  %x, %y:2, %z = foo.op
+  ///
+  ///    getResultName(0) == {"x", 0 }
+  ///    getResultName(1) == {"y", 0 }
+  ///    getResultName(2) == {"y", 1 }
+  ///    getResultName(3) == {"z", 0 }
+  std::pair<StringRef, unsigned>
+  getResultName(unsigned resultNo) const override {
+    // Scan for the resultID that contains this result number.
+    for (unsigned nameID = 0, e = resultIDs.size(); nameID != e; ++nameID) {
+      auto &entry = resultIDs[nameID];
+      if (resultNo < std::get<1>(entry)) {
+        // Don't pass on the leading %.
+        auto name = std::get<0>(entry).drop_front();
+        return {name, resultNo};
+      }
+    }
+
+    // Invalid result number.
+    return {"", ~0U};
+  }
+
   llvm::SMLoc getNameLoc() const override { return nameLoc; }
 
   //===--------------------------------------------------------------------===//
   // Token Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a `->` token.
   ParseResult parseArrow() override {
     return parser.parseToken(Token::arrow, "expected '->'");
   }
 
   /// Parses a `->` if present.
   ParseResult parseOptionalArrow() override {
     return success(parser.consumeIf(Token::arrow));
if (parseLParen())
       return failure();
     return parser.parseCommaSeparatedListUntil(Token::r_paren, parseElt);
   }
 
 private:
   /// The source location of the operation name.
   SMLoc nameLoc;
 
+  /// Information about the result name specifiers.
+  ArrayRef<OperationParser::ResultRecord> resultIDs;
+
   /// The abstract information of the operation.
   const AbstractOperation *opDefinition;
 
   /// The main operation parser.
   OperationParser &parser;
 
   /// A flag that indicates if any errors were emitted during parsing.
   bool emittedError = false;
 };
 } // end anonymous namespace.
 
-Operation *OperationParser::parseCustomOperation() {
+Operation *
+OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
   auto opLoc = getToken().getLoc();
   auto opName = getTokenSpelling();
 
   auto *opDefinition = AbstractOperation::lookup(opName, getContext());
   if (!opDefinition && !opName.contains('.')) {
     // If the operation name has no namespace prefix we treat it as a standard
     // operation and prefix it with "std".
     // TODO: Would it be better to just build a mapping of the registered
llvm::PrettyStackTraceFormat fmt("MLIR Parser: custom op parser '%s'",
                                    opNameStr.c_str());
 
   // Get location information for the operation.
   auto srcLocation = getEncodedSourceLocation(opLoc);
 
   // Have the op implementation take a crack and parsing this.
   OperationState opState(srcLocation, opDefinition->name);
   CleanupOpStateRegions guard{opState};
-  CustomOpAsmParser opAsmParser(opLoc, opDefinition, *this);
+  CustomOpAsmParser opAsmParser(opLoc, resultIDs, opDefinition, *this);
   if (opAsmParser.parseOperation(opState))
     return nullptr;
 
   // If it emitted an error, we failed.
   if (opAsmParser.didEmitError())
     return nullptr;
 
   // Parse a location if one is present.

mlir/test/IR/parser.mlir

func @custom_asm_names() -> (i32, i32, i32, i32, i32, i32, i32) {
   %3, %4 = "test.asm_interface_op"() : () -> (i32, i32)
 
   // CHECK: %[[RESULT:result.*]]
   %5 = "test.asm_dialect_interface_op"() : () -> (i32)
 
   // CHECK: return %[[FIRST]], %[[MIDDLE]]#0, %[[MIDDLE]]#1, %[[LAST]], %[[FIRST_2]], %[[LAST_2]]
   return %0, %1#0, %1#1, %2, %3, %4, %5 : i32, i32, i32, i32, i32, i32, i32
 }
+
+
+// CHECK-LABEL: func @pretty_names
+
+// This tests the behavior
+func @pretty_names() {
+  // Simple case, should parse and print as %x being an implied 'name'
+  // attribute.
+  %x = test.string_attr_pretty_name
+  // CHECK: %x = test.string_attr_pretty_name
+  // CHECK-NOT: attributes
+  
+  // This specifies an explicit name, which should override the result.
+  %YY = test.string_attr_pretty_name attributes { name = "y" }
+  // CHECK: %y = test.string_attr_pretty_name
+  // CHECK-NOT: attributes
+  
+  // Conflicts with the 'y' name, so need an explicit attribute.
+  %0 = "test.string_attr_pretty_name"() { name = "y"} : () -> i32
+  // CHECK: %y_0 = test.string_attr_pretty_name attributes {name = "y"}
+
+  // Name contains a space.
+  %1 = "test.string_attr_pretty_name"() { name = "space name"} : () -> i32
+  // CHECK: %space_name = test.string_attr_pretty_name attributes {name = "space name"}
+
+  "unknown.use"(%x, %YY, %0, %1) : (i32, i32, i32, i32) -> ()
+  return
+}
+

mlir/test/lib/TestDialect/TestDialect.cpp

for (Attribute element : effectsAttr) {
     if (effectElement.get("test_resource"))
       resource = TestResource::get();
 
     effects.emplace_back(effect, value, resource);
   }
 }
 
 //===----------------------------------------------------------------------===//
+// StringAttrPrettyNameOp
+//===----------------------------------------------------------------------===//
+
+// This op has fancy handling of its SSA result name.
+
+static ParseResult parseStringAttrPrettyNameOp(OpAsmParser &parser,
+                                               OperationState &result) {
+  result.addTypes(parser.getBuilder().getIntegerType(32));
+
+  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
+    return failure();
+
+  // If the attribute dictionary contains no 'name' attribute, infer it from
+  // the SSA name (if specified).
+  bool hadName = false;
+  for (auto &attr : result.attributes) {
+    if (attr.first.is("name")) {
+      hadName = true;
+      break;
+    }
+  }
+
+  // If there was no name specified, check to see if there was a useful name
+  // specified in the asm file.
+  if (!hadName) {
+    auto resultName = parser.getResultName(0);
+    if (resultName.second == 0 && !resultName.first.empty() &&
+        !isdigit(resultName.first[0])) {
+      auto *context = result.getContext();
+      auto nameAttr = StringAttr::get(resultName.first, context);
+      result.attributes.push_back({Identifier::get("name", context), nameAttr});
+    }
+  }
+
+  return success();
+}
+
+static void print(OpAsmPrinter &p, StringAttrPrettyNameOp op) {
+  p << "test.string_attr_pretty_name";
+
+  // Note that we only need to print the "name" attribute if the asmprinter
+  // result name disagrees with it.  This can happen in strange cases, e.g.
+  // when there are conflicts.
+  std::string resultNameStr;
+  llvm::raw_string_ostream tmpStream(resultNameStr);
+  p.printOperand(op.getResult(), tmpStream);
+
+  // If the name is the same as we would otherwise use, then we're good!
+  ArrayRef<StringRef> elidedAttrs;
+  if (StringRef(tmpStream.str()).drop_front() == op.name())
+    elidedAttrs = {"name"};
+
+  p.printOptionalAttrDictWithKeyword(op.getAttrs(), elidedAttrs);
+}
+
+// We set the SSA name in the asm syntax to the contents of the name attribute.
+void StringAttrPrettyNameOp::getAsmResultNames(
+    function_ref<void(Value, StringRef)> setNameFn) {
+
+  auto value = name();
+  if (!value.empty())
+    setNameFn(getResult(), value);
+}
+
+//===----------------------------------------------------------------------===//
 // Dialect Registration
 //===----------------------------------------------------------------------===//
 
 // Static initialization for Test dialect registration.
 static mlir::DialectRegistration<mlir::TestDialect> testDialect;
 
 #include "TestOpEnums.cpp.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.cpp.inc"

mlir/test/lib/TestDialect/TestOps.td

def AttrSizedResultOp : TEST_Op<"attr_sized_results",
   let results = (outs
     Variadic<I32>:$a,
     Variadic<I32>:$b,
     I32:$c,
     Variadic<I32>:$d
   );
 }
 
+// This is used to test encoding of a string attribute into an SSA name of a
+// pretty printed value name.
+def StringAttrPrettyNameOp
+ : TEST_Op<"string_attr_pretty_name",
+           [DeclareOpInterfaceMethods<OpAsmOpInterface>]> {
+  let arguments = (ins StrAttr:$name);
+  let results = (outs I32:$r);
+
+  let printer = [{ return ::print(p, *this); }];
+  let parser = [{ return ::parse$cppClass(parser, result); }];
+}
+
 //===----------------------------------------------------------------------===//
 // Test Patterns
 //===----------------------------------------------------------------------===//
 
 def OpA : TEST_Op<"op_a"> {
   let arguments = (ins I32, I32Attr:$attr);
   let results = (outs I32);
 }