[mlir][ods] Allow sharding of op definitions

[Mogball]

Summary

This patch adds an option to gen-op-defs and gen-op-decls called op-shard-count that, when set to a value greater than 1, splits the op class definitions into N shards and generates N registration functions, i.e. calls to addOperation<...>, guarded by the definition GET_OP_DEFS_<N>. It will also generate a top-level registration hook registerFooDialectOperations(FooDialect *) that calls each of the sharded registration hooks.

Sharding the op definitions and registration hooks can greatly speed up compilation time, especially for dialects with large opsets. To assist in sharding, a new buildrule is added (gentbl_sharded_ops in bazel and add_mlir_sharded_ops) which will copy a template source file and compile it multiple times with each shard.

Fixes #51789

[Mogball]

Add a build rule for automatically sharding ops.

Any feedback on the build setup is welcome, by the way.

[rriddle]

mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp

17

Can you drop this using?

80

Is the only thing this tool does? i.e. it just prints a #define GET_OP_DEFS_<N> line? I'm not really sure this tool pulls it's weight. Can you describe the file setup that you are trying to get from cmake?

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

3161–3172

Do we really need new generators? Could you not just detect and dispatch accordingly based on the shard count? i.e. if there is only one shard (or if the shard count option wasn't specified) do something unique?

3175

What is this used for?

mlir/tools/mlir-tblgen/OpGenHelpers.h

29–31

AFAICT you don't really need to copy into different std::vectors here, you're just chunking the input vector. It's much more optimal in that case to just use ArrayRef.

[Mogball]

mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp

80

Ideally, what I want is to compile the same source file with different definitions *in the same CMake/Bazel library*, i.e.

clang++ foo.cpp -DSHARD_0 -o foo0.o
clang++ foo.cpp -DSHARD_1 -o foo1.o
ld foodialect.o foo0.o. foo1.o

But I wasn't able to find a way to pull that off. CMake only allows setting per-file defines once per file and Bazel only allows local_defines per build target (I think... I'm not a build file wizard). Doing it with separate libraries would likely end up creating cyclic dependencies or otherwise make the dependency management of a dialect kind of tricky.

I originally used echo -e "#define ..\n${cat src.cpp}" but I'm not sure that works on windows.

[Mogball]

review comments

[rriddle]

mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp

80

If it's just about generating files, file(GENERATE or configure_file( should probably serve what you need. Worst case, when generating the file we would just do something like:

#define GET_OP_DEFS_${OP_DEF_N}
#include "src.cpp"

With ${OP_DEF_N} referencing a variable we set in cmake.

Another option of course, would be to consider what it would take to let tablegen generate multiple files...

[Mogball]

mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp

80

I'll give that a try. Thanks!

The problem with generating multiple files is that they all still the correct includes...

[Mogball]

review comments

[Mogball]

wrong diff

[Mogball]

wrong diff

mlir/CMakeLists.txt

 include_directories( ${MLIR_INCLUDE_DIR})
 
 # Adding tools/mlir-tblgen here as calling add_tablegen sets some variables like
 # MLIR_TABLEGEN_EXE in PARENT_SCOPE which gets lost if that folder is included
 # from another directory like tools
 add_subdirectory(tools/mlir-tblgen)
 add_subdirectory(tools/mlir-linalg-ods-gen)
 add_subdirectory(tools/mlir-pdll)
+add_subdirectory(tools/mlir-src-sharder)
 
 add_subdirectory(include/mlir)
 add_subdirectory(lib)
 # C API needs all dialects for registration, but should be built before tests.
 add_subdirectory(lib/CAPI)
 
 if (MLIR_INCLUDE_TESTS)
   add_definitions(-DMLIR_INCLUDE_TESTS)

mlir/cmake/modules/AddMLIR.cmake

 include(GNUInstallDirs)
 include(LLVMDistributionSupport)
 
 # Clear out any pre-existing compile_commands file before processing. This
 # allows for generating a clean compile_commands on each configure.
 file(REMOVE ${CMAKE_BINARY_DIR}/tablegen_compile_commands.yml)
 
 function(mlir_tablegen ofn)
   tablegen(MLIR ${ARGV})
   set(TABLEGEN_OUTPUT ${TABLEGEN_OUTPUT} ${CMAKE_CURRENT_BINARY_DIR}/${ofn}
       PARENT_SCOPE)
-      
   # Get the current set of include paths for this td file.
   cmake_parse_arguments(ARG "" "" "DEPENDS;EXTRA_INCLUDES" ${ARGN})
   get_directory_property(tblgen_includes INCLUDE_DIRECTORIES)
   list(APPEND tblgen_includes ${ARG_EXTRA_INCLUDES})
   # Filter out any empty include items.
   list(REMOVE_ITEM tblgen_includes "")
 
   # Build the absolute path for the current input file.
function(add_mlir_pdll_library target inputFile ofn)
 
   # Append the includes used for this file to the pdll_compilation_commands
   # file.
   file(APPEND ${CMAKE_BINARY_DIR}/pdll_compile_commands.yml
       "--- !FileInfo:\n"
       "  filepath: \"${LLVM_TARGET_DEFINITIONS_ABSOLUTE}\"\n"
       "  includes: \"${CMAKE_CURRENT_SOURCE_DIR};${tblgen_includes}\"\n"
   )
