[llvm] [AMDGPU] Split struct kernel arguments (PR #133786)
Yaxun Liu via llvm-commits
llvm-commits at lists.llvm.org
Fri Jan 23 11:21:31 PST 2026
https://github.com/yxsamliu updated https://github.com/llvm/llvm-project/pull/133786
>From 45a7c587725991dde1aff1cb2fa8200d2cc1aeed Mon Sep 17 00:00:00 2001
From: "Yaxun (Sam) Liu" <yaxun.liu at amd.com>
Date: Fri, 23 Jan 2026 09:45:50 -0500
Subject: [PATCH] [AMDGPU] Split struct kernel arguments for preloading
Integrate split kernel arguments logic into the preload pass. When a
struct kernel argument has fields that can be individually preloaded,
split the argument into separate scalar/vector arguments for more
efficient SGPR preloading.
Key features:
- Layout-preserving mode (default): Fields stay at original offsets,
uses backup declaration to preserve metadata for CLR compatibility
- Layout-changing mode (-amdgpu-kernarg-layout-change): Reorders fields
for optimal packing
- Only splits when target supports preload and args fit in SGPRs
- Handles nested structs, arrays with constant indices, vectors
- Variable-index arrays and address-taken args are not split
The backup declaration approach allows stock CLR to work without
patching by emitting original struct metadata.
---
.../AMDGPU/AMDGPUHSAMetadataStreamer.cpp | 53 +-
.../Target/AMDGPU/AMDGPUHSAMetadataStreamer.h | 4 +-
.../AMDGPU/AMDGPUPreloadKernelArguments.cpp | 520 +++++++++++++++++-
...mdgpu-split-kernel-args-preserve-layout.ll | 163 ++++++
.../AMDGPU/amdgpu-split-kernel-args.ll | 263 +++++++++
5 files changed, 995 insertions(+), 8 deletions(-)
create mode 100644 llvm/test/CodeGen/AMDGPU/amdgpu-split-kernel-args-preserve-layout.ll
create mode 100644 llvm/test/CodeGen/AMDGPU/amdgpu-split-kernel-args.ll
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
index 3c88d1b8214f7..37bc565f2ace1 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
@@ -15,6 +15,9 @@
#include "AMDGPUHSAMetadataStreamer.h"
#include "AMDGPU.h"
#include "GCNSubtarget.h"
+#include "llvm/Support/Debug.h"
+
+#define DEBUG_TYPE "amdgpu-hsa-metadata"
#include "MCTargetDesc/AMDGPUTargetStreamer.h"
#include "SIMachineFunctionInfo.h"
#include "SIProgramInfo.h"
@@ -286,7 +289,21 @@ void MetadataStreamerMsgPackV4::emitKernelArgs(const MachineFunction &MF,
auto &Func = MF.getFunction();
unsigned Offset = 0;
auto Args = HSAMetadataDoc->getArrayNode();
- for (auto &Arg : Func.args()) {
+
+ // Check if there's a backup declaration for the original kernel signature.
+ // If so, use the backup's args for metadata (layout-preserving split mode).
+ const Function *MetadataFunc = &Func;
+ Attribute OrigKernelAttr = Func.getFnAttribute("amdgpu-original-kernel");
+ if (OrigKernelAttr.isValid()) {
+ StringRef BackupName = OrigKernelAttr.getValueAsString();
+ if (const Function *BackupF = Func.getParent()->getFunction(BackupName)) {
+ MetadataFunc = BackupF;
+ LLVM_DEBUG(dbgs() << "Using backup declaration for metadata: "
+ << BackupName << '\n');
+ }
+ }
+
+ for (auto &Arg : MetadataFunc->args()) {
if (Arg.hasAttribute("amdgpu-hidden-argument"))
continue;
@@ -357,17 +374,39 @@ void MetadataStreamerMsgPackV4::emitKernelArg(const Argument &Arg,
Align ArgAlign;
std::tie(ArgTy, ArgAlign) = getArgumentTypeAlign(Arg, DL);
+ // Check if the argument was split from a struct-type argument.
+ // The "amdgpu-original-arg" attribute encodes original index and offset.
+ unsigned OriginalArgIndex = ~0U;
+ uint64_t OriginalArgOffset = 0;
+ Attribute Attr =
+ Func->getAttributes().getParamAttr(ArgNo, "amdgpu-original-arg");
+ if (Attr.isValid()) {
+ StringRef MappingStr = Attr.getValueAsString();
+ SmallVector<StringRef, 2> Elements;
+ MappingStr.split(Elements, ':');
+ bool Valid = Elements.size() == 2 &&
+ !Elements[0].getAsInteger(10, OriginalArgIndex) &&
+ !Elements[1].getAsInteger(10, OriginalArgOffset);
+ assert(Valid && "Invalid amdgpu-original-arg attribute format");
+ if (!Valid) {
+ // Invalid format, ignore the attribute.
+ OriginalArgIndex = ~0U;
+ OriginalArgOffset = 0;
+ }
+ }
+
emitKernelArg(DL, ArgTy, ArgAlign,
getValueKind(ArgTy, TypeQual, BaseTypeName), Offset, Args,
- PointeeAlign, Name, TypeName, BaseTypeName, ActAccQual,
- AccQual, TypeQual);
+ PointeeAlign, Name, TypeName, BaseTypeName, ActAccQual, AccQual,
+ TypeQual, OriginalArgIndex, OriginalArgOffset);
}
void MetadataStreamerMsgPackV4::emitKernelArg(
const DataLayout &DL, Type *Ty, Align Alignment, StringRef ValueKind,
unsigned &Offset, msgpack::ArrayDocNode Args, MaybeAlign PointeeAlign,
StringRef Name, StringRef TypeName, StringRef BaseTypeName,
- StringRef ActAccQual, StringRef AccQual, StringRef TypeQual) {
+ StringRef ActAccQual, StringRef AccQual, StringRef TypeQual,
+ unsigned OriginalArgIndex, uint64_t OriginalArgOffset) {
auto Arg = Args.getDocument()->getMapNode();
if (!Name.empty())
@@ -409,6 +448,12 @@ void MetadataStreamerMsgPackV4::emitKernelArg(
Arg[".is_pipe"] = Arg.getDocument()->getNode(true);
}
+ // Add original argument index and offset to the metadata
+ if (OriginalArgIndex != ~0U) {
+ Arg[".original_arg_index"] = Arg.getDocument()->getNode(OriginalArgIndex);
+ Arg[".original_arg_offset"] = Arg.getDocument()->getNode(OriginalArgOffset);
+ }
+
Args.push_back(Arg);
}
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h
index 6bad4dbdf5321..501b158ee17d9 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h
+++ b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.h
@@ -116,7 +116,9 @@ class LLVM_EXTERNAL_VISIBILITY MetadataStreamerMsgPackV4
MaybeAlign PointeeAlign = std::nullopt,
StringRef Name = "", StringRef TypeName = "",
StringRef BaseTypeName = "", StringRef ActAccQual = "",
- StringRef AccQual = "", StringRef TypeQual = "");
+ StringRef AccQual = "", StringRef TypeQual = "",
+ unsigned OriginalArgIndex = ~0U,
+ uint64_t OriginalArgOffset = 0);
void emitHiddenKernelArgs(const MachineFunction &MF, unsigned &Offset,
msgpack::ArrayDocNode Args) override;
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUPreloadKernelArguments.cpp b/llvm/lib/Target/AMDGPU/AMDGPUPreloadKernelArguments.cpp
index 7d6e3edc75e1f..c9c8965ca27cc 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUPreloadKernelArguments.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUPreloadKernelArguments.cpp
@@ -15,13 +15,23 @@
/// Additionally, hidden kernel arguments may be preloaded, in which case they
/// are appended to the kernel signature after explicit arguments. Preloaded
/// arguments will be marked with `inreg`.
+///
+/// This pass also handles splitting of byref struct kernel arguments into
+/// scalar arguments when doing so would allow them to be preloaded. The
+/// splitting only occurs if the split arguments can fit in available SGPRs.
//
//===----------------------------------------------------------------------===//
#include "AMDGPU.h"
#include "AMDGPUTargetMachine.h"
+#include "GCNSubtarget.h"
+#include "Utils/AMDGPUBaseInfo.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
#include "llvm/Analysis/ValueTracking.h"
#include "llvm/IR/Function.h"
+#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicsAMDGPU.h"
#include "llvm/IR/Module.h"
@@ -42,8 +52,496 @@ static cl::opt<bool>
cl::desc("Enable preload kernel arguments to SGPRs"),
cl::init(true));
+static cl::opt<bool> EnableKernargLayoutChange(
+ "amdgpu-kernarg-layout-change",
+ cl::desc("Allow changing kernel argument segment layout when splitting "
+ "byref structs (remove unused fields, reorder for packing). "
+ "When disabled (default), all struct fields are preserved in "
+ "their original order."),
+ cl::init(false));
+
namespace {
+//===----------------------------------------------------------------------===//
+// Kernel Argument Splitting Logic
+//
+// The following functions handle splitting of byref struct kernel arguments
+// into scalar arguments. This enables preloading of struct fields that would
+// otherwise not be preloadable due to the byref attribute.
+//===----------------------------------------------------------------------===//
+
+// Attribute name for tracking original argument index and offset
+static constexpr StringRef OriginalArgAttr = "amdgpu-original-arg";
+
+// Prefix for backup declaration of original kernel (used for metadata
+// generation)
+static constexpr StringRef OriginalKernelPrefix = "__amdgpu_orig_kernel_";
+
+// Attribute to store the name of the backup declaration
+static constexpr StringRef OriginalKernelAttr = "amdgpu-original-kernel";
+
+static bool parseOriginalArgAttribute(StringRef S, unsigned &RootIdx,
+ uint64_t &BaseOff) {
+ auto Parts = S.split(':');
+ if (Parts.second.empty())
+ return false;
+ if (Parts.first.getAsInteger(10, RootIdx))
+ return false;
+ if (Parts.second.getAsInteger(10, BaseOff))
+ return false;
+ return true;
+}
+
+/// Traverses all users of an argument to check if it's suitable for
+/// splitting. A suitable argument is only used by a chain of
+/// GEPs that terminate in LoadInsts.
+static bool
+areArgUsersValidForSplit(Argument &Arg, SmallVectorImpl<LoadInst *> &Loads,
+ SmallVectorImpl<GetElementPtrInst *> &GEPs) {
+ SmallVector<User *, 16> Worklist(Arg.user_begin(), Arg.user_end());
+ SetVector<User *> Visited;
+
+ while (!Worklist.empty()) {
+ User *U = Worklist.pop_back_val();
+ if (!Visited.insert(U))
+ continue;
+
+ if (auto *LI = dyn_cast<LoadInst>(U)) {
+ Loads.push_back(LI);
+ } else if (auto *GEP = dyn_cast<GetElementPtrInst>(U)) {
+ GEPs.push_back(GEP);
+ for (User *GEPUser : GEP->users())
+ Worklist.push_back(GEPUser);
+ } else {
+ return false;
+ }
+ }
+
+ const DataLayout &DL = Arg.getParent()->getParent()->getDataLayout();
+ for (const LoadInst *LI : Loads) {
+ APInt Offset(DL.getPointerSizeInBits(), 0);
+ const Value *Base =
+ LI->getPointerOperand()->stripAndAccumulateConstantOffsets(
+ DL, Offset, /*AllowNonInbounds=*/false);
+ if (Base != &Arg)
+ return false;
+ }
+
+ return true;
+}
+
+/// Information about a struct field to be flattened into a scalar argument.
+struct FieldInfo {
+ Type *Ty = nullptr;
+ uint64_t Offset = 0;
+ LoadInst *Load = nullptr; // nullptr if field is unused
+};
+
+/// Recursively collect all leaf (scalar) fields from a type with their offsets.
+/// This flattens nested structs and arrays into individual scalar fields.
+static void collectLeafFields(Type *Ty, const DataLayout &DL,
+ uint64_t BaseOffset,
+ SmallVectorImpl<FieldInfo> &Fields) {
+ if (auto *STy = dyn_cast<StructType>(Ty)) {
+ const StructLayout *SL = DL.getStructLayout(STy);
+ for (unsigned I = 0; I < STy->getNumElements(); ++I) {
+ Type *ElemTy = STy->getElementType(I);
+ uint64_t ElemOffset = BaseOffset + SL->getElementOffset(I);
+ collectLeafFields(ElemTy, DL, ElemOffset, Fields);
+ }
+ } else if (auto *ATy = dyn_cast<ArrayType>(Ty)) {
+ Type *ElemTy = ATy->getElementType();
+ uint64_t ElemSize = DL.getTypeAllocSize(ElemTy);
+ for (uint64_t I = 0; I < ATy->getNumElements(); ++I) {
+ collectLeafFields(ElemTy, DL, BaseOffset + I * ElemSize, Fields);
+ }
+ } else {
+ // Leaf type (scalar, vector, pointer, etc.)
+ Fields.push_back({Ty, BaseOffset, nullptr});
+ }
+}
+
+/// Check if split arguments can be preloaded into SGPRs.
+/// This calculates the new arg layout size after splitting and checks if it
+/// fits in available user SGPRs.
+static bool canPreloadSplitArgs(
+ Function &F, const GCNSubtarget &ST,
+ const DenseMap<Argument *, SmallVector<FieldInfo, 8>> &ArgToFieldsMap) {
+ GCNUserSGPRUsageInfo UserSGPRInfo(F, ST);
+ unsigned NumFreeUserSGPRs = UserSGPRInfo.getNumFreeUserSGPRs();
+ uint64_t AvailableBytes = NumFreeUserSGPRs * 4;
+
+ const DataLayout &DL = F.getParent()->getDataLayout();
+ uint64_t NewArgOffset = 0;
+
+ // Calculate the new arg layout size after splitting
+ for (Argument &Arg : F.args()) {
+ auto It = ArgToFieldsMap.find(&Arg);
+ if (It != ArgToFieldsMap.end()) {
+ // This arg will be split - add sizes of replacement scalar args
+ for (const FieldInfo &FI : It->second) {
+ Align ABITypeAlign = DL.getABITypeAlign(FI.Ty);
+ uint64_t AllocSize = DL.getTypeAllocSize(FI.Ty);
+ NewArgOffset = alignTo(NewArgOffset, ABITypeAlign) + AllocSize;
+ }
+ } else {
+ // This arg is not split - keep original size
+ Type *ArgTy = Arg.getType();
+ if (Arg.hasByRefAttr())
+ ArgTy = Arg.getParamByRefType();
+ Align ABITypeAlign = DL.getABITypeAlign(ArgTy);
+ uint64_t AllocSize = DL.getTypeAllocSize(ArgTy);
+ NewArgOffset = alignTo(NewArgOffset, ABITypeAlign) + AllocSize;
+ }
+ }
+
+ return NewArgOffset <= AvailableBytes;
+}
+
+/// Try to split byref struct kernel arguments into scalar arguments.
+/// Returns the new function with split arguments, or nullptr if no split
+/// was performed. If a new function is returned, the original function F
+/// has been erased and should not be used.
+///
+/// When EnableKernargLayoutChange is false (default), ALL struct fields are
+/// preserved in their original order (recursively flattening nested structs),
+/// maintaining the kernel argument segment layout. Unused fields become dead
+/// arguments.
+///
+/// When EnableKernargLayoutChange is true, only used fields are kept and
+/// the layout may change.
+static Function *trySplitKernelArguments(Function &F, const GCNSubtarget &ST) {
+ if (!ST.hasKernargPreload())
+ return nullptr;
+
+ if (F.isDeclaration() || !AMDGPU::isKernel(F.getCallingConv()) ||
+ F.arg_empty())
+ return nullptr;
+
+ const DataLayout &DL = F.getParent()->getDataLayout();
+
+ // Mappings from new arg index to original arg: (NewArgIdx, OrigArgIdx,
+ // Offset)
+ SmallVector<std::tuple<unsigned, unsigned, uint64_t>, 8> NewArgMappings;
+ DenseMap<Argument *, SmallVector<LoadInst *, 8>> ArgToLoadsMap;
+ DenseMap<Argument *, SmallVector<GetElementPtrInst *, 8>> ArgToGEPsMap;
+ // Maps struct arg to field info (type, offset, associated load if any)
+ DenseMap<Argument *, SmallVector<FieldInfo, 8>> ArgToFieldsMap;
+ SmallVector<Argument *, 8> StructArgs;
+ SmallVector<Type *, 8> NewArgTypes;
+
+ unsigned OriginalArgIndex = 0;
+ unsigned NewArgIndex = 0;
+ auto HandlePassthroughArg = [&](Argument &Arg) {
+ NewArgTypes.push_back(Arg.getType());
+ if (!Arg.hasAttribute(OriginalArgAttr) && NewArgIndex != OriginalArgIndex)
+ NewArgMappings.emplace_back(NewArgIndex, OriginalArgIndex, 0);
+ ++NewArgIndex;
+ ++OriginalArgIndex;
+ };
+
+ for (Argument &Arg : F.args()) {
+ PointerType *PT = dyn_cast<PointerType>(Arg.getType());
+ if (!PT || !Arg.hasByRefAttr()) {
+ HandlePassthroughArg(Arg);
+ continue;
+ }
+
+ StructType *STy = dyn_cast<StructType>(Arg.getParamByRefType());
+ if (!STy) {
+ HandlePassthroughArg(Arg);
+ continue;
+ }
+
+ // Collect loads from this struct argument
+ SmallVector<LoadInst *, 8> Loads;
+ SmallVector<GetElementPtrInst *, 8> GEPs;
+
+ // Check if all users are valid for splitting (GEPs + loads)
+ bool HasValidUsers =
+ Arg.use_empty() || areArgUsersValidForSplit(Arg, Loads, GEPs);
+ if (!HasValidUsers) {
+ HandlePassthroughArg(Arg);
+ continue;
+ }
+
+ // Helper to get load offset. Returns std::nullopt if offset can't be
+ // computed (e.g., variable-index GEP).
+ auto GetLoadOffset = [&](LoadInst *LI) -> std::optional<uint64_t> {
+ Value *Ptr = LI->getPointerOperand();
+ // Direct load from argument (offset 0)
+ if (Ptr == &Arg)
+ return 0;
+ if (auto *GEP = dyn_cast<GetElementPtrInst>(Ptr)) {
+ APInt OffsetAPInt(DL.getPointerSizeInBits(), 0);
+ if (GEP->accumulateConstantOffset(DL, OffsetAPInt))
+ return OffsetAPInt.getZExtValue();
+ }
+ return std::nullopt;
+ };
+
+ unsigned RootIdx = OriginalArgIndex;
+ uint64_t BaseOffset = 0;
+
+ if (Arg.hasAttribute(OriginalArgAttr)) {
+ Attribute Attr = F.getAttributeAtIndex(OriginalArgIndex, OriginalArgAttr);
+ (void)parseOriginalArgAttribute(Attr.getValueAsString(), RootIdx,
+ BaseOffset);
+ }
+
+ // Build maps from offset to load and load to offset. Skip splitting if any
+ // load has a variable-index GEP (can't compute constant offset).
+ DenseMap<uint64_t, LoadInst *> OffsetToLoad;
+ DenseMap<LoadInst *, uint64_t> LoadToOffset;
+ bool HasVariableIndexLoad = false;
+ for (LoadInst *LI : Loads) {
+ std::optional<uint64_t> Off = GetLoadOffset(LI);
+ if (!Off) {
+ HasVariableIndexLoad = true;
+ break;
+ }
+ OffsetToLoad[*Off] = LI;
+ LoadToOffset[LI] = *Off;
+ }
+
+ if (HasVariableIndexLoad) {
+ LLVM_DEBUG(dbgs() << "Skipping split for " << F.getName()
+ << ": load with variable-index GEP\n");
+ HandlePassthroughArg(Arg);
+ continue;
+ }
+
+ StructArgs.push_back(&Arg);
+ ArgToLoadsMap[&Arg] = Loads;
+ ArgToGEPsMap[&Arg] = GEPs;
+
+ SmallVector<FieldInfo, 8> Fields;
+
+ if (EnableKernargLayoutChange) {
+ // Layout change allowed: only keep used fields, sorted by offset
+ llvm::sort(Loads, [&](LoadInst *A, LoadInst *B) {
+ return LoadToOffset[A] < LoadToOffset[B];
+ });
+
+ for (LoadInst *LI : Loads) {
+ uint64_t LocalOff = LoadToOffset[LI];
+ Fields.push_back({LI->getType(), LocalOff, LI});
+ NewArgTypes.push_back(LI->getType());
+ uint64_t FinalOff = BaseOffset + LocalOff;
+ NewArgMappings.emplace_back(NewArgIndex, RootIdx, FinalOff);
+ ++NewArgIndex;
+ }
+ } else {
+ // Layout preserved: keep ALL leaf fields in original struct order
+ // Recursively flatten nested structs
+ collectLeafFields(STy, DL, /*BaseOffset=*/0, Fields);
+
+ // Build a map from (offset, type) to field index for matching
+ DenseMap<std::pair<uint64_t, Type *>, unsigned> OffsetTypeToField;
+ for (unsigned I = 0; I < Fields.size(); ++I)
+ OffsetTypeToField[{Fields[I].Offset, Fields[I].Ty}] = I;
+
+ // Verify all loads can be matched to a leaf field
+ bool AllLoadsMatch = true;
+ for (LoadInst *LI : Loads) {
+ uint64_t Off = LoadToOffset[LI];
+ auto Key = std::make_pair(Off, LI->getType());
+ if (!OffsetTypeToField.count(Key)) {
+ AllLoadsMatch = false;
+ break;
+ }
+ }
+
+ if (!AllLoadsMatch) {
+ LLVM_DEBUG(dbgs() << "Skipping split for " << F.getName()
+ << ": load type doesn't match leaf field type\n");
+ // Undo: remove from StructArgs, restore passthrough
+ StructArgs.pop_back();
+ ArgToLoadsMap.erase(&Arg);
+ ArgToGEPsMap.erase(&Arg);
+ HandlePassthroughArg(Arg);
+ continue;
+ }
+
+ // Associate loads with their corresponding fields
+ for (LoadInst *LI : Loads) {
+ uint64_t Off = LoadToOffset[LI];
+ auto Key = std::make_pair(Off, LI->getType());
+ unsigned FieldIdx = OffsetTypeToField[Key];
+ Fields[FieldIdx].Load = LI;
+ }
+
+ for (const FieldInfo &FI : Fields) {
+ NewArgTypes.push_back(FI.Ty);
+ uint64_t FinalOff = BaseOffset + FI.Offset;
+ NewArgMappings.emplace_back(NewArgIndex, RootIdx, FinalOff);
+ ++NewArgIndex;
+ }
+ }
+
+ ArgToFieldsMap[&Arg] = Fields;
+ ++OriginalArgIndex;
+ }
+
+ if (StructArgs.empty())
+ return nullptr;
+
+ if (!canPreloadSplitArgs(F, ST, ArgToFieldsMap)) {
+ LLVM_DEBUG(dbgs() << "Skipping split for " << F.getName()
+ << ": split args would not fit in preload SGPRs\n");
+ return nullptr;
+ }
+
+ LLVM_DEBUG(dbgs() << "Splitting kernel arguments for " << F.getName()
+ << "\n");
+
+ // Create new function
+ AttributeList OldAttrs = F.getAttributes();
+ AttributeSet FnAttrs = OldAttrs.getFnAttrs();
+ AttributeSet RetAttrs = OldAttrs.getRetAttrs();
+
+ FunctionType *NewFT =
+ FunctionType::get(F.getReturnType(), NewArgTypes, F.isVarArg());
+ Function *NewF =
+ Function::Create(NewFT, F.getLinkage(), F.getAddressSpace(), F.getName());
+ F.getParent()->getFunctionList().insert(F.getIterator(), NewF);
+ NewF->takeName(&F);
+ NewF->setVisibility(F.getVisibility());
+ if (F.hasComdat())
+ NewF->setComdat(F.getComdat());
+ NewF->setDSOLocal(F.isDSOLocal());
+ NewF->setUnnamedAddr(F.getUnnamedAddr());
+ NewF->setCallingConv(F.getCallingConv());
+
+ SmallVector<AttributeSet, 8> NewArgAttrSets;
+ NewArgIndex = 0;
+ for (Argument &Arg : F.args()) {
+ if (ArgToFieldsMap.count(&Arg)) {
+ for ([[maybe_unused]] const FieldInfo &FI : ArgToFieldsMap[&Arg]) {
+ NewArgAttrSets.push_back(AttributeSet());
+ ++NewArgIndex;
+ }
+ } else {
+ AttributeSet ArgAttrs = OldAttrs.getParamAttrs(Arg.getArgNo());
+ NewArgAttrSets.push_back(ArgAttrs);
+ ++NewArgIndex;
+ }
+ }
+
+ AttributeList NewAttrList =
+ AttributeList::get(F.getContext(), FnAttrs, RetAttrs, NewArgAttrSets);
+ NewF->setAttributes(NewAttrList);
+
+ // In layout-changing mode, add original-arg attributes so CLR can map
+ // split args back. In layout-preserving mode, we use a backup declaration
+ // instead, so no per-arg attributes needed.
+ if (EnableKernargLayoutChange) {
+ for (const auto &Info : NewArgMappings) {
+ unsigned NewArgIdx, RootArgIdx;
+ uint64_t Offset;
+ std::tie(NewArgIdx, RootArgIdx, Offset) = Info;
+ NewF->addParamAttr(
+ NewArgIdx,
+ Attribute::get(NewF->getContext(), OriginalArgAttr,
+ (Twine(RootArgIdx) + ":" + Twine(Offset)).str()));
+ }
+ }
+
+ LLVM_DEBUG(dbgs() << "New function signature:\n" << *NewF << '\n');
+
+ NewF->splice(NewF->begin(), &F);
+
+ DenseMap<Value *, Value *> VMap;
+ auto NewArgIt = NewF->arg_begin();
+ for (Argument &Arg : F.args()) {
+ if (ArgToFieldsMap.contains(&Arg)) {
+ for (const FieldInfo &FI : ArgToFieldsMap[&Arg]) {
+ Value *NewArg = &*NewArgIt++;
+ if (FI.Load) {
+ // This field has an associated load - map it
+ NewArg->takeName(FI.Load);
+ if (isa<PointerType>(NewArg->getType()) &&
+ isa<PointerType>(FI.Load->getType())) {
+ IRBuilder<> Builder(FI.Load);
+ Value *CastedArg = Builder.CreatePointerBitCastOrAddrSpaceCast(
+ NewArg, FI.Load->getType());
+ VMap[FI.Load] = CastedArg;
+ } else {
+ VMap[FI.Load] = NewArg;
+ }
+ }
+ // If FI.Load is null, this is an unused field - arg exists but unused
+ }
+ PoisonValue *PoisonArg = PoisonValue::get(Arg.getType());
+ Arg.replaceAllUsesWith(PoisonArg);
+ } else {
+ NewArgIt->takeName(&Arg);
+ Value *NewArg = &*NewArgIt;
+ if (isa<PointerType>(NewArg->getType()) &&
+ isa<PointerType>(Arg.getType())) {
+ IRBuilder<> Builder(&*NewF->begin()->begin());
+ Value *CastedArg =
+ Builder.CreatePointerBitCastOrAddrSpaceCast(NewArg, Arg.getType());
+ Arg.replaceAllUsesWith(CastedArg);
+ } else {
+ Arg.replaceAllUsesWith(NewArg);
+ }
+ ++NewArgIt;
+ }
+ }
+
+ for (auto &Entry : ArgToLoadsMap) {
+ for (LoadInst *LI : Entry.second) {
+ Value *NewArg = VMap.lookup(LI);
+ assert(NewArg && "Load not mapped to new argument - did we miss a "
+ "variable-index GEP check?");
+ LI->replaceAllUsesWith(NewArg);
+ LI->eraseFromParent();
+ }
+ }
+
+ for (auto &Entry : ArgToGEPsMap) {
+ for (GetElementPtrInst *GEP : Entry.second) {
+ GEP->replaceAllUsesWith(PoisonValue::get(GEP->getType()));
+ GEP->eraseFromParent();
+ }
+ }
+
+ LLVM_DEBUG(dbgs() << "Function after splitting:\n" << *NewF << '\n');
+
+ // In layout-preserving mode, create a backup declaration of the original
+ // function. This backup is used by the metadata streamer to emit the
+ // original kernel argument metadata (so CLR doesn't need patching).
+ if (!EnableKernargLayoutChange) {
+ std::string BackupName =
+ (Twine(OriginalKernelPrefix) + NewF->getName()).str();
+ // Create a declaration with the original signature (no body needed).
+ // This is used purely for metadata extraction, not code generation.
+ Function *BackupF =
+ Function::Create(F.getFunctionType(), GlobalValue::ExternalLinkage,
+ F.getAddressSpace(), BackupName, F.getParent());
+ // Copy function attributes for metadata extraction
+ BackupF->setAttributes(F.getAttributes());
+ BackupF->copyMetadata(&F, 0);
+ // Use C calling convention so it's not processed as a kernel
+ BackupF->setCallingConv(CallingConv::C);
+
+ // Store the backup name in the new function
+ NewF->addFnAttr(OriginalKernelAttr, BackupName);
+
+ LLVM_DEBUG(dbgs() << "Created backup declaration: " << BackupName << '\n');
+ }
+
+ F.replaceAllUsesWith(NewF);
+ F.eraseFromParent();
+
+ return NewF;
+}
+
+//===----------------------------------------------------------------------===//
+// Preload Kernel Arguments Logic
+//===----------------------------------------------------------------------===//
+
class AMDGPUPreloadKernelArgumentsLegacy : public ModulePass {
const GCNTargetMachine *TM;
@@ -285,11 +783,28 @@ static bool markKernelArgsAsInreg(Module &M, const TargetMachine &TM) {
SmallVector<Function *, 4> FunctionsToErase;
bool Changed = false;
+
+ // First, collect functions to process (split may modify the function list)
+ SmallVector<Function *, 16> FunctionsToProcess;
for (auto &F : M) {
const GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(F);
if (!ST.hasKernargPreload() ||
F.getCallingConv() != CallingConv::AMDGPU_KERNEL)
continue;
+ FunctionsToProcess.push_back(&F);
+ }
+
+ for (Function *FPtr : FunctionsToProcess) {
+ const GCNSubtarget &ST = TM.getSubtarget<GCNSubtarget>(*FPtr);
+
+ // Try to split byref struct arguments first. This may create a new
+ // function and erase the old one.
+ if (Function *NewF = trySplitKernelArguments(*FPtr, ST)) {
+ FPtr = NewF;
+ Changed = true;
+ }
+
+ Function &F = *FPtr;
PreloadKernelArgInfo PreloadInfo(F, ST);
uint64_t ExplicitArgOffset = 0;
@@ -299,9 +814,8 @@ static bool markKernelArgsAsInreg(Module &M, const TargetMachine &TM) {
unsigned NumPreloadedExplicitArgs = 0;
for (Argument &Arg : F.args()) {
// Avoid incompatible attributes and guard against running this pass
- // twice.
- //
- // TODO: Preload byref kernel arguments
+ // twice. Note: byref struct arguments are handled by splitting them
+ // into scalar arguments above via trySplitKernelArguments().
if (Arg.hasByRefAttr() || Arg.hasNestAttr() ||
Arg.hasAttribute("amdgpu-hidden-argument"))
break;
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-split-kernel-args-preserve-layout.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-split-kernel-args-preserve-layout.ll
new file mode 100644
index 0000000000000..dd9318ce5205d
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-split-kernel-args-preserve-layout.ll
@@ -0,0 +1,163 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --check-globals all --version 5
+; Test default behavior: layout preserved (all struct fields kept in original order)
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1250 -amdgpu-kernarg-preload-count=16 -passes=amdgpu-preload-kernel-arguments < %s | FileCheck %s
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1250 -amdgpu-kernarg-preload-count=16 -passes=amdgpu-preload-kernel-arguments < %s | llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1250 | FileCheck %s --check-prefix=ASM
+
+; Simple struct with 4 fields - all should be flattened even if only some are used
+%struct.S = type { i32, i8, i64, ptr }
+
+; Test: Only field at offset 8 (i64) is used, but ALL 4 fields should be preserved
+define amdgpu_kernel void @test_partial_use(ptr addrspace(4) byref(%struct.S) %s, ptr addrspace(1) %out) {
+; CHECK-LABEL: define amdgpu_kernel void @test_partial_use(
+; CHECK-SAME: i32 inreg [[TMP0:%.*]], i8 inreg [[TMP1:%.*]], i64 inreg [[S_L:%.*]], ptr inreg [[TMP2:%.*]], ptr addrspace(1) inreg [[OUT:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT: store i64 [[S_L]], ptr addrspace(1) [[OUT]], align 8
+; CHECK-NEXT: ret void
+;
+ %s.l.addr = getelementptr inbounds i8, ptr addrspace(4) %s, i64 8
+ %s.l = load i64, ptr addrspace(4) %s.l.addr, align 8
+ store i64 %s.l, ptr addrspace(1) %out, align 8
+ ret void
+}
+
+; Test: No fields used - all 4 fields should still be preserved as dead args
+define amdgpu_kernel void @test_unused_struct(ptr addrspace(4) byref(%struct.S) %unused) {
+; CHECK-LABEL: define amdgpu_kernel void @test_unused_struct(
+; CHECK-SAME: i32 inreg [[TMP0:%.*]], i8 inreg [[TMP1:%.*]], i64 inreg [[TMP2:%.*]], ptr inreg [[TMP3:%.*]]) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; Test: Multiple fields used - all fields preserved in original order
+define amdgpu_kernel void @test_multiple_fields(ptr addrspace(4) byref(%struct.S) %s, ptr addrspace(1) %out) {
+; CHECK-LABEL: define amdgpu_kernel void @test_multiple_fields(
+; CHECK-SAME: i32 inreg [[S_I:%.*]], i8 inreg [[TMP0:%.*]], i64 inreg [[S_L:%.*]], ptr inreg [[TMP1:%.*]], ptr addrspace(1) inreg [[OUT:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT: [[S_I_ZEXT:%.*]] = zext i32 [[S_I]] to i64
+; CHECK-NEXT: [[SUM:%.*]] = add i64 [[S_I_ZEXT]], [[S_L]]
+; CHECK-NEXT: store i64 [[SUM]], ptr addrspace(1) [[OUT]], align 8
+; CHECK-NEXT: ret void
+;
+ %s.i.addr = getelementptr inbounds i8, ptr addrspace(4) %s, i64 0
+ %s.i = load i32, ptr addrspace(4) %s.i.addr, align 4
+ %s.l.addr = getelementptr inbounds i8, ptr addrspace(4) %s, i64 8
+ %s.l = load i64, ptr addrspace(4) %s.l.addr, align 8
+ %s.i.zext = zext i32 %s.i to i64
+ %sum = add i64 %s.i.zext, %s.l
+ store i64 %sum, ptr addrspace(1) %out, align 8
+ ret void
+}
+
+; Nested struct: inner struct should be recursively flattened
+%struct.Inner = type { i32, i64 }
+%struct.Outer = type { i8, %struct.Inner, i32 }
+
+; Test: Access nested field - all leaf fields preserved
+define amdgpu_kernel void @test_nested_struct(ptr addrspace(4) byref(%struct.Outer) %o, ptr addrspace(1) %out) {
+; CHECK-LABEL: define amdgpu_kernel void @test_nested_struct(
+; CHECK-SAME: i8 inreg [[TMP0:%.*]], i32 inreg [[TMP1:%.*]], i64 inreg [[O_INNER_L:%.*]], i32 inreg [[TMP2:%.*]], ptr addrspace(1) inreg [[OUT:%.*]]) #[[ATTR3:[0-9]+]] {
+; CHECK-NEXT: store i64 [[O_INNER_L]], ptr addrspace(1) [[OUT]], align 8
+; CHECK-NEXT: ret void
+;
+ ; Access inner.l at offset 16 (outer.c=0, padding to 8, inner.i=8, inner.l=16)
+ %o.inner.l.addr = getelementptr inbounds i8, ptr addrspace(4) %o, i64 16
+ %o.inner.l = load i64, ptr addrspace(4) %o.inner.l.addr, align 8
+ store i64 %o.inner.l, ptr addrspace(1) %out, align 8
+ ret void
+}
+
+; Struct with vector field - vectors are kept as leaf types, not decomposed
+%struct.WithVector = type { i32, <4 x float>, i64 }
+
+; Test: Vector field should be kept as single vector argument
+define amdgpu_kernel void @test_vector_field(ptr addrspace(4) byref(%struct.WithVector) %s, ptr addrspace(1) %out) {
+ ; Access vector at offset 16
+; CHECK-LABEL: define amdgpu_kernel void @test_vector_field(
+; CHECK-SAME: i32 inreg [[TMP0:%.*]], <4 x float> inreg [[VEC:%.*]], i64 inreg [[TMP1:%.*]], ptr addrspace(1) inreg [[OUT:%.*]]) #[[ATTR4:[0-9]+]] {
+; CHECK-NEXT: store <4 x float> [[VEC]], ptr addrspace(1) [[OUT]], align 16
+; CHECK-NEXT: ret void
+;
+ %vec.addr = getelementptr inbounds i8, ptr addrspace(4) %s, i64 16
+ %vec = load <4 x float>, ptr addrspace(4) %vec.addr, align 16
+ store <4 x float> %vec, ptr addrspace(1) %out, align 16
+ ret void
+}
+
+; Struct with array field
+%struct.WithArray = type { i32, [4 x i32], i64 }
+
+; Test: Array field with constant index access - should be split, all elements flattened
+define amdgpu_kernel void @test_array_const_index(ptr addrspace(4) byref(%struct.WithArray) %s, ptr addrspace(1) %out) {
+ ; Access arr[2] at offset 4 + 2*4 = 12
+; CHECK-LABEL: define amdgpu_kernel void @test_array_const_index(
+; CHECK-SAME: i32 inreg [[TMP0:%.*]], i32 inreg [[TMP1:%.*]], i32 inreg [[TMP2:%.*]], i32 inreg [[ARR2:%.*]], i32 inreg [[TMP3:%.*]], i64 inreg [[TMP4:%.*]], ptr addrspace(1) inreg [[OUT:%.*]]) #[[ATTR5:[0-9]+]] {
+; CHECK-NEXT: store i32 [[ARR2]], ptr addrspace(1) [[OUT]], align 4
+; CHECK-NEXT: ret void
+;
+ %arr2.addr = getelementptr inbounds i8, ptr addrspace(4) %s, i64 12
+ %arr2 = load i32, ptr addrspace(4) %arr2.addr, align 4
+ store i32 %arr2, ptr addrspace(1) %out, align 4
+ ret void
+}
+
+; Test: Array field with variable index access - should NOT be split
+define amdgpu_kernel void @test_array_variable_index(ptr addrspace(4) byref(%struct.WithArray) %s, i32 %idx, ptr addrspace(1) %out) {
+ ; Access arr[idx] - non-constant index prevents splitting
+; CHECK-LABEL: define amdgpu_kernel void @test_array_variable_index(
+; CHECK-SAME: ptr addrspace(4) byref([[STRUCT_WITHARRAY:%.*]]) [[S:%.*]], i32 [[IDX:%.*]], ptr addrspace(1) [[OUT:%.*]]) #[[ATTR6:[0-9]+]] {
+; CHECK-NEXT: [[ARR_ADDR:%.*]] = getelementptr inbounds i8, ptr addrspace(4) [[S]], i64 4
+; CHECK-NEXT: [[ELEM_ADDR:%.*]] = getelementptr inbounds i32, ptr addrspace(4) [[ARR_ADDR]], i32 [[IDX]]
+; CHECK-NEXT: [[ELEM:%.*]] = load i32, ptr addrspace(4) [[ELEM_ADDR]], align 4
+; CHECK-NEXT: store i32 [[ELEM]], ptr addrspace(1) [[OUT]], align 4
+; CHECK-NEXT: ret void
+;
+ %arr.addr = getelementptr inbounds i8, ptr addrspace(4) %s, i64 4
+ %elem.addr = getelementptr inbounds i32, ptr addrspace(4) %arr.addr, i32 %idx
+ %elem = load i32, ptr addrspace(4) %elem.addr, align 4
+ store i32 %elem, ptr addrspace(1) %out, align 4
+ ret void
+}
+
+attributes #0 = { nounwind }
+
+; ASM-LABEL: .amdhsa_kernel test_partial_use
+; ASM: .amdhsa_user_sgpr_kernarg_preload_length 8
+; ASM: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+; ASM-LABEL: .amdhsa_kernel test_unused_struct
+; ASM: .amdhsa_user_sgpr_kernarg_preload_length 6
+; ASM: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+; ASM-LABEL: .amdhsa_kernel test_multiple_fields
+; ASM: .amdhsa_user_sgpr_kernarg_preload_length 8
+; ASM: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+; ASM-LABEL: .amdhsa_kernel test_nested_struct
+; ASM: .amdhsa_user_sgpr_kernarg_preload_length 8
+; ASM: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+; Vector field kept as single argument (not decomposed)
+; ASM-LABEL: .amdhsa_kernel test_vector_field
+; ASM: .amdhsa_user_sgpr_kernarg_preload_length 12
+; ASM: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+; ASM-LABEL: .amdhsa_kernel test_array_const_index
+; ASM: .amdhsa_user_sgpr_kernarg_preload_length 10
+; ASM: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+; Variable index prevents splitting - no preloading
+; ASM-LABEL: .amdhsa_kernel test_array_variable_index
+; ASM: .amdhsa_user_sgpr_kernarg_preload_length 0
+; ASM: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+; Verify that backup declarations do NOT generate any code or metadata
+; ASM-NOT: __amdgpu_orig_kernel
+; ASM-NOT: .amdhsa_kernel __amdgpu_orig_kernel
+;.
+; CHECK: attributes #[[ATTR0]] = { "amdgpu-original-kernel"="__amdgpu_orig_kernel_test_partial_use" "target-cpu"="gfx1250" }
+; CHECK: attributes #[[ATTR1]] = { "amdgpu-original-kernel"="__amdgpu_orig_kernel_test_unused_struct" "target-cpu"="gfx1250" }
+; CHECK: attributes #[[ATTR2]] = { "amdgpu-original-kernel"="__amdgpu_orig_kernel_test_multiple_fields" "target-cpu"="gfx1250" }
+; CHECK: attributes #[[ATTR3]] = { "amdgpu-original-kernel"="__amdgpu_orig_kernel_test_nested_struct" "target-cpu"="gfx1250" }
+; CHECK: attributes #[[ATTR4]] = { "amdgpu-original-kernel"="__amdgpu_orig_kernel_test_vector_field" "target-cpu"="gfx1250" }
+; CHECK: attributes #[[ATTR5]] = { "amdgpu-original-kernel"="__amdgpu_orig_kernel_test_array_const_index" "target-cpu"="gfx1250" }
+; CHECK: attributes #[[ATTR6]] = { "target-cpu"="gfx1250" }
+;.
diff --git a/llvm/test/CodeGen/AMDGPU/amdgpu-split-kernel-args.ll b/llvm/test/CodeGen/AMDGPU/amdgpu-split-kernel-args.ll
new file mode 100644
index 0000000000000..5381198d2ddd2
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/amdgpu-split-kernel-args.ll
@@ -0,0 +1,263 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-attributes --check-globals all --version 5
+; Test with layout change enabled (only used fields kept)
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1250 -amdgpu-kernarg-preload-count=16 -amdgpu-kernarg-layout-change -passes=amdgpu-preload-kernel-arguments < %s | FileCheck %s
+; RUN: opt -S -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1250 -amdgpu-kernarg-preload-count=16 -amdgpu-kernarg-layout-change -passes=amdgpu-preload-kernel-arguments < %s > %t.ll
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx1250 < %t.ll | FileCheck --check-prefix=GCN %s
+;
+; The LLVM IR is from the following HIP program:
+;
+; struct A {
+; int i;
+; char c;
+; long l;
+; int *p;
+; };
+
+; struct B {
+; char c;
+; A a1;
+; int i;
+; A a2;
+; };
+;
+; __global__ void test(int *out, int i, A a, char c, B b) {
+; *out = i + a.l + c + a.l + b.a1.c;
+; b.a2.p[2] = a.l + b.a2.c;
+;}
+;
+%struct.A = type { i32, i8, i64, ptr }
+%struct.B = type { i8, %struct.A, i32, %struct.A }
+
+; The "amdgpu-original-arg" function parameter attribute encodes how is the
+; argument split from the original kernel argument.
+;
+; Format: "amdgpu-original-arg"="OrigIndex:OrigOffset"
+; - OrigIndex: Index of the original kernel argument before splitting
+; - OrigOffset: Byte offset within the original argument
+
+;--- Main test case for successful split ---
+
+define amdgpu_kernel void @_Z4testPii1Ac1B(
+; CHECK-LABEL: define amdgpu_kernel void @_Z4testPii1Ac1B(
+; CHECK-SAME: ptr addrspace(1) inreg noundef writeonly captures(none) initializes((0, 4)) [[OUT:%.*]], i32 inreg noundef [[I:%.*]], i64 inreg "amdgpu-original-arg"="2:8" [[A_L:%.*]], i8 inreg noundef [[C:%.*]], i8 inreg "amdgpu-original-arg"="4:12" [[B_A1_C:%.*]], i8 inreg "amdgpu-original-arg"="4:44" [[B_A2_C:%.*]], ptr inreg "amdgpu-original-arg"="4:56" [[B_A2_P:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP1:%.*]] = addrspacecast ptr [[B_A2_P]] to ptr addrspace(1)
+; CHECK-NEXT: [[CONV:%.*]] = zext i32 [[I]] to i64
+; CHECK-NEXT: [[CONV3:%.*]] = sext i8 [[C]] to i64
+; CHECK-NEXT: [[CONV8:%.*]] = sext i8 [[B_A1_C]] to i64
+; CHECK-NEXT: [[FACTOR:%.*]] = shl i64 [[A_L]], 1
+; CHECK-NEXT: [[ADD4:%.*]] = add nsw i64 [[CONV]], [[CONV3]]
+; CHECK-NEXT: [[ADD6:%.*]] = add i64 [[ADD4]], [[FACTOR]]
+; CHECK-NEXT: [[ADD9:%.*]] = add i64 [[ADD6]], [[CONV8]]
+; CHECK-NEXT: [[CONV10:%.*]] = trunc i64 [[ADD9]] to i32
+; CHECK-NEXT: store i32 [[CONV10]], ptr addrspace(1) [[OUT]], align 4
+; CHECK-NEXT: [[B_A2_C_SEXT:%.*]] = sext i8 [[B_A2_C]] to i64
+; CHECK-NEXT: [[ADD14:%.*]] = add nsw i64 [[A_L]], [[B_A2_C_SEXT]]
+; CHECK-NEXT: [[CONV15:%.*]] = trunc i64 [[ADD14]] to i32
+; CHECK-NEXT: [[ARRAYIDX:%.*]] = getelementptr inbounds i32, ptr addrspace(1) [[TMP1]], i64 2
+; CHECK-NEXT: store i32 [[CONV15]], ptr addrspace(1) [[ARRAYIDX]], align 4
+; CHECK-NEXT: ret void
+;
+ ptr addrspace(1) noundef writeonly captures(none) initializes((0, 4)) %out,
+ i32 noundef %i,
+ ptr addrspace(4) noundef readonly byref(%struct.A) align 8 captures(none) %a,
+ i8 noundef %c,
+ ptr addrspace(4) noundef readonly byref(%struct.B) align 8 captures(none) %b
+) {
+entry:
+ ; Load a.l from struct A
+ %a.l.addr = getelementptr inbounds nuw i8, ptr addrspace(4) %a, i64 8
+ %a.l = load i64, ptr addrspace(4) %a.l.addr, align 8
+
+ ; Load b.a1.c from struct B
+ %b.a1.c.addr = getelementptr inbounds nuw i8, ptr addrspace(4) %b, i64 12
+ %b.a1.c = load i8, ptr addrspace(4) %b.a1.c.addr, align 4
+
+ ; Load b.a2.c from struct B
+ %b.a2.c.addr = getelementptr inbounds nuw i8, ptr addrspace(4) %b, i64 44
+ %b.a2.c = load i8, ptr addrspace(4) %b.a2.c.addr, align 4
+
+ ; Load b.a2.p from struct B
+ %b.a2.p.addr = getelementptr inbounds nuw i8, ptr addrspace(4) %b, i64 56
+ %b.a2.p = load ptr, ptr addrspace(4) %b.a2.p.addr, align 8
+
+ ; Cast b.a2.p to global address space
+ %b.a2.p.global = addrspacecast ptr %b.a2.p to ptr addrspace(1)
+
+ ; Compute i + a.l + c + a.l + b.a1.c
+ %i.zext = zext i32 %i to i64
+ %c.sext = sext i8 %c to i64
+ %b.a1.c.sext = sext i8 %b.a1.c to i64
+ %a.l.x2 = shl i64 %a.l, 1
+
+ %tmp_sum1 = add nsw i64 %i.zext, %c.sext
+ %tmp_sum2 = add i64 %tmp_sum1, %a.l.x2
+ %tmp_sum3 = add i64 %tmp_sum2, %b.a1.c.sext
+ %result = trunc i64 %tmp_sum3 to i32
+ store i32 %result, ptr addrspace(1) %out, align 4
+
+ ; Compute a.l + b.a2.c and store to b.a2.p[2]
+ %b.a2.c.sext = sext i8 %b.a2.c to i64
+ %sum_store = add nsw i64 %a.l, %b.a2.c.sext
+ %store_val = trunc i64 %sum_store to i32
+
+ %b.a2.p.elem2 = getelementptr inbounds i32, ptr addrspace(1) %b.a2.p.global, i64 2
+ store i32 %store_val, ptr addrspace(1) %b.a2.p.elem2, align 4
+
+ ret void
+}
+
+; --- Re-split test: arg #0 and #1 were split previously, arg #2 is byref struct ---
+; The second run of the pass must flatten arg #2 into three new parameters.
+;
+define amdgpu_kernel void @test_resplit(
+; CHECK-LABEL: define amdgpu_kernel void @test_resplit(
+; CHECK-SAME: i32 inreg "amdgpu-original-arg"="1:0" [[A_I:%.*]], i64 inreg "amdgpu-original-arg"="1:8" [[A_L:%.*]], i8 inreg "amdgpu-original-arg"="2:12" [[B_A1_C:%.*]], i8 inreg "amdgpu-original-arg"="2:44" [[B_A2_C:%.*]], ptr addrspace(1) inreg noundef "amdgpu-original-arg"="3:0" [[DST:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[A_I_ZEXT:%.*]] = zext i32 [[A_I]] to i64
+; CHECK-NEXT: [[B_A1_C_SEXT:%.*]] = sext i8 [[B_A1_C]] to i64
+; CHECK-NEXT: [[B_A2_C_SEXT:%.*]] = sext i8 [[B_A2_C]] to i64
+; CHECK-NEXT: [[SUM1:%.*]] = add i64 [[A_I_ZEXT]], [[A_L]]
+; CHECK-NEXT: [[SUM2:%.*]] = add i64 [[SUM1]], [[B_A1_C_SEXT]]
+; CHECK-NEXT: [[SUM3:%.*]] = add i64 [[SUM2]], [[B_A2_C_SEXT]]
+; CHECK-NEXT: [[TRUNC:%.*]] = trunc i64 [[SUM3]] to i32
+; CHECK-NEXT: store i32 [[TRUNC]], ptr addrspace(1) [[DST]], align 4
+; CHECK-NEXT: ret void
+;
+ i32 "amdgpu-original-arg"="1:0" %a.i, ; piece of original arg #1
+ i64 "amdgpu-original-arg"="1:8" %a.l, ; piece of original arg #1
+ ptr addrspace(4) noundef readonly
+ byref(%struct.B) align 8 %b, ; original arg #2 (to split)
+ ptr addrspace(1) noundef %dst) { ; ordinary output pointer
+entry:
+ ; load b.a1.c (offset 12)
+ %b.a1.c.addr = getelementptr inbounds i8, ptr addrspace(4) %b, i64 12
+ %b.a1.c = load i8, ptr addrspace(4) %b.a1.c.addr, align 4
+
+ ; load b.a2.c (offset 44)
+ %b.a2.c.addr = getelementptr inbounds i8, ptr addrspace(4) %b, i64 44
+ %b.a2.c = load i8, ptr addrspace(4) %b.a2.c.addr, align 4
+
+ ; sum up and store to the separate dst pointer
+ %a.i.zext = zext i32 %a.i to i64
+ %b.a1.c.sext = sext i8 %b.a1.c to i64
+ %b.a2.c.sext = sext i8 %b.a2.c to i64
+ %sum1 = add i64 %a.i.zext, %a.l
+ %sum2 = add i64 %sum1, %b.a1.c.sext
+ %sum3 = add i64 %sum2, %b.a2.c.sext
+ %sum.trunc = trunc i64 %sum3 to i32
+ store i32 %sum.trunc, ptr addrspace(1) %dst, align 4
+ ret void
+}
+
+; --- Additional test cases for passthrough logic ---
+
+; Test case for a struct argument that is never used.
+; Should not be split.
+define amdgpu_kernel void @test_unused_arg(ptr addrspace(4) byref(%struct.A) %unused_arg) {
+; CHECK-LABEL: define amdgpu_kernel void @test_unused_arg(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT: ret void
+;
+ ret void
+}
+
+; Test case for a pointer argument that does not have the 'byref' attribute.
+; Should not be split.
+define amdgpu_kernel void @test_no_byref_arg(ptr %ptr) {
+; CHECK-LABEL: define amdgpu_kernel void @test_no_byref_arg(
+; CHECK-SAME: ptr inreg [[PTR:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[VAL:%.*]] = load i32, ptr [[PTR]], align 4
+; CHECK-NEXT: ret void
+;
+ %val = load i32, ptr %ptr, align 4
+ ret void
+}
+
+; Test case for a 'byref' argument that points to a non-struct type.
+; Should not be split.
+define amdgpu_kernel void @test_byref_non_struct_arg(ptr byref(i32) %ptr) {
+; CHECK-LABEL: define amdgpu_kernel void @test_byref_non_struct_arg(
+; CHECK-SAME: ptr byref(i32) [[PTR:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[VAL:%.*]] = load i32, ptr [[PTR]], align 4
+; CHECK-NEXT: ret void
+;
+ %val = load i32, ptr %ptr, align 4
+ ret void
+}
+
+; Test case for an argument that is used by an unsupported instruction (a store).
+; Should not be split.
+define amdgpu_kernel void @test_unsupported_user(ptr byref(%struct.A) %a) {
+; CHECK-LABEL: define amdgpu_kernel void @test_unsupported_user(
+; CHECK-SAME: ptr byref([[STRUCT_A:%.*]]) [[A:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: store ptr null, ptr [[A]], align 8
+; CHECK-NEXT: ret void
+;
+ store ptr null, ptr %a, align 8
+ ret void
+}
+
+; Test case for a load from a GEP with a variable, non-constant offset.
+; Should not be split.
+define amdgpu_kernel void @test_variable_offset(ptr byref(%struct.A) %a, i32 %idx) {
+ ; GEP into the 'p' field (a ptr) with a variable index.
+; CHECK-LABEL: define amdgpu_kernel void @test_variable_offset(
+; CHECK-SAME: ptr byref([[STRUCT_A:%.*]]) [[A:%.*]], i32 [[IDX:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[P_FIELD_PTR:%.*]] = getelementptr inbounds [[STRUCT_A]], ptr [[A]], i32 [[IDX]]
+; CHECK-NEXT: [[P_FIELD_VAL:%.*]] = load ptr, ptr [[P_FIELD_PTR]], align 8
+; CHECK-NEXT: ret void
+;
+ %p_field_ptr = getelementptr inbounds %struct.A, ptr %a, i32 %idx
+ %p_field_val = load ptr, ptr %p_field_ptr, align 8
+ ret void
+}
+
+; Check kernel descriptor preload fields (comes first in assembly output)
+; GCN-LABEL: .amdhsa_kernel _Z4testPii1Ac1B
+; GCN: .amdhsa_user_sgpr_kernarg_preload_length 10
+; GCN: .amdhsa_user_sgpr_kernarg_preload_offset 0
+
+; Check HSA metadata for the main test kernel (comes later in assembly output)
+; GCN-LABEL: amdhsa.kernels:
+; GCN: - .address_space: global
+; GCN-NEXT: .name: out
+; GCN-NEXT: .offset: 0
+; GCN-NEXT: .size: 8
+; GCN-NEXT: .value_kind: global_buffer
+; GCN-NEXT: - .name: i
+; GCN-NEXT: .offset: 8
+; GCN-NEXT: .size: 4
+; GCN-NEXT: .value_kind: by_value
+; GCN-NEXT: - .name: a.l
+; GCN-NEXT: .offset: 16
+; GCN-NEXT: .original_arg_index: 2
+; GCN-NEXT: .original_arg_offset: 8
+; GCN-NEXT: .size: 8
+; GCN-NEXT: .value_kind: by_value
+; GCN-NEXT: - .name: c
+; GCN-NEXT: .offset: 24
+; GCN-NEXT: .size: 1
+; GCN-NEXT: .value_kind: by_value
+; GCN-NEXT: - .name: b.a1.c
+; GCN-NEXT: .offset: 25
+; GCN-NEXT: .original_arg_index: 4
+; GCN-NEXT: .original_arg_offset: 12
+; GCN-NEXT: .size: 1
+; GCN-NEXT: .value_kind: by_value
+; GCN-NEXT: - .name: b.a2.c
+; GCN-NEXT: .offset: 26
+; GCN-NEXT: .original_arg_index: 4
+; GCN-NEXT: .original_arg_offset: 44
+; GCN-NEXT: .size: 1
+; GCN-NEXT: .value_kind: by_value
+; GCN-NEXT: - .address_space: generic
+; GCN-NEXT: .name: b.a2.p
+; GCN-NEXT: .offset: 32
+; GCN-NEXT: .original_arg_index: 4
+; GCN-NEXT: .original_arg_offset: 56
+; GCN-NEXT: .size: 8
+; GCN-NEXT: .value_kind: global_buffer
+;.
+; CHECK: attributes #[[ATTR0]] = { "target-cpu"="gfx1250" }
+;.
More information about the llvm-commits
mailing list