https://github.com/yxsamliu updated 
https://github.com/llvm/llvm-project/pull/211059

>From 1bdc18cf9273ab2be826e941d2bbbd655ac0403f Mon Sep 17 00:00:00 2001
From: "Yaxun (Sam) Liu" <[email protected]>
Date: Tue, 21 Jul 2026 09:22:40 -0400
Subject: [PATCH 1/2] [AMDGPU] Add per-kernel kernarg preload count

The backend option sets one kernarg preload count for all kernels. Some
kernels need a smaller limit, and some kernels need preload disabled.

Add `amdgpu_kernarg_preload_count(N)` for AMDGPU kernels. Clang lowers
it to the `"amdgpu-kernarg-preload-count"` IR attribute. The AMDGPU
preload pass uses this value when it is present. A value of `0` disables
preload for that kernel. Kernels without the attribute keep using the
global option.
---
 clang/docs/ReleaseNotes.md                    |  2 +
 clang/include/clang/Basic/Attr.td             |  7 ++++
 clang/include/clang/Basic/AttrDocs.td         | 40 +++++++++++++++++++
 clang/include/clang/Sema/SemaAMDGPU.h         |  1 +
 clang/lib/CodeGen/Targets/AMDGPU.cpp          |  4 ++
 clang/lib/Sema/SemaAMDGPU.cpp                 | 11 +++++
 clang/lib/Sema/SemaDeclAttr.cpp               |  7 ++++
 clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu |  6 +++
 ...mdgpu-kernarg-preload-count-save-temps.hip | 25 ++++++++++++
 ...a-attribute-supported-attributes-list.test |  1 +
 clang/test/SemaOpenCL/amdgpu-attrs.cl         | 10 +++++
 .../AMDGPU/AMDGPUPreloadKernelArguments.cpp   | 15 ++++++-
 .../AMDGPU/preload-kernargs-count-attr.ll     | 22 ++++++++++
 13 files changed, 150 insertions(+), 1 deletion(-)
 create mode 100644 
clang/test/CodeGenHIP/amdgpu-kernarg-preload-count-save-temps.hip
 create mode 100644 llvm/test/CodeGen/AMDGPU/preload-kernargs-count-attr.ll

diff --git a/clang/docs/ReleaseNotes.md b/clang/docs/ReleaseNotes.md
index 80e770a40838a..7a80cd2746337 100644
--- a/clang/docs/ReleaseNotes.md
+++ b/clang/docs/ReleaseNotes.md
@@ -139,6 +139,8 @@ features cannot lower the translation-unit ABI level;
 ### Attribute Changes in Clang
 
 - Clang now properly propagates attributes on class and variable templates to 
their redeclarations, which will result in redeclarations not interfering with 
diagnostics. (#GH209812)
+- Added the `amdgpu_kernarg_preload_count` attribute for AMDGPU kernels to
+  control kernel argument preloading per kernel.
 
 ### Improvements to Clang's diagnostics
 
diff --git a/clang/include/clang/Basic/Attr.td 
b/clang/include/clang/Basic/Attr.td
index 1a5bd2301dfc8..9cb175fb2afbc 100644
--- a/clang/include/clang/Basic/Attr.td
+++ b/clang/include/clang/Basic/Attr.td
@@ -2525,6 +2525,13 @@ def AMDGPUNumVGPR : InheritableAttr {
   let Subjects = SubjectList<[Function], ErrorDiag, "kernel functions">;
 }
 
+def AMDGPUKernargPreloadCount : InheritableAttr {
+  let Spellings = [Clang<"amdgpu_kernarg_preload_count", 0>];
+  let Args = [UnsignedArgument<"Count">];
+  let Documentation = [AMDGPUKernargPreloadCountDocs];
+  let Subjects = SubjectList<[Function], ErrorDiag, "kernel functions">;
+}
+
 def AMDGPUMaxNumWorkGroups : InheritableAttr {
   let Spellings = [Clang<"amdgpu_max_num_work_groups", 0>];
   let Args = [ExprArgument<"MaxNumWorkGroupsX">, 
ExprArgument<"MaxNumWorkGroupsY", 1>, ExprArgument<"MaxNumWorkGroupsZ", 1>];
diff --git a/clang/include/clang/Basic/AttrDocs.td 
b/clang/include/clang/Basic/AttrDocs.td
index 05e4cb0870652..f752d75e3ab54 100644
--- a/clang/include/clang/Basic/AttrDocs.td
+++ b/clang/include/clang/Basic/AttrDocs.td
@@ -3574,6 +3574,46 @@ An error will be given if:
     attributes;
   - The AMDGPU target backend is unable to create machine code that can meet 
the
     request.
+}];
+}
+
+def AMDGPUKernargPreloadCountDocs : Documentation {
+  let Category = DocCatAMDGPUAttributes;
+  let Content = [{
+Clang supports the
+``__attribute__((amdgpu_kernarg_preload_count(<count>)))`` attribute on AMDGPU
+kernel functions.
+
+Kernel argument preloading can reduce kernel launch latency. Without
+preloading, the kernel may need to load explicit kernel arguments from kernel
+argument segment memory. On targets that support kernel argument preloading,
+the hardware can copy part of the kernel argument segment into user SGPRs
+before the kernel starts, avoiding those initial kernel argument memory loads.
+
+The attribute gives a kernel its own kernel argument preload limit. This is
+useful when a translation unit contains kernels that need different preload
+behavior. For example, a small hot kernel may benefit from preloading a few
+explicit kernel arguments, while another kernel in the same translation unit
+may need preload disabled to save user SGPRs.
+
+The ``<count>`` value is the maximum number of explicit kernel arguments that
+the backend should try to preload. The backend considers explicit kernel
+arguments from the start of the kernel argument list and stops after
+``<count>`` arguments. It may preload fewer arguments if an argument is not
+supported, if the arguments do not fit in the available user SGPRs, or if
+another target limit is reached.
+
+Kernel argument preloading is only supported on AMDGPU targets with the
+``kernarg-preload`` target feature. In the current backend, the known
+processors with this feature are ``gfx90a``, ``gfx942``, ``gfx950``,
+``gfx1250``, and ``gfx1251``. The generic aliases ``gfx9-4-generic`` and
+``gfx12-5-generic`` also include the feature. Future processors may also
+support it if their feature set includes ``FeatureKernargPreload`` in
+``llvm/lib/Target/AMDGPU/AMDGPU.td``. On targets that do not support kernel
+argument preloading, this attribute has no effect.
+
+Passing ``0`` disables kernel argument preloading for the annotated kernel.
+Omitting the attribute keeps the default command-line behavior.
   }];
 }
 
