[IR] Add "Large Data Threshold" module metadata

[aeubanks]

This allows us to not have to pass -mllvm flags to set the large data threshold for (in-LLD/not-distributed) ThinLTO.

Follows https://reviews.llvm.org/D52322, which did the same for the code model.

Since the large data threshold is tied to the code model and we disallow mixing different code models, do the same for the large data threshold.

[llvmbot]

@llvm/pr-subscribers-lto

@llvm/pr-subscribers-llvm-ir

Changes

This allows us to not have to pass -mllvm flags to set the large data threshold for (in-LLD/not-distributed) ThinLTO.

Follows https://reviews.llvm.org/D52322, which did the same for the code model.

Since the large data threshold is tied to the code model and we disallow mixing different code models, do the same for the large data threshold.

---

Full diff: https://github.com/llvm/llvm-project/pull/66797.diff

7 Files Affected:

• (modified) llvm/include/llvm/IR/Module.h (+11)
• (modified) llvm/lib/IR/Module.cpp (+17)
• (modified) llvm/lib/LTO/LTOBackend.cpp (+5)
• (added) llvm/test/LTO/X86/Inputs/largedatathreshold.ll (+10)
• (added) llvm/test/LTO/X86/largedatathreshold-1.ll (+21)
• (added) llvm/test/LTO/X86/largedatathreshold-2.ll (+20)
• (added) llvm/test/LTO/X86/largedatathreshold-3.ll (+20)

