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

[rriddle]

Summary

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%).

[jpienaar]

And this doesn't affect the error cases? E.g., split is often used when testing for errors reported, so using 1 context doesn't affect that

[rriddle]

In D107488#2926295, @jpienaar wrote:

And this doesn't affect the error cases? E.g., split is often used when testing for errors reported, so using 1 context doesn't affect that

Yeah, having one context doesn't affect diagnostics or anything related to split really (if it does, that's probably a bug somewhere). The context only really holds immutable state, so everything should still work the same before.

mlir/lib/Tools/mlir-lsp-server/MLIRServer.cpp

 //===----------------------------------------------------------------------===//
 // MLIRDocument
 //===----------------------------------------------------------------------===//
 
 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;
 
   //===--------------------------------------------------------------------===//
   // Definitions and References
   //===--------------------------------------------------------------------===//
 
   void getLocationsOf(const lsp::URIForFile &uri, const lsp::Position &defPos,
struct MLIRDocument {
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
   void findDocumentSymbols(Operation *op,
                            std::vector<lsp::DocumentSymbol> &symbols);
 
   //===--------------------------------------------------------------------===//
   // 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;
 
   /// The container for the IR parsed from the input file.
   Block parsedIR;
 
   /// The source manager containing the contents of the input file.
   llvm::SourceMgr sourceMgr;
 };
 } // 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));
   });
 
   // Try to parsed the given IR string.
   auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
   if (!memBuffer) {
     lsp::Logger::error("Failed to create memory buffer for file", uri.file());

 
 //===----------------------------------------------------------------------===//
 // MLIRTextFileChunk
 //===----------------------------------------------------------------------===//
 
 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.
   void adjustLocForChunkOffset(lsp::Range &range) {
     adjustLocForChunkOffset(range.start);
     adjustLocForChunkOffset(range.end);
   }
   /// Adjust the line number of the given position to anchor at the beginning of
public:
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
 
 private:
   /// Find the MLIR document that contains the given position, and update the
   /// position to be anchored at the start of the found chunk instead of the
   /// 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;
 
   /// The version of this file.
   int64_t version;
 
   /// The number of lines in the file.
   int64_t totalNumLines;
 
   /// The chunks of this file. The order of these chunks is the order in which
   /// they appear in the text file.
   std::vector<std::unique_ptr<MLIRTextFileChunk>> chunks;
 };
 } // namespace
 
 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
   // marker doesn't go out of sync.
   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
     // beginning of the file.
     for (lsp::Diagnostic &diag :
          llvm::drop_begin(diagnostics, currentNumDiags)) {
       chunk->adjustLocForChunkOffset(diag.range);