diff --git a/clang/include/clang/Sema/SemaAMDGPU.h 
b/clang/include/clang/Sema/SemaAMDGPU.h
index a6205534e0de3..603b2f80c5f21 100644
--- a/clang/include/clang/Sema/SemaAMDGPU.h
+++ b/clang/include/clang/Sema/SemaAMDGPU.h
@@ -77,6 +77,7 @@ class SemaAMDGPU : public SemaBase {
   void handleAMDGPUWavesPerEUAttr(Decl *D, const ParsedAttr &AL);
   void handleAMDGPUNumSGPRAttr(Decl *D, const ParsedAttr &AL);
   void handleAMDGPUNumVGPRAttr(Decl *D, const ParsedAttr &AL);
+  void handleAMDGPUKernargPreloadCountAttr(Decl *D, const ParsedAttr &AL);
   void handleAMDGPUMaxNumWorkGroupsAttr(Decl *D, const ParsedAttr &AL);
   void handleAMDGPUFlatWorkGroupSizeAttr(Decl *D, const ParsedAttr &AL);
 
diff --git a/clang/lib/CodeGen/Targets/AMDGPU.cpp 
b/clang/lib/CodeGen/Targets/AMDGPU.cpp
index 7b37f3f7f9b6e..0e2f8c5aa6a72 100644
--- a/clang/lib/CodeGen/Targets/AMDGPU.cpp
+++ b/clang/lib/CodeGen/Targets/AMDGPU.cpp
@@ -377,6 +377,10 @@ void AMDGPUTargetCodeGenInfo::setFunctionDeclAttributes(
       F->addFnAttr("amdgpu-num-vgpr", llvm::utostr(NumVGPR));
   }
 
+  if (const auto *Attr = FD->getAttr<AMDGPUKernargPreloadCountAttr>())
+    F->addFnAttr("amdgpu-kernarg-preload-count",
+                 llvm::utostr(Attr->getCount()));
+
   if (const auto *Attr = FD->getAttr<AMDGPUMaxNumWorkGroupsAttr>()) {
     uint32_t X = Attr->getMaxNumWorkGroupsX()
                      ->EvaluateKnownConstInt(M.getContext())
diff --git a/clang/lib/Sema/SemaAMDGPU.cpp b/clang/lib/Sema/SemaAMDGPU.cpp
index 48230fa262d5c..6fd7e711ed542 100644
--- a/clang/lib/Sema/SemaAMDGPU.cpp
+++ b/clang/lib/Sema/SemaAMDGPU.cpp
@@ -729,6 +729,17 @@ void SemaAMDGPU::handleAMDGPUNumVGPRAttr(Decl *D, const 
ParsedAttr &AL) {
                  AMDGPUNumVGPRAttr(getASTContext(), AL, NumVGPR));
 }
 
+void SemaAMDGPU::handleAMDGPUKernargPreloadCountAttr(Decl *D,
+                                                     const ParsedAttr &AL) {
+  uint32_t Count = 0;
+  Expr *CountExpr = AL.getArgAsExpr(0);
+  if (!SemaRef.checkUInt32Argument(AL, CountExpr, Count))
+    return;
+
+  D->addAttr(::new (getASTContext())
+                 AMDGPUKernargPreloadCountAttr(getASTContext(), AL, Count));
+}
+
 static bool
 checkAMDGPUMaxNumWorkGroupsArguments(Sema &S, Expr *XExpr, Expr *YExpr,
                                      Expr *ZExpr,
diff --git a/clang/lib/Sema/SemaDeclAttr.cpp b/clang/lib/Sema/SemaDeclAttr.cpp
index 1b272b5416860..11e5483471e40 100644
--- a/clang/lib/Sema/SemaDeclAttr.cpp
+++ b/clang/lib/Sema/SemaDeclAttr.cpp
@@ -7717,6 +7717,9 @@ ProcessDeclAttribute(Sema &S, Scope *scope, Decl *D, 
const ParsedAttr &AL,
   case ParsedAttr::AT_AMDGPUNumVGPR:
     S.AMDGPU().handleAMDGPUNumVGPRAttr(D, AL);
     break;
+  case ParsedAttr::AT_AMDGPUKernargPreloadCount:
+    S.AMDGPU().handleAMDGPUKernargPreloadCountAttr(D, AL);
+    break;
   case ParsedAttr::AT_AMDGPUMaxNumWorkGroups:
     S.AMDGPU().handleAMDGPUMaxNumWorkGroupsAttr(D, AL);
     break;
@@ -8607,6 +8610,10 @@ void Sema::ProcessDeclAttributeList(
       Diag(D->getLocation(), diag::err_attribute_wrong_decl_type)
           << A << A->isRegularKeywordAttribute() << ExpectedKernelFunction;
       D->setInvalidDecl();
+    } else if (const auto *A = D->getAttr<AMDGPUKernargPreloadCountAttr>()) {
+      Diag(D->getLocation(), diag::err_attribute_wrong_decl_type)
+          << A << A->isRegularKeywordAttribute() << ExpectedKernelFunction;
+      D->setInvalidDecl();
     }
   }
   checkAMDGPUReqdWorkGroupSize(*this, D);
diff --git a/clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu 
b/clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu
index cb0989dc8344a..f02be68dcdc61 100644
--- a/clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu
+++ b/clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu
@@ -45,6 +45,10 @@ __attribute__((amdgpu_num_vgpr(64))) // 
expected-no-diagnostics
 __global__ void num_vgpr_64() {
 // CHECK: define{{.*}} amdgpu_kernel void @_Z11num_vgpr_64v() 
[[NUM_VGPR_64:#[0-9]+]]
 }
+__attribute__((amdgpu_kernarg_preload_count(2))) // expected-no-diagnostics
+__global__ void kernarg_preload_count_2() {
+// CHECK: define{{.*}} amdgpu_kernel void @_Z23kernarg_preload_count_2v() 
[[KERNARG_PRELOAD_COUNT_2:#[0-9]+]]
+}
 __attribute__((amdgpu_max_num_work_groups(32, 4, 2))) // 
expected-no-diagnostics
 __global__ void max_num_work_groups_32_4_2() {
 // CHECK: define{{.*}} amdgpu_kernel void @_Z26max_num_work_groups_32_4_2v() 
[[MAX_NUM_WORK_GROUPS_32_4_2:#[0-9]+]]
@@ -101,6 +105,7 @@ template __global__ void 
template_a_b_c_max_num_work_groups<32, 4, 2>();
 // NAMD-NOT: "amdgpu-waves-per-eu"
 // NAMD-NOT: "amdgpu-num-vgpr"
 // NAMD-NOT: "amdgpu-num-sgpr"
+// NAMD-NOT: "amdgpu-kernarg-preload-count"
 // NAMD-NOT: "amdgpu-max-num-work-groups"
 
 // DEFAULT-DAG: attributes [[FLAT_WORK_GROUP_SIZE_DEFAULT]] = 
{{.*}}"amdgpu-flat-work-group-size"="1,1024"{{.*}}"uniform-work-group-size"
@@ -111,6 +116,7 @@ template __global__ void 
template_a_b_c_max_num_work_groups<32, 4, 2>();
 // CHECK-DAG: attributes [[WAVES_PER_EU_2]] = {{.*}}"amdgpu-waves-per-eu"="2"
 // CHECK-DAG: attributes [[NUM_SGPR_32]] = {{.*}}"amdgpu-num-sgpr"="32"
 // CHECK-DAG: attributes [[NUM_VGPR_64]] = {{.*}}"amdgpu-num-vgpr"="64"
+// CHECK-DAG: attributes [[KERNARG_PRELOAD_COUNT_2]] = 
{{.*}}"amdgpu-kernarg-preload-count"="2"
 // CHECK-DAG: attributes [[MAX_NUM_WORK_GROUPS_32_4_2]] = 
{{.*}}"amdgpu-max-num-workgroups"="32,4,2"
 // CHECK-DAG: attributes [[MAX_NUM_WORK_GROUPS_32_1_1]] = 
{{.*}}"amdgpu-max-num-workgroups"="32,1,1"
 
diff --git a/clang/test/CodeGenHIP/amdgpu-kernarg-preload-count-save-temps.hip 
b/clang/test/CodeGenHIP/amdgpu-kernarg-preload-count-save-temps.hip
new file mode 100644
index 0000000000000..4d809bbeca1d3
--- /dev/null
+++ b/clang/test/CodeGenHIP/amdgpu-kernarg-preload-count-save-temps.hip
@@ -0,0 +1,25 @@
+// REQUIRES: amdgpu-registered-target
+
+// RUN: rm -rf %t && mkdir -p %t
+// RUN: cd %t && %clang --target=x86_64-unknown-linux-gnu 
--offload-arch=gfx942 \
+// RUN:   --cuda-device-only -nogpulib -nogpuinc -save-temps -O2 -c -x hip %s \
+// RUN:   -o out.o
+// RUN: FileCheck %s 
--input-file=%t/amdgpu-kernarg-preload-count-save-temps-hip-amdgcn-amd-amdhsa-gfx942.s
+
+extern "C" __attribute__((global, amdgpu_kernarg_preload_count(2)))
+void preload_count_2(int *out, int a, int b) {
+  *out = a + b;
+}
+
+extern "C" __attribute__((global, amdgpu_kernarg_preload_count(0)))
+void preload_count_0(int *out, int a, int b) {
+  *out = a + b;
+}
+
+// CHECK-LABEL: .amdhsa_kernel preload_count_2
+// CHECK: .amdhsa_user_sgpr_kernarg_preload_length 3
+// CHECK: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+// CHECK-LABEL: .amdhsa_kernel preload_count_0
+// CHECK: .amdhsa_user_sgpr_kernarg_preload_length 0
+// CHECK: .amdhsa_user_sgpr_kernarg_preload_offset 0
diff --git a/clang/test/Misc/pragma-attribute-supported-attributes-list.test 
b/clang/test/Misc/pragma-attribute-supported-attributes-list.test
index 8bca68e2119e7..c9f16d9aace3c 100644
--- a/clang/test/Misc/pragma-attribute-supported-attributes-list.test
+++ b/clang/test/Misc/pragma-attribute-supported-attributes-list.test
@@ -4,6 +4,7 @@
 
 // CHECK: #pragma clang attribute supports the following attributes:
 // CHECK-NEXT: AMDGPUFlatWorkGroupSize (SubjectMatchRule_function)
+// CHECK-NEXT: AMDGPUKernargPreloadCount (SubjectMatchRule_function)
 // CHECK-NEXT: AMDGPUMaxNumWorkGroups (SubjectMatchRule_function)
 // CHECK-NEXT: AMDGPUNumSGPR (SubjectMatchRule_function)
 // CHECK-NEXT: AMDGPUNumVGPR (SubjectMatchRule_function)
diff --git a/clang/test/SemaOpenCL/amdgpu-attrs.cl 
b/clang/test/SemaOpenCL/amdgpu-attrs.cl
index 9321c2f83e01c..1933655499e16 100644
--- a/clang/test/SemaOpenCL/amdgpu-attrs.cl
+++ b/clang/test/SemaOpenCL/amdgpu-attrs.cl
@@ -20,12 +20,17 @@ typedef __attribute__((amdgpu_num_vgpr(64))) struct 
struct_num_vgpr_64 { // expe
   int x;
   float y;
 } struct_num_vgpr_64;
+typedef __attribute__((amdgpu_kernarg_preload_count(2))) struct 
struct_kernarg_preload_count_2 { // expected-error 
{{'amdgpu_kernarg_preload_count' attribute only applies to kernel functions}}
+  int x;
+  float y;
+} struct_kernarg_preload_count_2;
 
 __attribute__((amdgpu_flat_work_group_size(32, 64))) void 
func_flat_work_group_size_32_64() {} // expected-error 
{{'amdgpu_flat_work_group_size' attribute only applies to kernel functions}}
 __attribute__((amdgpu_waves_per_eu(2))) void func_waves_per_eu_2() {} // 
expected-error {{'amdgpu_waves_per_eu' attribute only applies to kernel 
functions}}
 __attribute__((amdgpu_waves_per_eu(2, 4))) void func_waves_per_eu_2_4() {} // 
expected-error {{'amdgpu_waves_per_eu' attribute only applies to kernel 
functions}}
 __attribute__((amdgpu_num_sgpr(32))) void func_num_sgpr_32() {} // 
expected-error {{'amdgpu_num_sgpr' attribute only applies to kernel functions}}
 __attribute__((amdgpu_num_vgpr(64))) void func_num_vgpr_64() {} // 
expected-error {{'amdgpu_num_vgpr' attribute only applies to kernel functions}}
+__attribute__((amdgpu_kernarg_preload_count(2))) void 
func_kernarg_preload_count_2() {} // expected-error 
{{'amdgpu_kernarg_preload_count' attribute only applies to kernel functions}}
 
 __attribute__((amdgpu_flat_work_group_size("ABC", "ABC"))) kernel void 
kernel_flat_work_group_size_ABC_ABC() {} // expected-error 
{{'amdgpu_flat_work_group_size' attribute requires parameter 0 to be an integer 
constant}}
 __attribute__((amdgpu_flat_work_group_size(32, "ABC"))) kernel void 
kernel_flat_work_group_size_32_ABC() {} // expected-error 
{{'amdgpu_flat_work_group_size' attribute requires parameter 1 to be an integer 
constant}}
@@ -35,6 +40,7 @@ __attribute__((amdgpu_waves_per_eu(2, "ABC"))) kernel void 
kernel_waves_per_eu_2
 __attribute__((amdgpu_waves_per_eu("ABC", 4))) kernel void 
kernel_waves_per_eu_ABC_4() {} // expected-error {{'amdgpu_waves_per_eu' 
attribute requires parameter 0 to be an integer constant}}
 __attribute__((amdgpu_num_sgpr("ABC"))) kernel void kernel_num_sgpr_ABC() {} 
// expected-error {{'amdgpu_num_sgpr' attribute requires an integer constant}}
 __attribute__((amdgpu_num_vgpr("ABC"))) kernel void kernel_num_vgpr_ABC() {} 
// expected-error {{'amdgpu_num_vgpr' attribute requires an integer constant}}
+__attribute__((amdgpu_kernarg_preload_count("ABC"))) kernel void 
kernel_kernarg_preload_count_ABC() {} // expected-error 
{{'amdgpu_kernarg_preload_count' attribute requires an integer constant}}
 
 __attribute__((amdgpu_flat_work_group_size(4294967296, 4294967296))) kernel 
void kernel_flat_work_group_size_L_L() {} // expected-error {{integer constant 
expression evaluates to value 4294967296 that cannot be represented in a 32-bit 
unsigned integer type}}
 __attribute__((amdgpu_flat_work_group_size(32, 4294967296))) kernel void 
kernel_flat_work_group_size_32_L() {} // expected-error {{integer constant 
expression evaluates to value 4294967296 that cannot be represented in a 32-bit 
unsigned integer type}}
@@ -44,6 +50,7 @@ __attribute__((amdgpu_waves_per_eu(2, 4294967296))) kernel 
void kernel_waves_per
 __attribute__((amdgpu_waves_per_eu(4294967296, 4))) kernel void 
kernel_waves_per_eu_L_4() {} // expected-error {{integer constant expression 
evaluates to value 4294967296 that cannot be represented in a 32-bit unsigned 
integer type}}
 __attribute__((amdgpu_num_sgpr(4294967296))) kernel void kernel_num_sgpr_L() 
{} // expected-error {{integer constant expression evaluates to value 
4294967296 that cannot be represented in a 32-bit unsigned integer type}}
 __attribute__((amdgpu_num_vgpr(4294967296))) kernel void kernel_num_vgpr_L() 
{} // expected-error {{integer constant expression evaluates to value 
4294967296 that cannot be represented in a 32-bit unsigned integer type}}
+__attribute__((amdgpu_kernarg_preload_count(4294967296))) kernel void 
kernel_kernarg_preload_count_L() {} // expected-error {{integer constant 
expression evaluates to value 4294967296 that cannot be represented in a 32-bit 
unsigned integer type}}
 
 __attribute__((amdgpu_flat_work_group_size(0, 64))) kernel void 
