[mlir] Expand error message to include unregistered dialects. (#158028)

It is possible to load unregistered dialects, this can result in a
confusing error message. Mark unregistered but loaded dialects.

This is currently done as a merged list as that is most concise but
requires some additional preprocessing (did merge sort given the other
two lists are, could do it shorter and probably at not too much extra
cost if I just used SetVectors - so alternative which uses less code and
may be preferred as performance not critical here).

---------

Co-authored-by: Mehdi Amini <[email protected]>

Author: jpienaar

mlir/lib/AsmParser/Parser.cpp

@@ -2076,10 +2076,46 @@ OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
       if (originalOpName != opName)
         diag << " (tried '" << opName << "' as well)";
       auto &note = diag.attachNote();
-      note << "Registered dialects: ";
-      llvm::interleaveComma(getContext()->getAvailableDialects(), note,
-                            [&](StringRef dialect) { note << dialect; });
-      note << " ; for more info on dialect registration see "
+      note << "Available dialects: ";
+      std::vector<StringRef> registered = getContext()->getAvailableDialects();
+      auto loaded = getContext()->getLoadedDialects();
+
+      // Merge the sorted lists of registered and loaded dialects.
+      SmallVector<std::pair<StringRef, bool>> mergedDialects;
+      auto regIt = registered.begin(), regEnd = registered.end();
+      auto loadIt = loaded.rbegin(), loadEnd = loaded.rend();
+      bool isRegistered = false;
+      bool isOnlyLoaded = true;
+      while (regIt != regEnd && loadIt != loadEnd) {
+        StringRef reg = *regIt;
+        StringRef load = (*loadIt)->getNamespace();
+        if (load < reg) {
+          mergedDialects.emplace_back(load, isOnlyLoaded);
+          ++loadIt;
+        } else {
+          mergedDialects.emplace_back(reg, isRegistered);
+          ++regIt;
+          if (reg == load)
+            ++loadIt;
+        }
+      }
+      for (; regIt != regEnd; ++regIt)
+        mergedDialects.emplace_back(*regIt, isRegistered);
+      for (; loadIt != loadEnd; ++loadIt)
+        mergedDialects.emplace_back((*loadIt)->getNamespace(), isOnlyLoaded);
+
+      bool loadedUnregistered = false;
+      llvm::interleaveComma(mergedDialects, note, [&](auto &pair) {
+        note << pair.first;
+        if (pair.second) {
+          loadedUnregistered = true;
+          note << " (*)";
+        }
+      });
+      note << " ";
+      if (loadedUnregistered)
+        note << "(* corresponding to loaded but unregistered dialects)";
+      note << "; for more info on dialect registration see "
               "https://mlir.llvm.org/getting_started/Faq/"
               "#registered-loaded-dependent-whats-up-with-dialects-management";
       return nullptr;

mlir/test/IR/invalid.mlir

@@ -114,7 +114,7 @@ func.func @non_operation() {
 
 func.func @unknown_dialect_operation() {
   // expected-error@below {{Dialect `foo' not found for custom op 'foo.asd'}}
-  // expected-note-re@below {{Registered dialects:{{.*}} test{{.*}}}}
+  // expected-note-re@below {{Available dialects:{{.*}} test{{.*}}}}
   foo.asd
 }