[mlir] Refactoring a few Parser APIs

[sdasgup3]

Summary

Refactored two new parser APIs parseGenericOperationAfterOperands and
parseCustomOperationName out of parseGenericOperation and parseCustomOperation.

Motivation: Sometimes an op can be printed in a special way if certain criteria
is met. While parsing, we need to handle all the versions.
parseGenericOperationAfterOperands is handy in situation where we already
parsed the operands and decide to fall back to default parsing.

parseCustomOperationName is useful when we need to know details (dialect,
operation name etc.) about a parsed token meant to be an mlir operation.

[rriddle]

mlir/lib/Parser/Parser.cpp

313–316

This feels like much too specific of an API, what is the exact contract you are going for here? Trying to inject pre-parsed bits into the generic operation parser?

351–356

Why does this need to expose the dialect and abstract operation?

[sdasgup3]

1. Updating D113719: [mlir] Refactoring a few Parser APIs #

Fixed an issue with the API CustomOpAsmParser::parseCustomOperationName exposing
the AbstractOperation.

[sdasgup3]

mlir/lib/Parser/Parser.cpp

313–316

Thanks for the comments!

Most of the custom printers prints the op-name and operands differently than the generic printer. e.g. the op-name w/o double quotes and operands w/o parenthesis. The op specific custom-parsers parses those accordingly. The above API is useful to parse beyond the operands if the op-specific parser decides to fall back to generic parsing.

One such scenario is provided in the example TestDialect, where after the operands are parsed, the custom-op parser can fallback to generic parsing and calls the above API. That way we can avoid calling the individual parsing methods like parseRegion, resolveSSAUse etc. and instead reuse a part of parseGenericOperation.

I find this API handy while implementing the "pretty-printing for the mhlo reduce-op".

Please let me know your thought?

351–356

Thanks for pointing out!
I agree that it does not make sense for OpAsmParser::parseCustomOperationName to expose the internal artifacts,
Although OperationParser::parseCustomOperationName needs to expose such internal details as it's caller OperationParser::parseCustomOperation needs that.

[sdasgup3]

1. Updating D113719: [mlir] Refactoring a few Parser APIs #

Fixing build failures.

[rriddle]

mlir/lib/Parser/Parser.cpp

313–316

Right, I think passing in pre-parsed things makes sense, but I would like to avoid having N parseGenericOperationAfter* methods for each combination of components. Could we just expose one method that allows for optionally taking all of the components, and we just parse the things not explicitly passed in?

351–356

Right, what we expose here doesn't need to be the same as what the internal parser needs. Here we care more about the API being exposed. Like, I think here we could even return FailureOr<OperationName> if that was desirable.

[sdasgup3]

1. Updating D113719: [mlir] Refactoring a few Parser APIs #

Refactored the APIs as per the following review comments.

1. The API parseGenericOperationAfterName allows for optionally taking all of

the components, and we just parse the things not explicitly passed in.

2. parseCustomOperationName returns FailureOr<std::pair<std::string, Dialect

*>> to accomodate the failure scenario.

[rriddle]

mlir/include/mlir/IR/OpImplementation.h

913

Can you switch this to OperationName now? A recent refactoring has made that a better choice.

933–937

Can you add llvm::None to reduce the number of forced parameters? (We could also possibly switch to a struct if it gets too crazy in practice, but this is fine for now).

mlir/lib/Parser/Parser.cpp

320–324

Can we add default params(llvm::None) for these?

1005

You shouldn't need the .hasValue() on any of these.

1060–1094

Can we just assign parsedFnType in the first if, and sink the rest of the logic out of the conditional? (i.e. similar to operandUseInfo above) e.g.

if (!parsedFnType) {
   ...
   parsedFnType = fnType;
}

result.addTypes(parsedFnType->getResults());

// Check that we have the right number of types for the operands.
...

1226–1233

Can you add a comment here? Realistically we shouldn't need this, we should make SSAUseInfo a typedef of OperandType. (You don't need to do that, I can do that in a followup).

1666

Can you add a move here to avoid copying the string? Either that or switch this to StringRef?

[mehdi_amini]

mlir/lib/Parser/Parser.cpp

351–353

The doc does not match the implementation: this isn't a pair but a tuple.

(I suspect that returning an FailureOr<OperationName> as River mentioned should be enough anyway)

1216