kernel_flat_work_group_size_0_64() {} // expected-error 
{{'amdgpu_flat_work_group_size' attribute argument is invalid: max must be 0 
since min is 0}}
 __attribute__((amdgpu_waves_per_eu(0, 4))) kernel void 
kernel_waves_per_eu_0_4() {} // expected-error {{'amdgpu_waves_per_eu' 
attribute argument is invalid: max must be 0 since min is 0}}
@@ -58,6 +65,8 @@ __attribute__((amdgpu_waves_per_eu(2, 4, 8))) kernel void 
kernel_waves_per_eu_2_
 
 __attribute__((amdgpu_flat_work_group_size(0, 0))) kernel void 
kernel_flat_work_group_size_0_0() {}
 __attribute__((amdgpu_waves_per_eu(0))) kernel void kernel_waves_per_eu_0() {}
+__attribute__((amdgpu_kernarg_preload_count(0))) kernel void 
kernel_kernarg_preload_count_0() {}
+__attribute__((amdgpu_kernarg_preload_count(2))) kernel void 
kernel_kernarg_preload_count_2() {}
 __attribute__((amdgpu_waves_per_eu(0, 0))) kernel void 
kernel_waves_per_eu_0_0() {}
 __attribute__((amdgpu_num_sgpr(0))) kernel void kernel_num_sgpr_0() {}
 __attribute__((amdgpu_num_vgpr(0))) kernel void kernel_num_vgpr_0() {}
