[mlir-lsp-server] Only use one MLIRContext per MLIRTextFile

A text file may be comprised of many different "chunks", when
the input file contains the `// -----` split markers. We don't
need to use a unique MLIRContext per chunk, as having
separate contexts is intended to allow for easy unloading of
unused data and all chunks have the same lifetime (tied to the
input file). This commit uses one context for the entire file,
greatly reducing memory consumption in certain situations (up
to 70%).

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

diff --git a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
index 35ada85..9f5a5d3 100644 (file)
--- a/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
+++ b/mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp
@@ -264,9 +264,8 @@ namespace {
 /// 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 @@ struct MLIRDocument {
   // 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 @@ struct MLIRDocument {
 };
 } // 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 @@ void MLIRDocument::findDocumentSymbols(
 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 @@ private:
   /// 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 @@ private:
 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 @@ MLIRTextFile::MLIRTextFile(const lsp::URIForFile &uri, StringRef fileContents,
   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