diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
--- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
+++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
@@ -264,9 +264,8 @@
 /// This class represents all of the information pertaining to a specific MLIR
 /// document.
 struct MLIRDocument {
-  MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
-               DialectRegistry &registry,
-               std::vector<lsp::Diagnostic> &diagnostics);
+  MLIRDocument(MLIRContext &context, const lsp::URIForFile &uri,
+               StringRef contents, std::vector<lsp::Diagnostic> &diagnostics);
   MLIRDocument(const MLIRDocument &) = delete;
   MLIRDocument &operator=(const MLIRDocument &) = delete;
 
@@ -310,9 +309,6 @@
   // Fields
   //===--------------------------------------------------------------------===//
 
-  /// The context used to hold the state contained by the parsed document.
-  MLIRContext context;
-
   /// The high level parser state used to find definitions and references within
   /// the source file.
   AsmParserState asmState;
@@ -325,11 +321,9 @@
 };
 } // namespace
 
-MLIRDocument::MLIRDocument(const lsp::URIForFile &uri, StringRef contents,
-                           DialectRegistry &registry,
-                           std::vector<lsp::Diagnostic> &diagnostics)
-    : context(registry, MLIRContext::Threading::DISABLED) {
-  context.allowUnregisteredDialects();
+MLIRDocument::MLIRDocument(MLIRContext &context, const lsp::URIForFile &uri,
+                           StringRef contents,
+                           std::vector<lsp::Diagnostic> &diagnostics) {
   ScopedDiagnosticHandler handler(&context, [&](Diagnostic &diag) {
     diagnostics.push_back(getLspDiagnoticFromDiag(sourceMgr, diag, uri));
   });
@@ -657,11 +651,10 @@
 namespace {
 /// This class represents a single chunk of an MLIR text file.
 struct MLIRTextFileChunk {
-  MLIRTextFileChunk(uint64_t lineOffset, const lsp::URIForFile &uri,
-                    StringRef contents, DialectRegistry &registry,
+  MLIRTextFileChunk(MLIRContext &context, uint64_t lineOffset,
+                    const lsp::URIForFile &uri, StringRef contents,
                     std::vector<lsp::Diagnostic> &diagnostics)
-      : lineOffset(lineOffset), document(uri, contents, registry, diagnostics) {
-  }
+      : lineOffset(lineOffset), document(context, uri, contents, diagnostics) {}
 
   /// Adjust the line number of the given range to anchor at the beginning of
   /// the file, instead of the beginning of this chunk.
@@ -713,6 +706,9 @@
   /// beginning of the file.
   MLIRTextFileChunk &getChunkFor(lsp::Position &pos);
 
+  /// The context used to hold the state contained by the parsed document.
+  MLIRContext context;
+
   /// The full string contents of the file.
   std::string contents;
 
@@ -731,7 +727,10 @@
 MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
                            int64_t version, DialectRegistry &registry,
                            std::vector<lsp::Diagnostic> &diagnostics)
-    : contents(fileContents.str()), version(version), totalNumLines(0) {
+    : context(registry, MLIRContext::Threading::DISABLED),
+      contents(fileContents.str()), version(version), totalNumLines(0) {
+  context.allowUnregisteredDialects();
+
   // Split the file into separate MLIR documents.
   // TODO: Find a way to share the split file marker with other tools. We don't
   // want to use `splitAndProcessBuffer` here, but we do want to make sure this
@@ -739,13 +738,13 @@
   SmallVector<StringRef, 8> subContents;
   StringRef(contents).split(subContents, "// -----");
   chunks.emplace_back(std::make_unique<MLIRTextFileChunk>(
-      /*lineOffset=*/0, uri, subContents.front(), registry, diagnostics));
+      context, /*lineOffset=*/0, uri, subContents.front(), diagnostics));
 
   uint64_t lineOffset = subContents.front().count('\n');
   for (StringRef docContents : llvm::drop_begin(subContents)) {
     unsigned currentNumDiags = diagnostics.size();
-    auto chunk = std::make_unique<MLIRTextFileChunk>(
-        lineOffset, uri, docContents, registry, diagnostics);
+    auto chunk = std::make_unique<MLIRTextFileChunk>(context, lineOffset, uri,
+                                                     docContents, diagnostics);
     lineOffset += docContents.count('\n');
 
     // Adjust locations used in diagnostics to account for the offset from the