[llvm] [llvm-split] Generalize split-by-category to accept an arbitrary attribute name (PR #197338)
Yury Plyakhin via llvm-commits
llvm-commits at lists.llvm.org
Tue May 12 18:00:40 PDT 2026
https://github.com/YuriPlyakhin created https://github.com/llvm/llvm-project/pull/197338
Rename the `module-id` split-by-category mode to `attribute` and add a `--category-attribute=<name>` flag that specifies which function attribute to group by. This decouples `llvm-split` from the hardcoded "module-id" attribute name, making the tool reusable for any attribute-based splitting.
Also fix the error reporting path to properly consume the llvm::Error (use `toString(std::move(E))`) instead of streaming it and letting the unchecked Error destructor abort.
Co-Authored-By: Claude
>From d54cf149282f50606b6cf0ceaf00d44b2c2907f3 Mon Sep 17 00:00:00 2001
From: "Plyakhin, Yury" <yury.plyakhin at intel.com>
Date: Wed, 13 May 2026 02:17:55 +0200
Subject: [PATCH 1/2] [llvm-split] Generalize split-by-category to accept an
arbitrary attribute name
Rename the `module-id` split-by-category mode to `attribute` and add a
`--category-attribute=<name>` flag that specifies which function attribute
to group by. This decouples `llvm-split` from the hardcoded "module-id"
attribute name, making the tool reusable for any attribute-based splitting.
Also fix the error reporting path to properly consume the llvm::Error
(use `toString(std::move(E))`) instead of streaming it and letting the
unchecked Error destructor abort.
Co-Authored-By: Claude
---
.../complex-indirect-call-chain1.ll | 2 +-
.../complex-indirect-call-chain2.ll | 2 +-
.../missing-attribute-value.ll | 6 +++
.../SplitByCategory/module-split-func-ptr.ll | 2 +-
.../SplitByCategory/split-by-source.ll | 2 +-
.../split-with-kernel-declarations.ll | 2 +-
llvm/tools/llvm-split/llvm-split.cpp | 50 ++++++++++++-------
7 files changed, 42 insertions(+), 24 deletions(-)
create mode 100644 llvm/test/tools/llvm-split/SplitByCategory/missing-attribute-value.ll
diff --git a/llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain1.ll b/llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain1.ll
index 80123d4dd8fb7..50e08cc093d83 100644
--- a/llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain1.ll
+++ b/llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain1.ll
@@ -1,7 +1,7 @@
; Check that Module splitting can trace through more complex call stacks
; involving several nested indirect calls.
-; RUN: llvm-split -split-by-category=module-id -S < %s -o %t
+; RUN: llvm-split -split-by-category=attribute --category-attribute=module-id -S < %s -o %t
; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix CHECK0 \
; RUN: --implicit-check-not @foo --implicit-check-not @kernel_A \
; RUN: --implicit-check-not @kernel_B --implicit-check-not @baz
diff --git a/llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain2.ll b/llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain2.ll
index 0c80602f99eef..aa84c5fbf904a 100644
--- a/llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain2.ll
+++ b/llvm/test/tools/llvm-split/SplitByCategory/complex-indirect-call-chain2.ll
@@ -1,6 +1,6 @@
; Check that Module splitting can trace indirect calls through signatures.
-; RUN: llvm-split -split-by-category=module-id -S < %s -o %t
+; RUN: llvm-split -split-by-category=attribute --category-attribute=module-id -S < %s -o %t
; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix CHECK0 \
; RUN: --implicit-check-not @kernel_A --implicit-check-not @bbb
; RUN: FileCheck %s -input-file=%t_1.ll --check-prefix CHECK1 \
diff --git a/llvm/test/tools/llvm-split/SplitByCategory/missing-attribute-value.ll b/llvm/test/tools/llvm-split/SplitByCategory/missing-attribute-value.ll
new file mode 100644
index 0000000000000..998810b16f264
--- /dev/null
+++ b/llvm/test/tools/llvm-split/SplitByCategory/missing-attribute-value.ll
@@ -0,0 +1,6 @@
+; Test that -split-by-category=attribute without --category-attribute produces an error.
+
+; RUN: not llvm-split -split-by-category=attribute -S < %s -o %t 2>&1 \
+; RUN: | FileCheck %s
+
+; CHECK: -split-by-category=attribute requires --category-attribute=<name>
diff --git a/llvm/test/tools/llvm-split/SplitByCategory/module-split-func-ptr.ll b/llvm/test/tools/llvm-split/SplitByCategory/module-split-func-ptr.ll
index 316500a4c7611..ee263fc38a893 100644
--- a/llvm/test/tools/llvm-split/SplitByCategory/module-split-func-ptr.ll
+++ b/llvm/test/tools/llvm-split/SplitByCategory/module-split-func-ptr.ll
@@ -1,7 +1,7 @@
; This test checks that Module splitting can properly perform device code split by tracking
; all uses of functions (not only direct calls).
-; RUN: llvm-split -split-by-category=module-id -S < %s -o %t
+; RUN: llvm-split -split-by-category=attribute --category-attribute=module-id -S < %s -o %t
; RUN: FileCheck %s -input-file=%t_0.ll --check-prefix=CHECK-IR0
; RUN: FileCheck %s -input-file=%t_1.ll --check-prefix=CHECK-IR1
diff --git a/llvm/test/tools/llvm-split/SplitByCategory/split-by-source.ll b/llvm/test/tools/llvm-split/SplitByCategory/split-by-source.ll
index 54485b7b7f348..dc0cc292f50fe 100644
--- a/llvm/test/tools/llvm-split/SplitByCategory/split-by-source.ll
+++ b/llvm/test/tools/llvm-split/SplitByCategory/split-by-source.ll
@@ -1,7 +1,7 @@
; Test checks that kernels are being split by attached module-id metadata and
; used functions are being moved with kernels that use them.
-; RUN: llvm-split -split-by-category=module-id -S < %s -o %t
+; RUN: llvm-split -split-by-category=attribute --category-attribute=module-id -S < %s -o %t
; RUN: FileCheck %s -input-file=%t_0.ll --check-prefixes CHECK-TU0,CHECK
; RUN: FileCheck %s -input-file=%t_1.ll --check-prefixes CHECK-TU1,CHECK
diff --git a/llvm/test/tools/llvm-split/SplitByCategory/split-with-kernel-declarations.ll b/llvm/test/tools/llvm-split/SplitByCategory/split-with-kernel-declarations.ll
index 0c1bd8b5c5fba..59a7a95761d9d 100644
--- a/llvm/test/tools/llvm-split/SplitByCategory/split-with-kernel-declarations.ll
+++ b/llvm/test/tools/llvm-split/SplitByCategory/split-with-kernel-declarations.ll
@@ -1,6 +1,6 @@
; The test checks that Module splitting does not treat declarations as entry points.
-; RUN: llvm-split -split-by-category=module-id -S < %s -o %t1
+; RUN: llvm-split -split-by-category=attribute --category-attribute=module-id -S < %s -o %t1
; RUN: FileCheck %s -input-file=%t1_0.ll --check-prefix CHECK-MODULE-ID0
; RUN: FileCheck %s -input-file=%t1_1.ll --check-prefix CHECK-MODULE-ID1
diff --git a/llvm/tools/llvm-split/llvm-split.cpp b/llvm/tools/llvm-split/llvm-split.cpp
index a987b8c1b3eb4..e0ff8906796f7 100644
--- a/llvm/tools/llvm-split/llvm-split.cpp
+++ b/llvm/tools/llvm-split/llvm-split.cpp
@@ -78,7 +78,7 @@ static cl::opt<std::string>
cl::value_desc("cpu"), cl::cat(SplitCategory));
enum class SplitByCategoryType {
- SBCT_ByModuleId,
+ SBCT_ByAttribute,
SBCT_ByKernel,
SBCT_None,
};
@@ -88,13 +88,19 @@ static cl::opt<SplitByCategoryType> SplitByCategory(
cl::desc("Split by category. If present, splitting by category is used "
"with the specified categorization type."),
cl::Optional, cl::init(SplitByCategoryType::SBCT_None),
- cl::values(clEnumValN(SplitByCategoryType::SBCT_ByModuleId, "module-id",
- "one output module per translation unit marked with "
- "\"module-id\" attribute"),
+ cl::values(clEnumValN(SplitByCategoryType::SBCT_ByAttribute, "attribute",
+ "one output module per unique value of the function "
+ "attribute named by --category-attribute"),
clEnumValN(SplitByCategoryType::SBCT_ByKernel, "kernel",
"one output module per kernel")),
cl::cat(SplitCategory));
+static cl::opt<std::string>
+ CategoryAttribute("category-attribute",
+ cl::desc("Function attribute name to use when splitting "
+ "with -split-by-category=attribute"),
+ cl::value_desc("name"), cl::cat(SplitCategory));
+
static cl::opt<bool> OutputAssembly{
"S", cl::desc("Write output as LLVM assembly"), cl::cat(SplitCategory)};
@@ -125,15 +131,16 @@ void writeModuleToFile(const Module &M, StringRef Path, bool OutputAssembly) {
WriteBitcodeToFile(M, OS);
}
-/// EntryPointCategorizer is used for splitting by category either by module-id
-/// or by kernels. It doesn't provide categories for functions other than
-/// kernels. Categorizer computes a string key for the given Function and
-/// records the association between the string key and an integer category. If a
-/// string key is already belongs to some category than the corresponding
-/// integer category is returned.
+/// EntryPointCategorizer is used for splitting by category either by a named
+/// function attribute or by kernels. It doesn't provide categories for
+/// functions other than kernels. Categorizer computes a string key for the
+/// given Function and records the association between the string key and an
+/// integer category. If a string key already belongs to some category then the
+/// corresponding integer category is returned.
class EntryPointCategorizer {
public:
- EntryPointCategorizer(SplitByCategoryType Type) : Type(Type) {}
+ EntryPointCategorizer(SplitByCategoryType Type, StringRef AttributeName)
+ : Type(Type), AttributeName(AttributeName) {}
EntryPointCategorizer() = delete;
EntryPointCategorizer(EntryPointCategorizer &) = delete;
@@ -163,16 +170,15 @@ class EntryPointCategorizer {
return F.hasKernelCallingConv();
}
- static SmallString<0> computeFunctionCategory(SplitByCategoryType Type,
- const Function &F) {
- static constexpr char ATTR_MODULE_ID[] = "module-id";
+ SmallString<0> computeFunctionCategory(SplitByCategoryType Type,
+ const Function &F) {
SmallString<0> Key;
switch (Type) {
case SplitByCategoryType::SBCT_ByKernel:
Key = F.getName().str();
break;
- case SplitByCategoryType::SBCT_ByModuleId:
- Key = F.getFnAttribute(ATTR_MODULE_ID).getValueAsString().str();
+ case SplitByCategoryType::SBCT_ByAttribute:
+ Key = F.getFnAttribute(AttributeName).getValueAsString().str();
break;
default:
llvm_unreachable("unexpected mode.");
@@ -197,6 +203,7 @@ class EntryPointCategorizer {
};
SplitByCategoryType Type;
+ std::string AttributeName;
DenseMap<SmallString<0>, int, KeyInfo> StrKeyToID;
};
@@ -209,6 +216,12 @@ void cleanupModule(Module &M) {
}
Error runSplitModuleByCategory(std::unique_ptr<Module> M) {
+ if (SplitByCategory == SplitByCategoryType::SBCT_ByAttribute &&
+ CategoryAttribute.empty())
+ return createStringError(
+ inconvertibleErrorCode(),
+ "-split-by-category=attribute requires --category-attribute=<name>");
+
size_t OutputID = 0;
auto PostSplitCallback = [&](std::unique_ptr<Module> MPart) -> Error {
if (verifyModule(*MPart)) {
@@ -228,7 +241,7 @@ Error runSplitModuleByCategory(std::unique_ptr<Module> M) {
return Error::success();
};
- auto Categorizer = EntryPointCategorizer(SplitByCategory);
+ auto Categorizer = EntryPointCategorizer(SplitByCategory, CategoryAttribute);
return splitModuleTransitiveFromEntryPoints(std::move(M), Categorizer,
PostSplitCallback);
}
@@ -291,8 +304,7 @@ int main(int argc, char **argv) {
if (SplitByCategory != SplitByCategoryType::SBCT_None) {
auto E = runSplitModuleByCategory(std::move(M));
if (E) {
- errs() << E << "\n";
- Err.print(argv[0], errs());
+ errs() << toString(std::move(E)) << "\n";
return 1;
}
>From e9dfddef6bc5a6db68fbd70d0b0772c6764eabfa Mon Sep 17 00:00:00 2001
From: "Plyakhin, Yury" <yury.plyakhin at intel.com>
Date: Wed, 13 May 2026 02:58:06 +0200
Subject: [PATCH 2/2] clean up
---
.../llvm-split/SplitByCategory/missing-attribute-value.ll | 2 +-
llvm/tools/llvm-split/llvm-split.cpp | 3 +--
2 files changed, 2 insertions(+), 3 deletions(-)
diff --git a/llvm/test/tools/llvm-split/SplitByCategory/missing-attribute-value.ll b/llvm/test/tools/llvm-split/SplitByCategory/missing-attribute-value.ll
index 998810b16f264..e05320584a372 100644
--- a/llvm/test/tools/llvm-split/SplitByCategory/missing-attribute-value.ll
+++ b/llvm/test/tools/llvm-split/SplitByCategory/missing-attribute-value.ll
@@ -3,4 +3,4 @@
; RUN: not llvm-split -split-by-category=attribute -S < %s -o %t 2>&1 \
; RUN: | FileCheck %s
-; CHECK: -split-by-category=attribute requires --category-attribute=<name>
+; CHECK: error: -split-by-category=attribute requires --category-attribute=<name>
diff --git a/llvm/tools/llvm-split/llvm-split.cpp b/llvm/tools/llvm-split/llvm-split.cpp
index e0ff8906796f7..4cc4fd945fc53 100644
--- a/llvm/tools/llvm-split/llvm-split.cpp
+++ b/llvm/tools/llvm-split/llvm-split.cpp
@@ -219,7 +219,6 @@ Error runSplitModuleByCategory(std::unique_ptr<Module> M) {
if (SplitByCategory == SplitByCategoryType::SBCT_ByAttribute &&
CategoryAttribute.empty())
return createStringError(
- inconvertibleErrorCode(),
"-split-by-category=attribute requires --category-attribute=<name>");
size_t OutputID = 0;
@@ -304,7 +303,7 @@ int main(int argc, char **argv) {
if (SplitByCategory != SplitByCategoryType::SBCT_None) {
auto E = runSplitModuleByCategory(std::move(M));
if (E) {
- errs() << toString(std::move(E)) << "\n";
+ errs() << "error: " << toString(std::move(E)) << "\n";
return 1;
}
More information about the llvm-commits
mailing list