[mlir:vscode] Add support for viewing and editing a bytecode file as .mlir

[rriddle]

Summary

This commit adds support for interacting with a (valid) bytecode file in the same
way as .mlir. This allows editing, using all of the traditional LSP features, etc. but
still using bytecode as the on-disk serialization format. Loading a bytecode file this
way will fail if the bytecode is invalid, and saving will fail if the edited .mlir is invalid.

[jpienaar]

Nice! The typescript I peripherally followed only. I was expecting this to be worse. Is period in schema the preferred form? (I saw some with colons but wasn't sure what is the defacto reference)

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

320

I thought ModuleOp was intended to not be required here?

612

Was this just discovered with this change or why is changed here along with bytecode?

1279

Could we avoid saving in this case to not make it accidentally lossy viewing?

[rriddle]

In D132970#3762811, @jpienaar wrote:

Nice! The typescript I peripherally followed only. I was expecting this to be worse. Is period in schema the preferred form? (I saw some with colons but wasn't sure what is the defacto reference)

period/dash/etc are all fine, colon separates the scheme from the rest of the content.

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

320

Thanks for the comment here. It was in my head that we need this, but the real constraint we have for bytecode right now is single op, restructured this to just block loading/saving on that (with TODOs that we can fix that if/when necessary).

612

Unnecessary now.

1279

We can't even detect this right now. I think we should have some mechanism for parsing in unknown external resources opaquely, which would allow for us to round trip. The producer case I'm not sure how we want to handle. If the user saves, technically the producer of the bytecode will be different. It will also be interesting in how the support here interacts with versioning and auto-upgrade whenever that happens.

mlir/lib/Tools/lsp-server-support/Protocol.h

enum class ErrorCode {
   InternalError = -32603,
 
   ServerNotInitialized = -32002,
   UnknownErrorCode = -32001,
 
   // Defined by the protocol.
   RequestCancelled = -32800,
   ContentModified = -32801,
+  RequestFailed = -32803,
 };
 
 /// Defines how the host (editor) should sync document changes to the language
 /// server.
 enum class TextDocumentSyncKind {
   /// Documents should not be synced at all.
   None = 0,
 

 /// URI in "file" scheme for a file.
 class URIForFile {
 public:
   URIForFile() = default;
 
   /// Try to build a URIForFile from the given URI string.
   static llvm::Expected<URIForFile> fromURI(StringRef uri);
 
-  /// Try to build a URIForFile from the given absolute file path.
-  static llvm::Expected<URIForFile> fromFile(StringRef absoluteFilepath);
+  /// Try to build a URIForFile from the given absolute file path and optional
+  /// scheme.
+  static llvm::Expected<URIForFile> fromFile(StringRef absoluteFilepath,
+                                             StringRef scheme = "file");
 
   /// Returns the absolute path to the file.
   StringRef file() const { return filePath; }
 
   /// Returns the original uri of the file.
   StringRef uri() const { return uriStr; }
 
+  /// Return the scheme of the uri.
+  StringRef scheme() const;
+
   explicit operator bool() const { return !filePath.empty(); }
 
   friend bool operator==(const URIForFile &lhs, const URIForFile &rhs) {
     return lhs.filePath == rhs.filePath;
   }
   friend bool operator!=(const URIForFile &lhs, const URIForFile &rhs) {
     return !(lhs == rhs);
   }
   friend bool operator<(const URIForFile &lhs, const URIForFile &rhs) {
     return lhs.filePath < rhs.filePath;
   }
 
+  /// Register a supported URI scheme. The protocol supports `file` by default,
+  /// so this is only necessary for any additional schemes that a server wants
+  /// to support.
+  static void registerSupportedScheme(StringRef scheme);
+
 private:
   explicit URIForFile(std::string &&filePath, std::string &&uriStr)
       : filePath(std::move(filePath)), uriStr(uriStr) {}
 
   std::string filePath;
   std::string uriStr;
 };
 

mlir/lib/Tools/lsp-server-support/Protocol.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "Protocol.h"
 #include "Logging.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/StringSwitch.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/Format.h"
 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/Path.h"
 #include "llvm/Support/raw_ostream.h"
if (*i == '%' && i + 2 < content.end() && llvm::isHexDigit(*(i + 1)) &&
       i += 2;
     } else {
       result.push_back(*i);
     }
   }
   return result;
 }
 
-static bool isValidScheme(StringRef scheme) {
+/// Return the set containing the supported URI schemes.
+static StringSet<> &getSupportedSchemes() {
+  static StringSet<> schemes({"file", "test"});
+  return schemes;
+}
+
+/// Returns true if the given scheme is structurally valid, i.e. it does not
+/// contain any invalid scheme characters. This does not check that the scheme
+/// is actually supported.
+static bool isStructurallyValidScheme(StringRef scheme) {
   if (scheme.empty())
     return false;
   if (!llvm::isAlpha(scheme[0]))
     return false;
   return llvm::all_of(llvm::drop_begin(scheme), [](char c) {
     return llvm::isAlnum(c) || c == '+' || c == '.' || c == '-';
   });
 }
 
-static llvm::Expected<std::string> uriFromAbsolutePath(StringRef absolutePath) {
+static llvm::Expected<std::string> uriFromAbsolutePath(StringRef absolutePath,
+                                                       StringRef scheme) {
   std::string body;
   StringRef authority;
   StringRef root = llvm::sys::path::root_name(absolutePath);
   if (isNetworkPath(root)) {
     // Windows UNC paths e.g. \\server\share => file://server/share
     authority = root.drop_front(2);
     absolutePath.consume_front(root);
   } else if (isWindowsPath(root)) {
     // Windows paths e.g. X:\path => file:///X:/path
     body = "/";
   }
   body += llvm::sys::path::convert_to_slash(absolutePath);
 
-  std::string uri = "file:";
+  std::string uri = scheme.str() + ":";
   if (authority.empty() && body.empty())
     return uri;
 
   // If authority if empty, we only print body if it starts with "/"; otherwise,
   // the URI is invalid.
   if (!authority.empty() || StringRef(body).startswith("/")) {
     uri.append("//");
     percentEncode(authority, uri);
static llvm::Expected<std::string> parseFilePathFromURI(StringRef origUri) {
   // Decode the scheme of the URI.
   size_t pos = uri.find(':');
   if (pos == StringRef::npos)
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "Scheme must be provided in URI: " +
                                        origUri);
   StringRef schemeStr = uri.substr(0, pos);
   std::string uriScheme = percentDecode(schemeStr);
-  if (!isValidScheme(uriScheme))
+  if (!isStructurallyValidScheme(uriScheme))
     return llvm::createStringError(llvm::inconvertibleErrorCode(),
                                    "Invalid scheme: " + schemeStr +
                                        " (decoded: " + uriScheme + ")");
   uri = uri.substr(pos + 1);
 
   // Decode the authority of the URI.
   std::string uriAuthority;
   if (uri.consume_front("//")) {
     pos = uri.find('/');
     uriAuthority = percentDecode(uri.substr(0, pos));
     uri = uri.substr(pos);
   }
 
   // Decode the body of the URI.
   std::string uriBody = percentDecode(uri);
 
   // Compute the absolute path for this uri.
-  if (uriScheme != "file" && uriScheme != "test") {
-    return llvm::createStringError(
-        llvm::inconvertibleErrorCode(),
-        "mlir-lsp-server only supports 'file' URI scheme for workspace files");
+  if (!getSupportedSchemes().contains(uriScheme)) {
+    return llvm::createStringError(llvm::inconvertibleErrorCode(),
+                                   "unsupported URI scheme `" + uriScheme +
+                                       "' for workspace files");
   }
   return getAbsolutePath(uriAuthority, uriBody);
 }
 
 llvm::Expected<URIForFile> URIForFile::fromURI(StringRef uri) {
   llvm::Expected<std::string> filePath = parseFilePathFromURI(uri);
   if (!filePath)
     return filePath.takeError();
   return URIForFile(std::move(*filePath), uri.str());
 }
 
-llvm::Expected<URIForFile> URIForFile::fromFile(StringRef absoluteFilepath) {
-  llvm::Expected<std::string> uri = uriFromAbsolutePath(absoluteFilepath);
+llvm::Expected<URIForFile> URIForFile::fromFile(StringRef absoluteFilepath,
+                                                StringRef scheme) {
+  llvm::Expected<std::string> uri =
+      uriFromAbsolutePath(absoluteFilepath, scheme);
   if (!uri)
     return uri.takeError();
   return fromURI(*uri);
 }
 
+StringRef URIForFile::scheme() const { return uri().split(':').first; }
+
+void URIForFile::registerSupportedScheme(StringRef scheme) {
+  getSupportedSchemes().insert(scheme);
+}
+
 bool mlir::lsp::fromJSON(const llvm::json::Value &value, URIForFile &result,
                          llvm::json::Path path) {
   if (Optional<StringRef> str = value.getAsString()) {
     llvm::Expected<URIForFile> expectedURI = URIForFile::fromURI(*str);
     if (!expectedURI) {
       path.report("unresolvable URI");
       consumeError(expectedURI.takeError());
       return false;

mlir/lib/Tools/mlir-lsp-server/CMakeLists.txt

 add_mlir_library(MLIRLspServerLib
   LSPServer.cpp
   MLIRServer.cpp
   MlirLspServerMain.cpp
+  Protocol.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/mlir-lsp-server
 
   LINK_LIBS PUBLIC
+  MLIRBytecodeWriter
   MLIRIR
   MLIRLspServerSupportLib
   MLIRParser
   )

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

 //===- LSPServer.cpp - MLIR Language Server -------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "LSPServer.h"
 #include "../lsp-server-support/Logging.h"
-#include "../lsp-server-support/Protocol.h"
 #include "../lsp-server-support/Transport.h"
 #include "MLIRServer.h"
+#include "Protocol.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/ADT/StringMap.h"
 
 #define DEBUG_TYPE "mlir-lsp-server"
 
 using namespace mlir;
 using namespace mlir::lsp;
 
struct LSPServer {
 
   //===--------------------------------------------------------------------===//
   // Code Action
 
   void onCodeAction(const CodeActionParams &params,
                     Callback<llvm::json::Value> reply);
 
   //===--------------------------------------------------------------------===//
+  // Bytecode
+
+  void onConvertFromBytecode(const MLIRConvertBytecodeParams &params,
+                             Callback<MLIRConvertBytecodeResult> reply);
+  void onConvertToBytecode(const MLIRConvertBytecodeParams &params,
+                           Callback<MLIRConvertBytecodeResult> reply);
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   MLIRServer &server;
 
   /// An outgoing notification used to send diagnostics to the client when they
   /// are ready to be processed.
   OutgoingNotification<PublishDiagnosticsParams> publishDiagnostics;
void LSPServer::onCodeAction(const CodeActionParams &params,
   // We provide a code action for fixes on the specified diagnostics.
   std::vector<CodeAction> actions;
   if (isKindAllowed(CodeAction::kQuickFix))
     server.getCodeActions(uri, params.range.start, params.context, actions);
   reply(std::move(actions));
 }
 
 //===----------------------------------------------------------------------===//
+// Bytecode
+
+void LSPServer::onConvertFromBytecode(
+    const MLIRConvertBytecodeParams &params,
+    Callback<MLIRConvertBytecodeResult> reply) {
+  reply(server.convertFromBytecode(params.uri));
+}
+
+void LSPServer::onConvertToBytecode(const MLIRConvertBytecodeParams &params,
+                                    Callback<MLIRConvertBytecodeResult> reply) {
+  reply(server.convertToBytecode(params.uri));
+}
+
+//===----------------------------------------------------------------------===//
 // Entry point
 //===----------------------------------------------------------------------===//
 
 LogicalResult lsp::runMlirLSPServer(MLIRServer &server,
                                     JSONTransport &transport) {
   LSPServer lspServer(server);
   MessageHandler messageHandler(transport);
 
LogicalResult lsp::runMlirLSPServer(MLIRServer &server,
   // Code Completion
   messageHandler.method("textDocument/completion", &lspServer,
                         &LSPServer::onCompletion);
 
   // Code Action
   messageHandler.method("textDocument/codeAction", &lspServer,
                         &LSPServer::onCodeAction);
 
+  // Bytecode
+  messageHandler.method("mlir/convertFromBytecode", &lspServer,
+                        &LSPServer::onConvertFromBytecode);
+  messageHandler.method("mlir/convertToBytecode", &lspServer,
+                        &LSPServer::onConvertToBytecode);
+
   // Diagnostics
   lspServer.publishDiagnostics =
       messageHandler.outgoingNotification<PublishDiagnosticsParams>(
           "textDocument/publishDiagnostics");
 
   // Run the main loop of the transport.
   LogicalResult result = success();
   if (llvm::Error error = transport.run(messageHandler)) {
     Logger::error("Transport error: {0}", error);
     llvm::consumeError(std::move(error));
     result = failure();
   } else {
     result = success(lspServer.shutdownRequestReceived);
   }
   return result;
 }

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

 //===- MLIRServer.h - MLIR General Language Server --------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LIB_MLIR_TOOLS_MLIRLSPSERVER_SERVER_H_
 #define LIB_MLIR_TOOLS_MLIRLSPSERVER_SERVER_H_
 
 #include "mlir/Support/LLVM.h"
+#include "llvm/Support/Error.h"
 #include <memory>
 
 namespace mlir {
 class DialectRegistry;
 
 namespace lsp {
 struct CodeAction;
 struct CodeActionContext;
 struct CompletionList;
 struct Diagnostic;
 struct DocumentSymbol;
 struct Hover;
 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:
   CompletionList getCodeCompletion(const URIForFile &uri,
                                    const Position &completePos);
 
   /// Get the set of code actions within the file.
   void getCodeActions(const URIForFile &uri, const Range &pos,
                       const CodeActionContext &context,
                       std::vector<CodeAction> &actions);
 
+  /// Convert the given bytecode file to the textual format.
+  llvm::Expected<MLIRConvertBytecodeResult>
+  convertFromBytecode(const URIForFile &uri);
+
+  /// Convert the given textual file to the bytecode format.
+  llvm::Expected<MLIRConvertBytecodeResult>
+  convertToBytecode(const URIForFile &uri);
+
 private:
   struct Impl;
 
   std::unique_ptr<Impl> impl;
 };
 
 } // namespace lsp
 } // namespace mlir
 
 #endif // LIB_MLIR_TOOLS_MLIRLSPSERVER_SERVER_H_

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

 //===- MLIRServer.cpp - MLIR Generic Language Server ----------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "MLIRServer.h"
 #include "../lsp-server-support/Logging.h"
-#include "../lsp-server-support/Protocol.h"
 #include "../lsp-server-support/SourceMgrUtils.h"
+#include "Protocol.h"
 #include "mlir/AsmParser/AsmParser.h"
 #include "mlir/AsmParser/AsmParserState.h"
 #include "mlir/AsmParser/CodeComplete.h"
+#include "mlir/Bytecode/BytecodeWriter.h"
 #include "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/Operation.h"
+#include "mlir/Parser/Parser.h"
+#include "llvm/Support/Base64.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 
 /// Returns a language server location from the given MLIR file location.
-static Optional<lsp::Location> getLocationFromLoc(FileLineColLoc loc) {
+/// `uriScheme` is the scheme to use when building new uris.
+static Optional<lsp::Location> getLocationFromLoc(StringRef uriScheme,
+                                                  FileLineColLoc loc) {
   llvm::Expected<lsp::URIForFile> sourceURI =
-      lsp::URIForFile::fromFile(loc.getFilename());
+      lsp::URIForFile::fromFile(loc.getFilename(), uriScheme);
   if (!sourceURI) {
     lsp::Logger::error("Failed to create URI for file `{0}`: {1}",
                        loc.getFilename(),
                        llvm::toString(sourceURI.takeError()));
     return llvm::None;
   }
 
   lsp::Position position;
   position.line = loc.getLine() - 1;
   position.character = loc.getColumn() ? loc.getColumn() - 1 : 0;
   return lsp::Location{*sourceURI, lsp::Range(position)};
 }
 
 /// Returns a language server location from the given MLIR location, or None if
-/// one couldn't be created. `uri` is an optional additional filter that, when
-/// present, is used to filter sub locations that do not share the same uri.
+/// one couldn't be created. `uriScheme` is the scheme to use when building new
+/// uris. `uri` is an optional additional filter that, when present, is used to
+/// filter sub locations that do not share the same uri.
 static Optional<lsp::Location>
 getLocationFromLoc(llvm::SourceMgr &sourceMgr, Location loc,
-                   const lsp::URIForFile *uri = nullptr) {
+                   StringRef uriScheme, const lsp::URIForFile *uri = nullptr) {
   Optional<lsp::Location> location;
   loc->walk([&](Location nestedLoc) {
     FileLineColLoc fileLoc = nestedLoc.dyn_cast<FileLineColLoc>();
     if (!fileLoc)
       return WalkResult::advance();
 
-    Optional<lsp::Location> sourceLoc = getLocationFromLoc(fileLoc);
+    Optional<lsp::Location> sourceLoc = getLocationFromLoc(uriScheme, fileLoc);
     if (sourceLoc && (!uri || sourceLoc->uri == *uri)) {
       location = *sourceLoc;
       SMLoc loc = sourceMgr.FindLocForLineAndColumn(
           sourceMgr.getMainFileID(), fileLoc.getLine(), fileLoc.getColumn());
 
       // Use range of potential identifier starting at location, else length 1
       // range.
       location->range.end.character += 1;
static void collectLocationsFromLoc(Location loc,
                                     std::vector<lsp::Location> &locations,
                                     const lsp::URIForFile &uri) {
   SetVector<Location> visitedLocs;
   loc->walk([&](Location nestedLoc) {
     FileLineColLoc fileLoc = nestedLoc.dyn_cast<FileLineColLoc>();
     if (!fileLoc || !visitedLocs.insert(nestedLoc))
       return WalkResult::advance();
 
-    Optional<lsp::Location> sourceLoc = getLocationFromLoc(fileLoc);
+    Optional<lsp::Location> sourceLoc =
+        getLocationFromLoc(uri.scheme(), fileLoc);
     if (sourceLoc && sourceLoc->uri != uri)
       locations.push_back(*sourceLoc);
     return WalkResult::advance();
   });
 }
 
 /// Returns true if the given range contains the given source location. Note
 /// that this has slightly different behavior than SMRange because it is
static lsp::Diagnostic getLspDiagnoticFromDiag(llvm::SourceMgr &sourceMgr,
 
   // Note: Right now all of the diagnostics are treated as parser issues, but
   // some are parser and some are verifier.
   lspDiag.category = "Parse Error";
 
   // Try to grab a file location for this diagnostic.
   // TODO: For simplicity, we just grab the first one. It may be likely that we
   // will need a more interesting heuristic here.'
+  StringRef uriScheme = uri.scheme();
   Optional<lsp::Location> lspLocation =
-      getLocationFromLoc(sourceMgr, diag.getLocation(), &uri);
+      getLocationFromLoc(sourceMgr, diag.getLocation(), uriScheme, &uri);
   if (lspLocation)
     lspDiag.range = lspLocation->range;
 
   // Convert the severity for the diagnostic.
   switch (diag.getSeverity()) {
   case DiagnosticSeverity::Note:
     llvm_unreachable("expected notes to be handled separately");
   case DiagnosticSeverity::Warning:
     lspDiag.severity = lsp::DiagnosticSeverity::Warning;
     break;
   case DiagnosticSeverity::Error:
     lspDiag.severity = lsp::DiagnosticSeverity::Error;
     break;
   case DiagnosticSeverity::Remark:
     lspDiag.severity = lsp::DiagnosticSeverity::Information;
     break;
   }
   lspDiag.message = diag.str();
 
   // Attach any notes to the main diagnostic as related information.
   std::vector<lsp::DiagnosticRelatedInformation> relatedDiags;
   for (Diagnostic &note : diag.getNotes()) {
     lsp::Location noteLoc;
     if (Optional<lsp::Location> loc =
-            getLocationFromLoc(sourceMgr, note.getLocation()))
+            getLocationFromLoc(sourceMgr, note.getLocation(), uriScheme))
       noteLoc = *loc;
     else
       noteLoc.uri = uri;
     relatedDiags.emplace_back(noteLoc, note.str());
   }
   if (!relatedDiags.empty())
     lspDiag.relatedInformation = std::move(relatedDiags);
 
struct MLIRDocument {
   //===--------------------------------------------------------------------===//
 
   void getCodeActionForDiagnostic(const lsp::URIForFile &uri,
                                   lsp::Position &pos, StringRef severity,
                                   StringRef message,
                                   std::vector<lsp::TextEdit> &edits);
 
   //===--------------------------------------------------------------------===//
+  // Bytecode
+  //===--------------------------------------------------------------------===//
+
+  llvm::Expected<lsp::MLIRConvertBytecodeResult> convertToBytecode();
+
+  //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   /// 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(MLIRContext &context, const lsp::URIForFile &uri,
                            StringRef contents,

 
 //===----------------------------------------------------------------------===//
 // MLIRDocument: Document Symbols
 //===----------------------------------------------------------------------===//
 
 void MLIRDocument::findDocumentSymbols(
     std::vector<lsp::DocumentSymbol> &symbols) {
   for (Operation &op : parsedIR)
     findDocumentSymbols(&op, symbols);
 }
 
 void MLIRDocument::findDocumentSymbols(
     Operation *op, std::vector<lsp::DocumentSymbol> &symbols) {
   std::vector<lsp::DocumentSymbol> *childSymbols = &symbols;
 
   // Check for the source information of this operation.
   if (const AsmParserState::OperationDefinition *def = asmState.getOpDef(op)) {
void MLIRDocument::getCodeActionForDiagnostic(
 
   edit.newText.append(indent, ' ');
   llvm::raw_string_ostream(edit.newText)
       << "// expected-" << severity << " @below {{" << message << "}}\n";
   edits.emplace_back(std::move(edit));
 }
 
 //===----------------------------------------------------------------------===//
+// MLIRDocument: Bytecode
+//===----------------------------------------------------------------------===//
+
+llvm::Expected<lsp::MLIRConvertBytecodeResult>
+MLIRDocument::convertToBytecode() {
+  // TODO: We currently require a single top-level operation, but this could
+  // conceptually be relaxed.
+  if (!llvm::hasSingleElement(parsedIR)) {
+    if (parsedIR.empty()) {
+      return llvm::make_error<lsp::LSPError>(
+          "expected a single and valid top-level operation, please ensure "
+          "there are no errors",
+          lsp::ErrorCode::RequestFailed);
+    }
+    return llvm::make_error<lsp::LSPError>(
+        "expected a single top-level operation", lsp::ErrorCode::RequestFailed);
+  }
+
+  lsp::MLIRConvertBytecodeResult result;
+  {
+    std::string rawBytecodeBuffer;
+    llvm::raw_string_ostream os(rawBytecodeBuffer);
+    writeBytecodeToFile(&parsedIR.front(), os);
+    result.output = llvm::encodeBase64(rawBytecodeBuffer);
+  }
+  return result;
+}
+
+//===----------------------------------------------------------------------===//
 // MLIRTextFileChunk
 //===----------------------------------------------------------------------===//
 
 namespace {
 /// This class represents a single chunk of an MLIR text file.
 struct MLIRTextFileChunk {
   MLIRTextFileChunk(MLIRContext &context, uint64_t lineOffset,
                     const lsp::URIForFile &uri, StringRef contents,
public:
   Optional<lsp::Hover> findHover(const lsp::URIForFile &uri,
                                  lsp::Position hoverPos);
   void findDocumentSymbols(std::vector<lsp::DocumentSymbol> &symbols);
   lsp::CompletionList getCodeCompletion(const lsp::URIForFile &uri,
                                         lsp::Position completePos);
   void getCodeActions(const lsp::URIForFile &uri, const lsp::Range &pos,
                       const lsp::CodeActionContext &context,
                       std::vector<lsp::CodeAction> &actions);
+  llvm::Expected<lsp::MLIRConvertBytecodeResult> convertToBytecode();
 
 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.
for (auto &diag : context.diagnostics) {
     action.edit.emplace();
     action.edit->changes[uri.uri().str()] = std::move(edits);
     action.diagnostics = {diag};
 
     actions.emplace_back(std::move(action));
   }
 }
 
+llvm::Expected<lsp::MLIRConvertBytecodeResult>
+MLIRTextFile::convertToBytecode() {
+  // Bail out if there is more than one chunk, bytecode wants a single module.
+  if (chunks.size() != 1) {
+    return llvm::make_error<lsp::LSPError>(
+        "unexpected split file, please remove all `// -----`",
+        lsp::ErrorCode::RequestFailed);
+  }
+  return chunks.front()->document.convertToBytecode();
+}
+
 MLIRTextFileChunk &MLIRTextFile::getChunkFor(lsp::Position &pos) {
   if (chunks.size() == 1)
     return *chunks.front();
 
   // Search for the first chunk with a greater line offset, the previous chunk
   // is the one that contains `pos`.
   auto it = llvm::upper_bound(
       chunks, pos, [](const lsp::Position &pos, const auto &chunk) {

 
 void lsp::MLIRServer::getCodeActions(const URIForFile &uri, const Range &pos,
                                      const CodeActionContext &context,
                                      std::vector<CodeAction> &actions) {
   auto fileIt = impl->files.find(uri.file());
   if (fileIt != impl->files.end())
     fileIt->second->getCodeActions(uri, pos, context, actions);
 }
+
+llvm::Expected<lsp::MLIRConvertBytecodeResult>
+lsp::MLIRServer::convertFromBytecode(const URIForFile &uri) {
+  MLIRContext tempContext(impl->registry);
+  tempContext.allowUnregisteredDialects();
+
+  // Collect any errors during parsing.
+  std::string errorMsg;
+  ScopedDiagnosticHandler diagHandler(
+      &tempContext,
+      [&](mlir::Diagnostic &diag) { errorMsg += diag.str() + "\n"; });
+
+  // Try to parse the given source file.
+  // TODO: This won't preserve external resources or the producer, we should try
+  // to fix this.
+  Block parsedBlock;
+  if (failed(parseSourceFile(uri.file(), &parsedBlock, &tempContext))) {
+    return llvm::make_error<lsp::LSPError>(
+        "failed to parse bytecode source file: " + errorMsg,
+        lsp::ErrorCode::RequestFailed);
+  }
+
+  // TODO: We currently expect a single top-level operation, but this could
+  // conceptually be relaxed.
+  if (!llvm::hasSingleElement(parsedBlock)) {
+    return llvm::make_error<lsp::LSPError>(
+        "expected bytecode to contain a single top-level operation",
+        lsp::ErrorCode::RequestFailed);
+  }
+
+  // Print the module to a buffer.
+  lsp::MLIRConvertBytecodeResult result;
+  {
+    // Extract the top-level op so that aliases get printed.
+    // FIXME: We should be able to enable aliases without having to do this!
+    OwningOpRef<Operation *> topOp = &parsedBlock.front();
+    (*topOp)->remove();
+
+    llvm::raw_string_ostream os(result.output);
+    (*topOp)->print(os, OpPrintingFlags().enableDebugInfo().assumeVerified());
+  }
+  return std::move(result);
+}
+
+llvm::Expected<lsp::MLIRConvertBytecodeResult>
+lsp::MLIRServer::convertToBytecode(const URIForFile &uri) {
+  auto fileIt = impl->files.find(uri.file());
+  if (fileIt == impl->files.end()) {
+    return llvm::make_error<lsp::LSPError>(
+        "language server does not contain an entry for this source file",
+        lsp::ErrorCode::RequestFailed);
+  }
+  return fileIt->second->convertToBytecode();
+}

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

LogicalResult mlir::MlirLspServerMain(int argc, char **argv,
 
   // Configure the logger.
   Logger::setLogLevel(logLevel);
 
   // Configure the transport used for communication.
   llvm::sys::ChangeStdinToBinary();
   JSONTransport transport(stdin, llvm::outs(), inputStyle, prettyPrint);
 
+  // Register the additionally supported URI schemes for the MLIR server.
+  URIForFile::registerSupportedScheme("mlir.bytecode-mlir");
+
   // Configure the servers and start the main language server.
   MLIRServer server(registry);
   return runMlirLSPServer(server, transport);
 }

mlir/lib/Tools/mlir-lsp-server/Protocol.h

+//===--- Protocol.h - Language Server Protocol Implementation ---*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains structs for LSP commands that are specific to the MLIR
+// server.
+//
+// Each struct has a toJSON and fromJSON function, that converts between
+// the struct and a JSON representation. (See JSON.h)
+//
+// Some structs also have operator<< serialization. This is for debugging and
+// tests, and is not generally machine-readable.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LIB_MLIR_TOOLS_MLIRLSPSERVER_PROTOCOL_H_
+#define LIB_MLIR_TOOLS_MLIRLSPSERVER_PROTOCOL_H_
+
+#include "../lsp-server-support/Protocol.h"
+
+namespace mlir {
+namespace lsp {
+//===----------------------------------------------------------------------===//
+// MLIRConvertBytecodeParams
+//===----------------------------------------------------------------------===//
+
+/// This class represents the parameters used when converting between MLIR's
+/// bytecode and textual format.
+struct MLIRConvertBytecodeParams {
+  /// The input file containing the bytecode or textual format.
+  URIForFile uri;
+};
+
+/// Add support for JSON serialization.
+bool fromJSON(const llvm::json::Value &value, MLIRConvertBytecodeParams &result,
+              llvm::json::Path path);
+
+//===----------------------------------------------------------------------===//
+// MLIRConvertBytecodeResult
+//===----------------------------------------------------------------------===//
+
+/// This class represents the result of converting between MLIR's bytecode and
+/// textual format.
+struct MLIRConvertBytecodeResult {
+  /// The resultant output of the conversion.
+  std::string output;
+};
+
+/// Add support for JSON serialization.
+llvm::json::Value toJSON(const MLIRConvertBytecodeResult &value);
+
+} // namespace lsp
+} // namespace mlir
+
+#endif

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

+//===--- Protocol.cpp - Language Server Protocol Implementation -----------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains the serialization code for the MLIR specific LSP structs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Protocol.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringSwitch.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/Format.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/JSON.h"
+#include "llvm/Support/Path.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+using namespace mlir::lsp;
+
+//===----------------------------------------------------------------------===//
+// MLIRConvertBytecodeParams
+//===----------------------------------------------------------------------===//
+
+bool mlir::lsp::fromJSON(const llvm::json::Value &value,
+                         MLIRConvertBytecodeParams &result,
+                         llvm::json::Path path) {
+  llvm::json::ObjectMapper o(value, path);
+  return o && o.map("uri", result.uri);
+}
+
+//===----------------------------------------------------------------------===//
+// MLIRConvertBytecodeResult
+//===----------------------------------------------------------------------===//
+
+llvm::json::Value mlir::lsp::toJSON(const MLIRConvertBytecodeResult &value) {
+  return llvm::json::Object{{"output", value.output}};
+}

mlir/utils/vscode/package-lock.json

 {
   "name": "vscode-mlir",
-  "version": "0.0.9",
+  "version": "0.0.10",
   "lockfileVersion": 2,
   "requires": true,
   "packages": {
     "": {
       "name": "vscode-mlir",
-      "version": "0.0.9",
+      "version": "0.0.10",
       "dependencies": {
+        "base64-js": "^1.5.1",
         "chokidar": "3.5.2",
         "vscode-languageclient": "^8.0.2-next.5"
       },
       "devDependencies": {
         "@types/mocha": "^7.0.2",
         "@types/node": "^14.17.0",
         "@types/vscode": "~1.67.0",
         "clang-format": "^1.8.0",
"node_modules/balanced-match": {
       "version": "1.0.2",
       "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz",
       "integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw=="
     },
     "node_modules/base64-js": {
       "version": "1.5.1",
       "resolved": "https://registry.npmjs.org/base64-js/-/base64-js-1.5.1.tgz",
       "integrity": "sha512-AKpaYlHn8t4SVbOHCy+b5+KKgvR4vrsD8vbvrbiQJps7fKDTkjkDry6ji0rUJjC0kzbNePLwzxq8iypo41qeWA==",
-      "dev": true,
       "funding": [
         {
           "type": "github",
           "url": "https://github.com/sponsors/feross"
         },
         {
           "type": "patreon",
           "url": "https://www.patreon.com/feross"
"dependencies": {
     "balanced-match": {
       "version": "1.0.2",
       "resolved": "https://registry.npmjs.org/balanced-match/-/balanced-match-1.0.2.tgz",
       "integrity": "sha512-3oSeUO0TMV67hN1AmbXsK4yaqU7tjiHlbxRDZOpH0KW9+CeX4bRAaX0Anxt0tx2MrpRpWwQaPwIlISEJhYU5Pw=="
     },
     "base64-js": {
       "version": "1.5.1",
       "resolved": "https://registry.npmjs.org/base64-js/-/base64-js-1.5.1.tgz",
-      "integrity": "sha512-AKpaYlHn8t4SVbOHCy+b5+KKgvR4vrsD8vbvrbiQJps7fKDTkjkDry6ji0rUJjC0kzbNePLwzxq8iypo41qeWA==",
-      "dev": true
+      "integrity": "sha512-AKpaYlHn8t4SVbOHCy+b5+KKgvR4vrsD8vbvrbiQJps7fKDTkjkDry6ji0rUJjC0kzbNePLwzxq8iypo41qeWA=="
     },
     "big-integer": {
       "version": "1.6.48",
       "resolved": "https://registry.npmjs.org/big-integer/-/big-integer-1.6.48.tgz",
       "integrity": "sha512-j51egjPa7/i+RdiRuJbPdJ2FIUYYPhvYLjzoYbcMMm62ooO6F94fETG4MTs46zPAF9Brs04OajboA/qTGuz78w==",
       "dev": true
     },
     "binary": {

mlir/utils/vscode/package.json

"keywords": [
     "LLVM",
     "MLIR",
     "PDLL",
     "TableGen",
     "tblgen",
     "tablegen"
   ],
   "activationEvents": [
+    "onFileSystem:mlir.bytecode-mlir",
+    "onCustomEditor:mlir.bytecode",
     "onLanguage:mlir",
     "onLanguage:pdll",
     "onLanguage:tablegen"
   ],
   "main": "./out/extension",
   "scripts": {
     "vscode:prepublish": "tsc -p ./",
     "compile": "tsc -watch -p ./",
     "format": "clang-format -i --glob=\"{src,test}/*.ts\"",
     "package": "vsce package",
     "publish": "vsce publish",
     "git-clang-format": "git-clang-format"
   },
   "dependencies": {
+    "base64-js": "^1.5.1",
     "chokidar": "3.5.2",
     "vscode-languageclient": "^8.0.2-next.5"
   },
   "devDependencies": {
     "@types/mocha": "^7.0.2",
     "@types/node": "^14.17.0",
     "@types/vscode": "~1.67.0",
     "clang-format": "^1.8.0",
     "typescript": "^4.6.4",
     "vsce": "^2.7.0",
     "vscode-test": "^1.3.0"
   },
   "repository": {
     "type": "git",
     "url": "https://github.com/llvm/vscode-mlir.git"
   },
   "contributes": {
+    "customEditors": [
+      {
+        "viewType": "mlir.bytecode",
+        "displayName": "MLIR Bytecode",
+        "priority": "default",
+        "selector": [
+          {
+            "filenamePattern": "*.mlirbc"
+          }
+        ]
+      }
+    ],
     "languages": [
       {
         "id": "mlir",
         "aliases": [
           "MLIR",
           "mlir"
         ],
         "extensions": [
-          ".mlir"
+          ".mlir",
+          ".mlirbc"
         ],
         "configuration": "./language-configuration.json"
       },
       {
         "id": "mlir-injection"
       },
       {
         "id": "pdll",

mlir/utils/vscode/src/MLIR/bytecodeProvider.ts

+import * as base64 from 'base64-js'
+import * as vscode from 'vscode'
+
+import {MLIRContext} from '../mlirContext';
+
+/**
+ * The parameters to the mlir/convert(To|From)Bytecode commands. These
+ * parameters are:
+ * - `uri`: The URI of the file to convert.
+ */
+type ConvertBytecodeParams = Partial<{uri : string}>;
+
+/**
+ * The output of the mlir/convert(To|From)Bytecode commands:
+ * - `output`: The output buffer of the command, e.g. a .mlir or bytecode
+ *             buffer.
+ */
+type ConvertBytecodeResult = Partial<{output : string}>;
+
+/**
+ * A custom filesystem that is used to convert MLIR bytecode files to text for
+ * use in the editor, but still use bytecode on disk.
+ */
+class BytecodeFS implements vscode.FileSystemProvider {
+  mlirContext: MLIRContext;
+
+  constructor(mlirContext: MLIRContext) { this.mlirContext = mlirContext; }
+
+  /*
+   * Forward to the default filesystem for the various methods that don't need
+   * to understand the bytecode <-> text translation.
+   */
+  readDirectory(uri: vscode.Uri): Thenable<[ string, vscode.FileType ][]> {
+    return vscode.workspace.fs.readDirectory(uri);
+  }
+  delete(uri: vscode.Uri): void {
+    vscode.workspace.fs.delete(uri.with({scheme : "file"}));
+  }
+  stat(uri: vscode.Uri): Thenable<vscode.FileStat> {
+    return vscode.workspace.fs.stat(uri.with({scheme : "file"}));
+  }
+  rename(oldUri: vscode.Uri, newUri: vscode.Uri,
+         options: {overwrite: boolean}): void {
+    vscode.workspace.fs.rename(oldUri.with({scheme : "file"}),
+                               newUri.with({scheme : "file"}), options);
+  }
+  createDirectory(uri: vscode.Uri): void {
+    vscode.workspace.fs.createDirectory(uri.with({scheme : "file"}));
+  }
+  watch(_uri: vscode.Uri, _options: {
+    readonly recursive: boolean; readonly excludes : readonly string[]
+  }): vscode.Disposable {
+    return new vscode.Disposable(() => {});
+  }
+
+  private _emitter = new vscode.EventEmitter<vscode.FileChangeEvent[]>();
+  readonly onDidChangeFile: vscode.Event<vscode.FileChangeEvent[]> =
+      this._emitter.event;
+
+  /*
+   * Read in a bytecode file, converting it to text before returning it to the
+   * caller.
+   */
+  async readFile(uri: vscode.Uri): Promise<Uint8Array> {
+    // Try to start a language client for this file so that we can parse
+    // it.
+    const client =
+        await this.mlirContext.getOrActivateLanguageClient(uri, 'mlir');
+    if (!client) {
+      throw new Error(
+          'Failed to activate mlir language server to read bytecode');
+    }
+    // Ask the client to do the conversion.
+    let convertParams: ConvertBytecodeParams = {uri : uri.toString()};
+    try {
+      const result: ConvertBytecodeResult =
+          await client.sendRequest('mlir/convertFromBytecode', convertParams);
+      return new TextEncoder().encode(result.output);
+    } catch (e) {
+      vscode.window.showErrorMessage(e.message);
+      throw new Error(`Failed to read bytecode file: ${e}`);
+    }
+  }
+
+  /*
+   * Save the provided content, which contains MLIR text, as bytecode.
+   */
+  async writeFile(uri: vscode.Uri, content: Uint8Array,
+                  _options: {create: boolean, overwrite: boolean}) {
+    // Get the language client managing this file.
+    let client = this.mlirContext.getLanguageClient(uri, 'mlir');
+    if (!client) {
+      throw new Error(
+          'Failed to activate mlir language server to write bytecode');
+    }
+
+    // Ask the client to do the conversion.
+    let convertParams: ConvertBytecodeParams = {
+      uri : uri.toString(),
+    };
+    const result: ConvertBytecodeResult =
+        await client.sendRequest('mlir/convertToBytecode', convertParams);
+    await vscode.workspace.fs.writeFile(uri.with({scheme : "file"}),
+                                        base64.toByteArray(result.output));
+  }
+}
+
+/**
+ * A custom bytecode document for use by the custom editor provider below.
+ */
+class BytecodeDocument implements vscode.CustomDocument {
+  readonly uri: vscode.Uri;
+
+  constructor(uri: vscode.Uri) { this.uri = uri; }
+  dispose(): void {}
+}
+
+/**
+ * A custom editor provider for MLIR bytecode that allows for non-binary
+ * interpretation.
+ */
+class BytecodeEditorProvider implements
+    vscode.CustomReadonlyEditorProvider<BytecodeDocument> {
+  public async openCustomDocument(uri: vscode.Uri, _openContext: any,
+                                  _token: vscode.CancellationToken):
+      Promise<BytecodeDocument> {
+    return new BytecodeDocument(uri);
+  }
+
+  public async resolveCustomEditor(document: BytecodeDocument,
+                                   _webviewPanel: vscode.WebviewPanel,
+                                   _token: vscode.CancellationToken):
+      Promise<void> {
+    // Ask the user for the desired view type.
+    const editType = await vscode.window.showQuickPick(
+        [ {label : '.mlir', description : "Edit as a .mlir text file"} ],
+        {title : 'Select an editor for the bytecode.'},
+    );
+
+    // If we don't have a valid view type, just bail.
+    if (!editType) {
+      await vscode.commands.executeCommand(
+          'workbench.action.closeActiveEditor');
+      return;
+    }
+
+    // TODO: We should also provide a non-`.mlir` way of viewing the
+    // bytecode, which should also ideally have some support for invalid
+    // bytecode files.
+
+    // Close the active editor given that we aren't using it.
+    await vscode.commands.executeCommand('workbench.action.closeActiveEditor');
+
+    // Display the file using a .mlir format.
+    await vscode.window.showTextDocument(
+        document.uri.with({scheme : "mlir.bytecode-mlir"}),
+        {preview : true, preserveFocus : false});
+  }
+}
+
+/**
+ *  Register the necessary providers for supporting MLIR bytecode.
+ */
+export function registerMLIRBytecodeExtensions(context: vscode.ExtensionContext,
+                                               mlirContext: MLIRContext) {
+  vscode.workspace.registerFileSystemProvider("mlir.bytecode-mlir",
+                                              new BytecodeFS(mlirContext));
+  vscode.window.registerCustomEditorProvider('mlir.bytecode',
+                                             new BytecodeEditorProvider());
+}

mlir/utils/vscode/src/MLIR/mlir.ts

+import * as vscode from 'vscode';
+
+import {MLIRContext} from '../mlirContext';
+import {registerMLIRBytecodeExtensions} from './bytecodeProvider';
+
+/**
+ *  Register the necessary extensions for supporting MLIR.
+ */
+export function registerMLIRExtensions(context: vscode.ExtensionContext,
+                                       mlirContext: MLIRContext) {
+  registerMLIRBytecodeExtensions(context, mlirContext);
+}

mlir/utils/vscode/src/PDLL/commands/viewOutput.ts

export class ViewPDLLCommand extends Command {
   constructor(context: MLIRContext) { super('mlir.viewPDLLOutput', context); }
 
   async execute() {
     const editor = vscode.window.activeTextEditor;
     if (editor.document.languageId != 'pdll')
       return;
 
     // Check to see if a language client is active for this document.
-    const workspaceFolder =
-        vscode.workspace.getWorkspaceFolder(editor.document.uri);
-    const pdllClient = this.context.getLanguageClient(workspaceFolder, "pdll");
+    const pdllClient =
+        this.context.getLanguageClient(editor.document.uri, "pdll");
     if (!pdllClient) {
       return;
     }
 
     // Ask the user for the desired output type.
     const outputType =
         await vscode.window.showQuickPick([ 'ast', 'mlir', 'cpp' ]);
     if (!outputType) {

mlir/utils/vscode/src/PDLL/pdll.ts

 import * as vscode from 'vscode';
 
 import {MLIRContext} from '../mlirContext';
 import {ViewPDLLCommand} from './commands/viewOutput';
 
 /**
- *  Register the necessary context and commands for PDLL.
+ *  Register the necessary extensions for supporting PDLL.
  */
-export function registerPDLLCommands(context: vscode.ExtensionContext,
+export function registerPDLLExtensions(context: vscode.ExtensionContext,
                                        mlirContext: MLIRContext) {
   context.subscriptions.push(new ViewPDLLCommand(mlirContext));
 }

mlir/utils/vscode/src/extension.ts

 import * as vscode from 'vscode';
 
+import {registerMLIRExtensions} from './MLIR/mlir';
 import {MLIRContext} from './mlirContext';
-import {registerPDLLCommands} from './PDLL/pdll';
+import {registerPDLLExtensions} from './PDLL/pdll';
 
 /**
  *  This method is called when the extension is activated. The extension is
  *  activated the very first time a command is executed.
  */
 export function activate(context: vscode.ExtensionContext) {
   const outputChannel = vscode.window.createOutputChannel('MLIR');
   context.subscriptions.push(outputChannel);
 
   const mlirContext = new MLIRContext();
   context.subscriptions.push(mlirContext);
 
   // Initialize the commands of the extension.
   context.subscriptions.push(
       vscode.commands.registerCommand('mlir.restart', async () => {
         // Dispose and reactivate the context.
         mlirContext.dispose();
         await mlirContext.activate(outputChannel);
       }));
-  registerPDLLCommands(context, mlirContext);
+  registerMLIRExtensions(context, mlirContext);
+  registerPDLLExtensions(context, mlirContext);
 
   mlirContext.activate(outputChannel);
 }

mlir/utils/vscode/src/mlirContext.ts

 
 /**
  *  This class manages all of the MLIR extension state,
  *  including the language client.
  */
 export class MLIRContext implements vscode.Disposable {
   subscriptions: vscode.Disposable[] = [];
   workspaceFolders: Map<string, WorkspaceFolderContext> = new Map();
+  outputChannel: vscode.OutputChannel;
 
   /**
    *  Activate the MLIR context, and start the language clients.
    */
   async activate(outputChannel: vscode.OutputChannel) {
+    this.outputChannel = outputChannel;
+
     // This lambda is used to lazily start language clients for the given
     // document. It removes the need to pro-actively start language clients for
     // every folder within the workspace and every language type we provide.
     const startClientOnOpenDocument = async (document: vscode.TextDocument) => {
-      if (document.uri.scheme !== 'file') {
-        return;
+      await this.getOrActivateLanguageClient(document.uri, document.languageId);
+    };
+    // Process any existing documents.
+    for (const textDoc of vscode.workspace.textDocuments) {
+      await startClientOnOpenDocument(textDoc);
+    }
+
+    // Watch any new documents to spawn servers when necessary.
+    this.subscriptions.push(
+        vscode.workspace.onDidOpenTextDocument(startClientOnOpenDocument));
+    this.subscriptions.push(
+        vscode.workspace.onDidChangeWorkspaceFolders((event) => {
+          for (const folder of event.removed) {
+            const client = this.workspaceFolders.get(folder.uri.toString());
+            if (client) {
+              client.dispose();
+              this.workspaceFolders.delete(folder.uri.toString());
+            }
           }
-      let serverSettingName: string;
-      if (document.languageId === 'mlir') {
-        serverSettingName = 'server_path';
-      } else if (document.languageId === 'pdll') {
-        serverSettingName = 'pdll_server_path';
-      } else if (document.languageId === 'tablegen') {
-        serverSettingName = 'tablegen_server_path';
-      } else {
-        return;
-      }
+        }));
+  }
+
+  /**
+   * Open or return a language server for the given uri and language.
+   */
+  async getOrActivateLanguageClient(uri: vscode.Uri, languageId: string):
+      Promise<vscodelc.LanguageClient> {
+    if (languageId === 'mlir') {
+    } else if (languageId === 'pdll') {
+    } else if (languageId === 'tablegen') {
+      return null;
+    }
+
+    // Check the scheme of the uri.
+    let validSchemes = [ 'file', 'mlir.bytecode-mlir' ];
+    if (!validSchemes.includes(uri.scheme)) {
+      return null;
 
-      // Resolve the workspace folder if this document is in one. We use the
-      // workspace folder when determining if a server needs to be started.
-      const uri = document.uri;
-      let workspaceFolder = vscode.workspace.getWorkspaceFolder(uri);
-      let workspaceFolderStr =
-          workspaceFolder ? workspaceFolder.uri.toString() : "";
 
-      // Get or create a client context for this folder.
-      let folderContext = this.workspaceFolders.get(workspaceFolderStr);
-      if (!folderContext) {
-        folderContext = new WorkspaceFolderContext();
-        this.workspaceFolders.set(workspaceFolderStr, folderContext);
-      }
-      // Start the client for this language if necessary.
-      if (!folderContext.clients.has(document.languageId)) {
-        let client = await this.activateWorkspaceFolder(
-            workspaceFolder, serverSettingName, document.languageId,
-            outputChannel);
-        folderContext.clients.set(document.languageId, client);
-      }
-    };
-    // Process any existing documents.
-    for (const textDoc of vscode.workspace.textDocuments) {
-      await startClientOnOpenDocument(textDoc);
-    }
-
-    // Watch any new documents to spawn servers when necessary.
-    this.subscriptions.push(
-        vscode.workspace.onDidOpenTextDocument(startClientOnOpenDocument));
-    this.subscriptions.push(
-        vscode.workspace.onDidChangeWorkspaceFolders((event) => {
-          for (const folder of event.removed) {
-            const client = this.workspaceFolders.get(folder.uri.toString());
-            if (client) {
-              client.dispose();
-              this.workspaceFolders.delete(folder.uri.toString());
-            }
-          }
-        }));
+    let client = folderContext.clients.get(languageId);
+    if (!client) {
+      client = await this.activateWorkspaceFolder(
+          workspaceFolder, serverSettingName, languageId, this.outputChannel);
+      folderContext.clients.set(languageId, client);
+    return client;
   }
 
   /**
    *  Prepare a compilation database option for a server.
    */
   async prepareCompilationDatabaseServerOptions(
       languageName: string, workspaceFolder: vscode.WorkspaceFolder,
       configsToWatch: string[], pathsToWatch: string[],
if (!workspaceFolder) {
           return Promise.resolve();
         }
       };
     }
 
     // Configure the client options.
     const clientOptions: vscodelc.LanguageClientOptions = {
       documentSelector : [
-        {scheme : 'file', language : languageName, pattern : selectorPattern}
+        {language : languageName, pattern : selectorPattern},
       ],
       synchronize : {
         // Notify the server about file changes to language files contained in
         // the workspace.
         fileEvents : vscode.workspace.createFileSystemWatcher(filePattern)
       },
       outputChannel : outputChannel,
       workspaceFolder : workspaceFolder,
async resolveServerPath(serverSettingName: string,
                           workspaceFolder: vscode.WorkspaceFolder):
       Promise<string> {
     const serverPath = config.get<string>(serverSettingName, workspaceFolder);
     const defaultPath = MLIRContext.getDefaultServerFilename(serverSettingName);
     return this.resolvePath(serverPath, defaultPath, workspaceFolder);
   }
 
   /**
-   * Return the language client for the given language and workspace folder, or
-   * null if no client is active.
+   * Return the language client for the given language and uri, or null if no
+   * client is active.
    */
-  getLanguageClient(workspaceFolder: vscode.WorkspaceFolder,
+  getLanguageClient(uri: vscode.Uri,
                     languageName: string): vscodelc.LanguageClient {
+    let workspaceFolder = vscode.workspace.getWorkspaceFolder(uri);
     let workspaceFolderStr =
         workspaceFolder ? workspaceFolder.uri.toString() : "";
     let folderContext = this.workspaceFolders.get(workspaceFolderStr);
     if (!folderContext) {
       return null;
     }
     return folderContext.clients.get(languageName);
   }
 
   dispose() {
     this.subscriptions.forEach((d) => { d.dispose(); });
     this.subscriptions = [];
     this.workspaceFolders.forEach((d) => { d.dispose(); });
     this.workspaceFolders.clear();
   }
 }