[llvm] [mlir] [mlir][lsp] Enable registering dialects based on URI. (PR #141331)

Jacques Pienaar via llvm-commits llvm-commits at lists.llvm.org
Sat May 24 22:55:34 PDT 2025 

https://github.com/jpienaar updated https://github.com/llvm/llvm-project/pull/141331

>From bcc3878265d5fdca8a3ca6ee7be9404825a0affe Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Sun, 25 May 2025 05:50:42 +0000
Subject: [PATCH 1/3] [mlir][lsp] Enable registering dialects based on URI.

Previously the dialects registered were fixed per LSP binary. This works
as long as all the dialects of interest from the different projects
across which one uses the LSP, are disjoint. This expands this to
support cases where there are dialects that overlap in dialect name but
usage of these are separate wrt projects. The alternative is multiple
binaries and switching LSP used in editor per project (there is some
extra complexity in hosted instances).

This handles a simple (I believe common case) where one can determine
based on path and have single binary - the cost of dynamically doing so
based on path would be either keeping different registries to return or
repopulating dialect & extension maps.
---
 .../mlir-lsp-server/MlirLspRegistryFunction.h | 29 ++++++++++++++
 .../Tools/mlir-lsp-server/MlirLspServerMain.h | 11 ++++-
 mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp | 22 +++++-----
 mlir/lib/Tools/mlir-lsp-server/MLIRServer.h   |  6 +--
 .../mlir-lsp-server/MlirLspServerMain.cpp     | 13 +++++-
 .../uri-based-registration.test               | 22 ++++++++++
 .../tools/mlir-lsp-server/mlir-lsp-server.cpp | 19 +++++++--
 .../llvm-project-overlay/mlir/BUILD.bazel     | 40 +++++++++++++++++++
 8 files changed, 141 insertions(+), 21 deletions(-)
 create mode 100644 mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h
 create mode 100644 mlir/test/mlir-lsp-server/uri-based-registration.test

diff --git a/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h
new file mode 100644
index 0000000000000..c83b361a0e51c
--- /dev/null
+++ b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h
@@ -0,0 +1,29 @@
+//===- MlirLspRegistryFunction.h - LSP registry functions -------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Registry function types for MLIR LSP.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPREGISTRYFUNCTION_H
+#define MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPREGISTRYFUNCTION_H
+
+namespace llvm {
+template <typename Fn> class function_ref;
+} // namespace llvm
+
+namespace mlir {
+class DialectRegistry;
+namespace lsp {
+class URIForFile;
+using DialectRegistryFn =
+    llvm::function_ref<DialectRegistry &(const URIForFile &uri)>;
+} // namespace lsp
+} // namespace mlir
+
+#endif // MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPREGISTRYFUNCTION_H
diff --git a/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h
index 66d5d40a6d28d..a461fc4702946 100644
--- a/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h
+++ b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspServerMain.h
@@ -12,20 +12,27 @@
 
 #ifndef MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPSERVERMAIN_H
 #define MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPSERVERMAIN_H
+#include "mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h"
 
 namespace llvm {
 struct LogicalResult;
 } // namespace llvm
 
 namespace mlir {
-class DialectRegistry;
 
 /// Implementation for tools like `mlir-lsp-server`.
 /// - registry should contain all the dialects that can be parsed in source IR
-/// passed to the server.
+///   passed to the server.
 llvm::LogicalResult MlirLspServerMain(int argc, char **argv,
                                       DialectRegistry &registry);
 
+/// Implementation for tools like `mlir-lsp-server`.
+/// - registry should contain all the dialects that can be parsed in source IR
+///   passed to the server and may register different dialects depending on the
+///   input URI.
+llvm::LogicalResult MlirLspServerMain(int argc, char **argv,
+                                      lsp::DialectRegistryFn registry_fn);
+
 } // namespace mlir
 
 #endif // MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPSERVERMAIN_H
diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
index 4e19274c3da40..61987525a5ca5 100644
--- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
+++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
@@ -997,7 +997,7 @@ namespace {
 class MLIRTextFile {
 public:
   MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
-               int64_t version, DialectRegistry &registry,
+               int64_t version, lsp::DialectRegistryFn registry_fn,
                std::vector<lsp::Diagnostic> &diagnostics);
 
   /// Return the current version of this text file.
@@ -1046,9 +1046,9 @@ class MLIRTextFile {
 } // namespace
 
 MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
-                           int64_t version, DialectRegistry &registry,
+                           int64_t version, lsp::DialectRegistryFn registry_fn,
                            std::vector<lsp::Diagnostic> &diagnostics)
-    : context(registry, MLIRContext::Threading::DISABLED),
+    : context(registry_fn(uri), MLIRContext::Threading::DISABLED),
       contents(fileContents.str()), version(version) {
   context.allowUnregisteredDialects();
 
@@ -1263,11 +1263,11 @@ MLIRTextFileChunk &MLIRTextFile::getChunkFor(lsp::Position &pos) {
 //===----------------------------------------------------------------------===//
 
 struct lsp::MLIRServer::Impl {
-  Impl(DialectRegistry &registry) : registry(registry) {}
+  Impl(lsp::DialectRegistryFn registry_fn) : registry_fn(registry_fn) {}
 
-  /// The registry containing dialects that can be recognized in parsed .mlir
-  /// files.
-  DialectRegistry ®istry;
+  /// The registry factory for containing dialects that can be recognized in
+  /// parsed .mlir files.
+  lsp::DialectRegistryFn registry_fn;
 
   /// The files held by the server, mapped by their URI file name.
   llvm::StringMap<std::unique_ptr<MLIRTextFile>> files;
@@ -1277,15 +1277,15 @@ struct lsp::MLIRServer::Impl {
 // MLIRServer
 //===----------------------------------------------------------------------===//
 
-lsp::MLIRServer::MLIRServer(DialectRegistry &registry)
-    : impl(std::make_unique<Impl>(registry)) {}
+lsp::MLIRServer::MLIRServer(lsp::DialectRegistryFn registry_fn)
+    : impl(std::make_unique<Impl>(registry_fn)) {}
 lsp::MLIRServer::~MLIRServer() = default;
 
 void lsp::MLIRServer::addOrUpdateDocument(
     const URIForFile &uri, StringRef contents, int64_t version,
     std::vector<Diagnostic> &diagnostics) {
   impl->files[uri.file()] = std::make_unique<MLIRTextFile>(
-      uri, contents, version, impl->registry, diagnostics);
+      uri, contents, version, impl->registry_fn, diagnostics);
 }
 
 std::optional<int64_t> lsp::MLIRServer::removeDocument(const URIForFile &uri) {
@@ -1348,7 +1348,7 @@ void lsp::MLIRServer::getCodeActions(const URIForFile &uri, const Range &pos,
 
 llvm::Expected<lsp::MLIRConvertBytecodeResult>
 lsp::MLIRServer::convertFromBytecode(const URIForFile &uri) {
-  MLIRContext tempContext(impl->registry);
+  MLIRContext tempContext(impl->registry_fn(uri));
   tempContext.allowUnregisteredDialects();
 
   // Collect any errors during parsing.
diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.h b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.h
index 979be615b82cc..85e69e69f6631 100644
--- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.h
+++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.h
@@ -10,6 +10,7 @@
 #define LIB_MLIR_TOOLS_MLIRLSPSERVER_SERVER_H_
 
 #include "mlir/Support/LLVM.h"
+#include "mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h"
 #include "llvm/Support/Error.h"
 #include <memory>
 #include <optional>
@@ -28,15 +29,14 @@ struct Location;
 struct MLIRConvertBytecodeResult;
 struct Position;
 struct Range;
-class URIForFile;
 
 /// This class implements all of the MLIR related functionality necessary for a
 /// language server. This class allows for keeping the MLIR specific logic
 /// separate from the logic that involves LSP server/client communication.
 class MLIRServer {
 public:
-  /// Construct a new server with the given dialect regitstry.
-  MLIRServer(DialectRegistry &registry);
+  /// Construct a new server with the given dialect registry function.
+  MLIRServer(DialectRegistryFn registry_fn);
   ~MLIRServer();
 
   /// Add or update the document, with the provided `version`, at the given URI.
diff --git a/mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp b/mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp
index 259bd2613a6cc..b1bbf98ce769e 100644
--- a/mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp
+++ b/mlir/lib/Tools/mlir-lsp-server/MlirLspServerMain.cpp
@@ -19,7 +19,7 @@ using namespace mlir;
 using namespace mlir::lsp;
 
 LogicalResult mlir::MlirLspServerMain(int argc, char **argv,
-                                      DialectRegistry &registry) {
+                                      DialectRegistryFn registry_fn) {
   llvm::cl::opt<JSONStreamStyle> inputStyle{
       "input-style",
       llvm::cl::desc("Input JSON stream encoding"),
@@ -72,6 +72,15 @@ LogicalResult mlir::MlirLspServerMain(int argc, char **argv,
   URIForFile::registerSupportedScheme("mlir.bytecode-mlir");
 
   // Configure the servers and start the main language server.
-  MLIRServer server(registry);
+  MLIRServer server(registry_fn);
   return runMlirLSPServer(server, transport);
 }
+
+llvm::LogicalResult mlir::MlirLspServerMain(int argc, char **argv,
+                                            DialectRegistry &registry) {
+  auto registry_fn =
+      [&registry](const lsp::URIForFile &uri) -> DialectRegistry & {
+    return registry;
+  };
+  return MlirLspServerMain(argc, argv, registry_fn);
+}
diff --git a/mlir/test/mlir-lsp-server/uri-based-registration.test b/mlir/test/mlir-lsp-server/uri-based-registration.test
new file mode 100644
index 0000000000000..52886b3ef04ee
--- /dev/null
+++ b/mlir/test/mlir-lsp-server/uri-based-registration.test
@@ -0,0 +1,22 @@
+// RUN: not mlir-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"mlir","capabilities":{},"trace":"off"}}
+// -----
+// Just regular parse, successful.
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo-regular-registration.mlir",
+  "languageId":"mlir",
+  "version":1,
+  "text":"func.func @fail_with_empty_registry() { return }"
+}}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK:     "diagnostics": []
+// -----
+// Just regular parse, successful.
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo-disable-lsp-registration.mlir",
+  "languageId":"mlir",
+  "version":1,
+  "text":"func.func @fail_with_empty_registry() { return }"
+}}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK: "message": "Dialect `func' not found for custom op 'func.func'
\ No newline at end of file
diff --git a/mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp b/mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp
index f0ecc5adc68b3..3e57ec6046564 100644
--- a/mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp
+++ b/mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp
@@ -6,10 +6,10 @@
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/IR/Dialect.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
 #include "mlir/InitAllExtensions.h"
+#include "mlir/Tools/lsp-server-support/Protocol.h"
 #include "mlir/Tools/mlir-lsp-server/MlirLspServerMain.h"
 
 using namespace mlir;
@@ -23,7 +23,7 @@ void registerTestTransformDialectExtension(DialectRegistry &);
 #endif
 
 int main(int argc, char **argv) {
-  DialectRegistry registry;
+  DialectRegistry registry, empty;
   registerAllDialects(registry);
   registerAllExtensions(registry);
 
@@ -32,5 +32,18 @@ int main(int argc, char **argv) {
   ::test::registerTestTransformDialectExtension(registry);
   ::test::registerTestDynDialect(registry);
 #endif
-  return failed(MlirLspServerMain(argc, argv, registry));
+
+  // Returns the registry, except in testing mode when the URI contains
+  // "_disable_lsp_registration".
+  auto registryFn = [&registry,
+                     &empty](const lsp::URIForFile &uri) -> DialectRegistry & {
+    (void)empty;
+#ifdef MLIR_INCLUDE_TESTS
+    llvm::errs() << "uri: " << uri.uri() << "\n";
+    if (uri.uri().contains("-disable-lsp-registration"))
+      return empty;
+#endif
+    return registry;
+  };
+  return failed(MlirLspServerMain(argc, argv, registryFn));
 }
diff --git a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
index 1cbc632952c4d..8b1c76717e2f2 100644
--- a/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/BUILD.bazel
@@ -8835,11 +8835,51 @@ cc_binary(
     name = "mlir-lsp-server",
     srcs = ["tools/mlir-lsp-server/mlir-lsp-server.cpp"],
     includes = ["include"],
+    local_defines = ["MLIR_INCLUDE_TESTS"],
     deps = [
         ":AllExtensions",
         ":AllPassesAndDialects",
         ":IR",
         ":MlirLspServerLib",
+        ":MlirLspServerSupportLib",
+        "//mlir/test:TestAffine",
+        "//mlir/test:TestAnalysis",
+        "//mlir/test:TestArith",
+        "//mlir/test:TestArmNeon",
+        "//mlir/test:TestArmSME",
+        "//mlir/test:TestBufferization",
+        "//mlir/test:TestControlFlow",
+        "//mlir/test:TestConvertToSPIRV",
+        "//mlir/test:TestDLTI",
+        "//mlir/test:TestDialect",
+        "//mlir/test:TestFunc",
+        "//mlir/test:TestFuncToLLVM",
+        "//mlir/test:TestGPU",
+        "//mlir/test:TestIR",
+        "//mlir/test:TestLLVM",
+        "//mlir/test:TestLinalg",
+        "//mlir/test:TestLoopLikeInterface",
+        "//mlir/test:TestMath",
+        "//mlir/test:TestMathToVCIX",
+        "//mlir/test:TestMemRef",
+        "//mlir/test:TestMesh",
+        "//mlir/test:TestNVGPU",
+        "//mlir/test:TestPDLL",
+        "//mlir/test:TestPass",
+        "//mlir/test:TestReducer",
+        "//mlir/test:TestRewrite",
+        "//mlir/test:TestSCF",
+        "//mlir/test:TestSPIRV",
+        "//mlir/test:TestShapeDialect",
+        "//mlir/test:TestTensor",
+        "//mlir/test:TestTestDynDialect",
+        "//mlir/test:TestTilingInterface",
+        "//mlir/test:TestTosaDialect",
+        "//mlir/test:TestTransformDialect",
+        "//mlir/test:TestTransforms",
+        "//mlir/test:TestVector",
+        "//mlir/test:TestVectorToSPIRV",
+        "//mlir/test:TestXeGPU",
     ],
 )
 

>From a789598c54d4574b7e639a1a900dc1f18238caf4 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Sun, 25 May 2025 05:54:16 +0000
Subject: [PATCH 2/3] Newline at EOF

---
 .../mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h       | 3 ++-
 mlir/test/mlir-lsp-server/uri-based-registration.test          | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h
index c83b361a0e51c..4811ecb5e92b7 100644
--- a/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h
+++ b/mlir/include/mlir/Tools/mlir-lsp-server/MlirLspRegistryFunction.h
@@ -14,7 +14,8 @@
 #define MLIR_TOOLS_MLIR_LSP_SERVER_MLIRLSPREGISTRYFUNCTION_H
 
 namespace llvm {
-template <typename Fn> class function_ref;
+template <typename Fn>
+class function_ref;
 } // namespace llvm
 
 namespace mlir {
diff --git a/mlir/test/mlir-lsp-server/uri-based-registration.test b/mlir/test/mlir-lsp-server/uri-based-registration.test
index 52886b3ef04ee..d6d06692a8d7a 100644
--- a/mlir/test/mlir-lsp-server/uri-based-registration.test
+++ b/mlir/test/mlir-lsp-server/uri-based-registration.test
@@ -19,4 +19,5 @@
   "text":"func.func @fail_with_empty_registry() { return }"
 }}}
 // CHECK: "method": "textDocument/publishDiagnostics",
-// CHECK: "message": "Dialect `func' not found for custom op 'func.func'
\ No newline at end of file
+// CHECK: "message": "Dialect `func' not found for custom op 'func.func'
+

>From 0fd04d8cc37863004fed85bf9ac48e9c964bde17 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Sun, 25 May 2025 05:55:22 +0000
Subject: [PATCH 3/3] Remove debugging

---
 mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp b/mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp
index 3e57ec6046564..6a759d9e0d60f 100644
--- a/mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp
+++ b/mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp
@@ -34,12 +34,12 @@ int main(int argc, char **argv) {
 #endif
 
   // Returns the registry, except in testing mode when the URI contains
-  // "_disable_lsp_registration".
+  // "-disable-lsp-registration". Testing for/example of registering dialects
+  // based on URI.
   auto registryFn = [&registry,
                      &empty](const lsp::URIForFile &uri) -> DialectRegistry & {
     (void)empty;
 #ifdef MLIR_INCLUDE_TESTS
-    llvm::errs() << "uri: " << uri.uri() << "\n";
     if (uri.uri().contains("-disable-lsp-registration"))
       return empty;
 #endif

More information about the llvm-commits mailing list