[mlir][lsp] Register all extensions and TestDynDialect in main.

The main function of the LSP server needs to load the dialects and
similar that the server should be able to understand. When extensions
where introduced, the loading of the extensions was apparently not added
to its main functions, so ops from extensions were previously not
recognized by the server. This patch registers all extensions through
the existing convenience function, and also registers the
TestDynDialect, which `mlir-opt`s main function also registers.

Reviewed By: springerm

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

Author: ingomueller-net

mlir/tools/mlir-lsp-server/CMakeLists.txt

@@ -17,6 +17,7 @@ if(MLIR_INCLUDE_TESTS)
     MLIRSPIRVTestPasses
     MLIRTestAnalysis
     MLIRTestDialect
+    MLIRTestDynDialect
     MLIRTestIR
     MLIRTestPass
     MLIRTestReducer
@@ -27,12 +28,15 @@ if(MLIR_INCLUDE_TESTS)
 endif()
 
 set(LIBS
-  ${dialect_libs}
   ${conversion_libs}
+  ${dialect_libs}
+  ${extension_libs}
   ${test_libs}
+
   MLIRAffineAnalysis
   MLIRAnalysis
   MLIRDialect
+  MLIRFuncAllExtensions
   MLIRLspServerLib
   MLIRParser
   MLIRPass

mlir/tools/mlir-lsp-server/mlir-lsp-server.cpp

@@ -9,23 +9,28 @@
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/InitAllDialects.h"
+#include "mlir/InitAllExtensions.h"
 #include "mlir/Tools/mlir-lsp-server/MlirLspServerMain.h"
 
 using namespace mlir;
 
 #ifdef MLIR_INCLUDE_TESTS
 namespace test {
 void registerTestDialect(DialectRegistry &);
+void registerTestDynDialect(DialectRegistry &);
 void registerTestTransformDialectExtension(DialectRegistry &);
 } // namespace test
 #endif
 
 int main(int argc, char **argv) {
   DialectRegistry registry;
   registerAllDialects(registry);
+  registerAllExtensions(registry);
+
 #ifdef MLIR_INCLUDE_TESTS
   ::test::registerTestDialect(registry);
   ::test::registerTestTransformDialectExtension(registry);
+  ::test::registerTestDynDialect(registry);
 #endif
   return failed(MlirLspServerMain(argc, argv, registry));
 }