[mlir][Tablegen-LSP] Add support for a basic TableGen language server

[rriddle]

Summary

This follows the same general structure of the MLIR and PDLL language
servers. This commits adds the basic functionality for setting up the server,
and initially only supports providing diagnostics. Followon commits will
build out more comprehensive behavior.

Realistically this should eventually live in llvm/, but building in MLIR is an easier
initial step given that:

• All of the necessary LSP functionality is already here
• It allows for proving out useful language features (e.g. compilation databases) without affecting wider scale tablegen users
• MLIR has a vscode extension that can immediately take advantage of it

[jpienaar]

I'll let the argument about the location of this slide as I know you'll follow up ;-)

Looks good as mostly close to PDLL one, couple of places it seemed leftovers which need update/may not apply, but overall LGTM.

mlir/include/mlir/Tools/mlir-tblgen-lsp-server/TableGenLspServerMain.h

20 ↗

(On Diff #428844)

Shall we just call this tblgen-lsp-server? Given this isn't mlir specific, should move eventually and doesn't cause confusion (there isn't another TableGen LSP to differentiate with)

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

15

Could we skip lex? E.g., why not just "find end of string that starts at curptr"?

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

24

Is there any useful expansion to what the heuristic is that could be added here?

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

16 ↗

(On Diff #428844)

So this is the main reason for this being inside mlir/? So if it moves to like tools or what not, this could move to be adjacent to tblgen binary itself?

mlir/lib/Tools/mlir-tblgen-lsp-server/LSPServer.h

23 ↗

(On Diff #428844)

Pdll -> Tablegen?

mlir/lib/Tools/mlir-tblgen-lsp-server/TableGenServer.cpp

57 ↗

(On Diff #428844)

Is the syntax registered as tblgen or tablegen? (Doc side https://github.com/llvm/llvm-project/blob/main/mlir/docs/OpDefinitions.md we have been using Tablegen, the tool is tblgen)

60 ↗

(On Diff #428844)

Is there any verifier errors in Tablegen?

mlir/lib/Tools/mlir-tblgen-lsp-server/TableGenServer.h

1 ↗

(On Diff #428844)

TablegenLspGenServer?

[DavidSpickett]

This is failing to build on one of our flang bots: https://lab.llvm.org/buildbot/#/builders/177/builds/5120

This bot is building llvm as (or with, I'm not 100% sure yet myself) shared libraries. I'll see if I can find the fix myself (it's a missing library include...somewhere) but if you get to it first let me know.

[DavidSpickett]

Turns out we specifically don't include tablegen in the llvm shared library https://github.com/llvm/llvm-project/blob/main/llvm/tools/llvm-shlib/CMakeLists.txt#L20.

If I do it will link but you get:

FAILED: tools/clang/include/clang/Sema/AttrParsedAttrList.inc
<...>
: CommandLine Error: Option 'o' registered more than once!
LLVM ERROR: inconsistency in registered CommandLine options

Given that this is a language server would it make sense to disable building it if LLVM_LINK_LLVM_DYLIB is on?

[rriddle]

In D125440#3518288, @DavidSpickett wrote:

Turns out we specifically don't include tablegen in the llvm shared library https://github.com/llvm/llvm-project/blob/main/llvm/tools/llvm-shlib/CMakeLists.txt#L20.

If I do it will link but you get:

FAILED: tools/clang/include/clang/Sema/AttrParsedAttrList.inc
<...>
: CommandLine Error: Option 'o' registered more than once!
LLVM ERROR: inconsistency in registered CommandLine options

Given that this is a language server would it make sense to disable building it if LLVM_LINK_LLVM_DYLIB is on?

Dang, thanks for the report. I think we can just do what we do for the PDLL parser, maybe something like:

llvm_add_library(TableGenLspServerLib STATIC
  LSPServer.cpp
  TableGenServer.cpp
  TableGenLspServerMain.cpp

  ADDITIONAL_HEADER_DIRS
  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/tblgen-lsp-server

  DISABLE_LLVM_LINK_LLVM_DYLIB

  LINK_COMPONENTS
  Support
  TableGen

  LINK_LIBS PUBLIC
  MLIRLspServerSupportLib
  MLIRSupport
  )

Could you give that a try?

mlir/include/mlir/Tools/mlir-tblgen-lsp-server/TableGenLspServerMain.h

20 ↗

(On Diff #428844)

Yeah, agreed there. Meant to rename this before, thanks for pointing this out!

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

16 ↗

(On Diff #428844)

Yep, we rely heavily on the LSPServer support library. If we can move that to a more common place, we could move the tblgen server out of MLIR.

mlir/lib/Tools/mlir-tblgen-lsp-server/LSPServer.h

23 ↗

(On Diff #428844)

Oops, thanks for the catch!

mlir/lib/Tools/mlir-tblgen-lsp-server/TableGenServer.cpp

57 ↗

(On Diff #428844)

The syntax is registered as tablegen IIRC. Let me fix the references to tblgen here.

60 ↗

(On Diff #428844)

Let me drop this for now, but I did initially have this comment purposefully. It's a little more nuanced than from MLIR/PDLL. Right now we only capture parser errors, but realistically a lot of warnings/errors/notes/etc. come during the execution of the tablegen backend. It may be interesting in the future to try and run backends on files we know they should run on, and try to capture errors. That would expose more problems to users that would otherwise only pop up during build time. I'm not sure right now how we should do that/if we should do that, so dropped the comment for now.

[DavidSpickett]

Could you give that a try?

I'll do that, in fact I was looking for an existing lib that did the same thing thanks for the pointer.

[DavidSpickett]

Fixed in 9e469ced42cd9335baffd906bed17fd81ce4e6e8, all good now.

[rriddle]

In D125440#3519346, @DavidSpickett wrote:

Fixed in 9e469ced42cd9335baffd906bed17fd81ce4e6e8, all good now.

Thanks for the fix! Really appreciated.

mlir/docs/Tools/MLIRLSP.md

 reviews sent to
 [phabricator](https://llvm.org/docs/Contributing.html#how-to-submit-a-patch).
 
 When developing or deploying this extension within the LLVM monorepo, a few
 extra setup steps are required:
 
 *   Copy `mlir/utils/textmate/mlir.json` to the extension directory and rename
     to `grammar.json`.
+*   Copy `llvm/utils/textmate/tablegen.json` to the extension directory and rename
+    to `tablegen-grammar.json`.
 *   Copy
     `https://mlir.llvm.org//LogoAssets/logo/PNG/full_color/mlir-identity-03.png`
     to the extension directory and rename to `icon.png`.
 
 Please follow the existing code style when contributing to the extension, we
 recommend to run `npm run format` before sending a patch.

mlir/include/mlir/Tools/tblgen-lsp-server/TableGenLspServerMain.h

+//===- TableGenLSPServerMain.h - TableGen Language Server main --*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Main entry function for tblgen-lsp-server when built as standalone binary.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENLSPSERVERMAIN_H
+#define MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENLSPSERVERMAIN_H
+
+namespace mlir {
+struct LogicalResult;
+
+/// Implementation for tools like `tblgen-lsp-server`.
+LogicalResult TableGenLspServerMain(int argc, char **argv);
+
+} // namespace mlir
+
+#endif // MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENLSPSERVERMAIN_H

mlir/lib/Tools/CMakeLists.txt

 add_subdirectory(lsp-server-support)
 add_subdirectory(mlir-lsp-server)
 add_subdirectory(mlir-opt)
 add_subdirectory(mlir-pdll-lsp-server)
 add_subdirectory(mlir-reduce)
 add_subdirectory(mlir-translate)
 add_subdirectory(PDLL)
+add_subdirectory(tblgen-lsp-server)

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

 add_mlir_library(MLIRLspServerSupportLib
   Logging.cpp
   Protocol.cpp
+  SourceMgrUtils.cpp
   Transport.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/lsp-server
 
   LINK_LIBS PUBLIC
   MLIRSupport
   )

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

+//===--- SourceMgrUtils.h - SourceMgr LSP Utils -----------------*- 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 an array of generally useful SourceMgr utilities for
+// interacting with LSP components.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LIB_MLIR_TOOLS_LSPSERVERSUPPORT_TRANSPORT_H_
+#define LIB_MLIR_TOOLS_LSPSERVERSUPPORT_TRANSPORT_H_
+
+#include "Protocol.h"
+#include "llvm/Support/SourceMgr.h"
+
+namespace mlir {
+namespace lsp {
+
+/// Returns the range of a lexical token given a SMLoc corresponding to the
+/// start of an token location. The range is computed heuristically, and
+/// supports identifier-like tokens, strings, etc.
+SMRange convertTokenLocToRange(SMLoc loc);
+
+} // namespace lsp
+} // namespace mlir
+
+#endif

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

+//===--- SourceMgrUtils.cpp - SourceMgr LSP Utils -------------------------===//
+//
+// 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 "SourceMgrUtils.h"
+
+using namespace mlir;
+using namespace mlir::lsp;
+
+/// Find the end of a string whose contents start at the given `curPtr`. Returns
+/// the position at the end of the string, after a terminal or invalid character
+/// (e.g. `"` or `\0`).
+static const char *lexLocStringTok(const char *curPtr) {
+  while (char c = *curPtr++) {
+    // Check for various terminal characters.
+    if (StringRef("\"\n\v\f").contains(c))
+      return curPtr;
+
+    // Check for escape sequences.
+    if (c == '\\') {
+      // Check a few known escapes and \xx hex digits.
+      if (*curPtr == '"' || *curPtr == '\\' || *curPtr == 'n' || *curPtr == 't')
+        ++curPtr;
+      else if (llvm::isHexDigit(*curPtr) && llvm::isHexDigit(curPtr[1]))
+        curPtr += 2;
+      else
+        return curPtr;
+    }
+  }
+
+  // If we hit this point, we've reached the end of the buffer. Update the end
+  // pointer to not point past the buffer.
+  return curPtr - 1;
+}
+
+SMRange lsp::convertTokenLocToRange(SMLoc loc) {
+  if (!loc.isValid())
+    return SMRange();
+  const char *curPtr = loc.getPointer();
+
+  // Check if this is a string token.
+  if (*curPtr == '"') {
+    curPtr = lexLocStringTok(curPtr + 1);
+
+    // Otherwise, default to handling an identifier.
+  } else {
+    // Return if the given character is a valid identifier character.
+    auto isIdentifierChar = [](char c) {
+      return isalnum(c) || c == '$' || c == '.' || c == '_' || c == '-';
+    };
+
+    while (*curPtr && isIdentifierChar(*(++curPtr)))
+      continue;
+  }
+
+  return SMRange(loc, SMLoc::getFromPointer(curPtr));
+}

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 "mlir/IR/FunctionInterfaces.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Parser/AsmParserState.h"
 #include "mlir/Parser/Parser.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 
loc->walk([&](Location nestedLoc) {
     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;
-      if (Optional<SMRange> range = AsmParserState::convertIdLocToRange(loc)) {
+      if (Optional<SMRange> range = lsp::convertTokenLocToRange(loc)) {
         auto lineCol = sourceMgr.getLineAndColumn(range->End);
         location->range.end.character =
             std::max(fileLoc.getColumn() + 1, lineCol.second - 1);
       }
       return WalkResult::interrupt();
     }
     return WalkResult::advance();
   });

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

+set(LLVM_LINK_COMPONENTS
+  Demangle
+  Support
+  TableGen
+)
+
+llvm_add_library(TableGenLspServerLib
+  LSPServer.cpp
+  TableGenServer.cpp
+  TableGenLspServerMain.cpp
+
+  ADDITIONAL_HEADER_DIRS
+  ${MLIR_MAIN_INCLUDE_DIR}/mlir/Tools/tblgen-lsp-server
+
+  LINK_LIBS PUBLIC
+  MLIRLspServerSupportLib
+  MLIRSupport
+  )

mlir/lib/Tools/tblgen-lsp-server/LSPServer.h

+//===- LSPServer.h - TableGen LSP 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_TBLGENLSPSERVER_LSPSERVER_H
+#define LIB_MLIR_TOOLS_TBLGENLSPSERVER_LSPSERVER_H
+
+#include <memory>
+
+namespace mlir {
+struct LogicalResult;
+
+namespace lsp {
+class JSONTransport;
+class TableGenServer;
+
+/// Run the main loop of the LSP server using the given TableGen server and
+/// transport.
+LogicalResult runTableGenLSPServer(TableGenServer &server,
+                                   JSONTransport &transport);
+
+} // namespace lsp
+} // namespace mlir
+
+#endif // LIB_MLIR_TOOLS_TBLGENLSPSERVER_LSPSERVER_H

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

+//===- LSPServer.cpp - TableGen 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 "TableGenServer.h"
+#include "llvm/ADT/FunctionExtras.h"
+#include "llvm/ADT/StringMap.h"
+
+using namespace mlir;
+using namespace mlir::lsp;
+
+//===----------------------------------------------------------------------===//
+// LSPServer
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct LSPServer {
+  LSPServer(TableGenServer &server, JSONTransport &transport)
+      : server(server), transport(transport) {}
+
+  //===--------------------------------------------------------------------===//
+  // Initialization
+
+  void onInitialize(const InitializeParams &params,
+                    Callback<llvm::json::Value> reply);
+  void onInitialized(const InitializedParams &params);
+  void onShutdown(const NoParams &params, Callback<std::nullptr_t> reply);
+
+  //===--------------------------------------------------------------------===//
+  // Document Change
+
+  void onDocumentDidOpen(const DidOpenTextDocumentParams &params);
+  void onDocumentDidClose(const DidCloseTextDocumentParams &params);
+  void onDocumentDidChange(const DidChangeTextDocumentParams &params);
+
+  //===--------------------------------------------------------------------===//
+  // Fields
+  //===--------------------------------------------------------------------===//
+
+  TableGenServer &server;
+  JSONTransport &transport;
+
+  /// An outgoing notification used to send diagnostics to the client when they
+  /// are ready to be processed.
+  OutgoingNotification<PublishDiagnosticsParams> publishDiagnostics;
+
+  /// Used to indicate that the 'shutdown' request was received from the
+  /// Language Server client.
+  bool shutdownRequestReceived = false;
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Initialization
+
+void LSPServer::onInitialize(const InitializeParams &params,
+                             Callback<llvm::json::Value> reply) {
+  // Send a response with the capabilities of this server.
+  llvm::json::Object serverCaps{
+      {"textDocumentSync",
+       llvm::json::Object{
+           {"openClose", true},
+           {"change", (int)TextDocumentSyncKind::Full},
+           {"save", true},
+       }},
+  };
+
+  llvm::json::Object result{
+      {{"serverInfo", llvm::json::Object{{"name", "tblgen-lsp-server"},
+                                         {"version", "0.0.1"}}},
+       {"capabilities", std::move(serverCaps)}}};
+  reply(std::move(result));
+}
+void LSPServer::onInitialized(const InitializedParams &) {}
+void LSPServer::onShutdown(const NoParams &, Callback<std::nullptr_t> reply) {
+  shutdownRequestReceived = true;
+  reply(nullptr);
+}
+
+//===----------------------------------------------------------------------===//
+// Document Change
+
+void LSPServer::onDocumentDidOpen(const DidOpenTextDocumentParams &params) {
+  PublishDiagnosticsParams diagParams(params.textDocument.uri,
+                                      params.textDocument.version);
+  server.addOrUpdateDocument(params.textDocument.uri, params.textDocument.text,
+                             params.textDocument.version,
+                             diagParams.diagnostics);
+
+  // Publish any recorded diagnostics.
+  publishDiagnostics(diagParams);
+}
+void LSPServer::onDocumentDidClose(const DidCloseTextDocumentParams &params) {
+  Optional<int64_t> version = server.removeDocument(params.textDocument.uri);
+  if (!version)
+    return;
+
+  // Empty out the diagnostics shown for this document. This will clear out
+  // anything currently displayed by the client for this document (e.g. in the
+  // "Problems" pane of VSCode).
+  publishDiagnostics(
+      PublishDiagnosticsParams(params.textDocument.uri, *version));
+}
+void LSPServer::onDocumentDidChange(const DidChangeTextDocumentParams &params) {
+  // TODO: We currently only support full document updates, we should refactor
+  // to avoid this.
+  if (params.contentChanges.size() != 1)
+    return;
+  PublishDiagnosticsParams diagParams(params.textDocument.uri,
+                                      params.textDocument.version);
+  server.addOrUpdateDocument(
+      params.textDocument.uri, params.contentChanges.front().text,
+      params.textDocument.version, diagParams.diagnostics);
+
+  // Publish any recorded diagnostics.
+  publishDiagnostics(diagParams);
+}
+
+//===----------------------------------------------------------------------===//
+// Entry Point
+//===----------------------------------------------------------------------===//
+
+LogicalResult mlir::lsp::runTableGenLSPServer(TableGenServer &server,
+                                              JSONTransport &transport) {
+  LSPServer lspServer(server, transport);
+  MessageHandler messageHandler(transport);
+
+  // Initialization
+  messageHandler.method("initialize", &lspServer, &LSPServer::onInitialize);
+  messageHandler.notification("initialized", &lspServer,
+                              &LSPServer::onInitialized);
+  messageHandler.method("shutdown", &lspServer, &LSPServer::onShutdown);
+
+  // Document Changes
+  messageHandler.notification("textDocument/didOpen", &lspServer,
+                              &LSPServer::onDocumentDidOpen);
+  messageHandler.notification("textDocument/didClose", &lspServer,
+                              &LSPServer::onDocumentDidClose);
+  messageHandler.notification("textDocument/didChange", &lspServer,
+                              &LSPServer::onDocumentDidChange);
+
+  // Diagnostics
+  lspServer.publishDiagnostics =
+      messageHandler.outgoingNotification<PublishDiagnosticsParams>(
+          "textDocument/publishDiagnostics");
+
+  // Run the main loop of the transport.
+  if (llvm::Error error = transport.run(messageHandler)) {
+    Logger::error("Transport error: {0}", error);
+    llvm::consumeError(std::move(error));
+    return failure();
+  }
+  return success(lspServer.shutdownRequestReceived);
+}

mlir/lib/Tools/tblgen-lsp-server/TableGenLspServerMain.cpp

+//===- TableGenLspServerMain.cpp - TableGen Language Server main ----------===//
+//
+// 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 "mlir/Tools/tblgen-lsp-server/TableGenLspServerMain.h"
+#include "../lsp-server-support/Logging.h"
+#include "../lsp-server-support/Transport.h"
+#include "LSPServer.h"
+#include "TableGenServer.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Program.h"
+
+using namespace mlir;
+using namespace mlir::lsp;
+
+LogicalResult mlir::TableGenLspServerMain(int argc, char **argv) {
+  llvm::cl::opt<JSONStreamStyle> inputStyle{
+      "input-style",
+      llvm::cl::desc("Input JSON stream encoding"),
+      llvm::cl::values(clEnumValN(JSONStreamStyle::Standard, "standard",
+                                  "usual LSP protocol"),
+                       clEnumValN(JSONStreamStyle::Delimited, "delimited",
+                                  "messages delimited by `// -----` lines, "
+                                  "with // comment support")),
+      llvm::cl::init(JSONStreamStyle::Standard),
+      llvm::cl::Hidden,
+  };
+  llvm::cl::opt<bool> litTest{
+      "lit-test",
+      llvm::cl::desc(
+          "Abbreviation for -input-style=delimited -pretty -log=verbose. "
+          "Intended to simplify lit tests"),
+      llvm::cl::init(false),
+  };
+  llvm::cl::opt<Logger::Level> logLevel{
+      "log",
+      llvm::cl::desc("Verbosity of log messages written to stderr"),
+      llvm::cl::values(
+          clEnumValN(Logger::Level::Error, "error", "Error messages only"),
+          clEnumValN(Logger::Level::Info, "info",
+                     "High level execution tracing"),
+          clEnumValN(Logger::Level::Debug, "verbose", "Low level details")),
+      llvm::cl::init(Logger::Level::Info),
+  };
+  llvm::cl::opt<bool> prettyPrint{
+      "pretty",
+      llvm::cl::desc("Pretty-print JSON output"),
+      llvm::cl::init(false),
+  };
+
+  llvm::cl::ParseCommandLineOptions(argc, argv, "TableGen LSP Language Server");
+
+  if (litTest) {
+    inputStyle = JSONStreamStyle::Delimited;
+    logLevel = Logger::Level::Debug;
+    prettyPrint = true;
+  }
+
+  // Configure the logger.
+  Logger::setLogLevel(logLevel);
+
+  // Configure the transport used for communication.
+  llvm::sys::ChangeStdinToBinary();
+  JSONTransport transport(stdin, llvm::outs(), inputStyle, prettyPrint);
+
+  // Configure the servers and start the main language server.
+  TableGenServer server;
+  return runTableGenLSPServer(server, transport);
+}

mlir/lib/Tools/tblgen-lsp-server/TableGenServer.h

+//===- TableGenServer.h - TableGen 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_TBLGENLSPSERVER_TABLEGENSERVER_H_
+#define LIB_MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENSERVER_H_
+
+#include "mlir/Support/LLVM.h"
+#include "llvm/ADT/StringRef.h"
+#include <memory>
+#include <string>
+
+namespace mlir {
+namespace lsp {
+struct Diagnostic;
+class URIForFile;
+
+/// This class implements all of the TableGen related functionality necessary
+/// for a language server. This class allows for keeping the TableGen specific
+/// logic separate from the logic that involves LSP server/client communication.
+class TableGenServer {
+public:
+  TableGenServer();
+  ~TableGenServer();
+
+  /// Add or update the document, with the provided `version`, at the given URI.
+  /// Any diagnostics emitted for this document should be added to
+  /// `diagnostics`.
+  void addOrUpdateDocument(const URIForFile &uri, StringRef contents,
+                           int64_t version,
+                           std::vector<Diagnostic> &diagnostics);
+
+  /// Remove the document with the given uri. Returns the version of the removed
+  /// document, or None if the uri did not have a corresponding document within
+  /// the server.
+  Optional<int64_t> removeDocument(const URIForFile &uri);
+
+private:
+  struct Impl;
+  std::unique_ptr<Impl> impl;
+};
+
+} // namespace lsp
+} // namespace mlir
+
+#endif // LIB_MLIR_TOOLS_TBLGENLSPSERVER_TABLEGENSERVER_H_

mlir/lib/Tools/tblgen-lsp-server/TableGenServer.cpp

+//===- TableGenServer.cpp - TableGen 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 "TableGenServer.h"
+
+#include "../lsp-server-support/Logging.h"
+#include "../lsp-server-support/Protocol.h"
+#include "../lsp-server-support/SourceMgrUtils.h"
+#include "llvm/ADT/IntervalMap.h"
+#include "llvm/ADT/StringMap.h"
+#include "llvm/ADT/StringSet.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/Path.h"
+#include "llvm/TableGen/Parser.h"
+#include "llvm/TableGen/Record.h"
+
+using namespace mlir;
+
+/// Returns a language server uri for the given source location. `mainFileURI`
+/// corresponds to the uri for the main file of the source manager.
+static lsp::URIForFile getURIFromLoc(const llvm::SourceMgr &mgr, SMLoc loc,
+                                     const lsp::URIForFile &mainFileURI) {
+  int bufferId = mgr.FindBufferContainingLoc(loc);
+  if (bufferId == 0 || bufferId == static_cast<int>(mgr.getMainFileID()))
+    return mainFileURI;
+  llvm::Expected<lsp::URIForFile> fileForLoc = lsp::URIForFile::fromFile(
+      mgr.getBufferInfo(bufferId).Buffer->getBufferIdentifier());
+  if (fileForLoc)
+    return *fileForLoc;
+  lsp::Logger::error("Failed to create URI for include file: {0}",
+                     llvm::toString(fileForLoc.takeError()));
+  return mainFileURI;
+}
+
+/// Returns a language server location from the given source range.
+static lsp::Location getLocationFromLoc(llvm::SourceMgr &mgr, SMLoc loc,
+                                        const lsp::URIForFile &uri) {
+  return lsp::Location(getURIFromLoc(mgr, loc, uri),
+                       lsp::Range(mgr, lsp::convertTokenLocToRange(loc)));
+}
+
+/// Convert the given TableGen diagnostic to the LSP form.
+static Optional<lsp::Diagnostic>
+getLspDiagnoticFromDiag(const llvm::SMDiagnostic &diag,
+                        const lsp::URIForFile &uri) {
+  auto *sourceMgr = const_cast<llvm::SourceMgr *>(diag.getSourceMgr());
+  if (!sourceMgr || !diag.getLoc().isValid())
+    return llvm::None;
+
+  lsp::Diagnostic lspDiag;
+  lspDiag.source = "tablegen";
+  lspDiag.category = "Parse Error";
+
+  // Try to grab a file location for this diagnostic.
+  lsp::Location loc = getLocationFromLoc(*sourceMgr, diag.getLoc(), uri);
+  lspDiag.range = loc.range;
+
+  // Skip diagnostics that weren't emitted within the main file.
+  if (loc.uri != uri)
+    return llvm::None;
+
+  // Convert the severity for the diagnostic.
+  switch (diag.getKind()) {
+  case llvm::SourceMgr::DK_Warning:
+    lspDiag.severity = lsp::DiagnosticSeverity::Warning;
+    break;
+  case llvm::SourceMgr::DK_Error:
+    lspDiag.severity = lsp::DiagnosticSeverity::Error;
+    break;
+  case llvm::SourceMgr::DK_Note:
+    // Notes are emitted separately from the main diagnostic, so we just treat
+    // them as remarks given that we can't determine the diagnostic to relate
+    // them to.
+  case llvm::SourceMgr::DK_Remark:
+    lspDiag.severity = lsp::DiagnosticSeverity::Information;
+    break;
+  }
+  lspDiag.message = diag.getMessage().str();
+
+  return lspDiag;
+}
+
+//===----------------------------------------------------------------------===//
+// TableGenTextFile
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// This class represents a text file containing one or more TableGen documents.
+class TableGenTextFile {
+public:
+  TableGenTextFile(const lsp::URIForFile &uri, StringRef fileContents,
+                   int64_t version, std::vector<lsp::Diagnostic> &diagnostics);
+
+  /// Return the current version of this text file.
+  int64_t getVersion() const { return version; }
+
+private:
+  /// The full string contents of the file.
+  std::string contents;
+
+  /// The version of this file.
+  int64_t version;
+
+  /// The include directories for this file.
+  std::vector<std::string> includeDirs;
+
+  /// The source manager containing the contents of the input file.
+  llvm::SourceMgr sourceMgr;
+
+  /// The record keeper containing the parsed tablegen constructs.
+  llvm::RecordKeeper recordKeeper;
+};
+} // namespace
+
+TableGenTextFile::TableGenTextFile(const lsp::URIForFile &uri,
+                                   StringRef fileContents, int64_t version,
+                                   std::vector<lsp::Diagnostic> &diagnostics)
+    : contents(fileContents.str()), version(version) {
+  auto memBuffer = llvm::MemoryBuffer::getMemBufferCopy(contents, uri.file());
+  if (!memBuffer) {
+    lsp::Logger::error("Failed to create memory buffer for file", uri.file());
+    return;
+  }
+
+  // Build the set of include directories for this file.
+  // TODO: Setup external include directories.
+  llvm::SmallString<32> uriDirectory(uri.file());
+  llvm::sys::path::remove_filename(uriDirectory);
+  includeDirs.push_back(uriDirectory.str().str());
+
+  sourceMgr.setIncludeDirs(includeDirs);
+  sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
+
+  // This class provides a context argument for the llvm::SourceMgr diagnostic
+  // handler.
+  struct DiagHandlerContext {
+    std::vector<lsp::Diagnostic> &diagnostics;
+    const lsp::URIForFile &uri;
+  } handlerContext{diagnostics, uri};
+
+  // Set the diagnostic handler for the tablegen source manager.
+  sourceMgr.setDiagHandler(
+      [](const llvm::SMDiagnostic &diag, void *rawHandlerContext) {
+        auto *ctx = reinterpret_cast<DiagHandlerContext *>(rawHandlerContext);
+        if (auto lspDiag = getLspDiagnoticFromDiag(diag, ctx->uri))
+          ctx->diagnostics.push_back(*lspDiag);
+      },
+      &handlerContext);
+  if (llvm::TableGenParseFile(sourceMgr, recordKeeper))
+    return;
+}
+
+//===----------------------------------------------------------------------===//
+// TableGenServer::Impl
+//===----------------------------------------------------------------------===//
+
+struct lsp::TableGenServer::Impl {
+  /// The files held by the server, mapped by their URI file name.
+  llvm::StringMap<std::unique_ptr<TableGenTextFile>> files;
+};
+
+//===----------------------------------------------------------------------===//
+// TableGenServer
+//===----------------------------------------------------------------------===//
+
+lsp::TableGenServer::TableGenServer() : impl(std::make_unique<Impl>()) {}
+lsp::TableGenServer::~TableGenServer() = default;
+
+void lsp::TableGenServer::addOrUpdateDocument(
+    const URIForFile &uri, StringRef contents, int64_t version,
+    std::vector<Diagnostic> &diagnostics) {
+  impl->files[uri.file()] =
+      std::make_unique<TableGenTextFile>(uri, contents, version, diagnostics);
+}
+
+Optional<int64_t> lsp::TableGenServer::removeDocument(const URIForFile &uri) {
+  auto it = impl->files.find(uri.file());
+  if (it == impl->files.end())
+    return llvm::None;
+
+  int64_t version = it->second->getVersion();
+  impl->files.erase(it);
+  return version;
+}

mlir/test/CMakeLists.txt

set(MLIR_TEST_DEPENDS
   mlir-linalg-ods-yaml-gen
   mlir-lsp-server
   mlir-pdll-lsp-server
   mlir-opt
   mlir-pdll
   mlir-reduce
   mlir-tblgen
   mlir-translate
+  tblgen-lsp-server
   )
 
 # The native target may not be enabled, in this case we won't
 # run tests that involves executing on the host: do not build
 # useless binaries.
 if(TARGET ${LLVM_NATIVE_ARCH})
   list(APPEND MLIR_TEST_DEPENDS
     mlir-cpu-runner

mlir/test/tblgen-lsp-server/diagnostics.test

+// RUN: tblgen-lsp-server -lit-test < %s | FileCheck -strict-whitespace %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"tablegen","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"textDocument/didOpen","params":{"textDocument":{
+  "uri":"test:///foo.td",
+  "languageId":"tablegen",
+  "version":1,
+  "text":"class Foo<>;"
+}}}
+// CHECK: "method": "textDocument/publishDiagnostics",
+// CHECK-NEXT: "params": {
+// CHECK-NEXT:     "diagnostics": [
+// CHECK-NEXT:       {
+// CHECK-NEXT:         "category": "Parse Error",
+// CHECK-NEXT:         "message": "Unknown token when expecting a type",
+// CHECK-NEXT:         "range": {
+// CHECK-NEXT:           "end": {
+// CHECK-NEXT:             "character": 11,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           },
+// CHECK-NEXT:           "start": {
+// CHECK-NEXT:             "character": 10,
+// CHECK-NEXT:             "line": 0
+// CHECK-NEXT:           }
+// CHECK-NEXT:         },
+// CHECK-NEXT:         "severity": 1,
+// CHECK-NEXT:         "source": "tablegen"
+// CHECK-NEXT:       }
+// CHECK-NEXT:     ],
+// CHECK-NEXT:     "uri": "test:///foo.td",
+// CHECK-NEXT:     "version": 1
+// CHECK-NEXT:   }
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/test/tblgen-lsp-server/exit-eof.test

+// RUN: not tblgen-lsp-server < %s 2> %t.err
+// RUN: FileCheck %s < %t.err
+//
+// No LSP messages here, just let tblgen-lsp-server see the end-of-file
+// CHECK: Transport error:
+// (Typically "Transport error: Input/output error" but platform-dependent).
+

mlir/test/tblgen-lsp-server/exit-with-shutdown.test

+// RUN: tblgen-lsp-server -lit-test < %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"tablegen","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/test/tblgen-lsp-server/exit-without-shutdown.test

+// RUN: not tblgen-lsp-server -lit-test < %s
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootPath":"tablegen","capabilities":{},"trace":"off"}}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/test/tblgen-lsp-server/initialize-params-invalid.test

+// RUN: tblgen-lsp-server -lit-test < %s | FileCheck %s
+// Test with invalid initialize request parameters
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":"","rootUri":"test:///workspace","capabilities":{},"trace":"verbose"}}
+//      CHECK:  "id": 0,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "capabilities": {
+// ...
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/test/tblgen-lsp-server/initialize-params.test

+// RUN: tblgen-lsp-server -lit-test < %s | FileCheck %s
+// Test initialize request parameters with rootUri
+{"jsonrpc":"2.0","id":0,"method":"initialize","params":{"processId":123,"rootUri":"test:///workspace","capabilities":{},"trace":"off"}}
+//      CHECK:  "id": 0,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": {
+// CHECK-NEXT:    "capabilities": {
+// CHECK-NEXT:      "textDocumentSync": {
+// CHECK-NEXT:        "change": 1,
+// CHECK-NEXT:        "openClose": true,
+// CHECK-NEXT:        "save": true
+// CHECK-NEXT:      }
+// CHECK-NEXT:    },
+// CHECK-NEXT:    "serverInfo": {
+// CHECK-NEXT:      "name": "tblgen-lsp-server",
+// CHECK-NEXT:      "version": "{{.*}}"
+// CHECK-NEXT:    }
+// CHECK-NEXT:  }
+// -----
+{"jsonrpc":"2.0","id":3,"method":"shutdown"}
+//      CHECK:  "id": 3,
+// CHECK-NEXT:  "jsonrpc": "2.0",
+// CHECK-NEXT:  "result": null
+// -----
+{"jsonrpc":"2.0","method":"exit"}

mlir/tools/CMakeLists.txt

 add_subdirectory(mlir-lsp-server)
 add_subdirectory(mlir-opt)
 add_subdirectory(mlir-parser-fuzzer)
 add_subdirectory(mlir-pdll-lsp-server)
 add_subdirectory(mlir-reduce)
 add_subdirectory(mlir-shlib)
 add_subdirectory(mlir-spirv-cpu-runner)
 add_subdirectory(mlir-translate)
 add_subdirectory(mlir-vulkan-runner)
+add_subdirectory(tblgen-lsp-server)
 
 # mlir-cpu-runner requires ExecutionEngine which is only built
 # when the native target is configured in.
 if(TARGET ${LLVM_NATIVE_ARCH})
   add_subdirectory(mlir-cpu-runner)
 endif()

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

+set(LIBS
+  TableGenLspServerLib
+  )
+
+add_llvm_tool(tblgen-lsp-server
+  tblgen-lsp-server.cpp
+
+  DEPENDS
+  ${LIBS}
+  )
+
+target_link_libraries(tblgen-lsp-server PRIVATE ${LIBS})
+llvm_update_compile_flags(tblgen-lsp-server)
+
+mlir_check_all_link_libraries(tblgen-lsp-server)

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

+//===- tblgen-lsp-server.cpp - TableGen Language Server main --------------===//
+//
+// 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 "mlir/Support/LogicalResult.h"
+#include "mlir/Tools/tblgen-lsp-server/TableGenLspServerMain.h"
+
+using namespace mlir;
+
+int main(int argc, char **argv) {
+  return failed(TableGenLspServerMain(argc, argv));
+}

mlir/utils/vscode/.gitignore

 out
 node_modules
 .vscode-test
 *.vsix
 grammar.json
+tablegen-grammar.json
 
 !.vscode
 No newline at end of file

mlir/utils/vscode/package.json

"engines": {
     "vscode": "^1.52.0"
   },
   "categories": [
     "Programming Languages"
   ],
   "keywords": [
     "LLVM",
     "MLIR",
-    "PDLL"
+    "PDLL",
+    "TableGen",
+    "tblgen",
+    "tablegen"
   ],
   "activationEvents": [
     "onLanguage:mlir",
-    "onLanguage:pdll"
+    "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",
"languages": [
         "aliases": [
           "PDLL",
           "pdll"
         ],
         "extensions": [
           ".pdll"
         ],
         "configuration": "./pdll-language-configuration.json"
+      },
+      {
+        "id": "tablegen",
+        "aliases": [
+          "TableGen",
+          "tblgen"
+        ],
+        "extensions": [
+          ".td"
+        ],
+        "configuration": "./tablegen-language-configuration.json"
       }
     ],
     "grammars": [
       {
         "language": "mlir",
         "scopeName": "source.mlir",
         "path": "./grammar.json"
       },
"grammars": [
         "embeddedLanguages": {
           "source.mlir": "mlir"
         }
       },
       {
         "language": "pdll",
         "scopeName": "source.pdll",
         "path": "./pdll-grammar.json"
+      },
+      {
+        "language": "tablegen",
+        "scopeName": "source.tablegen",
+        "path": "./tablegen-grammar.json"
       }
     ],
     "configuration": {
       "type": "object",
       "title": "MLIR",
       "properties": {
         "mlir.server_path": {
           "scope": "resource",
           "type": "string",
           "description": "The file path of the mlir-lsp-server executable."
         },
         "mlir.pdll_server_path": {
           "scope": "resource",
           "type": "string",
           "description": "The file path of the mlir-pdll-lsp-server executable."
         },
         "mlir.pdll_compilation_databases": {
           "scope": "resource",
           "type": "array",
           "description": "A list of `pdll_compile_commands.yml` database files containing information about .pdll files processed by the server."
         },
+        "mlir.tablegen_server_path": {
+          "scope": "resource",
+          "type": "string",
+          "description": "The file path of the tblgen-lsp-server executable."
+        },
         "mlir.onSettingsChanged": {
           "type": "string",
           "default": "prompt",
           "description": "Action taken when a setting change requires a server restart to take effect.",
           "enum": [
             "prompt",
             "restart",
             "ignore"

mlir/utils/vscode/src/mlirContext.ts

const startClientOnOpenDocument = async (document: vscode.TextDocument) => {
       if (document.uri.scheme !== 'file') {
         return;
       }
       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;
       }
 
       // 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);
export class MLIRContext implements vscode.Disposable {
    */
   static getDefaultServerFilename(serverSettingName: string): string {
     if (serverSettingName === 'pdll_server_path') {
       return 'mlir-pdll-lsp-server';
     }
     if (serverSettingName === 'server_path') {
       return 'mlir-lsp-server';
     }
+    if (serverSettingName === 'tablegen_server_path') {
+      return 'tblgen-lsp-server';
+    }
     return '';
   }
 
   /**
    * Try to resolve the given path, or the default path, with an optional
    * workspace folder. If a path could not be resolved, just returns the
    * input filePath.
    */

mlir/utils/vscode/tablegen-language-configuration.json

+{
+  "comments": {
+    "lineComment": "//",
+    "blockComment": [
+      "/*",
+      "*/"
+    ]
+  },
+  "brackets": [
+    [
+      "{",
+      "}"
+    ],
+    [
+      "[",
+      "]"
+    ],
+    [
+      "(",
+      ")"
+    ],
+    [
+      "<",
+      ">"
+    ]
+  ],
+  "autoClosingPairs": [
+    [
+      "{",
+      "}"
+    ],
+    [
+      "[",
+      "]"
+    ],
+    [
+      "(",
+      ")"
+    ],
+    [
+      "<",
+      ">"
+    ],
+    [
+      "\"",
+      "\""
+    ]
+  ],
+  "surroundingPairs": [
+    [
+      "{",
+      "}"
+    ],
+    [
+      "[",
+      "]"
+    ],
+    [
+      "(",
+      ")"
+    ],
+    [
+      "<",
+      ">"
+    ],
+    [
+      "\"",
+      "\""
+    ]
+  ]
+}
 No newline at end of file