Missing std::move here as well to avoid string copy

[sdasgup3]

1. Updating D113719: [mlir] Refactoring a few Parser APIs #

Addressed all the comments except the one related to leveraging a recent change
https://reviews.llvm.org/D114049. Waiting of the change to sync in the
repository I am working on.

[sdasgup3]

mlir/lib/Parser/Parser.cpp

1060–1094

Thanks for pointing this out as I am also not happy with this piece of code.

As you found that in both if else, we have some repeated code after the comment "// Check that we have the right number of types for the operands." .
The only difference is the call to emitError: In the if part, we are leveraging the extra location information typeLoc, which is not available if the argument parsedFnType is explicitly provided.
I wish if I could use some location info to use in the else part as well. That way we can refactor-out the then-becoming-common-code out of the conditional.

Let me know about any suggestions.

1226–1233

I agree with you and added the comment.
Actually these the two types: one internal to parser implementation (OperationParser::SSAUseInfo {name, number, loc}) and the other exposed to clients of parser (OpAsmParser::OperandType {location, name , number}), both have the same members but in different order. As you said, It would be much cleaner if one could be the alias of other.

[rriddle]

mlir/lib/Parser/Parser.cpp

1060–1094

I think for now you could just use result.location as the location in the else.

[sdasgup3]

1. Updating D113719: [mlir] Refactoring a few Parser APIs #

Changed the return type of pasreCustomOperationName to OperationName.

mlir/lib/Parser/Parser.cpp

1060–1094

Done.
Just for FYI:

I used mlir::emitError(mlir::Location, ...) version instead of the Parser::emitError(llvm::SMLoc, ...) to accommodate the fact that result.location is of type mlir::Location. The decision is made because

1. There is no conversion routine available from llvm::SMLoc to mlir::Location and
2. There exist no version like Parser::emitError(mlir::Location. ...)

[sdasgup3]

mlir/lib/Parser/Parser.cpp

1216

This is not applicable after the recent code change.

[rriddle]

mlir/include/mlir/IR/OpImplementation.h

910–912

nit: I don't think the description of OperationName is necessary here.

mlir/lib/Parser/Parser.cpp

1701

opNameInfo.getIdentifier().c_str()?

That would avoid a string copy.

1705

