[Mlir-commits] [mlir] 1155c1f - [mlir:Parser] Emit a better diagnostic when a custom operation is unknown

[River Riddle]

Author: River Riddle
Date: 2022-05-11T22:54:44-07:00
New Revision: 1155c1fe6589492fa3b3cf4d871c0749736b40c0

URL: https://github.com/llvm/llvm-project/commit/1155c1fe6589492fa3b3cf4d871c0749736b40c0
DIFF: https://github.com/llvm/llvm-project/commit/1155c1fe6589492fa3b3cf4d871c0749736b40c0.diff

LOG: [mlir:Parser] Emit a better diagnostic when a custom operation is unknown

When a custom operation is unknown and does not have a dialect prefix, we currently
emit an error using the name of the operation with the default dialect prefix. This
leads to a confusing error message, especially when operations get moved between dialects.
For example, `func` was recently moved out of `builtin` and to the `func` dialect. The current
error message we get is:

```
func @foo()
^ custom op 'builtin.func' is unknown
```

This could lead users to believe that there is supposed to be a `builtin.func`,
because there used to be. This commit adds a better error message that does
not assume that the operation is supposed to be in the default dialect:

```
func @foo()
^ custom op 'func' is unknown (tried 'builtin.func' as well)
```

Differential Revision: https://reviews.llvm.org/D125351

Added: 
    

Modified: 
    mlir/lib/Parser/Parser.cpp
    mlir/test/IR/invalid.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Parser/Parser.cpp b/mlir/lib/Parser/Parser.cpp
index 05daf32eb264..1f88353c1bc2 100644
--- a/mlir/lib/Parser/Parser.cpp
+++ b/mlir/lib/Parser/Parser.cpp
@@ -1641,52 +1641,38 @@ FailureOr<OperationName> OperationParser::parseCustomOperationName() {
   std::string opName = getTokenSpelling().str();
   if (opName.empty())
     return (emitError("empty operation name is invalid"), failure());
-
   consumeToken();
 
+  // Check to see if this operation name is already registered.
   Optional<RegisteredOperationName> opInfo =
       RegisteredOperationName::lookup(opName, getContext());
-  StringRef defaultDialect = getState().defaultDialectStack.back();
-  Dialect *dialect = nullptr;
-  if (opInfo) {
-    dialect = &opInfo->getDialect();
-  } else {
-    if (StringRef(opName).contains('.')) {
-      // This op has a dialect, we try to check if we can register it in the
-      // context on the fly.
-      StringRef dialectName = StringRef(opName).split('.').first;
-      dialect = getContext()->getLoadedDialect(dialectName);
-      if (!dialect && (dialect = getContext()->getOrLoadDialect(dialectName)))
-        opInfo = RegisteredOperationName::lookup(opName, getContext());
-    } else {
-      // If the operation name has no namespace prefix we lookup the current
-      // default dialect (set through OpAsmOpInterface).
-      opInfo = RegisteredOperationName::lookup(
-          Twine(defaultDialect + "." + opName).str(), getContext());
-      if (opInfo) {
-        dialect = &opInfo->getDialect();
-        opName = opInfo->getStringRef().str();
-      } else if (!defaultDialect.empty()) {
-        dialect = getContext()->getOrLoadDialect(defaultDialect);
-        opName = (defaultDialect + "." + opName).str();
-      }
-    }
+  if (opInfo)
+    return *opInfo;
+
+  // If the operation doesn't have a dialect prefix try using the default
+  // dialect.
+  auto opNameSplit = StringRef(opName).split('.');
+  StringRef dialectName = opNameSplit.first;
+  if (opNameSplit.second.empty()) {
+    dialectName = getState().defaultDialectStack.back();
+    opName = (dialectName + "." + opName).str();
   }
 
+  // Try to load the dialect before returning the operation name to make sure
+  // the operation has a chance to be registered.
+  getContext()->getOrLoadDialect(dialectName);
   return OperationName(opName, getContext());
 }
 
 Operation *
 OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
   SMLoc opLoc = getToken().getLoc();
+  StringRef originalOpName = getTokenSpelling();
 
   FailureOr<OperationName> opNameInfo = parseCustomOperationName();
   if (failed(opNameInfo))
     return nullptr;
-
   StringRef opName = opNameInfo->getStringRef();
-  Dialect *dialect = opNameInfo->getDialect();
-  Optional<RegisteredOperationName> opInfo = opNameInfo->getRegisteredInfo();
 
   // This is the actual hook for the custom op parsing, usually implemented by
   // the op itself (`Op::parse()`). We retrieve it either from the
@@ -1695,7 +1681,7 @@ OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
   bool isIsolatedFromAbove = false;
 
   StringRef defaultDialect = "";
-  if (opInfo) {
+  if (auto opInfo = opNameInfo->getRegisteredInfo()) {
     parseAssemblyFn = opInfo->getParseAssemblyFn();
     isIsolatedFromAbove = opInfo->hasTrait<OpTrait::IsIsolatedFromAbove>();
     auto *iface = opInfo->getInterface<OpAsmOpInterface>();
@@ -1703,10 +1689,13 @@ OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
       defaultDialect = iface->getDefaultDialect();
   } else {
     Optional<Dialect::ParseOpHook> dialectHook;
-    if (dialect)
+    if (Dialect *dialect = opNameInfo->getDialect())
       dialectHook = dialect->getParseOperationHook(opName);
     if (!dialectHook.hasValue()) {
-      emitError(opLoc) << "custom op '" << opName << "' is unknown";
+      InFlightDiagnostic diag =
+          emitError(opLoc) << "custom op '" << originalOpName << "' is unknown";
+      if (originalOpName != opName)
+        diag << " (tried '" << opName << "' as well)";
       return nullptr;
     }
     parseAssemblyFn = *dialectHook;

diff  --git a/mlir/test/IR/invalid.mlir b/mlir/test/IR/invalid.mlir
index 59fe5ef6c23d..7c802954f43b 100644
--- a/mlir/test/IR/invalid.mlir
+++ b/mlir/test/IR/invalid.mlir
@@ -277,6 +277,13 @@ func.func @non_operation() {
 
 // -----
 
+func.func @non_operation() {
+  // expected-error at +1 {{custom op 'asd' is unknown (tried 'func.asd' as well)}}
+  asd
+}
+
+// -----
+
 func.func @invalid_if_conditional2() {
   affine.for %i = 1 to 10 {
     affine.if affine_set<(i)[N] : (i >= )>  // expected-error {{expected '== 0' or '>= 0' at end of affine constraint}}
@@ -826,7 +833,7 @@ func.func @mixed_named_arguments(f32,
 func.func @f(f32) {
 ^bb0(%a : f32):
   %18 = arith.cmpi slt, %idx, %idx : index
-  tensor<42 x index  // expected-error {{custom op 'func.tensor' is unknown}}
+  tensor<42 x index  // expected-error {{custom op 'tensor' is unknown (tried 'func.tensor' as well)}}
   return
 }