[mlir] do not elide dialect prefix for ops with dots in the name

[ftynse]

Summary

For the hypothetical "a.b.c" op printed within a region that declares "a" as
the default dialect, MLIR would currently elide the "a." prefix and only print
"b.c". However, this becomes ambiguous while parsing as "b.c" may be exist as
the "c" op in the "b" dialect. If it does not, the parsing currently fails. Do
not elide the default dialect if the op name contains further dots to avoid the
ambiguity.

See https://discourse.llvm.org/t/dropping-dialect-prefix-for-ops-with-multiple-dots-in-the-name/62562

[rriddle]

LG, thanks!

mlir/lib/IR/AsmPrinter.cpp

if (auto opInfo = op->getRegisteredInfo()) {
       opInfo->printAssembly(op, *this, defaultDialectStack.back());
       return;
     }
     // Otherwise try to dispatch to the dialect, if available.
     if (Dialect *dialect = op->getDialect()) {
       if (auto opPrinter = dialect->getOperationPrinter(op)) {
         // Print the op name first.
         StringRef name = op->getName().getStringRef();
+        // Only drop the default dialect prefix when it cannot lead to
+        // ambiguities.
+        if (name.count('.') == 1)
           name.consume_front((defaultDialectStack.back() + ".").str());
         printEscapedString(name, os);
         // Print the rest of the op now.
         opPrinter(op, *this);
         return;
       }
     }
   }
 

mlir/lib/IR/Operation.cpp

void OpState::print(Operation *op, OpAsmPrinter &p, StringRef defaultDialect) {
   if (auto printFn = op->getDialect()->getOperationPrinter(op)) {
     printOpName(op, p, defaultDialect);
     printFn(op, p);
   } else {
     p.printGenericOp(op);
   }
 }
 
-/// Print an operation name, eliding the dialect prefix if necessary.
+/// Print an operation name, eliding the dialect prefix if necessary and doesn't
+/// lead to ambiguities.
 void OpState::printOpName(Operation *op, OpAsmPrinter &p,
                           StringRef defaultDialect) {
   StringRef name = op->getName().getStringRef();
-  if (name.startswith((defaultDialect + ".").str()))
+  if (name.startswith((defaultDialect + ".").str()) && name.count('.') == 1)
     name = name.drop_front(defaultDialect.size() + 1);
   p.getStream() << name;
 }
 
 /// Emit an error about fatal conditions with this operation, reporting up to
 /// any diagnostic handlers that may be listening.
 InFlightDiagnostic OpState::emitError(const Twine &message) {
   return getOperation()->emitError(message);

mlir/test/IR/parser.mlir

test.default_dialect {
     // CHECK-NOT: test.parse_integer_literal
     parse_integer_literal : 5
     // CHECK: parse_integer_literal : 6
     test.parse_integer_literal : 6
     // Verify that only an op prefix is stripped, not an attribute value for
     // example.
     // CHECK:  "test.op_with_attr"() {test.attr = "test.value"} : () -> ()
     "test.op_with_attr"() {test.attr = "test.value"} : () -> ()
+    // Verify that the prefix is not stripped when it can lead to ambiguity.
+    // CHECK: test.op.with_dot_in_name
+    test.op.with_dot_in_name
+    // This is an unregistered operation, the printing/parsing is handled by the
+    // dialect, and the dialect prefix should not be stripped while printing
+    // because of potential ambiguity.
+    // CHECK: test.dialect_custom_printer.with.dot
+    test.dialect_custom_printer.with.dot
     "test.terminator"() : ()->()
   }
   // The same operation outside of the region does not have an func. prefix.
   // CHECK: return
   func.return
 }
 
 // CHECK-LABEL: func @unreachable_dominance_violation_ok

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

return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {
       return parser.parseKeyword("custom_format");
     }};
   }
   if (opName == "test.dialect_custom_format_fallback") {
     return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {
       return parser.parseKeyword("custom_format_fallback");
     }};
   }
+  if (opName == "test.dialect_custom_printer.with.dot") {
+    return ParseOpHook{[](OpAsmParser &parser, OperationState &state) {
+      return ParseResult::success();
+    }};
+  }
   return None;
 }
 
 llvm::unique_function<void(Operation *, OpAsmPrinter &)>
 TestDialect::getOperationPrinter(Operation *op) const {
   StringRef opName = op->getName().getStringRef();
   if (opName == "test.dialect_custom_printer") {
     return [](Operation *op, OpAsmPrinter &printer) {

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

let extraClassDeclaration = [{
     static ::llvm::StringRef getDefaultDialect() {
       return "test";
     }
     void getAsmResultNames(::llvm::function_ref<void(::mlir::Value, ::llvm::StringRef)> setNameFn) {}
   }];
   let assemblyFormat = "regions attr-dict-with-keyword";
 }
 
+// This is used to test that the default dialect is not elided when printing an
+// op with dots in the name to avoid parsing ambiguity.
+def OpWithDotInNameOp : TEST_Op<"op.with_dot_in_name"> {
+  let assemblyFormat = "attr-dict";
+}
+
 // This is used to test the OpAsmOpInterface::getAsmBlockName() feature:
 // blocks nested in a region under this op will have a name defined by the
 // interface.
 def AsmBlockNameOp : TEST_Op<"block_names", [OpAsmOpInterface]> {
  let regions = (region AnyRegion:$body);
   let extraClassDeclaration = [{
     void getAsmBlockNames(mlir::OpAsmSetBlockNameFn setNameFn) {
       std::string name;