Can we pass in the OperationName (opNameInfo) here instead of opName? (it's just going to re-lookup the OperationName otherwise).

[sdasgup3]

1. Updating D113719: [mlir] Refactoring a few Parser APIs #

Addressed the review comments.

mlir/lib/Parser/Parser.cpp

1701

Agreed! changed it to opNameInfo.getIdentifier().data()

[mehdi_amini]

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

747

Please avoid auto when the type isn't obvious from the context (or not spell-able/readable).

[sdasgup3]

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

747

Agreed. Thanks for the fix.

mlir/include/mlir/IR/OpImplementation.h

public:
   /// 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
   /// aren't mapped to any named value in the parser. Returns nullptr on
   /// failure.
   virtual Operation *parseGenericOperation(Block *insertBlock,
                                            Block::iterator insertPt) = 0;
 
+  /// Parse the name of an operation in the custom form. On success, return a
+  /// pair containing name of the operation and pointer, possibly nullptr, to
+  /// the corresponding dialect. Otherwise, failure is returned.
+  using ParsedOpNameAndDialect = std::pair<std::string, Dialect *>;
+  virtual FailureOr<ParsedOpNameAndDialect> parseCustomOperationName() = 0;
+
   //===--------------------------------------------------------------------===//
   // Operand Parsing
   //===--------------------------------------------------------------------===//
 
   /// This is the representation of an operand reference.
   struct OperandType {
     llvm::SMLoc location; // Location of the token.
     StringRef name;       // Value name, e.g. %42 or %abc
     unsigned number;      // Number, e.g. 12 for an operand like %xyz#12
   };
 
+  /// Parse different components, viz., use-info of operand(s), successor(s),
+  /// region(s), attribute(s) and function-type, of the generic form of an
+  /// operation instance and populate the input operation-state 'result' with
+  /// those components. If any of the components is explicitly provided, then
+  /// skip parsing that component.
+  virtual ParseResult parseGenericOperationAfterOpName(
+      OperationState &result, Optional<ArrayRef<OperandType>> parsedOperandType,
+      Optional<ArrayRef<Block *>> parsedSuccessors,
+      Optional<MutableArrayRef<std::unique_ptr<Region>>> parsedRegions,
+      Optional<ArrayRef<NamedAttribute>> parsedAttributes,
+      Optional<FunctionType> parsedFnType) = 0;
+
   /// Parse a single operand.
   virtual ParseResult parseOperand(OperandType &result) = 0;
 
   /// Parse a single operand if present.
   virtual OptionalParseResult parseOptionalOperand(OperandType &result) = 0;
 
   /// Parse zero or more SSA comma-separated operand references with a specified
   /// surrounding delimiter, and an optional required operand count.

mlir/lib/Parser/Parser.cpp

public:
   ParseResult parseSuccessor(Block *&dest);
 
   /// Parse a comma-separated list of operation successors in brackets.
   ParseResult parseSuccessors(SmallVectorImpl<Block *> &destinations);
 
   /// Parse an operation instance that is in the generic form.
   Operation *parseGenericOperation();
 
+  /// Parse different components, viz., use-info of operand(s), successor(s),
+  /// region(s), attribute(s) and function-type, of the generic form of an
+  /// operation instance and populate the input operation-state 'result' with
+  /// those components. If any of the components is explicitly provided, then
+  /// skip parsing that component.
+  ParseResult parseGenericOperationAfterOpName(
+      OperationState &result,
+      Optional<ArrayRef<SSAUseInfo>> parsedOperandUseInfo,
+      Optional<ArrayRef<Block *>> parsedSuccessors,
+      Optional<MutableArrayRef<std::unique_ptr<Region>>> parsedRegions,
+      Optional<ArrayRef<NamedAttribute>> parsedAttributes,
+      Optional<FunctionType> parsedFnType);
+
   /// 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 type is used to keep track of things that are either an Operation or
   /// a BlockArgument.  We cannot use Value for this, because not all Operations
   /// have results.
public:
   /// This is the structure of a result specifier in the assembly syntax,
   /// including the name, number of results, and location.
   using ResultRecord = std::tuple<StringRef, unsigned, SMLoc>;
 
   /// Parse an operation instance that is in the op-defined custom form.
   /// resultInfo specifies information about the "%name =" specifiers.
   Operation *parseCustomOperation(ArrayRef<ResultRecord> resultIDs);
 
+  /// Parse the name of an operation in the custom form. On success, return a
+  /// pair containing name of the operation and pointer, possibly nullptr, to
+  /// the corresponding dialect. Otherwise, failure is returned.
+  using parsedOpInfo =
+      std::tuple<std::string, Dialect *, const AbstractOperation *>;
+  FailureOr<parsedOpInfo> parseCustomOperationName();
+
   //===--------------------------------------------------------------------===//
   // 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.
   ParseResult parseRegion(Region &region,
for (auto &region : state.regions)
       if (region)
         for (auto &block : *region)
           block.dropAllDefinedValueUses();
   }
   OperationState &state;
 };
 } // namespace
 
