[Mlir-commits] [mlir] [MLIR][ROCDL] Shard generated operation definitions (NFC) (PR #221729)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Sep 7 06:38:15 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
@llvm/pr-subscribers-mlir-llvm
Author: Mehdi Amini (joker-eph)
<details>
<summary>Changes</summary>
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
---
Full diff: https://github.com/llvm/llvm-project/pull/221729.diff
6 Files Affected:
- (modified) mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt (+8-1)
- (modified) mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h (+2-2)
- (modified) mlir/lib/Dialect/LLVMIR/CMakeLists.txt (+15)
- (modified) mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp (+5-9)
- (added) mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.cpp (+16)
- (added) mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.h (+22)
``````````diff
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/221729
More information about the Mlir-commits
mailing list