Add a mechanism for Dialects to customize printing/parsing operations when they are unregistered

[rriddle]

mlir/include/mlir/IR/Dialect.h

46

Do we need the bool result here? This is only used for name shadowing (IIRC), so do we have an actual need to allow unregistered operations to name shadow? Also, I would expect that if an unregistered op does have that trait, we would develop a proper mechanism (e.g. via a dialect fallback) to query if the op has that trait. WDYT?

mlir/lib/IR/AsmPrinter.cpp

2412

mlir/lib/Parser/Parser.cpp

1624–1626

nit: Drop the llvm:: here.

[mehdi_amini]

mlir/include/mlir/IR/Dialect.h

46

Yeah that's a good point!

[mehdi_amini]

address comments

[jpienaar]

mlir/include/mlir/IR/Dialect.h

42

This one I struggle to read (perhaps as I keep on reading bookmark). Is isolated from above returned/required still?

mlir/lib/Parser/Parser.cpp

1674

So this is only place where it can be changed to true?

[mehdi_amini]

Fix doc

[rriddle]

mlir/include/mlir/IR/Dialect.h

46

Could we return function_ref (or llvm::unique_function) instead? I imagine that a majority of the functions passed back are going to be function-pointer-like(i.e. cheap to construct/copy).

mlir/lib/Parser/Parser.cpp

895

nit: Drop the llvm:: here.

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

212

nit: Add braces here.

[mehdi_amini]

mlir/include/mlir/IR/Dialect.h

42

Oh I forgot to update the doc when I changed the API...

[mehdi_amini]

Address River's most recent comments

mlir/lib/Parser/Parser.cpp

1674

Yes!
If/when we have Traits fallback to the dialect, we can also implement it in the else branch