-Operation *OperationParser::parseGenericOperation() {
-  // Get location information for the operation.
-  auto srcLocation = getEncodedSourceLocation(getToken().getLoc());
-
-  std::string name = getToken().getStringValue();
-  if (name.empty())
-    return (emitError("empty operation name is invalid"), nullptr);
-  if (name.find('\0') != StringRef::npos)
-    return (emitError("null character not allowed in operation name"), nullptr);
-
-  consumeToken(Token::string);
-
-  OperationState result(srcLocation, name);
-
-  // Lazy load dialects in the context as needed.
-  if (!result.name.getAbstractOperation()) {
-    StringRef dialectName = StringRef(name).split('.').first;
-    if (!getContext()->getLoadedDialect(dialectName)) {
-      if (getContext()->getOrLoadDialect(dialectName)) {
-        result.name = OperationName(name, getContext());
-      } else if (!getContext()->allowsUnregisteredDialects()) {
-        // Emit an error if the dialect couldn't be loaded (i.e., it was not
-        // registered) and unregistered dialects aren't allowed.
-        return emitError(
-                   "operation being parsed with an unregistered dialect. If "
-                   "this is intended, please use -allow-unregistered-dialect "
-                   "with the MLIR tool used"),
-               nullptr;
-      }
-    }
-  }
-
-  // If we are populating the parser state, start a new operation definition.
-  if (state.asmState)
-    state.asmState->startOperationDefinition(result.name);
-
-  // Parse the operand list.
-  SmallVector<SSAUseInfo, 8> operandInfos;
+ParseResult OperationParser::parseGenericOperationAfterOpName(
+    OperationState &result, Optional<ArrayRef<SSAUseInfo>> parsedOperandUseInfo,
+    Optional<ArrayRef<Block *>> parsedSuccessors,
+    Optional<MutableArrayRef<std::unique_ptr<Region>>> parsedRegions,
+    Optional<ArrayRef<NamedAttribute>> parsedAttributes,
+    Optional<FunctionType> parsedFnType) {
+
+  // Parse the operand list, if not explicitly provided.
+  SmallVector<SSAUseInfo, 8> opInfo;
+  ArrayRef<SSAUseInfo> operandUseInfo;
+  if (!parsedOperandUseInfo.hasValue()) {
     if (parseToken(Token::l_paren, "expected '(' to start operand list") ||
-      parseOptionalSSAUseList(operandInfos) ||
+        parseOptionalSSAUseList(opInfo) ||
         parseToken(Token::r_paren, "expected ')' to end operand list")) {
-    return nullptr;
+      return failure();
+    }
+    operandUseInfo = opInfo;
+  } else {
+    operandUseInfo = *parsedOperandUseInfo;
   }
 
-  // Parse the successor list.
+  // Parse the successor list, if not explicitly provided.
+  if (!parsedSuccessors.hasValue()) {
     if (getToken().is(Token::l_square)) {
       // Check if the operation is a known terminator.
       const AbstractOperation *abstractOp = result.name.getAbstractOperation();
       if (abstractOp && !abstractOp->hasTrait<OpTrait::IsTerminator>())
-      return emitError("successors in non-terminator"), nullptr;
+        return emitError("successors in non-terminator");
 
       SmallVector<Block *, 2> successors;
       if (parseSuccessors(successors))
-      return nullptr;
+        return failure();
       result.addSuccessors(successors);
     }
+  } else {
+    result.addSuccessors(*parsedSuccessors);
+  }
 
-  // Parse the region list.
-  CleanupOpStateRegions guard{result};
+  // Parse the region list, if not explicitly provided.
+  if (!parsedRegions.hasValue()) {
     if (consumeIf(Token::l_paren)) {
       do {
         // Create temporary regions with the top level region as parent.
         result.regions.emplace_back(new Region(topLevelOp));
         if (parseRegion(*result.regions.back(), /*entryArguments=*/{}))
-        return nullptr;
+          return failure();
       } while (consumeIf(Token::comma));
       if (parseToken(Token::r_paren, "expected ')' to end region list"))
-      return nullptr;
+        return failure();
+    }
+  } else {
+    result.addRegions(*parsedRegions);
   }
 
+  // Parse the attributes, if not explicitly provided.
+  if (!parsedAttributes.hasValue()) {
     if (getToken().is(Token::l_brace)) {
       if (parseAttributeDict(result.attributes))
-      return nullptr;
+        return failure();
+    }
+  } else {
+    result.addAttributes(*parsedAttributes);
   }
 
+  // Parse the operation type, if not explicitly provided.
+  ArrayRef<Type> operandTypes;
+  if (!parsedFnType.hasValue()) {
     if (parseToken(Token::colon, "expected ':' followed by operation type"))
-    return nullptr;
+      return failure();
 
     auto typeLoc = getToken().getLoc();
     auto type = parseType();
     if (!type)
-    return nullptr;
+      return failure();
     auto fnType = type.dyn_cast<FunctionType>();
     if (!fnType)
-    return (emitError(typeLoc, "expected function type"), nullptr);
+      return emitError(typeLoc, "expected function type");
 
     result.addTypes(fnType.getResults());
 
     // Check that we have the right number of types for the operands.
-  auto operandTypes = fnType.getInputs();
-  if (operandTypes.size() != operandInfos.size()) {
-    auto plural = "s"[operandInfos.size() == 1];
-    return (emitError(typeLoc, "expected ")
-                << operandInfos.size() << " operand type" << plural
-                << " but had " << operandTypes.size(),
-            nullptr);
+    operandTypes = fnType.getInputs();
+    if (operandTypes.size() != operandUseInfo.size()) {
+      auto plural = "s"[operandUseInfo.size() == 1];
+      return emitError(typeLoc, "expected ")
+             << operandUseInfo.size() << " operand type" << plural
+             << " but had " << operandTypes.size();
+    }
+  } else {
+    result.addTypes(parsedFnType->getResults());
+
+    // Check that we have the right number of types for the operands.
+    operandTypes = parsedFnType->getInputs();
+    if (operandTypes.size() != operandUseInfo.size()) {
+      auto plural = "s"[operandUseInfo.size() == 1];
+      return emitError("expected ")
+             << operandUseInfo.size() << " operand type" << plural
+             << " but had " << operandTypes.size();
+    }
   }
 
   // Resolve all of the operands.
-  for (unsigned i = 0, e = operandInfos.size(); i != e; ++i) {
-    result.operands.push_back(resolveSSAUse(operandInfos[i], operandTypes[i]));
+  for (unsigned i = 0, e = operandUseInfo.size(); i != e; ++i) {
+    result.operands.push_back(
+        resolveSSAUse(operandUseInfo[i], operandTypes[i]));
     if (!result.operands.back())
-      return nullptr;
+      return failure();
+  }
+
+  return success();
+}
+
+Operation *OperationParser::parseGenericOperation() {
+  // Get location information for the operation.
+  auto srcLocation = getEncodedSourceLocation(getToken().getLoc());
+
+  std::string name = getToken().getStringValue();
+  if (name.empty())
+    return (emitError("empty operation name is invalid"), nullptr);
+  if (name.find('\0') != StringRef::npos)
+    return (emitError("null character not allowed in operation name"), nullptr);
+
+  consumeToken(Token::string);
+
+  OperationState result(srcLocation, name);
+  CleanupOpStateRegions guard{result};
+
+  // Lazy load dialects in the context as needed.
+  if (!result.name.getAbstractOperation()) {
+    StringRef dialectName = StringRef(name).split('.').first;
+    if (!getContext()->getLoadedDialect(dialectName)) {
+      if (getContext()->getOrLoadDialect(dialectName)) {
+        result.name = OperationName(name, getContext());
+      } else if (!getContext()->allowsUnregisteredDialects()) {
+        // Emit an error if the dialect couldn't be loaded (i.e., it was not
+        // registered) and unregistered dialects aren't allowed.
+        return emitError(
+                   "operation being parsed with an unregistered dialect. If "
+                   "this is intended, please use -allow-unregistered-dialect "
+                   "with the MLIR tool used"),
+               nullptr;
+      }
+    }
   }
 
+  // If we are populating the parser state, start a new operation definition.
+  if (state.asmState)
+    state.asmState->startOperationDefinition(result.name);
+
+  if (parseGenericOperationAfterOpName(
+          result, /*parsedOperandUseInfo=*/llvm::None,
+          /*parsedSuccessors=*/llvm::None, /*parsedRegions=*/llvm::None,
+          /*parsedAttributes=*/llvm::None, /*parsedFnType=*/llvm::None))
+    return nullptr;
+
   // Create the operation and try to parse a location for it.
   Operation *op = opBuilder.createOperation(result);
   if (parseTrailingLocationSpecifier(op))
     return nullptr;
   return op;
 }
 
 Operation *OperationParser::parseGenericOperation(Block *insertBlock,
ParseResult parseOperation(OperationState &opState) {
     return success();
   }
 
   Operation *parseGenericOperation(Block *insertBlock,
                                    Block::iterator insertPt) final {
     return parser.parseGenericOperation(insertBlock, insertPt);
   }
 
+  FailureOr<ParsedOpNameAndDialect> parseCustomOperationName() final {
+    auto opInfo = parser.parseCustomOperationName();
+    if (failed(opInfo))
+      return failure();
+
+    return ParsedOpNameAndDialect(std::get<0>(*opInfo), std::get<1>(*opInfo));
+  }
+
+  ParseResult parseGenericOperationAfterOpName(
+      OperationState &result,
+      Optional<ArrayRef<OperandType>> parsedOperandTypes,
+      Optional<ArrayRef<Block *>> parsedSuccessors,
+      Optional<MutableArrayRef<std::unique_ptr<Region>>> parsedRegions,
+      Optional<ArrayRef<NamedAttribute>> parsedAttributes,
+      Optional<FunctionType> parsedFnType) final {
+
+    SmallVector<OperationParser::SSAUseInfo> parsedOperandUseInfo;
+    if (parsedOperandTypes.hasValue()) {
+      for (auto &parsedOperandType : *parsedOperandTypes)
+        parsedOperandUseInfo.push_back({
+            parsedOperandType.name,
+            parsedOperandType.number,
+            parsedOperandType.location,
+        });
+    }
+
+    return parser.parseGenericOperationAfterOpName(
+        result,
+        parsedOperandTypes.hasValue() ? llvm::makeArrayRef(parsedOperandUseInfo)
+                                      : llvm::None,
+        parsedSuccessors, parsedRegions, parsedAttributes, parsedFnType);
+  }
   //===--------------------------------------------------------------------===//
   // Utilities
   //===--------------------------------------------------------------------===//
 
   /// 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
private:
   bool isIsolatedFromAbove;
   StringRef opName;
 
   /// The backing operation parser.
   OperationParser &parser;
 };
 } // end anonymous namespace.
 
-Operation *
-OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
-  llvm::SMLoc opLoc = getToken().getLoc();
+FailureOr<OperationParser::parsedOpInfo>
+OperationParser::parseCustomOperationName() {
   std::string opName = getTokenSpelling().str();
+  if (opName.empty())
+    return (emitError("empty operation name is invalid"), failure());
+
+  consumeToken();
+
   auto *opDefinition = AbstractOperation::lookup(opName, getContext());
   StringRef defaultDialect = getState().defaultDialectStack.back();
   Dialect *dialect = nullptr;
   if (opDefinition) {
     dialect = &opDefinition->dialect;
   } else {
     if (StringRef(opName).contains('.')) {
       // This op has a dialect, we try to check if we can register it in the
if (StringRef(opName).contains('.')) {
         opName = opDefinition->name.str();
       } else if (!defaultDialect.empty()) {
         dialect = getContext()->getOrLoadDialect(defaultDialect);
         opName = (defaultDialect + "." + opName).str();
       }
     }
   }
 
+  return parsedOpInfo(opName, dialect, opDefinition);
+}
+
+Operation *
+OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
+  llvm::SMLoc opLoc = getToken().getLoc();
+
+  auto opInfo = parseCustomOperationName();
+  if (failed(opInfo))
+    return nullptr;
+
+  std::string opName = std::get<0>(*opInfo);
+  Dialect *dialect = std::get<1>(*opInfo);
+  const AbstractOperation *opDefinition = std::get<2>(*opInfo);
+
   // This is the actual hook for the custom op parsing, usually implemented by
   // the op itself (`Op::parse()`). We retrieve it either from the
   // AbstractOperation or from the Dialect.
   function_ref<ParseResult(OpAsmParser &, OperationState &)> parseAssemblyFn;
   bool isIsolatedFromAbove = false;
 
-  defaultDialect = "";
+  StringRef defaultDialect = "";
   if (opDefinition) {
     parseAssemblyFn = opDefinition->getParseAssemblyFn();
     isIsolatedFromAbove =
         opDefinition->hasTrait<OpTrait::IsIsolatedFromAbove>();
     auto *iface = opDefinition->getInterface<OpAsmOpInterface>();
     if (iface && !iface->getDefaultDialect().empty())
       defaultDialect = iface->getDefaultDialect();
   } else {
     Optional<Dialect::ParseOpHook> dialectHook;
     if (dialect)
       dialectHook = dialect->getParseOperationHook(opName);
     if (!dialectHook.hasValue()) {
       emitError(opLoc) << "custom op '" << opName << "' is unknown";
       return nullptr;
     }
     parseAssemblyFn = *dialectHook;
   }
   getState().defaultDialectStack.push_back(defaultDialect);
   auto restoreDefaultDialect = llvm::make_scope_exit(
       [&]() { getState().defaultDialectStack.pop_back(); });
 
-  consumeToken();
-
   // If the custom op parser crashes, produce some indication to help
   // debugging.
   llvm::PrettyStackTraceFormat fmt("MLIR Parser: custom op parser '%s'",
                                    opName.c_str());
 
   // Get location information for the operation.
   auto srcLocation = getEncodedSourceLocation(opLoc);
   OperationState opState(srcLocation, opName);
 
   // If we are populating the parser state, start a new operation definition.
   if (state.asmState)
     state.asmState->startOperationDefinition(opState.name);
 
   // Have the op implementation take a crack and parsing this.
   CleanupOpStateRegions guard{opState};
   CustomOpAsmParser opAsmParser(opLoc, resultIDs, parseAssemblyFn,

mlir/test/IR/pretty_printed_region_op.mlir

+// RUN: mlir-opt -allow-unregistered-dialect -split-input-file  %s | FileCheck %s --check-prefixes=CHECK-CUSTOM,CHECK
+// RUN: mlir-opt -allow-unregistered-dialect -mlir-print-op-generic -split-input-file   %s | FileCheck %s --check-prefixes=CHECK,CHECK-GENERIC
+
+// -----
+
+func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) {
+// CHECK-CUSTOM:  test.pretty_printed_region %arg1, %arg0 start special.op end : (f32, f32) -> f32
+// CHECK-GENERIC: "test.pretty_printed_region"(%arg1, %arg0)
+// CHECK-GENERIC:   ^bb0(%arg[[x:[0-9]+]]: f32, %arg[[y:[0-9]+]]: f32
+// CHECK-GENERIC:     %[[RES:.*]] = "special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32
+// CHECK-GENERIC:     "test.return"(%[[RES]]) : (f32) -> ()
+// CHECK-GENERIC:  : (f32, f32) -> f32
+
+  %res = test.pretty_printed_region %arg1, %arg0 start special.op end : (f32, f32) -> (f32) loc("some_NameLoc")
+  return %res : f32
+}
+
+// -----
+
+func @pretty_printed_region_op(%arg0 : f32, %arg1 : f32) -> (f32) {
+// CHECK-CUSTOM:  test.pretty_printed_region %arg1, %arg0
+// CHECK-GENERIC: "test.pretty_printed_region"(%arg1, %arg0)
+// CHECK:          ^bb0(%arg[[x:[0-9]+]]: f32, %arg[[y:[0-9]+]]: f32):
+// CHECK:            %[[RES:.*]] = "non.special.op"(%arg[[x]], %arg[[y]]) : (f32, f32) -> f32
+// CHECK:            "test.return"(%[[RES]]) : (f32) -> ()
+// CHECK:          : (f32, f32) -> f32
+
+  %0 = test.pretty_printed_region %arg1, %arg0 ( {
+    ^bb0(%arg2: f32, %arg3: f32):
+      %1 = "non.special.op"(%arg2, %arg3) : (f32, f32) -> f32
+      "test.return"(%1) : (f32) -> ()
+    }) : (f32, f32) -> f32
+    return %0 : f32
+}
+

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

 }
 
 static void print(OpAsmPrinter &p, WrappingRegionOp op) {
   p << " wraps ";
   p.printGenericOp(&op.getRegion().front().front());
 }
 
 //===----------------------------------------------------------------------===//
+// Test PrettyPrintedRegionOp -  exercising the following parser APIs
+//   parseGenericOperationAfterOpName
+//   parseCustomOperationName
+//===----------------------------------------------------------------------===//
+
+static ParseResult parsePrettyPrintedRegionOp(OpAsmParser &parser,
+                                              OperationState &result) {
+
+  llvm::SMLoc loc = parser.getCurrentLocation();
+  Location currLocation = parser.getEncodedSourceLoc(loc);
+
+  // Parse the operands.
+  SmallVector<OpAsmParser::OperandType, 2> operands;
+  if (parser.parseOperandList(operands))
+    return failure();
+
+  // Check if we are parsing the pretty-printed version
+  //  test.pretty_printed_region start <inner-op> end : <functional-type>
+  // Else fallback to parsing the "non pretty-printed" version.
+  if (!succeeded(parser.parseOptionalKeyword("start")))
+    return parser.parseGenericOperationAfterOpName(
+        result, llvm::makeArrayRef(operands), /*parsedSuccessors=*/llvm::None,
+        /*parsedRegions=*/llvm::None, /*parsedAttributes=*/llvm::None,
+        /*parsedFnType=*/llvm::None);
+
+  auto parseOpNameResult = parser.parseCustomOperationName();
+  if (failed(parseOpNameResult))
+    return failure();
+
+  StringRef innerOpName = std::get<0>(*parseOpNameResult);
+
+  FunctionType opFntype;
+  Optional<Location> explicitLoc;
+  if (parser.parseKeyword("end") || parser.parseColon() ||
+      parser.parseType(opFntype) ||
+      parser.parseOptionalLocationSpecifier(explicitLoc))
+    return failure();
+
+  // If location of the op is explicitly provided, then use it; Else use
+  // the parser's current location.
+  Location opLoc = explicitLoc.getValueOr(currLocation);
+
+  // Derive the SSA-values for op's operands.
+  if (parser.resolveOperands(operands, opFntype.getInputs(), loc,
+                             result.operands))
+    return failure();
+
+  // Add a region for op.
+  Region &region = *result.addRegion();
+
+  // Create a basic-block inside op's region.
+  Block &block = region.emplaceBlock();
+
+  // Create and insert an "inner-op" operation in the block.
+  // Just for testing purposes, we can assume that inner op is a binary op with
+  // result and operand types all same as the test-op's first operand.
+  auto innerOpType = opFntype.getInput(0);
+  auto lhs = block.addArgument(innerOpType, opLoc);
+  auto rhs = block.addArgument(innerOpType, opLoc);
+
+  OpBuilder builder(parser.getBuilder().getContext());
+  builder.setInsertionPointToStart(&block);
+
+  OperationState innerOpState(opLoc, innerOpName);
+  innerOpState.operands.push_back(lhs);
+  innerOpState.operands.push_back(rhs);
+  innerOpState.addTypes(innerOpType);
+
+  Operation *innerOp = builder.createOperation(innerOpState);
+
+  // Insert a return statement in the block returning the inner-op's result.
+  builder.create<TestReturnOp>(innerOp->getLoc(), innerOp->getResults());
+
+  // Populate the op operation-state with result-type and location.
+  result.addTypes(opFntype.getResults());
+  result.location = innerOp->getLoc();
+
+  return success();
+}
+
+static void print(OpAsmPrinter &p, PrettyPrintedRegionOp op) {
+  p << ' ';
+  p.printOperands(op.getOperands());
+
+  Operation &innerOp = op.getRegion().front().front();
+  // Assuming that region has a single non-terminator inner-op, if the inner-op
+  // meets some criteria (which in this case is a simple one  based on the name
+  // of inner-op), then we can print the entire region in a succinct way.
+  // Here we assume that the prototype of "special.op" can be trivially derived
+  // while parsing it back.
+  if (innerOp.getName().getStringRef().equals("special.op")) {
+    p << " start special.op end";
+  } else {
+    p << " (";
+    p.printRegion(op.getRegion());
+    p << ")";
+  }
+
+  p << " : ";
+  p.printFunctionalType(op);
+}
+
+//===----------------------------------------------------------------------===//
 // Test PolyForOp - parse list of region arguments.
 //===----------------------------------------------------------------------===//
 
 static ParseResult parsePolyForOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 4> ivsInfo;
   // Parse list of region arguments without a delimiter.
   if (parser.parseRegionArgumentList(ivsInfo))
     return failure();

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

def WrappingRegionOp : TEST_Op<"wrapping_region",
   }];
 
   let results = (outs Variadic<AnyType>);
   let regions = (region SizedRegion<1>:$region);
   let parser = [{ return ::parse$cppClass(parser, result); }];
   let printer = [{ return ::print(p, *this); }];
 }
 
+def PrettyPrintedRegionOp : TEST_Op<"pretty_printed_region",
+    [SingleBlockImplicitTerminator<"TestReturnOp">]> {
+  let summary =  "pretty_printed_region operation";
+  let description = [{
+    Test-op can be printed either in a "pretty" or "non-pretty" way based on
+    some criteria. The custom parser parsers both the versions while testing
+    APIs: parseCustomOperationName & parseGenericOperationAfterOpName.
+  }];
+  let arguments = (ins
+    AnyType:$input1,
+    AnyType:$input2
+  );
+
+  let results = (outs AnyType);
+  let regions = (region SizedRegion<1>:$region);
+  let parser = [{ return ::parse$cppClass(parser, result); }];
+  let printer = [{ return ::print(p, *this); }];
+}
+
 def PolyForOp : TEST_Op<"polyfor">
 {
   let summary =  "polyfor operation";
   let description = [{
     Test op with multiple region arguments, each argument of index type.
   }];
 
   let regions = (region SizedRegion<1>:$region);