-  
   add_public_tablegen_target(${target})
 endfunction()
 
 # Declare a dialect in the include directory
 function(add_mlir_dialect dialect dialect_namespace)
   set(LLVM_TARGET_DEFINITIONS ${dialect}.td)
   mlir_tablegen(${dialect}.h.inc -gen-op-decls)
   mlir_tablegen(${dialect}.cpp.inc -gen-op-defs)
   mlir_tablegen(${dialect}Types.h.inc -gen-typedef-decls -typedefs-dialect=${dialect_namespace})
   mlir_tablegen(${dialect}Types.cpp.inc -gen-typedef-defs -typedefs-dialect=${dialect_namespace})
   mlir_tablegen(${dialect}Dialect.h.inc -gen-dialect-decls -dialect=${dialect_namespace})
   mlir_tablegen(${dialect}Dialect.cpp.inc -gen-dialect-defs -dialect=${dialect_namespace})
   add_public_tablegen_target(MLIR${dialect}IncGen)
   add_dependencies(mlir-headers MLIR${dialect}IncGen)
 endfunction()
 
+# Declare sharded dialect operation declarations and definitions
+function(add_sharded_ops ops_target shard_count)
+  set(LLVM_TARGET_DEFINITIONS ${ops_target}.td)
+  mlir_tablegen(${ops_target}.h.inc -gen-sharded-op-decls -op-shard-count=${shard_count})
+  mlir_tablegen(${ops_target}.cpp.inc -gen-sharded-op-defs -op-shard-count=${shard_count})
+  set(LLVM_TARGET_DEFINITIONS ${ops_target}.cpp)
+  foreach(index RANGE ${shard_count})
+    set(SHARDED_SRC ${ops_target}.${index}.cpp)
+    list(APPEND SHARDED_SRCS ${SHARDED_SRC})
+    tablegen(MLIR_SRC_SHARDER ${SHARDED_SRC} -op-shard-index=${index})
+    set(TABLEGEN_OUTPUT ${TABLEGEN_OUTPUT} ${CMAKE_CURRENT_BINARY_DIR}/${SHARDED_SRC})
+  endforeach()
+  add_public_tablegen_target(MLIR${ops_target}ShardGen)
+  set(SHARDED_SRCS ${SHARDED_SRCS} PARENT_SCOPE)
+endfunction()
+
 # Declare a dialect in the include directory
 function(add_mlir_interface interface)
   set(LLVM_TARGET_DEFINITIONS ${interface}.td)
   mlir_tablegen(${interface}.h.inc -gen-op-interface-decls)
   mlir_tablegen(${interface}.cpp.inc -gen-op-interface-defs)
   add_public_tablegen_target(MLIR${interface}IncGen)
   add_dependencies(mlir-generic-headers MLIR${interface}IncGen)
 endfunction()

mlir/cmake/modules/CMakeLists.txt

 set(MLIR_CONFIG_INCLUDE_EXPORTS "include(\"\${MLIR_CMAKE_DIR}/MLIRTargets.cmake\")")
 set(MLIR_CONFIG_INCLUDE_DIRS
   "${MLIR_SOURCE_DIR}/include"
   "${MLIR_BINARY_DIR}/include"
   )
 # Refer to the best host mlir-tbgen, which might be a host-optimized version
 set(MLIR_CONFIG_TABLEGEN_EXE "${MLIR_TABLEGEN_EXE}")
 set(MLIR_CONFIG_PDLL_TABLEGEN_EXE "${MLIR_PDLL_TABLEGEN_EXE}")
+set(MLIR_CONFIG_SRC_SHARDER_TABLEGEN_EXE "${MLIR_SRC_SHARDER_TABLEGEN_EXE}")
 
 configure_file(
   ${CMAKE_CURRENT_SOURCE_DIR}/MLIRConfig.cmake.in
   ${mlir_cmake_builddir}/MLIRConfig.cmake
   @ONLY)
 set(MLIR_CONFIG_CMAKE_DIR)
 set(MLIR_CONFIG_LLVM_CMAKE_DIR)
 set(MLIR_CONFIG_INCLUDE_DIRS)

 extend_path(base_includedir "\${MLIR_INSTALL_PREFIX}" "${CMAKE_INSTALL_INCLUDEDIR}")
 set(MLIR_CONFIG_INCLUDE_DIRS
   "${base_includedir}"
   )
 # Ensure that we are using the installed mlir-tblgen. This might not be MLIR_TABLEGEN_EXE
 # if we're building with a host-optimized mlir-tblgen (with LLVM_OPTIMIZED_TABLEGEN).
 set(MLIR_CONFIG_TABLEGEN_EXE mlir-tblgen)
 set(MLIR_CONFIG_PDLL_TABLEGEN_EXE mlir-pdll)
