[Mlir-commits] [mlir] a760450 - [MLIR][ROCDL] Shard generated operation definitions (NFC) (#221729)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue Sep 8 06:51:04 PDT 2026
Author: Mehdi Amini
Date: 2026-09-08T13:50:59Z
New Revision: a76045040a1297217256d917bec9e02f9723804c
URL: https://github.com/llvm/llvm-project/commit/a76045040a1297217256d917bec9e02f9723804c
DIFF: https://github.com/llvm/llvm-project/commit/a76045040a1297217256d917bec9e02f9723804c.diff
LOG: [MLIR][ROCDL] Shard generated operation definitions (NFC) (#221729)
ROCDL generated operation definitions is amongst 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
Added:
mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.cpp
mlir/lib/Dialect/LLVMIR/IR/ROCDLOps.h
Modified:
mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h
mlir/lib/Dialect/LLVMIR/CMakeLists.txt
mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt b/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
index 452af9683ceb8..a81288495f625 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
@@ -79,7 +79,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 799e04360cd2b..18057048a509a 100644
--- a/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
+++ b/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
@@ -2,6 +2,7 @@ add_subdirectory(Transforms)
set(LLVM_OPTIONAL_SOURCES
IR/LLVMOps.cpp
+ IR/ROCDLOps.cpp
)
set(LLVM_TARGET_DEFINITIONS IR/LLVMOps.cpp)
@@ -87,8 +88,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
@@ -96,6 +107,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