[Mlir-commits] [mlir] [MLIR][ROCDL] Shard generated operation definitions (NFC) (PR #221729)

Mehdi Amini llvmlistbot at llvm.org
Mon Sep 7 06:37:35 PDT 2026 

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/221729

ROCDL generated operation definitions dominate the slowest MLIR translation unit. Generate eight definition shards, compile them as separate sources, and use the generated sharded registration hook while keeping custom parser and printer helpers in one implementation unit.

Assisted-by: Codex

>From 71607826a225efc04fb8513a3002268d18e603e5 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Mon, 7 Sep 2026 05:18:46 -0700
Subject: [PATCH] [mlir][rocdl] Shard generated operation definitions

ROCDL generated operation definitions dominate the slowest MLIR translation
unit. Generate eight definition shards, compile them as separate sources, and
use the generated sharded registration hook while keeping custom parser and
printer helpers in one implementation unit.

The clean all build completed 5,525 Ninja steps in 7m17.17s. It used
11,387.70s user time and 694.49s system time at 2,763% effective CPU, with
3,007,068 KiB peak RSS and zero swaps. The build produced 1,377 MLIR C++ TU
traces totaling 9,710.73s of ExecuteCompiler time; the build tree was 7.8 GiB
and MLIR trace JSON files occupied 1.1 GiB.

Inclusive compiler phase totals were 7,366.12s frontend (75.9%), 5,480.06s
source processing (56.4%), 4,712.96s class parsing (48.5%), 2,707.24s class
instantiation (27.9%), 2,227.97s backend (22.9%), 1,902.95s function
instantiation (19.6%), and 1,467.13s optimizer (15.1%). These nested phase
percentages are not additive.

The ten slowest TUs were ROCDLDialect.cpp at 156.19s (51.50s frontend,
104.12s backend), LLVMDialect.cpp at 136.03s (43.84s, 91.70s),
NVVMDialect.cpp at 111.68s (40.40s, 70.80s), SPIRVOpDefinition.cpp at 87.30s
(25.79s, 61.23s), TilingInterfaceImpl.cpp at 86.20s (31.28s, 54.44s),
SPIRVDialect.cpp at 83.92s (38.29s, 45.00s), TosaProfileCompliance.cpp at
65.86s (7.28s, 58.48s), OpenACC.cpp at 60.99s (20.44s, 40.31s),
OpenMPDialect.cpp at 60.07s (21.39s, 38.28s), and LinalgDialect.cpp at 53.58s
(19.11s, 34.08s).

The top 10 TUs contributed 901.82s (9.3%), and the top 100 contributed
2,832.26s (29.2%). There were 49 TUs over 20s and 274 over 10s; median TU time
was 4.90s. Related registration instantiations took 24.39s for SPIR-V, 16.52s
for ROCDL, and 16.45s for LLVM. The Linalg registerAll instantiation took
19.21s.

Against the 156.19s monolithic compile (156.188s ExecuteCompiler, 51.50s
frontend, 104.12s backend, and 156.890s Ninja edge), the eight shards plus
dialect TU totaled 115.32s ExecuteCompiler (115.323s unrounded), 63.74s
frontend, 50.73s backend, and 115.868s summed Ninja time. This is a 26.2%
aggregate compiler reduction, a 23.8% frontend increase, and a 51.3% backend
reduction. The longest compile edge fell 89.7% to 16.10s (16.096s unrounded).

The final affected-target run took 25.65s wall time, including 3.5s of
automatic CMake configuration and 4.2s of generation (7.7s total), leaving
about 17.9s for shard generation, compilation, and archive linking. It used
118.11s user time and 7.22s system time at 488% CPU, with 1,072,364 KiB peak
RSS.

The static archive grew from 35,489,242 to 36,254,922 bytes (2.2%), while
summed ELF sections grew from 9,295,988 to 9,668,989 bytes (4.0%).

Validation: build MLIRROCDLDialect and mlir-opt with -Werror; pass 56 focused
ROCDL dialect, target, and conversion tests; pass check-mlir with 3,937 passed,
619 unsupported, and 1 expected failure; and run CodeRabbit over all six files
with no findings.

Assisted-by: Codex
---
 .../mlir/Dialect/LLVMIR/CMakeLists.txt        |  9 +++++++-
 .../mlir/Dialect/LLVMIR/ROCDLDialect.h        |  4 ++--
 mlir/lib/Dialect/LLVMIR/CMakeLists.txt        | 15 +++++++++++++
 mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp   | 14 +++++-------
 mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.cpp       | 16 ++++++++++++++
 mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.h         | 22 +++++++++++++++++++
 6 files changed, 68 insertions(+), 12 deletions(-)
 create mode 100644 mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.cpp
 create mode 100644 mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.h

diff --git a/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt b/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
index de9a3e1c93282..b822ced98a649 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
@@ -74,7 +74,14 @@ mlir_tablegen(NVVMOpsAttributes.h.inc -gen-attrdef-decls -attrdefs-dialect=nvvm)
 mlir_tablegen(NVVMOpsAttributes.cpp.inc -gen-attrdef-defs -attrdefs-dialect=nvvm)
 add_mlir_dialect_tablegen_target(MLIRNVVMConversionsIncGen)
 
-add_mlir_dialect(ROCDLOps rocdl)
+set(LLVM_TARGET_DEFINITIONS ROCDLOps.td)
+mlir_tablegen(ROCDLOps.h.inc -gen-op-decls -op-shard-count=8)
+mlir_tablegen(ROCDLOps.cpp.inc -gen-op-defs -op-shard-count=8)
+mlir_tablegen(ROCDLOpsTypes.h.inc -gen-typedef-decls -typedefs-dialect=rocdl)
+mlir_tablegen(ROCDLOpsTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=rocdl)
+mlir_tablegen(ROCDLOpsDialect.h.inc -gen-dialect-decls -dialect=rocdl)
+mlir_tablegen(ROCDLOpsDialect.cpp.inc -gen-dialect-defs -dialect=rocdl)
+add_mlir_dialect_tablegen_target(MLIRROCDLOpsIncGen)
 add_mlir_doc(ROCDLOps ROCDLDialect Dialects/ -gen-dialect-doc -dialect=rocdl)
 set(LLVM_TARGET_DEFINITIONS ROCDLOps.td)
 mlir_tablegen(ROCDLConversions.inc -gen-llvmir-conversions)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h b/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h
index 58d35a4373656..f89c8812ca4d4 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h
@@ -34,9 +34,9 @@
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.h.inc"
 
+#include "mlir/Dialect/LLVMIR/ROCDLOpsDialect.h.inc"
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/ROCDLOps.h.inc"
 
-#include "mlir/Dialect/LLVMIR/ROCDLOpsDialect.h.inc"
-
 #endif /* MLIR_DIALECT_LLVMIR_ROCDLDIALECT_H_ */
diff --git a/mlir/lib/Dialect/LLVMIR/CMakeLists.txt b/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
index 494a1e4c3ee97..2e4f916228bcf 100644
--- a/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
+++ b/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
@@ -1,5 +1,9 @@
 add_subdirectory(Transforms)
 
+set(LLVM_OPTIONAL_SOURCES
+  IR/ROCDLOps.cpp
+  )
+
 add_mlir_dialect_library(MLIRLLVMDialect
   IR/FunctionCallUtils.cpp
   IR/LLVMAttrs.cpp
@@ -71,8 +75,18 @@ add_mlir_dialect_library(MLIRNVVMDialect
   MLIRInferIntRangeInterface
   )
 
+set(LLVM_TARGET_DEFINITIONS IR/ROCDLOps.cpp)
+foreach(index RANGE 7)
+  set(ROCDL_SHARDED_SRC ROCDLOps.${index}.cpp)
+  list(APPEND ROCDL_SHARDED_SRCS ${ROCDL_SHARDED_SRC})
+  tablegen(MLIR_SRC_SHARDER ${ROCDL_SHARDED_SRC}
+    -op-shard-index=${index})
+endforeach()
+add_public_tablegen_target(MLIRROCDLOpsShardGen)
+
 add_mlir_dialect_library(MLIRROCDLDialect
   IR/ROCDLDialect.cpp
+  ${ROCDL_SHARDED_SRCS}
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LLVMIR
@@ -80,6 +94,7 @@ add_mlir_dialect_library(MLIRROCDLDialect
   DEPENDS
   MLIRGPUCompilationAttrInterfacesIncGen
   MLIRROCDLOpsIncGen
+  MLIRROCDLOpsShardGen
   MLIRROCDLConversionsIncGen
   intrinsics_gen
 
diff --git a/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
index 686b62f1dd80e..e5d2bed07d05c 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
@@ -16,6 +16,8 @@
 
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
 
+#include "IR/ROCDLOps.h"
+
 #include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Builders.h"
@@ -51,10 +53,7 @@ struct ROCDLInlinerInterface final : DialectInlinerInterface {
 
 // TODO: This should be the llvm.rocdl dialect once this is supported.
 void ROCDLDialect::initialize() {
-  addOperations<
-#define GET_OP_LIST
-#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
-      >();
+  registerROCDLDialectOperations(this);
 
   addAttributes<
 #define GET_ATTRDEF_LIST
@@ -125,7 +124,7 @@ static ParseResult parseCachePolicyEnum(OpAsmParser &parser,
   return success();
 }
 
-static ParseResult parseCachePolicy(OpAsmParser &parser,
+ParseResult ROCDL::parseCachePolicy(OpAsmParser &parser,
                                     Attribute &cachePolicy) {
   uint32_t rawValue;
   OptionalParseResult rawValueParseResult =
@@ -166,7 +165,7 @@ static void printCachePolicyEnum(OpAsmPrinter &printer, EnumAttrT cachePolicy,
   printer << family << "<" << cachePolicy.getValue() << ">";
 }
 
-static void printCachePolicy(OpAsmPrinter &printer, Operation *,
+void ROCDL::printCachePolicy(OpAsmPrinter &printer, Operation *,
                              Attribute cachePolicy) {
   llvm::TypeSwitch<Attribute>(cachePolicy)
       .Case<IntegerAttr>([&](IntegerAttr rawPolicy) {
@@ -220,8 +219,5 @@ ROCDLTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError,
   return success();
 }
 
-#define GET_OP_CLASSES
-#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
-
 #define GET_ATTRDEF_CLASSES
 #include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"
diff --git a/mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.cpp b/mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.cpp
new file mode 100644
index 0000000000000..c198b4a83a9ba
--- /dev/null
+++ b/mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.cpp
@@ -0,0 +1,16 @@
+//===- ROCDLOps.cpp - ROCDL operation definitions ------------------------===//
+//
+// 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/Dialect/LLVMIR/ROCDLDialect.h"
+
+#include "IR/ROCDLOps.h"
+
+using namespace mlir;
+using namespace ROCDL;
+
+#include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
diff --git a/mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.h b/mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.h
new file mode 100644
index 0000000000000..b77bf0b26d1f4
--- /dev/null
+++ b/mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.h
@@ -0,0 +1,22 @@
+//===- ROCDLOps.h - ROCDL operation implementation 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_LIB_DIALECT_LLVMIR_IR_ROCDLOPS_H
+#define MLIR_LIB_DIALECT_LLVMIR_IR_ROCDLOPS_H
+
+#include "mlir/IR/OpImplementation.h"
+
+namespace mlir::ROCDL {
+
+ParseResult parseCachePolicy(OpAsmParser &parser, Attribute &cachePolicy);
+void printCachePolicy(OpAsmPrinter &printer, Operation *,
+                      Attribute cachePolicy);
+
+} // namespace mlir::ROCDL
+
+#endif // MLIR_LIB_DIALECT_LLVMIR_IR_ROCDLOPS_H

More information about the Mlir-commits mailing list