[llvm-branch-commits] [llvm] [mlir] [mlir][test] Shard the Test Dialect (NFC) (PR #89628)

[Jeff Niu via llvm-branch-commits]

https://github.com/Mogball created https://github.com/llvm/llvm-project/pull/89628

[mlir][test] Shard the Test Dialect (NFC)

This PR uses the new op sharding mechanism in tablegen to shard the test
dialect's op definitions. This breaks the definition of ops into
multiple source files, speeding up compile time of the test dialect
dramatically. This improves developer cycle times when iterating on the
test dialect.

>From 498ea497b99e1328fc16cbc41ebf37575c1c4a59 Mon Sep 17 00:00:00 2001
From: Mogball <jeff at modular.com>
Date: Mon, 22 Apr 2024 16:39:58 +0000
Subject: [PATCH] [mlir][test] Shard the Test Dialect (NFC)

This PR uses the new op sharding mechanism in tablegen to shard the test
dialect's op definitions. This breaks the definition of ops into
multiple source files, speeding up compile time of the test dialect
dramatically. This improves developer cycle times when iterating on the
test dialect.
---
 mlir/test/lib/Dialect/Test/CMakeLists.txt     |  6 +++--
 mlir/test/lib/Dialect/Test/TestDialect.cpp    |  5 +---
 mlir/test/lib/Dialect/Test/TestOps.cpp        |  1 -
 .../mlir/test/BUILD.bazel                     | 25 +++++++++++--------
 4 files changed, 20 insertions(+), 17 deletions(-)

diff --git a/mlir/test/lib/Dialect/Test/CMakeLists.txt b/mlir/test/lib/Dialect/Test/CMakeLists.txt
index f63e4d330e6ac1..fab89378093326 100644
--- a/mlir/test/lib/Dialect/Test/CMakeLists.txt
+++ b/mlir/test/lib/Dialect/Test/CMakeLists.txt
@@ -31,8 +31,6 @@ mlir_tablegen(TestOpEnums.cpp.inc -gen-enum-defs)
 add_public_tablegen_target(MLIRTestEnumDefIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestOps.td)
-mlir_tablegen(TestOps.h.inc -gen-op-decls)
-mlir_tablegen(TestOps.cpp.inc -gen-op-defs)
 mlir_tablegen(TestOpsDialect.h.inc -gen-dialect-decls -dialect=test)
 mlir_tablegen(TestOpsDialect.cpp.inc -gen-dialect-defs -dialect=test)
 mlir_tablegen(TestPatterns.inc -gen-rewriters)
@@ -43,6 +41,8 @@ mlir_tablegen(TestOpsSyntax.h.inc -gen-op-decls)
 mlir_tablegen(TestOpsSyntax.cpp.inc -gen-op-defs)
 add_public_tablegen_target(MLIRTestOpsSyntaxIncGen)
 
+add_sharded_ops(TestOps 20)
+
 # Exclude tests from libMLIR.so
 add_mlir_library(MLIRTestDialect
   TestAttributes.cpp
@@ -56,6 +56,7 @@ add_mlir_library(MLIRTestDialect
   TestTypes.cpp
   TestOpsSyntax.cpp
   TestDialectInterfaces.cpp
+  ${SHARDED_SRCS}
 
   EXCLUDE_FROM_LIBMLIR
 
@@ -66,6 +67,7 @@ add_mlir_library(MLIRTestDialect
   MLIRTestTypeDefIncGen
   MLIRTestOpsIncGen
   MLIRTestOpsSyntaxIncGen
+  MLIRTestOpsShardGen
 
   LINK_LIBS PUBLIC
   MLIRControlFlowInterfaces
diff --git a/mlir/test/lib/Dialect/Test/TestDialect.cpp b/mlir/test/lib/Dialect/Test/TestDialect.cpp
index 77fd7e61bd3a06..bfb9592e638288 100644
--- a/mlir/test/lib/Dialect/Test/TestDialect.cpp
+++ b/mlir/test/lib/Dialect/Test/TestDialect.cpp
@@ -326,12 +326,9 @@ struct TestOpEffectInterfaceFallback
 void TestDialect::initialize() {
   registerAttributes();
   registerTypes();
-  addOperations<
-#define GET_OP_LIST
-#include "TestOps.cpp.inc"
-      >();
   registerOpsSyntax();
   addOperations<ManualCppOpWithFold>();
+  registerTestDialectOperations(this);
   registerDynamicOp(getDynamicGenericOp(this));
   registerDynamicOp(getDynamicOneOperandTwoResultsOp(this));
   registerDynamicOp(getDynamicCustomParserPrinterOp(this));
diff --git a/mlir/test/lib/Dialect/Test/TestOps.cpp b/mlir/test/lib/Dialect/Test/TestOps.cpp
index ce7e476be74e65..47d5b1b19121ef 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOps.cpp
@@ -14,5 +14,4 @@
 using namespace mlir;
 using namespace test;
 
-#define GET_OP_CLASSES
 #include "TestOps.cpp.inc"
diff --git a/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel b/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
index dc5f4047c286db..b98f7eb5613af4 100644
--- a/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
+++ b/utils/bazel/llvm-project-overlay/mlir/test/BUILD.bazel
@@ -4,7 +4,7 @@
 
 load("@bazel_skylib//rules:expand_template.bzl", "expand_template")
 load("//llvm:lit_test.bzl", "package_path")
-load("//mlir:tblgen.bzl", "gentbl_cc_library", "td_library")
+load("//mlir:tblgen.bzl", "gentbl_cc_library", "td_library", "gentbl_sharded_ops", "td_library")
 
 package(
     default_visibility = ["//visibility:public"],
@@ -151,14 +151,6 @@ gentbl_cc_library(
     name = "TestOpsIncGen",
     strip_include_prefix = "lib/Dialect/Test",
     tbl_outs = [
-        (
-            ["-gen-op-decls"],
-            "lib/Dialect/Test/TestOps.h.inc",
-        ),
-        (
-            ["-gen-op-defs"],
-            "lib/Dialect/Test/TestOps.cpp.inc",
-        ),
         (
             [
                 "-gen-dialect-decls",
@@ -370,12 +362,25 @@ cc_library(
     ],
 )
 
+gentbl_sharded_ops(
+    name = "TestDialectOpSrcs",
+    hdr_out = "lib/Dialect/Test/TestOps.h.inc",
+    shard_count = 20,
+    sharder = "//mlir:mlir-src-sharder",
+    src_file = "lib/Dialect/Test/TestOps.cpp",
+    src_out = "lib/Dialect/Test/TestOps.cpp.inc",
+    tblgen = "//mlir:mlir-tblgen",
+    td_file = "lib/Dialect/Test/TestOps.td",
+    test = True,
+    deps = [":TestOpTdFiles"],
+)
+
 cc_library(
     name = "TestDialect",
     srcs = glob(
         ["lib/Dialect/Test/*.cpp"],
         exclude = ["lib/Dialect/Test/TestToLLVMIRTranslation.cpp"],
-    ),
+    ) + [":TestDialectOpSrcs"],
     hdrs = glob(["lib/Dialect/Test/*.h"]),
     includes = [
         "lib/Dialect/Test",