mlir/include/mlir/IR/Dialect.h

 /// behavior associated with the entire group.  For example, hooks into other
 /// systems for constant folding, interfaces, default named types for asm
 /// printing, etc.
 ///
 /// Instances of the dialect object are loaded in a specific MLIRContext.
 ///
 class Dialect {
 public:
+  /// Type for a callback provided by the dialect to parse a custom operation.
+  /// This is used for the dialect to provide an alternative way to parse custom
+  /// operations, including unregistered ones.
+  using ParseOpHook =
+      std::function<ParseResult(OpAsmParser &parser, OperationState &result)>;
+
   virtual ~Dialect();
 
   /// Utility function that returns if the given string is a valid dialect
   /// namespace.
   static bool isValidNamespace(StringRef str);
 
   MLIRContext *getContext() const { return context; }
 
public:
   /// Parse a type registered to this dialect.
   virtual Type parseType(DialectAsmParser &parser) const;
 
   /// Print a type registered to this dialect.
   virtual void printType(Type, DialectAsmPrinter &) const {
     llvm_unreachable("dialect has no registered type printing hook");
   }
 
+  /// Return the hook to parse an operation registered to this dialect, if any.
+  /// By default this will lookup for registered operations and return the
+  /// `parse()` method registered on the AbstractOperation. Dialects can
+  /// override this behavior and handle unregistered operations as well.
+  virtual Optional<ParseOpHook> getParseOperationHook(StringRef opName) const;
+
+  /// Print an operation registered to this dialect.
+  /// This hook is invoked for registered operation which don't override the
+  /// `print()` method to define their own custom assembly.
+  virtual LogicalResult printOperation(Operation *op,
+                                       OpAsmPrinter &printer) const;
+
   //===--------------------------------------------------------------------===//
   // Verification Hooks
   //===--------------------------------------------------------------------===//
 
   /// Verify an attribute from this dialect on the argument at 'argIndex' for
   /// the region at 'regionIndex' on the given operation. Returns failure if
   /// the verification failed, success otherwise. This hook may optionally be
   /// invoked from any operation containing a region.

mlir/include/mlir/IR/OperationSupport.h

public:
   Dialect &dialect;
 
   /// The unique identifier of the derived Op class.
   TypeID typeID;
 
   /// Use the specified object to parse this ops custom assembly format.
   ParseResult parseAssembly(OpAsmParser &parser, OperationState &result) const;
 
+  /// Return the static hook for parsing this operation assembly.
+  ParseAssemblyFn getParseAssemblyFn() const { return parseAssemblyFn; }
+
   /// This hook implements the AsmPrinter for this operation.
   void printAssembly(Operation *op, OpAsmPrinter &p) const {
     return printAssemblyFn(op, p);
   }
 
   /// This hook implements the verifier for this operation.  It should emits an
   /// error message and returns failure if a problem is detected, or returns
   /// success if everything is ok.

mlir/lib/IR/AsmPrinter.cpp

void OperationPrinter::printOperation(Operation *op) {
   // If requested, always print the generic form.
   if (!printerFlags.shouldPrintGenericOpForm()) {
     // Check to see if this is a known operation.  If so, use the registered
     // custom printer hook.
     if (auto *opInfo = op->getAbstractOperation()) {
       opInfo->printAssembly(op, *this);
       return;
     }
+    // Otherwise try to dispatch to the dialect, if available.
+    if (Dialect *dialect = op->getDialect()) {
+      if (succeeded(dialect->printOperation(op, *this)))
+        return;
+    }
   }
 
   // Otherwise print with the generic assembly form.
   printGenericOp(op);
 }
 
 void OperationPrinter::printGenericOp(Operation *op) {
   os << '"';

mlir/lib/IR/Dialect.cpp

if (allowsUnknownTypes()) {
     return OpaqueType::get(ns, parser.getFullSymbolSpec());
   }
 
   parser.emitError(parser.getNameLoc())
       << "dialect '" << getNamespace() << "' provides no type parsing hook";
   return Type();
 }
 
+Optional<Dialect::ParseOpHook>
+Dialect::getParseOperationHook(StringRef opName) const {
+  return None;
+}
+
+LogicalResult Dialect::printOperation(Operation *op,
+                                      OpAsmPrinter &printer) const {
+  assert(op->getDialect() == this &&
+         "Dialect hook invoked on non-dialect owned operation");
+  return failure();
+}
+
 /// Utility function that returns if the given string is a valid dialect
 /// namespace.
 bool Dialect::isValidNamespace(StringRef str) {
   if (str.empty())
     return true;
   llvm::Regex dialectNameRegex("^[a-zA-Z_][a-zA-Z_0-9\\$]*$");
   return dialectNameRegex.match(str);
 }

mlir/lib/Parser/Parser.cpp

Operation *OperationParser::parseGenericOperation(Block *insertBlock,
   OpBuilder::InsertionGuard restoreInsertionPoint(opBuilder);
   opBuilder.setInsertionPoint(insertBlock, insertPt);
   return parseGenericOperation();
 }
 
 namespace {
 class CustomOpAsmParser : public OpAsmParser {
 public:
-  CustomOpAsmParser(SMLoc nameLoc,
-                    ArrayRef<OperationParser::ResultRecord> resultIDs,
-                    const AbstractOperation *opDefinition,
-                    OperationParser &parser)
-      : nameLoc(nameLoc), resultIDs(resultIDs), opDefinition(opDefinition),
+  CustomOpAsmParser(
+      SMLoc nameLoc, ArrayRef<OperationParser::ResultRecord> resultIDs,
+      llvm::function_ref<ParseResult(OpAsmParser &, OperationState &)>
+          parseAssembly,
+      bool isIsolatedFromAbove, StringRef opName, OperationParser &parser)
+      : nameLoc(nameLoc), resultIDs(resultIDs), parseAssembly(parseAssembly),
+        isIsolatedFromAbove(isIsolatedFromAbove), opName(opName),
         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))
+    if (parseAssembly(*this, opState))
       return failure();
     // Verify that the parsed attributes does not have duplicate attributes.
     // This can happen if an attribute set during parsing is also specified in
     // the attribute dictionary in the assembly, or the attribute is set
     // multiple during parsing.
     Optional<NamedAttribute> duplicate = opState.attributes.findDuplicate();
     if (duplicate)
       return emitError(getNameLoc(), "attribute '")
public:
   //===--------------------------------------------------------------------===//
 
   /// Return if any errors were emitted during parsing.
   bool didEmitError() const { return emittedError; }
 
   /// Emit a diagnostic at the specified location and return failure.
   InFlightDiagnostic emitError(llvm::SMLoc loc, const Twine &message) override {
     emittedError = true;
-    return parser.emitError(loc, "custom op '" + opDefinition->name.strref() +
-                                     "' " + message);
+    return parser.emitError(loc, "custom op '" + opName + "' " + message);
   }
 
   llvm::SMLoc getCurrentLocation() override {
     return parser.getToken().getLoc();
   }
 
   Builder &getBuilder() const override { return parser.builder; }
 
for (auto pair : llvm::zip(arguments, argTypes)) {
       const OperandType &operand = std::get<0>(pair);
       Type type = std::get<1>(pair);
       OperationParser::SSAUseInfo operandInfo = {operand.name, operand.number,
                                                  operand.location};
       regionArguments.emplace_back(operandInfo, type);
     }
 
     // Try to parse the region.
-    assert((!enableNameShadowing ||
-            opDefinition->hasTrait<OpTrait::IsIsolatedFromAbove>()) &&
+    assert((!enableNameShadowing || isIsolatedFromAbove) &&
            "name shadowing is only allowed on isolated regions");
     if (parser.parseRegion(region, regionArguments, enableNameShadowing))
       return failure();
     return success();
   }
 
   /// Parses a region if present.
   OptionalParseResult parseOptionalRegion(Region &region,

 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;
+  function_ref<ParseResult(OpAsmParser &, OperationState &)> parseAssembly;
+  bool isIsolatedFromAbove;
+  StringRef opName;
 
   /// 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(ArrayRef<ResultRecord> resultIDs) {
   llvm::SMLoc opLoc = getToken().getLoc();
   StringRef opName = getTokenSpelling();
-
   auto *opDefinition = AbstractOperation::lookup(opName, getContext());
-  if (!opDefinition) {
+  Dialect *dialect = nullptr;
+  if (opDefinition) {
+    dialect = &opDefinition->dialect;
+  } else {
     if (opName.contains('.')) {
       // This op has a dialect, we try to check if we can register it in the
       // context on the fly.
       StringRef dialectName = opName.split('.').first;
-      if (!getContext()->getLoadedDialect(dialectName) &&
-          getContext()->getOrLoadDialect(dialectName)) {
+      dialect = getContext()->getLoadedDialect(dialectName);
+      if (!dialect && (dialect = getContext()->getOrLoadDialect(dialectName)))
         opDefinition = AbstractOperation::lookup(opName, getContext());
-      }
     } else {
       // 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
       // operations in the standard dialect?
-      if (getContext()->getOrLoadDialect("std"))
+      if (getContext()->getOrLoadDialect("std")) {
         opDefinition = AbstractOperation::lookup(Twine("std." + opName).str(),
                                                  getContext());
+        if (opDefinition)
+          opName = opDefinition->name.strref();
       }
     }
+  }
+
+  // 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.
+  std::function<ParseResult(OpAsmParser &, OperationState &)> parseAssemblyFn;
+  bool isIsolatedFromAbove = false;
 
-  if (!opDefinition) {
+  if (opDefinition) {
+    parseAssemblyFn = opDefinition->getParseAssemblyFn();
+    isIsolatedFromAbove =
+        opDefinition->hasTrait<OpTrait::IsIsolatedFromAbove>();
+  } 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;
+  }
 
   consumeToken();
 
   // If the custom op parser crashes, produce some indication to help
   // debugging.
   std::string opNameStr = opName.str();
   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);
+  OperationState opState(srcLocation, opName);
   CleanupOpStateRegions guard{opState};
-  CustomOpAsmParser opAsmParser(opLoc, resultIDs, opDefinition, *this);
+  CustomOpAsmParser opAsmParser(opLoc, resultIDs, parseAssemblyFn,
+                                isIsolatedFromAbove, opName, *this);
   if (opAsmParser.parseOperation(opState))
     return nullptr;
 
   // If it emitted an error, we failed.
   if (opAsmParser.didEmitError())
     return nullptr;
 
   // Otherwise, create the operation and try to parse a location for it.

mlir/test/IR/parser.mlir

 // CHECK: "unregistered_func_might_have_graph_region"() ( {
 // CHECK: [[VAL1:%.*]] = "foo"([[VAL1]], [[VAL2:%.*]]) : (i64, i64) -> i64
 // CHECK: [[VAL2]] = "bar"([[VAL1]])
 "unregistered_func_might_have_graph_region"() ( {
   %1 = "foo"(%1, %2) : (i64, i64) -> i64
   %2 = "bar"(%1) : (i64) -> i64
   "unregistered_terminator"() : () -> ()
 }) {sym_name = "unregistered_op_dominance_violation_ok", type = () -> i1} : () -> ()
+
+// This is an unregister operation, the printing/parsing is handled by the dialect.
+// CHECK: test.dialect_custom_printer custom_format
+test.dialect_custom_printer custom_format

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

 TestDialect::verifyRegionResultAttribute(Operation *op, unsigned regionIndex,
                                          unsigned resultIndex,
                                          NamedAttribute namedAttr) {
   if (namedAttr.first == "test.invalid_attr")
     return op->emitError() << "invalid to use 'test.invalid_attr'";
   return success();
 }
 
+Optional<Dialect::ParseOpHook>
+TestDialect::getParseOperationHook(StringRef opName) const {
+  if (opName == "test.dialect_custom_printer")
+    return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {
+      return parser.parseKeyword("custom_format");
+    }};
+  return None;
+}
+
+LogicalResult TestDialect::printOperation(Operation *op,
+                                          OpAsmPrinter &printer) const {
+  StringRef opName = op->getName().getStringRef();
+  if (opName == "test.dialect_custom_printer") {
+    printer.getStream() << opName << " custom_format";
+    return success();
+  }
+  return failure();
+}
+
 //===----------------------------------------------------------------------===//
 // TestBranchOp
 //===----------------------------------------------------------------------===//
 
 Optional<MutableOperandRange>
 TestBranchOp::getMutableSuccessorOperands(unsigned index) {
   assert(index == 0 && "invalid successor index");
   return targetOperandsMutable();

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

def Test_Dialect : Dialect {
   let extraClassDeclaration = [{
     void registerAttributes();
     void registerTypes();
 
     Attribute parseAttribute(DialectAsmParser &parser,
                              Type type) const override;
     void printAttribute(Attribute attr,
                         DialectAsmPrinter &printer) const override;
+
+    // Provides a custom printing/parsing for some operations.
+    Optional<ParseOpHook>
+      getParseOperationHook(StringRef opName) const override;
+    LogicalResult printOperation(Operation *op,
+                                 OpAsmPrinter &printer) const override;
   }];
 }
 
 class TEST_Op<string mnemonic, list<OpTrait> traits = []> :
     Op<Test_Dialect, mnemonic, traits>;
 
 //===----------------------------------------------------------------------===//
 // Test Types