[Mlir-commits] [mlir] [MLIR][NVVM] Shard generated operation definitions (PR #221812)
Mehdi Amini
llvmlistbot at llvm.org
Mon Sep 7 12:48:38 PDT 2026
https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/221812
Generate the NVVM operation definitions in eight shards, use the generated registration hook, and expose only the four parser and range helpers referenced by generated code.
Assisted-by: Codex
>From ffac23eb3ef2c9e74f5c80bcd355719b13184f70 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Mon, 7 Sep 2026 06:09:15 -0700
Subject: [PATCH] [mlir][nvvm] Shard generated operation definitions
Generate the NVVM operation definitions in eight shards, use the generated
registration hook, and expose only the four parser and range helpers referenced
by generated code.
The baseline monolithic Ninja edge took 112.157s, with 111.676s of
ExecuteCompiler time. Its inclusive trace totals were 40.405s frontend,
70.799s backend, 42.292s optimizer, 28.417s code generation, 15.331s function
instantiation, 10.353s class instantiation, 22.391s source processing, and
10.271s class parsing.
The eight shards plus the remaining dialect TU totaled 110.072s of
ExecuteCompiler time. Their inclusive totals were 72.098s frontend, 36.952s
backend, 22.565s optimizer, 14.322s code generation, 17.093s function
instantiation, 22.340s class instantiation, 55.510s source processing, and
42.989s class parsing. The slowest new compile edge was 14.612s.
Aggregate compiler time fell 1.4%, backend time fell 47.8%, optimizer time
fell 46.6%, and code-generation time fell 49.6%. The longest compile edge fell
87.0%. Repeated work increased frontend time 78.4%, function-instantiation
time 11.5%, class-instantiation time 115.8%, source-processing time 147.9%,
and class-parsing time 318.6%.
The affected-target build ran 10 steps in 15.38s wall time, using 105.66s user
time and 5.49s system time at 722% CPU, with 1,184,916 KiB peak RSS and zero
swaps. The baseline monolithic object contained 6,342,618 text bytes; the
remaining dialect object plus eight shard objects contain 6,740,626 text
bytes, an increase of 6.3%. The static archive grew from 23,634,936 to
24,977,400 bytes, an increase of 5.7%.
Rebuilding mlir-opt and running six focused suites took 26.40s wall time,
367.48s user time, and 35.80s system time at 1,527% CPU, with 1,966,540 KiB
peak RSS and zero swaps. The suites ran 11 tests in 0.36s, 1 in 0.08s, 6 in
0.83s, 2 in 0.22s, 1 in 0.09s, and 124 in 0.22s; all 145 tests passed.
The full check-mlir invocation took 6.27s wall time, 1.52s user time, and
0.36s system time at 30% CPU, with 175,268 KiB peak RSS and zero swaps. Its
lit suite ran in 5.47s with 3,937 passed, 619 unsupported, and 1 expected
failure.
Assisted-by: Codex
---
.../mlir/Dialect/LLVMIR/CMakeLists.txt | 9 ++++-
.../include/mlir/Dialect/LLVMIR/NVVMDialect.h | 4 +--
mlir/lib/Dialect/LLVMIR/CMakeLists.txt | 15 +++++++++
mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp | 31 ++++++++---------
mlir/lib/Dialect/LLVMIR/IR/NVVMOps.cpp | 16 +++++++++
mlir/lib/Dialect/LLVMIR/IR/NVVMOps.h | 33 +++++++++++++++++++
6 files changed, 87 insertions(+), 21 deletions(-)
create mode 100644 mlir/lib/Dialect/LLVMIR/IR/NVVMOps.cpp
create mode 100644 mlir/lib/Dialect/LLVMIR/IR/NVVMOps.h
diff --git a/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt b/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
index de9a3e1c93282..3a39ebdf6137b 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
+++ b/mlir/include/mlir/Dialect/LLVMIR/CMakeLists.txt
@@ -62,7 +62,14 @@ mlir_tablegen(NVVMRequiresSMTraits.h.inc -gen-op-interface-decls)
mlir_tablegen(NVVMRequiresSMTraits.cpp.inc -gen-op-interface-defs)
add_mlir_dialect_tablegen_target(MLIRNVVMRequiresSMTraitsIncGen)
-add_mlir_dialect(NVVMOps nvvm)
+set(LLVM_TARGET_DEFINITIONS NVVMOps.td)
+mlir_tablegen(NVVMOps.h.inc -gen-op-decls -op-shard-count=8)
+mlir_tablegen(NVVMOps.cpp.inc -gen-op-defs -op-shard-count=8)
+mlir_tablegen(NVVMOpsTypes.h.inc -gen-typedef-decls -typedefs-dialect=nvvm)
+mlir_tablegen(NVVMOpsTypes.cpp.inc -gen-typedef-defs -typedefs-dialect=nvvm)
+mlir_tablegen(NVVMOpsDialect.h.inc -gen-dialect-decls -dialect=nvvm)
+mlir_tablegen(NVVMOpsDialect.cpp.inc -gen-dialect-defs -dialect=nvvm)
+add_mlir_dialect_tablegen_target(MLIRNVVMOpsIncGen)
add_mlir_doc(NVVMOps NVVMOps Dialects/ -gen-op-doc -dialect=nvvm)
set(LLVM_TARGET_DEFINITIONS NVVMOps.td)
mlir_tablegen(NVVMConversions.inc -gen-llvmir-conversions)
diff --git a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
index 6bd582d66ed25..551949b415db0 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
+++ b/mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h
@@ -67,9 +67,9 @@ std::pair<mlir::Type, unsigned> inferMMAType(mlir::NVVM::MMATypes type,
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.h.inc"
+#include "mlir/Dialect/LLVMIR/NVVMOpsDialect.h.inc"
+
#define GET_OP_CLASSES
#include "mlir/Dialect/LLVMIR/NVVMOps.h.inc"
-#include "mlir/Dialect/LLVMIR/NVVMOpsDialect.h.inc"
-
#endif /* MLIR_DIALECT_LLVMIR_NVVMDIALECT_H_ */
diff --git a/mlir/lib/Dialect/LLVMIR/CMakeLists.txt b/mlir/lib/Dialect/LLVMIR/CMakeLists.txt
index 494a1e4c3ee97..a93dc0635a9e5 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/NVVMOps.cpp
+ )
+
add_mlir_dialect_library(MLIRLLVMDialect
IR/FunctionCallUtils.cpp
IR/LLVMAttrs.cpp
@@ -41,10 +45,20 @@ add_mlir_dialect_library(MLIRLLVMDialect
MLIRSupport
)
+set(LLVM_TARGET_DEFINITIONS IR/NVVMOps.cpp)
+foreach(index RANGE 7)
+ set(NVVM_SHARDED_SRC NVVMOps.${index}.cpp)
+ list(APPEND NVVM_SHARDED_SRCS ${NVVM_SHARDED_SRC})
+ tablegen(MLIR_SRC_SHARDER ${NVVM_SHARDED_SRC}
+ -op-shard-index=${index})
+endforeach()
+add_public_tablegen_target(MLIRNVVMOpsShardGen)
+
add_mlir_dialect_library(MLIRNVVMDialect
IR/NVVMDialect.cpp
IR/BasicPtxBuilderInterface.cpp
IR/NVVMRequiresSMTraits.cpp
+ ${NVVM_SHARDED_SRCS}
ADDITIONAL_HEADER_DIRS
${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/LLVMIR
@@ -52,6 +66,7 @@ add_mlir_dialect_library(MLIRNVVMDialect
DEPENDS
MLIRGPUCompilationAttrInterfacesIncGen
MLIRNVVMOpsIncGen
+ MLIRNVVMOpsShardGen
MLIRNVVMConversionsIncGen
MLIRBasicPtxBuilderInterfaceIncGen
MLIRNVVMRequiresSMTraitsIncGen
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
index c6ed0a840a1b0..b201c09712810 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMDialect.cpp
@@ -16,6 +16,8 @@
#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+#include "IR/NVVMOps.h"
+
#include "mlir/Conversion/ConvertToLLVM/ToLLVMInterface.h"
#include "mlir/Dialect/GPU/IR/CompilationInterfaces.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
@@ -93,8 +95,8 @@ getNVVMCtaGroupKind(NVVM::CTAGroupKind ctaGroup) {
llvm_unreachable("unsupported cta_group value");
}
-static ParseResult parseCTAGroup(OpAsmParser &parser,
- NVVM::CTAGroupKindAttr &groupAttr) {
+ParseResult mlir::NVVM::parseCTAGroup(OpAsmParser &parser,
+ NVVM::CTAGroupKindAttr &groupAttr) {
StringRef keyword;
if (parser.parseKeyword(&keyword))
return failure();
@@ -106,8 +108,8 @@ static ParseResult parseCTAGroup(OpAsmParser &parser,
return success();
}
-static void printCTAGroup(OpAsmPrinter &printer, Operation *,
- NVVM::CTAGroupKindAttr groupAttr) {
+void mlir::NVVM::printCTAGroup(OpAsmPrinter &printer, Operation *,
+ NVVM::CTAGroupKindAttr groupAttr) {
printer << NVVM::stringifyCTAGroupKind(groupAttr.getValue());
}
@@ -5818,10 +5820,10 @@ LogicalResult Tcgen05StOp::verify() {
/// Infer the result ranges for the NVVM SpecialRangeableRegisterOp that might
/// have ConstantRangeAttr.
-static void nvvmInferResultRanges(std::optional<LLVM::ConstantRangeAttr> range,
- Value result,
- ArrayRef<::mlir::ConstantIntRanges> argRanges,
- SetIntRangeFn setResultRanges) {
+void mlir::NVVM::nvvmInferResultRanges(
+ std::optional<LLVM::ConstantRangeAttr> range, Value result,
+ ArrayRef<::mlir::ConstantIntRanges> argRanges,
+ SetIntRangeFn setResultRanges) {
if (range) {
setResultRanges(result, {range->getLower(), range->getUpper(),
range->getLower(), range->getUpper()});
@@ -5832,9 +5834,8 @@ static void nvvmInferResultRanges(std::optional<LLVM::ConstantRangeAttr> range,
/// Verify the range attribute satisfies LLVM ConstantRange constructor
/// requirements for NVVM SpecialRangeableRegisterOp.
-static LogicalResult
-verifyConstantRangeAttr(Operation *op,
- std::optional<LLVM::ConstantRangeAttr> rangeAttr) {
+LogicalResult mlir::NVVM::verifyConstantRangeAttr(
+ Operation *op, std::optional<LLVM::ConstantRangeAttr> rangeAttr) {
if (!rangeAttr)
return success();
@@ -6957,10 +6958,7 @@ struct NVVMInlinerInterface final : DialectInlinerInterface {
// TODO: This should be the llvm.nvvm dialect once this is supported.
void NVVMDialect::initialize() {
- addOperations<
-#define GET_OP_LIST
-#include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
- >();
+ registerNVVMDialectOperations(this);
addAttributes<
#define GET_ATTRDEF_LIST
#include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
@@ -7173,8 +7171,5 @@ LogicalResult NVVMTargetAttr::verifyTarget(Operation *gpuModule) {
return success();
}
-#define GET_OP_CLASSES
-#include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
-
#define GET_ATTRDEF_CLASSES
#include "mlir/Dialect/LLVMIR/NVVMOpsAttributes.cpp.inc"
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMOps.cpp b/mlir/lib/Dialect/LLVMIR/IR/NVVMOps.cpp
new file mode 100644
index 0000000000000..8bfa736721ca0
--- /dev/null
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMOps.cpp
@@ -0,0 +1,16 @@
+//===- NVVMOps.cpp - NVVM 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/NVVMDialect.h"
+
+#include "IR/NVVMOps.h"
+
+using namespace mlir;
+using namespace NVVM;
+
+#include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
diff --git a/mlir/lib/Dialect/LLVMIR/IR/NVVMOps.h b/mlir/lib/Dialect/LLVMIR/IR/NVVMOps.h
new file mode 100644
index 0000000000000..19ba277b8fcc0
--- /dev/null
+++ b/mlir/lib/Dialect/LLVMIR/IR/NVVMOps.h
@@ -0,0 +1,33 @@
+//===- NVVMOps.h - NVVM 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_NVVMOPS_H
+#define MLIR_LIB_DIALECT_LLVMIR_IR_NVVMOPS_H
+
+#include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+
+#include <optional>
+
+namespace mlir::NVVM {
+
+mlir::ParseResult parseCTAGroup(mlir::OpAsmParser &parser,
+ mlir::NVVM::CTAGroupKindAttr &groupAttr);
+void printCTAGroup(mlir::OpAsmPrinter &printer, mlir::Operation *,
+ mlir::NVVM::CTAGroupKindAttr groupAttr);
+
+void nvvmInferResultRanges(std::optional<mlir::LLVM::ConstantRangeAttr> range,
+ mlir::Value result,
+ mlir::ArrayRef<mlir::ConstantIntRanges> argRanges,
+ mlir::SetIntRangeFn setResultRanges);
+mlir::LogicalResult
+verifyConstantRangeAttr(mlir::Operation *op,
+ std::optional<mlir::LLVM::ConstantRangeAttr> rangeAttr);
+
+} // namespace mlir::NVVM
+
+#endif // MLIR_LIB_DIALECT_LLVMIR_IR_NVVMOPS_H
More information about the Mlir-commits
mailing list