diff --git a/llvm/include/llvm/IR/Module.h b/llvm/include/llvm/IR/Module.h
index 61b19409a96d0af..70beddddc1c1615 100644
--- a/llvm/include/llvm/IR/Module.h
+++ b/llvm/include/llvm/IR/Module.h
@@ -922,6 +922,17 @@ class LLVM_EXTERNAL_VISIBILITY Module {
   void setCodeModel(CodeModel::Model CL);
   /// @}
 
+  /// @}
+  /// @name Utility function for querying and setting the large data threshold
+  /// @{
+
+  /// Returns the code model (tiny, small, kernel, medium or large model)
+  std::optional<uint64_t> getLargeDataThreshold() const;
+
+  /// Set the code model (tiny, small, kernel, medium or large)
+  void setLargeDataThreshold(uint64_t Threshold);
+  /// @}
+
   /// @name Utility functions for querying and setting PGO summary
   /// @{
 
diff --git a/llvm/lib/IR/Module.cpp b/llvm/lib/IR/Module.cpp
index 5861bbd1f293e9b..dba660bbe5bafd3 100644
--- a/llvm/lib/IR/Module.cpp
+++ b/llvm/lib/IR/Module.cpp
@@ -631,6 +631,23 @@ void Module::setCodeModel(CodeModel::Model CL) {
   addModuleFlag(ModFlagBehavior::Error, "Code Model", CL);
 }
 
+std::optional<uint64_t> Module::getLargeDataThreshold() const {
+  auto *Val =
+      cast_or_null<ConstantAsMetadata>(getModuleFlag("Large Data Threshold"));
+
+  if (!Val)
+    return std::nullopt;
+
+  return cast<ConstantInt>(Val->getValue())->getZExtValue();
+}
+
+void Module::setLargeDataThreshold(uint64_t Threshold) {
+  // Since the large data threshold goes along with the code model, the merge
+  // behavior is the same.
+  addModuleFlag(ModFlagBehavior::Error, "Large Data Threshold",
+                ConstantInt::get(Type::getInt64Ty(Context), Threshold));
+}
+
 void Module::setProfileSummary(Metadata *M, ProfileSummary::Kind Kind) {
   if (Kind == ProfileSummary::PSK_CSInstr)
     setModuleFlag(ModFlagBehavior::Error, "CSProfileSummary", M);
diff --git a/llvm/lib/LTO/LTOBackend.cpp b/llvm/lib/LTO/LTOBackend.cpp
index 29e28876760816d..ccc4276e36dacf0 100644
--- a/llvm/lib/LTO/LTOBackend.cpp
+++ b/llvm/lib/LTO/LTOBackend.cpp
@@ -225,7 +225,12 @@ createTargetMachine(const Config &Conf, const Target *TheTarget, Module &M) {
   std::unique_ptr<TargetMachine> TM(TheTarget->createTargetMachine(
       TheTriple, Conf.CPU, Features.getString(), Conf.Options, RelocModel,
       CodeModel, Conf.CGOptLevel));
+
   assert(TM && "Failed to create target machine");
+
+  if (std::optional<uint64_t> LargeDataThreshold = M.getLargeDataThreshold())
+    TM->setLargeDataThreshold(*LargeDataThreshold);
+
   return TM;
 }
 
diff --git a/llvm/test/LTO/X86/Inputs/largedatathreshold.ll b/llvm/test/LTO/X86/Inputs/largedatathreshold.ll
new file mode 100644
index 000000000000000..34174deb78c0e4f
--- /dev/null
+++ b/llvm/test/LTO/X86/Inputs/largedatathreshold.ll
@@ -0,0 +1,10 @@
+target triple = "x86_64-unknown-linux-gnu"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+define void @bar() {
+  ret void
+}
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 1, !"Code Model", i32 3}
+!1 = !{i32 1, !"Large Data Threshold", i32 101}
diff --git a/llvm/test/LTO/X86/largedatathreshold-1.ll b/llvm/test/LTO/X86/largedatathreshold-1.ll
new file mode 100644
index 000000000000000..e3be5c11baaac25
--- /dev/null
+++ b/llvm/test/LTO/X86/largedatathreshold-1.ll
@@ -0,0 +1,21 @@
+; RUN: llvm-as %s -o %t.o
+; RUN: llvm-lto2 run -r %t.o,_start,px %t.o -o %t.s
+; RUN: llvm-objdump -d %t.s.0 | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+@data = internal constant [20 x i8] zeroinitializer
+
+define ptr @_start() {
+entry:
+; CHECK-LABEL:  <_start>:
+; CHECK: leaq    (%rip), %rax
+; CHECK-NOT: movabsq
+    ret ptr @data
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 1, !"Code Model", i32 3}
+!1 = !{i32 1, !"Large Data Threshold", i32 100}
diff --git a/llvm/test/LTO/X86/largedatathreshold-2.ll b/llvm/test/LTO/X86/largedatathreshold-2.ll
new file mode 100644
index 000000000000000..103c066b744d0fa
--- /dev/null
+++ b/llvm/test/LTO/X86/largedatathreshold-2.ll
@@ -0,0 +1,20 @@
+; RUN: llvm-as %s -o %t.o
+; RUN: llvm-lto2 run -r %t.o,_start,px %t.o -o %t.s
+; RUN: llvm-objdump -d %t.s.0 | FileCheck %s
+
+target triple = "x86_64-unknown-linux-gnu"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+@data = internal constant [20 x i8] zeroinitializer
+
+define ptr @_start() {
+entry:
+; CHECK-LABEL:  <_start>:
+; CHECK: movabsq $0x0, %rax
+    ret ptr @data
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 1, !"Code Model", i32 3}
+!1 = !{i32 1, !"Large Data Threshold", i32 10}
diff --git a/llvm/test/LTO/X86/largedatathreshold-3.ll b/llvm/test/LTO/X86/largedatathreshold-3.ll
new file mode 100644
index 000000000000000..3c0653db334d85f
--- /dev/null
+++ b/llvm/test/LTO/X86/largedatathreshold-3.ll
@@ -0,0 +1,20 @@
+; RUN: llvm-as %s -o %t0.o
+; RUN: llvm-as < %p/Inputs/largedatathreshold.ll > %t1.o
+; RUN: not llvm-lto2 run -r %t0.o,_start,px -r %t1.o,bar,px %t0.o %t1.o -o %t2.s 2>&1 | FileCheck %s
+
+; CHECK: 'Large Data Threshold': IDs have conflicting values
+
+target triple = "x86_64-unknown-linux-gnu"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+@data = internal constant [20 x i8] zeroinitializer
+
+define ptr @_start() {
+entry:
+  ret ptr @data
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 1, !"Code Model", i32 3}
+!1 = !{i32 1, !"Large Data Threshold", i32 100}

[rnk]

lgtm

[rnk]

I considered recommending that we fence this off as some x86-specific option, but nothing prevents any other target from adding similar code models, so I think its fine as is.