+set(MLIR_CONFIG_SRC_SHARDER_TABLEGEN_EXE mlir-src-sharder)
 
 configure_file(
   ${CMAKE_CURRENT_SOURCE_DIR}/MLIRConfig.cmake.in
   ${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/MLIRConfig.cmake
   @ONLY)
 set(MLIR_CONFIG_CODE)
 set(MLIR_CONFIG_CMAKE_DIR)
 set(MLIR_CONFIG_LLVM_CMAKE_DIR)

mlir/cmake/modules/MLIRConfig.cmake.in

 # This file allows users to call find_package(MLIR) and pick up our targets.
 
 @MLIR_CONFIG_CODE@
 
 find_package(LLVM REQUIRED CONFIG
              HINTS "@MLIR_CONFIG_LLVM_CMAKE_DIR@")
 
 set(MLIR_EXPORTED_TARGETS "@MLIR_EXPORTS@")
 set(MLIR_CMAKE_DIR "@MLIR_CONFIG_CMAKE_DIR@")
 set(MLIR_INCLUDE_DIRS "@MLIR_CONFIG_INCLUDE_DIRS@")
 set(MLIR_TABLEGEN_EXE "@MLIR_CONFIG_TABLEGEN_EXE@")
 set(MLIR_PDLL_TABLEGEN_EXE "@MLIR_CONFIG_PDLL_TABLEGEN_EXE@")
+set(MLIR_SRC_SHARDER_TABLEGEN_EXE "@MLIR_CONFIG_SRC_SHARDER_TABLEGEN_EXE@")
 set(MLIR_INSTALL_AGGREGATE_OBJECTS "@MLIR_INSTALL_AGGREGATE_OBJECTS@")
 set(MLIR_ENABLE_BINDINGS_PYTHON "@MLIR_ENABLE_BINDINGS_PYTHON@")
 
 # For mlir_tablegen()
 set(MLIR_INCLUDE_DIR "@MLIR_INCLUDE_DIR@")
 set(MLIR_MAIN_SRC_DIR "@MLIR_MAIN_SRC_DIR@")
 
 set_property(GLOBAL PROPERTY MLIR_ALL_LIBS "@MLIR_ALL_LIBS@")

mlir/include/mlir/TableGen/CodeGenHelpers.h

 /// LogicalResult OpB::verify(...) {
 ///  if (failed(localVerify(...)))
 ///    return failure();
 ///  ...
 /// }
 ///
 class StaticVerifierFunctionEmitter {
 public:
+  /// Create a constraint uniquer with a unique prefix derived from the record
+  /// keeper with an optional tag.
   StaticVerifierFunctionEmitter(raw_ostream &os,
-                                const llvm::RecordKeeper &records);
+                                const llvm::RecordKeeper &records,
+                                StringRef tag = "");
+
+  /// Collect and unique all the constraints used by operations.
+  void collectOpConstraints(ArrayRef<llvm::Record *> opDefs);
 
   /// Collect and unique all compatible type, attribute, successor, and region
   /// constraints from the operations in the file and emit them at the top of
   /// the generated file.
   ///
   /// Constraints that do not meet the restriction that they can only reference
   /// `$_self` and `$_op` are not uniqued.
-  void emitOpConstraints(ArrayRef<llvm::Record *> opDefs, bool emitDecl);
+  void emitOpConstraints(ArrayRef<llvm::Record *> opDefs);
 
   /// Unique all compatible type and attribute constraints from a pattern file
   /// and emit them at the top of the generated file.
   ///
   /// Constraints that do not meet the restriction that they can only reference
   /// `$_self`, `$_op`, and `$_builder` are not uniqued.
   void emitPatternConstraints(const ArrayRef<DagLeaf> constraints);
 
private:
   /// Emit static successor constraint functions.
   void emitSuccessorConstraints();
   /// Emit static region constraint functions.
   void emitRegionConstraints();
 
   /// Emit pattern constraints.
   void emitPatternConstraints();
 
-  /// Collect and unique all the constraints used by operations.
-  void collectOpConstraints(ArrayRef<llvm::Record *> opDefs);
   /// Collect and unique all pattern constraints.
   void collectPatternConstraints(ArrayRef<DagLeaf> constraints);
 
   /// The output stream.
   raw_ostream &os;
 
   /// A unique label for the file currently being generated. This is used to
   /// ensure that the static functions have a unique name.

 
 namespace detail {
 template <typename> struct stringifier {
   template <typename T> static std::string apply(T &&t) {
     return std::string(std::forward<T>(t));
   }
 };
 template <> struct stringifier<Twine> {
-  static std::string apply(const Twine &twine) {
-    return twine.str();
-  }
+  static std::string apply(const Twine &twine) { return twine.str(); }
 };
 template <typename OptionalT>
 struct stringifier<Optional<OptionalT>> {
   static std::string apply(Optional<OptionalT> optional) {
     return optional.hasValue() ? stringifier<OptionalT>::apply(*optional)
                                : std::string();
   }
 };

mlir/test/mlir-tblgen/shard-op-defs.td

+// RUN: mlir-tblgen -gen-sharded-op-defs -op-shard-count=2 -I %S/../../include %s | FileCheck %s --check-prefix=DEFS
+// RUN: mlir-tblgen -gen-sharded-op-decls -op-shard-count=2 -I %S/../../include %s | FileCheck %s --check-prefix=DECLS
+
+include "mlir/IR/OpBase.td"
+
+def Test_Dialect : Dialect {
+  let name = "test";
+  let cppNamespace = "test";
+}
+
+class Test_Op<string mnemonic, list<Trait> traits = []> 
+    : Op<Test_Dialect, mnemonic, traits>;
+
+def OpA : Test_Op<"a">;
+def OpB : Test_Op<"b">;
+def OpC : Test_Op<"c">;
+
+// DECLS: OpA
+// DECLS: OpB
+// DECLS: OpC
+// DECLS: registerTestDialectOperations(
+// DECLS: registerTestDialectOperations0(
+// DECLS: registerTestDialectOperations1(
+
+// DEFS-LABEL: GET_OP_DEFS_0
+// DEFS: void test::registerTestDialectOperations(
+// DEFS: void test::registerTestDialectOperations0(
+// DEFS: OpAAdaptor
+// DEFS: OpBAdaptor
+
+// DEFS-LABEL: GET_OP_DEFS_1
+// DEFS: void test::registerTestDialectOperations1(
+// DEFS: OpCAdaptor

mlir/tools/mlir-src-sharder/CMakeLists.txt

+set(LLVM_LINK_COMPONENTS Support)
+set(LIBS MLIRSupport)
+
+add_tablegen(mlir-src-sharder MLIR_SRC_SHARDER
+  mlir-src-sharder.cpp
+
+  DEPENDS
+  ${LIBS}
+  )
+
+set_target_properties(mlir-src-sharder PROPERTIES FOLDER "Tablegenning")
+target_link_libraries(mlir-src-sharder PRIVATE ${LIBS})
+
+mlir_check_all_link_libraries(mlir-src-sharder)

mlir/tools/mlir-src-sharder/mlir-src-sharder.cpp

+//===- mlir-src-sharder.cpp - A tool for sharder generated source files ---===//
+//
+// 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/FileUtilities.h"
+#include "mlir/Support/LogicalResult.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/FileSystem.h"
+#include "llvm/Support/InitLLVM.h"
+#include "llvm/Support/MemoryBuffer.h"
+#include "llvm/Support/ToolOutputFile.h"
+
+using namespace llvm;
+using namespace mlir;
+
+/// Create a dependency file for `-d` option.
+///
+/// This functionality is generally only for the benefit of the build system,
+/// and is modeled after the same option in TableGen.
+static LogicalResult createDependencyFile(StringRef outputFilename,
+                                          StringRef dependencyFile) {
+  if (outputFilename == "-") {
+    llvm::errs() << "error: the option -d must be used together with -o\n";
+    return failure();
+  }
+
+  std::string errorMessage;
+  std::unique_ptr<llvm::ToolOutputFile> outputFile =
+      openOutputFile(dependencyFile, &errorMessage);
+  if (!outputFile) {
+    llvm::errs() << errorMessage << "\n";
+    return failure();
+  }
+
+  outputFile->os() << outputFilename << ":\n";
+  outputFile->keep();
+  return success();
+}
+
+int main(int argc, char **argv) {
+  // FIXME: This is necessary because we link in TableGen, which defines its
+  // options as static variables.. some of which overlap with our options.
+  cl::ResetCommandLineParser();
+
+  cl::opt<unsigned> opShardIndex("op-shard-index",
+                                 cl::desc("The current shard index"));
+  cl::opt<std::string> inputFilename(cl::Positional, cl::desc("<input file>"),
+                                     cl::init("-"));
+  cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
+                                      cl::value_desc("filename"),
+                                      cl::init("-"));
+  cl::list<std::string> includeDirs("I", cl::desc("Directory of include files"),
+                                    cl::value_desc("directory"), cl::Prefix);
+  cl::opt<std::string> dependencyFilename("d", cl::desc("Dependency filename"),
+                                          cl::value_desc("filename"),
+                                          cl::init(""));
+  cl::opt<bool> writeIfChanged(
+      "write-if-changed",
+      cl::desc("Only write to the output file if it changed"));
+
+  InitLLVM y(argc, argv);
+  cl::ParseCommandLineOptions(argc, argv);
+
+  // Open the input file.
+  std::string errorMessage;
+  std::unique_ptr<llvm::MemoryBuffer> inputFile =
+      openInputFile(inputFilename, &errorMessage);
+  if (!inputFile) {
+    llvm::errs() << errorMessage << "\n";
+    return 1;
+  }
+
+  // Write the output to a buffer.
+  std::string outputStr;
+  raw_string_ostream os(outputStr);
+  os << "#define GET_OP_DEFS_" << opShardIndex << "\n"
+     << inputFile->getBuffer();
+
+  // Determine whether we need to write the output file.
+  bool shouldWriteOutput = true;
+  if (writeIfChanged) {
+    // Only update the real output file if there are any differences. This
+    // prevents recompilation of all the files depending on it if there aren't
+    // any.
+    if (auto existingOrErr =
+            MemoryBuffer::getFile(outputFilename, /*IsText=*/true))
+      if (std::move(existingOrErr.get())->getBuffer() == os.str())
+        shouldWriteOutput = false;
+  }
+
+  // Populate the output file if necessary.
+  if (shouldWriteOutput) {
+    std::unique_ptr<ToolOutputFile> outputFile =
+        openOutputFile(outputFilename, &errorMessage);
+    if (!outputFile) {
+      llvm::errs() << errorMessage << "\n";
+      return 1;
+    }
+    outputFile->os() << os.str();
+    outputFile->keep();
+  }
+
+  // Always write the depfile, even if the main output hasn't changed. If it's
+  // missing, Ninja considers the output dirty.
+  if (!dependencyFilename.empty())
+    if (failed(createDependencyFile(outputFilename, dependencyFilename)))
+      return 1;
+
+  return 0;
+}

mlir/tools/mlir-tblgen/CodeGenHelpers.cpp

 #include "llvm/TableGen/Record.h"
 
 using namespace llvm;
 using namespace mlir;
 using namespace mlir::tblgen;
 
 /// Generate a unique label based on the current file name to prevent name
 /// collisions if multiple generated files are included at once.
-static std::string getUniqueOutputLabel(const llvm::RecordKeeper &records) {
+static std::string getUniqueOutputLabel(const llvm::RecordKeeper &records,
+                                        StringRef tag) {
   // Use the input file name when generating a unique name.
   std::string inputFilename = records.getInputFilename();
 
   // Drop all but the base filename.
   StringRef nameRef = llvm::sys::path::filename(inputFilename);
   nameRef.consume_back(".td");
 
   // Sanitize any invalid characters.
-  std::string uniqueName;
+  std::string uniqueName(tag);
   for (char c : nameRef) {
     if (llvm::isAlnum(c) || c == '_')
       uniqueName.push_back(c);
     else
       uniqueName.append(llvm::utohexstr((unsigned char)c));
   }
   return uniqueName;
 }
 
 StaticVerifierFunctionEmitter::StaticVerifierFunctionEmitter(
-    raw_ostream &os, const llvm::RecordKeeper &records)
-    : os(os), uniqueOutputLabel(getUniqueOutputLabel(records)) {}
+    raw_ostream &os, const llvm::RecordKeeper &records, StringRef tag)
+    : os(os), uniqueOutputLabel(getUniqueOutputLabel(records, tag)) {}
 
 void StaticVerifierFunctionEmitter::emitOpConstraints(
-    ArrayRef<llvm::Record *> opDefs, bool emitDecl) {
-  collectOpConstraints(opDefs);
-  if (emitDecl)
-    return;
-
+    ArrayRef<llvm::Record *> opDefs) {
   NamespaceEmitter namespaceEmitter(os, Operator(*opDefs[0]).getCppNamespace());
   emitTypeConstraints();
   emitAttrConstraints();
   emitSuccessorConstraints();
   emitRegionConstraints();
 }
 
 void StaticVerifierFunctionEmitter::emitPatternConstraints(

mlir/tools/mlir-tblgen/OpDefinitionsGen.cpp

 
 void OpOperandAdaptorEmitter::emitDef(
     const Operator &op,
     const StaticVerifierFunctionEmitter &staticVerifierEmitter,
     raw_ostream &os) {
   OpOperandAdaptorEmitter(op, staticVerifierEmitter).adaptor.writeDefTo(os);
 }
 
-// Emits the opcode enum and op classes.
-static void emitOpClasses(const RecordKeeper &recordKeeper,
-                          const std::vector<Record *> &defs, raw_ostream &os,
-                          bool emitDecl) {
-  // First emit forward declaration for each class, this allows them to refer
-  // to each others in traits for example.
-  if (emitDecl) {
-    os << "#if defined(GET_OP_CLASSES) || defined(GET_OP_FWD_DEFINES)\n";
-    os << "#undef GET_OP_FWD_DEFINES\n";
-    for (auto *def : defs) {
-      Operator op(*def);
-      NamespaceEmitter emitter(os, op.getCppNamespace());
-      os << "class " << op.getCppClassName() << ";\n";
-    }
-    os << "#endif\n\n";
-  }
-
-  IfDefScope scope("GET_OP_CLASSES", os);
+/// Emit the class declarations or definitions for the given op defs.
+static void
+emitOpClasses(const RecordKeeper &recordKeeper,
+              const StaticVerifierFunctionEmitter &staticVerifierEmitter,
   if (defs.empty())
     return;
 
-  // Generate all of the locally instantiated methods first.
-  StaticVerifierFunctionEmitter staticVerifierEmitter(os, recordKeeper);
-  os << formatv(opCommentHeader, "Local Utility Method", "Definitions");
-  staticVerifierEmitter.emitOpConstraints(defs, emitDecl);
-
   for (auto *def : defs) {
     Operator op(*def);
     if (emitDecl) {
       {
         NamespaceEmitter emitter(os, op.getCppNamespace());
         os << formatv(opCommentHeader, op.getQualCppClassName(),
                       "declarations");
         OpOperandAdaptorEmitter::emitDecl(op, staticVerifierEmitter, os);
if (emitDecl) {
       // Emit the TypeID explicit specialization to have a single definition.
       if (!op.getCppNamespace().empty())
         os << "MLIR_DEFINE_EXPLICIT_TYPE_ID(" << op.getCppNamespace()
            << "::" << op.getCppClassName() << ")\n\n";
     }
   }
 }
 
-// Emits a comma-separated list of the ops.
-static void emitOpList(const std::vector<Record *> &defs, raw_ostream &os) {
-  IfDefScope scope("GET_OP_LIST", os);
+/// Emit the declarations for the provided op classes.
+static void emitOpClassDecls(const RecordKeeper &recordKeeper,
+                             const std::vector<Record *> &defs,
+                             raw_ostream &os) {
+  // First emit forward declaration for each class, this allows them to refer
+  // to each others in traits for example.
+  for (auto *def : defs) {
+    Operator op(*def);
+    NamespaceEmitter emitter(os, op.getCppNamespace());
+    os << "class " << op.getCppClassName() << ";\n";
+  }
 
-  interleave(
-      // TODO: We are constructing the Operator wrapper instance just for
-      // getting it's qualified class name here. Reduce the overhead by having a
-      // lightweight version of Operator class just for that purpose.
-      defs, [&os](Record *def) { os << Operator(def).getQualCppClassName(); },
-      [&os]() { os << ",\n"; });
+  // Emit the op class declarations.
+  IfDefScope scope("GET_OP_CLASSES", os);
+  if (defs.empty())
+    return;
+  StaticVerifierFunctionEmitter staticVerifierEmitter(os, recordKeeper);
+  staticVerifierEmitter.collectOpConstraints(defs);
+  emitOpClasses(recordKeeper, defs, os, staticVerifierEmitter,
+                /*emitDecl=*/true);
+}
+
+/// Emit the definitions for the provided op classes.
+static void emitOpClassDefs(const RecordKeeper &recordKeeper,
+                            ArrayRef<Record *> defs, raw_ostream &os,
+                            StringRef constraintPrefix = "") {
+  if (defs.empty())
+    return;
+
+  // Generate all of the locally instantiated methods first.
+  StaticVerifierFunctionEmitter staticVerifierEmitter(os, recordKeeper,
+                                                      constraintPrefix);
+  os << formatv(opCommentHeader, "Local Utility Method", "Definitions");
+  staticVerifierEmitter.collectOpConstraints(defs);
+  staticVerifierEmitter.emitOpConstraints(defs);
+
+  // Emit the classes.
+  emitOpClasses(recordKeeper, defs, os, staticVerifierEmitter,
+                /*emitDecl=*/false);
 }
 
+/// Emit op declarations for all op records.
 static bool emitOpDecls(const RecordKeeper &recordKeeper, raw_ostream &os) {
   emitSourceFileHeader("Op Declarations", os);
 
   std::vector<Record *> defs = getRequestedOpDefinitions(recordKeeper);
-  emitOpClasses(recordKeeper, defs, os, /*emitDecl=*/true);
+  emitOpClassDecls(recordKeeper, defs, os);
 
   return false;
 }
 
+/// Emit op definitions for all op records.
 static bool emitOpDefs(const RecordKeeper &recordKeeper, raw_ostream &os) {
   emitSourceFileHeader("Op Definitions", os);
 
   std::vector<Record *> defs = getRequestedOpDefinitions(recordKeeper);
-  emitOpList(defs, os);
-  emitOpClasses(recordKeeper, defs, os, /*emitDecl=*/false);
+  {
+    IfDefScope scope("GET_OP_LIST", os);
+    interleave(
+        defs, os,
+        [&](Record *def) { os << Operator(def).getQualCppClassName(); }, ",\n");
+  }
+  {
+    IfDefScope scope("GET_OP_CLASSES", os);
+    emitOpClassDefs(recordKeeper, defs, os);
+  }
+  return false;
+}
+
+/// Emit op declarations and sharded op registration hook declarations.
+static bool emitShardedOpDecls(const RecordKeeper &recordKeeper,
+                               raw_ostream &os) {
+  emitSourceFileHeader("Op Declarations", os);
+
+  std::vector<Record *> defs = getRequestedOpDefinitions(recordKeeper);
+  Dialect dialect = Operator(defs.front()).getDialect();
+  emitOpClassDecls(recordKeeper, defs, os);
+
+  // Emit the op registration hooks.
+  SmallVector<std::vector<Record *>, 4> shardedDefs;
+  shardOpDefinitions(std::move(defs), shardedDefs);
+  NamespaceEmitter ns(os, dialect);
+
+  const char *const opRegistrationHook =
+      "void register{0}Operations{1}({2}::{0} *dialect);\n";
+  os << formatv(opRegistrationHook, dialect.getCppClassName(), "",
+                dialect.getCppNamespace());
+  for (unsigned i = 0; i < shardedDefs.size(); ++i) {
+    os << formatv(opRegistrationHook, dialect.getCppClassName(), i,
+                  dialect.getCppNamespace());
+  }
+
+  return false;
+}
+
+/// Generate the dialect op registration hook and the op class definitions.
+static void emitOpDefShard(const RecordKeeper &recordKeeper,
+                           ArrayRef<Record *> defs, const Dialect &dialect,
+                           unsigned shardIndex, unsigned shardCount,
+                           raw_ostream &os) {
+  std::string shardGuard = "GET_OP_DEFS_";
+  std::string indexStr = std::to_string(shardIndex);
+  shardGuard += indexStr;
+  IfDefScope scope(shardGuard, os);
+
+  // Emit the op registration hook in the first shard.
+  const char *const opRegistrationHook =
+      "void {0}::register{1}Operations{2}({0}::{1} *dialect) {{\n";
+  if (shardIndex == 0) {
+    os << formatv(opRegistrationHook, dialect.getCppNamespace(),
+                  dialect.getCppClassName(), "");
+    for (unsigned i = 0; i < shardCount; ++i) {
+      os << formatv("  {0}::register{1}Operations{2}(dialect);\n",
+                    dialect.getCppNamespace(), dialect.getCppClassName(), i);
+    }
+    os << "}\n";
+  }
+
+  // Generate the per-shard op registration hook.
+  os << formatv(opCommentHeader, dialect.getCppClassName(),
+                "Op Registration Hook")
+     << formatv(opRegistrationHook, dialect.getCppNamespace(),
+                dialect.getCppClassName(), shardIndex);
+  for (Record *def : defs) {
+    os << formatv("  ::mlir::RegisteredOperationName::insert<{0}>(*dialect);\n",
+                  Operator(def).getQualCppClassName());
+  }
+  os << "}\n";
+
+  // Generate the per-shard op definitions.
+  emitOpClassDefs(recordKeeper, defs, os, indexStr);
+}
+
+/// Emit sharded op definitions and registration hooks.
+static bool emitShardedOpDefs(const RecordKeeper &recordKeeper,
+                              raw_ostream &os) {
+  std::vector<Record *> defs = getRequestedOpDefinitions(recordKeeper);
+  if (defs.empty())
+    return false;
+  Dialect dialect = Operator(defs.front()).getDialect();
+
+  SmallVector<std::vector<Record *>, 4> shardedDefs;
+  shardOpDefinitions(std::move(defs), shardedDefs);
+
+  for (auto &it : llvm::enumerate(shardedDefs)) {
+    emitOpDefShard(recordKeeper, it.value(), dialect, it.index(),
+                   shardedDefs.size(), os);
+  }
 
   return false;
 }
 
 static mlir::GenRegistration
     genOpDecls("gen-op-decls", "Generate op declarations",
                [](const RecordKeeper &records, raw_ostream &os) {
                  return emitOpDecls(records, os);
                });
 
 static mlir::GenRegistration genOpDefs("gen-op-defs", "Generate op definitions",
                                        [](const RecordKeeper &records,
                                           raw_ostream &os) {
                                          return emitOpDefs(records, os);
                                        });
+
+static mlir::GenRegistration
+    genShardedOpDecls("gen-sharded-op-decls",
+                      "Generate sharded op declarations",
+                      [](const RecordKeeper &records, raw_ostream &os) {
+                        return emitShardedOpDecls(records, os);
+                      });
+
+static mlir::GenRegistration
+    genShardedOpDefs("gen-sharded-op-defs", "Generate sharded op definitions",
+                     [](const RecordKeeper &records, raw_ostream &os) {
+                       return emitShardedOpDefs(records, os);
+                     });
+
+static mlir::GenRegistration genShardSourceFile(
+    "gen-shard-src", "Generate a sharded source file for op definitions",
+    [](const RecordKeeper &records, raw_ostream &os) { return false; });

mlir/tools/mlir-tblgen/OpGenHelpers.h

 //===- OpGenHelpers.h - MLIR operation generator helpers --------*- 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 defines helpers used in the op generators.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TOOLS_MLIRTBLGEN_OPGENHELPERS_H_
 #define MLIR_TOOLS_MLIRTBLGEN_OPGENHELPERS_H_
 
+#include "mlir/Support/LLVM.h"
 #include "llvm/TableGen/Record.h"
 #include <vector>
 
 namespace mlir {
 namespace tblgen {
 
 /// Returns all the op definitions filtered by the user. The filtering is via
 /// command-line option "op-include-regex" and "op-exclude-regex".
 std::vector<llvm::Record *>
 getRequestedOpDefinitions(const llvm::RecordKeeper &recordKeeper);
 
+/// Shard the op defintions into the number of shards set by "op-shard-count".
+void shardOpDefinitions(
+    std::vector<llvm::Record *> &&defs,
+    SmallVectorImpl<std::vector<llvm::Record *>> &shardedDefs);
+
 } // namespace tblgen
 } // namespace mlir
 
 #endif //  MLIR_TOOLS_MLIRTBLGEN_OPGENHELPERS_H_

mlir/tools/mlir-tblgen/OpGenHelpers.cpp

 static cl::opt<std::string> opIncFilter(
     "op-include-regex",
     cl::desc("Regex of name of op's to include (no filter if empty)"),
     cl::cat(opDefGenCat));
 static cl::opt<std::string> opExcFilter(
     "op-exclude-regex",
     cl::desc("Regex of name of op's to exclude (no filter if empty)"),
     cl::cat(opDefGenCat));
+static cl::opt<unsigned> opShardCount(
+    "op-shard-count",
+    cl::desc("The number of shards into which the op classes will be divided"),
+    cl::cat(opDefGenCat), cl::init(1));
 
 static std::string getOperationName(const Record &def) {
   auto prefix = def.getValueAsDef("opDialect")->getValueAsString("name");
   auto opName = def.getValueAsString("opName");
   if (prefix.empty())
     return std::string(opName);
   return std::string(llvm::formatv("{0}.{1}", prefix, opName));
 }
for (const auto &def : recordKeeper.getDefs()) {
     if (!opExcFilter.empty() &&
         excludeRegex.match(getOperationName(*def.second)))
       continue;
     defs.push_back(def.second.get());
   }
 
   return defs;
 }
+
+void mlir::tblgen::shardOpDefinitions(
+    std::vector<llvm::Record *> &&defs,
+    SmallVectorImpl<std::vector<llvm::Record *>> &shardedDefs) {
+  assert(opShardCount > 0 && "expected a positive shard count");
+  if (opShardCount == 1) {
+    shardedDefs.push_back(std::move(defs));
+    return;
+  }
+
+  unsigned minShardSize = defs.size() / opShardCount;
+  unsigned numMissing = defs.size() - minShardSize * opShardCount;
+  shardedDefs.assign(opShardCount, {});
+  auto it = defs.begin();
+  for (unsigned i = 0; i < opShardCount; ++i) {
+    std::vector<Record *> &shardDefs = shardedDefs[i];
+    for (unsigned opCount = 0; opCount < minShardSize + (i < numMissing);
+         ++opCount, ++it)
+      shardDefs.push_back(*it);
+  }
+}

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":TableGen",
         "//llvm:Support",
         "//llvm:TableGen",
         "//llvm:config",
     ],
 )
 
 cc_binary(
+    name = "mlir-src-sharder",
+    srcs = ["tools/mlir-src-sharder/mlir-src-sharder.cpp"],
+    deps = [
+        ":Support",
+        "//llvm:Support",
+    ],
+)
+
+cc_binary(
     name = "mlir-linalg-ods-yaml-gen",
     srcs = [
         "tools/mlir-linalg-ods-gen/mlir-linalg-ods-yaml-gen.cpp",
     ],
     deps = [
         ":IR",
         ":Parser",
         ":Support",

utils/bazel/llvm-project-overlay/mlir/tblgen.bzl

native.cc_library(
         name = name,
         # strip_include_prefix does not apply to textual_hdrs.
         # https://github.com/bazelbuild/bazel/issues/12424
         hdrs = [":" + filegroup_name] if strip_include_prefix else [],
         strip_include_prefix = strip_include_prefix,
         textual_hdrs = [":" + filegroup_name],
         **kwargs
     )
+
+def _gentbl_shard_impl(ctx):
+    args = ctx.actions.args()
+    args.add(ctx.file.src_file)
+    args.add("-op-shard-index", ctx.attr.index)
+    args.add("-o", ctx.outputs.out.path)
+    ctx.actions.run(
+        outputs = [ctx.outputs.out],
+        inputs = [ctx.file.src_file],
+        executable = ctx.executable.sharder,
+        arguments = [args],
+        use_default_shell_env = True,
+        mnemonic = "ShardGenerate",
+    )
+
+gentbl_shard_rule = rule(
+    _gentbl_shard_impl,
+    doc = "",
+    output_to_genfiles = True,
+    attrs = {
+        "index": attr.int(mandatory = True, doc = ""),
+        "sharder": attr.label(
+            doc = "",
+            executable = True,
+            cfg = "exec",
+        ),
+        "src_file": attr.label(
+            doc = "",
+            allow_single_file = True,
+            mandatory = True,
+        ),
+        "out": attr.output(
+            doc = "",
+            mandatory = True,
+        ),
+    },
+)
+
+def gentbl_sharded_ops(
+        name,
+        tblgen,
+        sharder,
+        td_file,
+        shard_count,
+        src_file,
+        src_out,
+        hdr_out,
+        test = False,
+        includes = [],
+        strip_include_prefix = None,
+        deps = []):
+    """Generate sharded op declarations and definitions.
+
+    This special build rule shards op definitions in a TableGen file and generates multiple copies
+    of a template source file for including and compiling each shard. The rule defines a filegroup
+    consisting of the source shards, the generated source file, and the generated header file.
+
+    Args:
+      name: The name of the filegroup.
+      tblgen: The binary used to produce the output.
+      sharder: The source file sharder to use.
+      td_file: The primary table definitions file.
+      shard_count: The number of op definition shards to produce.
+      src_file: The source file template.
+      src_out: The generated source file.
+      hdr_out: The generated header file.
+      test: Whether this is a test target.
+      includes: See gentbl_rule.includes
+      deps: See gentbl_rule.deps
+      strip_include_prefix: Attribute to pass through to cc_library.
+    """
+    cc_lib_name = name + "__gentbl_cc_lib"
+    gentbl_cc_library(
+        name = cc_lib_name,
+        strip_include_prefix = strip_include_prefix,
+        includes = includes,
+        tbl_outs = [
+            (
+                [
+                    "-gen-sharded-op-defs",
+                    "-op-shard-count=" + str(shard_count),
+                ],
+                src_out,
+            ),
+            (
+                [
+                    "-gen-sharded-op-decls",
+                    "-op-shard-count=" + str(shard_count),
+                ],
+                hdr_out,
+            ),
+        ],
+        tblgen = tblgen,
+        td_file = td_file,
+        test = test,
+        deps = deps,
+    )
+    all_files = [hdr_out, src_out]
+    for i in range(0, shard_count):
+        out_file = "shard_copy_" + str(i) + "_" + src_file
+        gentbl_shard_rule(
+            index = i,
+            name = name + "__src_shard" + str(i),
+            testonly = test,
+            out = out_file,
+            sharder = sharder,
+            src_file = src_file,
+        )
+        all_files.append(out_file)
+    native.filegroup(name = name, srcs = all_files)
+
+def gentbl_sharded_op_defs(name, source_file, shard_count):
+    """Generates multiple copies of a source file that includes sharded op definitions.
+
+    Args:
+      name: The name of the rule.
+      source_file: The source to copy.
+      shard_count: The number of shards.
+
+    Returns:
+      A list of the copied filenames to be included in the dialect library.
+    """
+    copies = []
+    for i in range(0, shard_count):
+        out_file = "shard_copy_" + str(i) + "_" + source_file
+        copies.append(out_file)
+        native.genrule(
+            name = name + "_shard_" + str(i),
+            srcs = [source_file],
+            outs = [out_file],
+            cmd = "echo -e \"#define GET_OP_DEFS_" + str(i) + "\n$$(cat $(SRCS))\" > $(OUTS)",
+        )
+    return copies