@@ -67,3 +76,4 @@ kernel __attribute__((amdgpu_waves_per_eu(2))) void 
kernel_waves_per_eu_2() {}
 kernel __attribute__((amdgpu_waves_per_eu(2, 4))) void 
kernel_waves_per_eu_2_4() {}
 kernel __attribute__((amdgpu_num_sgpr(32))) void kernel_num_sgpr_32() {}
 kernel __attribute__((amdgpu_num_vgpr(64))) void kernel_num_vgpr_64() {}
+kernel __attribute__((amdgpu_kernarg_preload_count(2))) void 
kernel_kernarg_preload_count_2_alt() {}
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUPreloadKernelArguments.cpp 
b/llvm/lib/Target/AMDGPU/AMDGPUPreloadKernelArguments.cpp
index 7d6e3edc75e1f..7702c1ae52df2 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUPreloadKernelArguments.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUPreloadKernelArguments.cpp
@@ -42,6 +42,9 @@ static cl::opt<bool>
                          cl::desc("Enable preload kernel arguments to SGPRs"),
                          cl::init(true));
 
+static constexpr StringRef KernargPreloadCountAttr =
+    "amdgpu-kernarg-preload-count";
+
 namespace {
 
 class AMDGPUPreloadKernelArgumentsLegacy : public ModulePass {
@@ -295,7 +298,14 @@ static bool markKernelArgsAsInreg(Module &M, const 
TargetMachine &TM) {
     uint64_t ExplicitArgOffset = 0;
     const DataLayout &DL = F.getDataLayout();
     const uint64_t BaseOffset = ST.getExplicitKernelArgOffset();
-    unsigned NumPreloadsRequested = KernargPreloadCount;
+    bool HasKernargPreloadCountAttr = 
F.hasFnAttribute(KernargPreloadCountAttr);
+    uint64_t NumPreloadsRequested =
+        HasKernargPreloadCountAttr
+            ? F.getFnAttributeAsParsedInteger(KernargPreloadCountAttr)
+            : KernargPreloadCount;
+    if (HasKernargPreloadCountAttr && NumPreloadsRequested == 0)
+      continue;
+
     unsigned NumPreloadedExplicitArgs = 0;
     for (Argument &Arg : F.args()) {
       // Avoid incompatible attributes and guard against running this pass
@@ -306,6 +316,9 @@ static bool markKernelArgsAsInreg(Module &M, const 
TargetMachine &TM) {
           Arg.hasAttribute("amdgpu-hidden-argument"))
         break;
 
+      if (HasKernargPreloadCountAttr && NumPreloadsRequested == 0)
+        break;
+
       // Inreg may be pre-existing on some arguments, try to preload these.
       if (NumPreloadsRequested == 0 && !Arg.hasInRegAttr())
         break;
diff --git a/llvm/test/CodeGen/AMDGPU/preload-kernargs-count-attr.ll 
b/llvm/test/CodeGen/AMDGPU/preload-kernargs-count-attr.ll
new file mode 100644
index 0000000000000..aa344e2174164
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/preload-kernargs-count-attr.ll
@@ -0,0 +1,22 @@
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx942 
-passes=amdgpu-preload-kernel-arguments -amdgpu-kernarg-preload-count=3 < %s | 
FileCheck %s
+
+define amdgpu_kernel void @global_default(ptr addrspace(1) %a, ptr 
addrspace(1) %b, ptr addrspace(1) %c) {
+; CHECK-LABEL: define amdgpu_kernel void @global_default(
+; CHECK-SAME: ptr addrspace(1) inreg %a, ptr addrspace(1) inreg %b, ptr 
addrspace(1) inreg %c)
+  ret void
+}
+
+define amdgpu_kernel void @attr_count_one(ptr addrspace(1) %a, ptr 
addrspace(1) %b, ptr addrspace(1) %c) #0 {
+; CHECK-LABEL: define amdgpu_kernel void @attr_count_one(
+; CHECK-SAME: ptr addrspace(1) inreg %a, ptr addrspace(1) %b, ptr addrspace(1) 
%c)
+  ret void
+}
+
+define amdgpu_kernel void @attr_count_zero(ptr addrspace(1) %a, ptr 
addrspace(1) %b, ptr addrspace(1) %c) #1 {
+; CHECK-LABEL: define amdgpu_kernel void @attr_count_zero(
+; CHECK-SAME: ptr addrspace(1) %a, ptr addrspace(1) %b, ptr addrspace(1) %c)
+  ret void
+}
+
+attributes #0 = { "amdgpu-kernarg-preload-count"="1" }
+attributes #1 = { "amdgpu-kernarg-preload-count"="0" }

>From e2414578034c5db272609b3591b0497018f1825d Mon Sep 17 00:00:00 2001
From: "Yaxun (Sam) Liu" <[email protected]>
Date: Thu, 23 Jul 2026 21:54:12 -0400
Subject: [PATCH 2/2] [AMDGPU] Support template kernarg preload counts

The attribute count must be a compile-time constant, but a template value is 
not known until instantiation.

Keep the expression until template instantiation and validate the substituted 
value then.
---
 clang/include/clang/Basic/Attr.td             |  2 +-
 clang/include/clang/Sema/SemaAMDGPU.h         |  7 ++++
 clang/lib/CodeGen/Targets/AMDGPU.cpp          |  8 ++--
 clang/lib/Sema/SemaAMDGPU.cpp                 | 39 +++++++++++++++----
 .../lib/Sema/SemaTemplateInstantiateDecl.cpp  | 20 ++++++++++
 clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu |  8 ++++
 clang/test/SemaCUDA/amdgpu-attrs.cu           |  8 +++-
 clang/test/SemaOpenCL/amdgpu-attrs.cl         |  2 +
 8 files changed, 82 insertions(+), 12 deletions(-)

diff --git a/clang/include/clang/Basic/Attr.td 
b/clang/include/clang/Basic/Attr.td
index 9cb175fb2afbc..88b2d42472f12 100644
--- a/clang/include/clang/Basic/Attr.td
+++ b/clang/include/clang/Basic/Attr.td
@@ -2527,7 +2527,7 @@ def AMDGPUNumVGPR : InheritableAttr {
 
 def AMDGPUKernargPreloadCount : InheritableAttr {
   let Spellings = [Clang<"amdgpu_kernarg_preload_count", 0>];
-  let Args = [UnsignedArgument<"Count">];
+  let Args = [ExprArgument<"Count">];
   let Documentation = [AMDGPUKernargPreloadCountDocs];
   let Subjects = SubjectList<[Function], ErrorDiag, "kernel functions">;
 }
diff --git a/clang/include/clang/Sema/SemaAMDGPU.h 
b/clang/include/clang/Sema/SemaAMDGPU.h
index 603b2f80c5f21..0f3b5ccf03aca 100644
--- a/clang/include/clang/Sema/SemaAMDGPU.h
+++ b/clang/include/clang/Sema/SemaAMDGPU.h
@@ -64,6 +64,13 @@ class SemaAMDGPU : public SemaBase {
   void addAMDGPUWavesPerEUAttr(Decl *D, const AttributeCommonInfo &CI,
                                Expr *Min, Expr *Max);
 
+  AMDGPUKernargPreloadCountAttr *
+  CreateAMDGPUKernargPreloadCountAttr(const AttributeCommonInfo &CI,
+                                      Expr *CountExpr);
+
+  void addAMDGPUKernargPreloadCountAttr(Decl *D, const AttributeCommonInfo &CI,
+                                        Expr *CountExpr);
+
   /// Create an AMDGPUMaxNumWorkGroupsAttr attribute.
   AMDGPUMaxNumWorkGroupsAttr *
   CreateAMDGPUMaxNumWorkGroupsAttr(const AttributeCommonInfo &CI, Expr *XExpr,
diff --git a/clang/lib/CodeGen/Targets/AMDGPU.cpp 
b/clang/lib/CodeGen/Targets/AMDGPU.cpp
index 0e2f8c5aa6a72..41700f55f4cb0 100644
--- a/clang/lib/CodeGen/Targets/AMDGPU.cpp
+++ b/clang/lib/CodeGen/Targets/AMDGPU.cpp
@@ -377,9 +377,11 @@ void AMDGPUTargetCodeGenInfo::setFunctionDeclAttributes(
       F->addFnAttr("amdgpu-num-vgpr", llvm::utostr(NumVGPR));
   }
 
-  if (const auto *Attr = FD->getAttr<AMDGPUKernargPreloadCountAttr>())
-    F->addFnAttr("amdgpu-kernarg-preload-count",
-                 llvm::utostr(Attr->getCount()));
+  if (const auto *Attr = FD->getAttr<AMDGPUKernargPreloadCountAttr>()) {
+    unsigned Count =
+        Attr->getCount()->EvaluateKnownConstInt(M.getContext()).getExtValue();
+    F->addFnAttr("amdgpu-kernarg-preload-count", llvm::utostr(Count));
+  }
 
   if (const auto *Attr = FD->getAttr<AMDGPUMaxNumWorkGroupsAttr>()) {
     uint32_t X = Attr->getMaxNumWorkGroupsX()
diff --git a/clang/lib/Sema/SemaAMDGPU.cpp b/clang/lib/Sema/SemaAMDGPU.cpp
index 6fd7e711ed542..1fcb480f620af 100644
--- a/clang/lib/Sema/SemaAMDGPU.cpp
+++ b/clang/lib/Sema/SemaAMDGPU.cpp
@@ -729,15 +729,40 @@ void SemaAMDGPU::handleAMDGPUNumVGPRAttr(Decl *D, const 
ParsedAttr &AL) {
                  AMDGPUNumVGPRAttr(getASTContext(), AL, NumVGPR));
 }
 
-void SemaAMDGPU::handleAMDGPUKernargPreloadCountAttr(Decl *D,
-                                                     const ParsedAttr &AL) {
+static bool checkAMDGPUKernargPreloadCountArgument(
+    Sema &S, Expr *CountExpr, const AMDGPUKernargPreloadCountAttr &Attr) {
+  if (S.DiagnoseUnexpandedParameterPack(CountExpr))
+    return true;
+
+  if (CountExpr->isValueDependent())
+    return false;
+
   uint32_t Count = 0;
-  Expr *CountExpr = AL.getArgAsExpr(0);
-  if (!SemaRef.checkUInt32Argument(AL, CountExpr, Count))
-    return;
+  return !S.checkUInt32Argument(Attr, CountExpr, Count);
+}
 
-  D->addAttr(::new (getASTContext())
-                 AMDGPUKernargPreloadCountAttr(getASTContext(), AL, Count));
+AMDGPUKernargPreloadCountAttr *
+SemaAMDGPU::CreateAMDGPUKernargPreloadCountAttr(const AttributeCommonInfo &CI,
+                                                Expr *CountExpr) {
+  ASTContext &Context = getASTContext();
+  AMDGPUKernargPreloadCountAttr TmpAttr(Context, CI, CountExpr);
+
+  if (checkAMDGPUKernargPreloadCountArgument(SemaRef, CountExpr, TmpAttr))
+    return nullptr;
+
+  return ::new (Context) AMDGPUKernargPreloadCountAttr(Context, CI, CountExpr);
+}
+
+void SemaAMDGPU::addAMDGPUKernargPreloadCountAttr(Decl *D,
+                                                  const AttributeCommonInfo 
&CI,
+                                                  Expr *CountExpr) {
+  if (auto *Attr = CreateAMDGPUKernargPreloadCountAttr(CI, CountExpr))
+    D->addAttr(Attr);
+}
+
+void SemaAMDGPU::handleAMDGPUKernargPreloadCountAttr(Decl *D,
+                                                     const ParsedAttr &AL) {
+  addAMDGPUKernargPreloadCountAttr(D, AL, AL.getArgAsExpr(0));
 }
 
 static bool
diff --git a/clang/lib/Sema/SemaTemplateInstantiateDecl.cpp 
b/clang/lib/Sema/SemaTemplateInstantiateDecl.cpp
index 921a9f965fb9c..bc00784fa5b8e 100644
--- a/clang/lib/Sema/SemaTemplateInstantiateDecl.cpp
+++ b/clang/lib/Sema/SemaTemplateInstantiateDecl.cpp
@@ -677,6 +677,19 @@ static void instantiateDependentAMDGPUWavesPerEUAttr(
   S.AMDGPU().addAMDGPUWavesPerEUAttr(New, Attr, MinExpr, MaxExpr);
 }
 
+static void instantiateDependentAMDGPUKernargPreloadCountAttr(
+    Sema &S, const MultiLevelTemplateArgumentList &TemplateArgs,
+    const AMDGPUKernargPreloadCountAttr &Attr, Decl *New) {
+  EnterExpressionEvaluationContext Unevaluated(
+      S, Sema::ExpressionEvaluationContext::ConstantEvaluated);
+
+  ExprResult Result = S.SubstExpr(Attr.getCount(), TemplateArgs);
+  if (Result.isInvalid())
+    return;
+
+  S.AMDGPU().addAMDGPUKernargPreloadCountAttr(New, Attr, Result.getAs<Expr>());
+}
+
 static void instantiateDependentAMDGPUMaxNumWorkGroupsAttr(
     Sema &S, const MultiLevelTemplateArgumentList &TemplateArgs,
     const AMDGPUMaxNumWorkGroupsAttr &Attr, Decl *New) {
@@ -963,6 +976,13 @@ void Sema::InstantiateAttrs(const 
MultiLevelTemplateArgumentList &TemplateArgs,
                                                *AMDGPUFlatWorkGroupSize, New);
     }
 
+    if (const auto *AMDGPUKernargPreloadCount =
+            dyn_cast<AMDGPUKernargPreloadCountAttr>(TmplAttr)) {
+      instantiateDependentAMDGPUKernargPreloadCountAttr(
+          *this, TemplateArgs, *AMDGPUKernargPreloadCount, New);
+      continue;
+    }
+
     if (const auto *AMDGPUMaxNumWorkGroups =
             dyn_cast<AMDGPUMaxNumWorkGroupsAttr>(TmplAttr)) {
       instantiateDependentAMDGPUMaxNumWorkGroupsAttr(
diff --git a/clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu 
b/clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu
index f02be68dcdc61..af1a9f49aa22e 100644
--- a/clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu
+++ b/clang/test/CodeGenCUDA/amdgpu-kernel-attrs.cu
@@ -49,6 +49,13 @@ __attribute__((amdgpu_kernarg_preload_count(2))) // 
expected-no-diagnostics
 __global__ void kernarg_preload_count_2() {
 // CHECK: define{{.*}} amdgpu_kernel void @_Z23kernarg_preload_count_2v() 
[[KERNARG_PRELOAD_COUNT_2:#[0-9]+]]
 }
+
+template<unsigned Count>
+__attribute__((amdgpu_kernarg_preload_count(Count)))
+__global__ void template_kernarg_preload_count() {}
+template __global__ void template_kernarg_preload_count<3>();
+// CHECK: define{{.*}} amdgpu_kernel void 
@_Z30template_kernarg_preload_countILj3EEvv() 
[[KERNARG_PRELOAD_COUNT_3:#[0-9]+]]
+
 __attribute__((amdgpu_max_num_work_groups(32, 4, 2))) // 
expected-no-diagnostics
 __global__ void max_num_work_groups_32_4_2() {
 // CHECK: define{{.*}} amdgpu_kernel void @_Z26max_num_work_groups_32_4_2v() 
[[MAX_NUM_WORK_GROUPS_32_4_2:#[0-9]+]]
@@ -117,6 +124,7 @@ template __global__ void 
template_a_b_c_max_num_work_groups<32, 4, 2>();
 // CHECK-DAG: attributes [[NUM_SGPR_32]] = {{.*}}"amdgpu-num-sgpr"="32"
 // CHECK-DAG: attributes [[NUM_VGPR_64]] = {{.*}}"amdgpu-num-vgpr"="64"
 // CHECK-DAG: attributes [[KERNARG_PRELOAD_COUNT_2]] = 
{{.*}}"amdgpu-kernarg-preload-count"="2"
+// CHECK-DAG: attributes [[KERNARG_PRELOAD_COUNT_3]] = 
{{.*}}"amdgpu-kernarg-preload-count"="3"
 // CHECK-DAG: attributes [[MAX_NUM_WORK_GROUPS_32_4_2]] = 
{{.*}}"amdgpu-max-num-workgroups"="32,4,2"
 // CHECK-DAG: attributes [[MAX_NUM_WORK_GROUPS_32_1_1]] = 
{{.*}}"amdgpu-max-num-workgroups"="32,1,1"
 
diff --git a/clang/test/SemaCUDA/amdgpu-attrs.cu 
b/clang/test/SemaCUDA/amdgpu-attrs.cu
index ee0219696bdf4..7ad89325f3e78 100644
--- a/clang/test/SemaCUDA/amdgpu-attrs.cu
+++ b/clang/test/SemaCUDA/amdgpu-attrs.cu
@@ -205,6 +205,13 @@ __global__ void non_cexpr_waves_per_eu_2() {}
 __attribute__((amdgpu_waves_per_eu(2, ipow2(2))))
 __global__ void non_cexpr_waves_per_eu_2_4() {}
 
+// expected-error@+3{{integer constant expression evaluates to value 
4294967296 that cannot be represented in a 32-bit unsigned integer type}}
+// expected-note@+4{{in instantiation of}}
+template<unsigned long long Count>
+__attribute__((amdgpu_kernarg_preload_count(Count)))
+__global__ void template_kernarg_preload_count_too_large() {}
+template __global__ void 
template_kernarg_preload_count_too_large<4294967296ULL>();
+
 __attribute__((amdgpu_max_num_work_groups(32)))
 __global__ void max_num_work_groups_32() {}
 
@@ -325,4 +332,3 @@ __attribute__((amdgpu_max_num_work_groups(32, 1, b)))
 __global__ void template_32_1_b_max_num_work_groups() {}
 template __global__ void template_32_1_b_max_num_work_groups<0>();
 
-
diff --git a/clang/test/SemaOpenCL/amdgpu-attrs.cl 
b/clang/test/SemaOpenCL/amdgpu-attrs.cl
index 1933655499e16..5875a0ac54156 100644
--- a/clang/test/SemaOpenCL/amdgpu-attrs.cl
+++ b/clang/test/SemaOpenCL/amdgpu-attrs.cl
@@ -41,6 +41,8 @@ __attribute__((amdgpu_waves_per_eu("ABC", 4))) kernel void 
kernel_waves_per_eu_A
 __attribute__((amdgpu_num_sgpr("ABC"))) kernel void kernel_num_sgpr_ABC() {} 
// expected-error {{'amdgpu_num_sgpr' attribute requires an integer constant}}
 __attribute__((amdgpu_num_vgpr("ABC"))) kernel void kernel_num_vgpr_ABC() {} 
// expected-error {{'amdgpu_num_vgpr' attribute requires an integer constant}}
 __attribute__((amdgpu_kernarg_preload_count("ABC"))) kernel void 
kernel_kernarg_preload_count_ABC() {} // expected-error 
{{'amdgpu_kernarg_preload_count' attribute requires an integer constant}}
+extern constant int kernarg_preload_count;
+__attribute__((amdgpu_kernarg_preload_count(kernarg_preload_count))) kernel 
void kernel_kernarg_preload_count_non_constant() {} // expected-error 
{{'amdgpu_kernarg_preload_count' attribute requires an integer constant}}
 
 __attribute__((amdgpu_flat_work_group_size(4294967296, 4294967296))) kernel 
void kernel_flat_work_group_size_L_L() {} // expected-error {{integer constant 
expression evaluates to value 4294967296 that cannot be represented in a 32-bit 
unsigned integer type}}
 __attribute__((amdgpu_flat_work_group_size(32, 4294967296))) kernel void 
kernel_flat_work_group_size_32_L() {} // expected-error {{integer constant 
expression evaluates to value 4294967296 that cannot be represented in a 32-bit 
unsigned integer type}}

_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to