[AMDGPU] Support preloading hidden kernel arguments #98861

kerbowa · 2024-07-15T06:42:59Z

Adds hidden kernel arguments to the function signature and marks them inreg if they should be preloaded into user SGPRs. The normal kernarg preloading logic then takes over with some additional checks for the correct implicitarg_ptr alignment.

Special care is needed so that metadata for the hidden arguments is not added twice when generating the code object.

llvmbot · 2024-07-15T06:43:31Z

@llvm/pr-subscribers-llvm-ir

@llvm/pr-subscribers-backend-amdgpu

Author: Austin Kerbow (kerbowa)

Changes

Adds hidden kernel arguments to the function signature and marks them inreg if they should be preloaded into user SGPRs. The normal kernarg preloading logic then takes over with some additional checks for the correct implicitarg_ptr alignment.

Special care is needed so that metadata for the hidden arguments is not added twice when generating the code object.

Patch is 64.19 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/98861.diff

8 Files Affected:

(modified) llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp (+3)
(modified) llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp (+178-2)
(modified) llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp (+2)
(modified) llvm/lib/Target/AMDGPU/SIISelLowering.cpp (+14)
(modified) llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp (+9)
(modified) llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h (+3)
(added) llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll (+698)
(modified) llvm/test/CodeGen/AMDGPU/preload-kernargs.ll (+37-39)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
index 1aaf514ae8f62..ef049505dc628 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp
@@ -283,6 +283,9 @@ void MetadataStreamerMsgPackV4::emitKernelArg(const Argument &Arg,
   else if (Arg.hasName())
     Name = Arg.getName();
 
+  if (Name.starts_with("_hidden"))
+    return;
+
   StringRef TypeName;
   Node = Func->getMetadata("kernel_arg_type");
   if (Node && ArgNo < Node->getNumOperands())
diff --git a/llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp b/llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp
index 83cce6021693a..e9ba41ef8681f 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp
@@ -13,6 +13,7 @@
 
 #include "AMDGPU.h"
 #include "GCNSubtarget.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/TargetPassConfig.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/IntrinsicsAMDGPU.h"
@@ -31,9 +32,88 @@ class PreloadKernelArgInfo {
   const GCNSubtarget &ST;
   unsigned NumFreeUserSGPRs;
 
-public:
-  SmallVector<llvm::Metadata *, 8> KernelArgMetadata;
+  enum HiddenArg : unsigned {
+    HIDDEN_BLOCK_COUNT_X,
+    HIDDEN_BLOCK_COUNT_Y,
+    HIDDEN_BLOCK_COUNT_Z,
+    HIDDEN_GROUP_SIZE_X,
+    HIDDEN_GROUP_SIZE_Y,
+    HIDDEN_GROUP_SIZE_Z,
+    HIDDEN_REMAINDER_X,
+    HIDDEN_REMAINDER_Y,
+    HIDDEN_REMAINDER_Z,
+    END_HIDDEN_ARGS
+  };
+
+  struct HiddenArgInfo {
+    unsigned Offset;
+    unsigned Size;
+    const char *Name;
+  };
+
+  static constexpr HiddenArgInfo HiddenArgs[END_HIDDEN_ARGS] = {
+      {0, 4, "_hidden_block_count_x"}, {4, 4, "_hidden_block_count_y"},
+      {8, 4, "_hidden_block_count_z"}, {12, 2, "_hidden_group_size_x"},
+      {14, 2, "_hidden_group_size_y"}, {16, 2, "_hidden_group_size_z"},
+      {18, 2, "_hidden_remainder_x"},  {20, 2, "_hidden_remainder_y"},
+      {22, 2, "_hidden_remainder_z"}};
+
+  static HiddenArg getHiddenArgIndexFromOffset(unsigned Offset) {
+    for (unsigned I = 0; I < END_HIDDEN_ARGS; ++I)
+      if (HiddenArgs[I].Offset == Offset)
+        return static_cast<HiddenArg>(I);
+
+    llvm_unreachable("Unexpected hidden argument offset.");
+  }
+
+  static Type *getHiddenArgType(LLVMContext &Ctx, HiddenArg HA) {
+    if (HA < END_HIDDEN_ARGS)
+      return Type::getIntNTy(Ctx, HiddenArgs[HA].Size * 8);
+
+    llvm_unreachable("Unexpected hidden argument.");
+  }
+
+  static const char *getHiddenArgName(HiddenArg HA) {
+    if (HA < END_HIDDEN_ARGS) {
+      return HiddenArgs[HA].Name;
+    }
+    llvm_unreachable("Unexpected hidden argument.");
+  }
 
+  Function *cloneFunctionWithPreloadImplicitArgs() {
+    FunctionType *FT = F.getFunctionType();
+    std::vector<Type *> FTypes(FT->param_begin(), FT->param_end());
+    for (unsigned I = 0; I < END_HIDDEN_ARGS; ++I)
+      FTypes.push_back(getHiddenArgType(F.getContext(), HiddenArg(I)));
+
+    FunctionType *NFT =
+        FunctionType::get(FT->getReturnType(), FTypes, FT->isVarArg());
+    Function *NF =
+        Function::Create(NFT, F.getLinkage(), F.getAddressSpace(), F.getName());
+
+    NF->copyAttributesFrom(&F);
+    NF->copyMetadata(&F, 0);
+    NF->setIsNewDbgInfoFormat(F.IsNewDbgInfoFormat);
+
+    F.getParent()->getFunctionList().insert(F.getIterator(), NF);
+    NF->takeName(&F);
+    assert(F.use_empty());
+    NF->splice(NF->begin(), &F);
+
+    Function::arg_iterator NFArg = NF->arg_begin();
+    for (Argument &Arg : F.args()) {
+      Arg.replaceAllUsesWith(&*NFArg);
+      NFArg->takeName(&Arg);
+      ++NFArg;
+    }
+
+    for (unsigned I = 0; I < END_HIDDEN_ARGS; ++I)
+      NFArg++->setName(getHiddenArgName(HiddenArg(I)));
+
+    return NF;
+  }
+
+public:
   PreloadKernelArgInfo(Function &F, const GCNSubtarget &ST) : F(F), ST(ST) {
     setInitialFreeUserSGPRsCount();
   }
@@ -64,6 +144,94 @@ class PreloadKernelArgInfo {
     NumFreeUserSGPRs -= (NumPreloadSGPRs + PaddingSGPRs);
     return true;
   }
+
+  // Try to allocate SGPRs to preload implicit kernel arguments.
+  void tryAllocImplicitArgPreloadSGPRs(uint64_t ImplicitArgsBaseOffset,
+                                       IRBuilder<> &Builder) {
+    StringRef Name = Intrinsic::getName(Intrinsic::amdgcn_implicitarg_ptr);
+    Function *ImplicitArgPtr = F.getParent()->getFunction(Name);
+    if (!ImplicitArgPtr)
+      return;
+
+    const DataLayout &DL = F.getParent()->getDataLayout();
+    // Pair is the load and the load offset.
+    SmallVector<std::pair<LoadInst *, unsigned>, 4> ImplicitArgLoads;
+    for (auto *U : ImplicitArgPtr->users()) {
+      Instruction *CI = dyn_cast<Instruction>(U);
+      if (!CI || CI->getParent()->getParent() != &F)
+        continue;
+
+      for (auto *U : CI->users()) {
+        int64_t Offset = 0;
+        auto *Load = dyn_cast<LoadInst>(U); // Load from ImplicitArgPtr?
+        if (!Load) {
+          if (GetPointerBaseWithConstantOffset(U, Offset, DL) != CI)
+            continue;
+
+          Load = dyn_cast<LoadInst>(*U->user_begin()); // Load from GEP?
+        }
+
+        if (!Load || !Load->isSimple())
+          continue;
+
+        // FIXME: Expand to handle 64-bit implicit args and large merged loads.
+        unsigned LoadSize = Load->getType()->getScalarSizeInBits();
+        if (LoadSize != 32 && LoadSize != 16)
+          continue;
+
+        ImplicitArgLoads.push_back(std::make_pair(Load, Offset));
+      }
+    }
+
+    if (ImplicitArgLoads.empty())
+      return;
+
+    // Allocate loads in order of offset. We need to be sure that the implicit
+    // argument can actually be preloaded.
+    std::sort(ImplicitArgLoads.begin(), ImplicitArgLoads.end(),
+              [](const std::pair<LoadInst *, unsigned> &A,
+                 const std::pair<LoadInst *, unsigned> &B) {
+                return A.second < B.second;
+              });
+
+    uint64_t LastExplicitArgOffset = ImplicitArgsBaseOffset;
+    bool AddedHiddenArgsToSignature = false;
+    Function *NF = nullptr;
+    unsigned LastPreloadIndex = 0;
+    for (const auto &Load : ImplicitArgLoads) {
+      LoadInst *LoadInst = Load.first;
+      Type *LoadType = LoadInst->getType();
+      auto LoadOffset = Load.second;
+      unsigned LoadSize = DL.getTypeStoreSize(LoadType);
+      // If we fail to preload any implicit argument we know we don't have SGPRs
+      // to preload any subsequent ones with larger offsets.
+      if (!tryAllocPreloadSGPRs(LoadSize, LoadOffset + ImplicitArgsBaseOffset,
+                                LastExplicitArgOffset))
+        break;
+
+      if (!AddedHiddenArgsToSignature) {
+        NF = cloneFunctionWithPreloadImplicitArgs();
+        AddedHiddenArgsToSignature = true;
+      }
+
+      LastExplicitArgOffset = LoadOffset + LoadSize;
+      unsigned HiddenArgIndex = getHiddenArgIndexFromOffset(LoadOffset);
+      assert(NF);
+      unsigned Index = NF->arg_size() - END_HIDDEN_ARGS + HiddenArgIndex;
+      Argument *Arg = NF->getArg(Index);
+      LoadInst->replaceAllUsesWith(Arg);
+      if (Index > HiddenArgIndex)
+        LastPreloadIndex = HiddenArgIndex;
+    }
+
+    // Ensure all hidden arguments up to the final preload are also
+    // preloaded, even if some are unused.
+    for (unsigned I = 0; I <= LastPreloadIndex; ++I)
+      NF->getArg(NF->arg_size() - END_HIDDEN_ARGS + I)
+          ->addAttr(Attribute::InReg);
+
+    F.removeFromParent();
+  }
 };
 
 class AMDGPULowerKernelArguments : public FunctionPass {
@@ -281,6 +449,14 @@ static bool lowerKernelArguments(Function &F, const TargetMachine &TM) {
   KernArgSegment->addRetAttr(
       Attribute::getWithAlignment(Ctx, std::max(KernArgBaseAlign, MaxAlign)));
 
+  if (InPreloadSequence) {
+    uint64_t ImplicitArgsBaseOffset =
+        alignTo(ExplicitArgOffset, ST.getAlignmentForImplicitArgPtr()) +
+        BaseOffset;
+    PreloadInfo.tryAllocImplicitArgPreloadSGPRs(ImplicitArgsBaseOffset,
+                                                Builder);
+  }
+
   return true;
 }
 
diff --git a/llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp b/llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp
index 21fe1bc31a27e..0b4565a8fdc2b 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp
@@ -601,6 +601,8 @@ uint64_t AMDGPUSubtarget::getExplicitKernArgSize(const Function &F,
   MaxAlign = Align(1);
 
   for (const Argument &Arg : F.args()) {
+    if (Arg.getName().starts_with("_hidden"))
+      continue;
     const bool IsByRef = Arg.hasByRefAttr();
     Type *ArgTy = IsByRef ? Arg.getParamByRefType() : Arg.getType();
     Align Alignment = DL.getValueOrABITypeAlignment(
diff --git a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
index a733295d2a511..954cab0702afb 100644
--- a/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
+++ b/llvm/lib/Target/AMDGPU/SIISelLowering.cpp
@@ -2496,6 +2496,7 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
   GCNUserSGPRUsageInfo &SGPRInfo = Info.getUserSGPRInfo();
   bool InPreloadSequence = true;
   unsigned InIdx = 0;
+  bool AlignedForImplictArgs = false;
   for (auto &Arg : F.args()) {
     if (!InPreloadSequence || !Arg.hasInRegAttr())
       break;
@@ -2518,6 +2519,19 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
       unsigned NumAllocSGPRs =
           alignTo(ArgLoc.getLocVT().getFixedSizeInBits(), 32) / 32;
 
+      if (!AlignedForImplictArgs && Arg.getName().starts_with("_hidden")) {
+        unsigned OffsetBefore = LastExplicitArgOffset;
+        LastExplicitArgOffset = alignTo(
+            LastExplicitArgOffset, Subtarget->getAlignmentForImplicitArgPtr());
+        if (OffsetBefore != LastExplicitArgOffset) {
+          unsigned PaddingSGPRs =
+              alignTo(LastExplicitArgOffset - OffsetBefore, 4) / 4;
+          Info.allocateUserSGPRs(PaddingSGPRs);
+          ArgOffset += PaddingSGPRs * 4;
+        }
+        AlignedForImplictArgs = true;
+      }
+
       // Arg is preloaded into the previous SGPR.
       if (ArgLoc.getLocVT().getStoreSize() < 4 && Alignment < 4) {
         Info.getArgInfo().PreloadKernArgs[InIdx].Regs.push_back(
diff --git a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp
index d9db0f7a4f531..35f41a3b87c3a 100644
--- a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp
+++ b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp
@@ -278,6 +278,15 @@ SmallVectorImpl<MCRegister> *SIMachineFunctionInfo::addPreloadedKernArg(
   return &ArgInfo.PreloadKernArgs[KernArgIdx].Regs;
 }
 
+bool SIMachineFunctionInfo::allocateUserSGPRs(
+    unsigned Number) {
+  if (Number <= getNumUserSGPRs())
+    return false;
+
+  NumUserSGPRs = Number;
+  return true;
+}
+
 void SIMachineFunctionInfo::allocateWWMSpill(MachineFunction &MF, Register VGPR,
                                              uint64_t Size, Align Alignment) {
   // Skip if it is an entry function or the register is already added.
diff --git a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
index 7af5e7388f841..f2ff919a4bc0e 100644
--- a/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
+++ b/llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.h
@@ -760,6 +760,9 @@ class SIMachineFunctionInfo final : public AMDGPUMachineFunction,
                       unsigned AllocSizeDWord, int KernArgIdx,
                       int PaddingSGPRs);
 
+  /// Reserve up to \p Number of user SGPRs.
+  bool allocateUserSGPRs(unsigned Number);
+
   /// Increment user SGPRs used for padding the argument list only.
   Register addReservedUserSGPR() {
     Register Next = getNextUserSGPR();
diff --git a/llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll b/llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll
new file mode 100644
index 0000000000000..456c9e3dc7ee5
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll
@@ -0,0 +1,698 @@
+; NOTE: Assertions have been autogenerated by utils/update_llc_test_checks.py UTC_ARGS: --version 4
+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx940 -verify-machineinstrs < %s | FileCheck -check-prefixes=GFX940-NO-PRELOAD %s
+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx940 -amdgpu-kernarg-preload-count=16 -verify-machineinstrs < %s | FileCheck -check-prefixes=GFX940-PRELOAD %s
+
+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx90a -verify-machineinstrs < %s | FileCheck -check-prefixes=GFX90a-NO-PRELOAD %s
+; RUN: llc -mtriple=amdgcn--amdhsa -mcpu=gfx90a -amdgpu-kernarg-preload-count=16 -verify-machineinstrs < %s | FileCheck -check-prefixes=GFX90a-PRELOAD %s
+
+define amdgpu_kernel void @preload_block_count_x(ptr addrspace(1) %out) {
+; GFX940-NO-PRELOAD-LABEL: preload_block_count_x:
+; GFX940-NO-PRELOAD:       ; %bb.0:
+; GFX940-NO-PRELOAD-NEXT:    s_load_dword s4, s[0:1], 0x8
+; GFX940-NO-PRELOAD-NEXT:    s_load_dwordx2 s[2:3], s[0:1], 0x0
+; GFX940-NO-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX940-NO-PRELOAD-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX940-NO-PRELOAD-NEXT:    v_mov_b32_e32 v1, s4
+; GFX940-NO-PRELOAD-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
+; GFX940-NO-PRELOAD-NEXT:    s_endpgm
+;
+; GFX940-PRELOAD-LABEL: preload_block_count_x:
+; GFX940-PRELOAD:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX940-PRELOAD-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX940-PRELOAD-NEXT:  ; %bb.0:
+; GFX940-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX940-PRELOAD-NEXT:    v_mov_b32_e32 v1, s4
+; GFX940-PRELOAD-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
+; GFX940-PRELOAD-NEXT:    s_endpgm
+;
+; GFX90a-NO-PRELOAD-LABEL: preload_block_count_x:
+; GFX90a-NO-PRELOAD:       ; %bb.0:
+; GFX90a-NO-PRELOAD-NEXT:    s_load_dword s2, s[4:5], 0x8
+; GFX90a-NO-PRELOAD-NEXT:    s_load_dwordx2 s[0:1], s[4:5], 0x0
+; GFX90a-NO-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX90a-NO-PRELOAD-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX90a-NO-PRELOAD-NEXT:    v_mov_b32_e32 v1, s2
+; GFX90a-NO-PRELOAD-NEXT:    global_store_dword v0, v1, s[0:1]
+; GFX90a-NO-PRELOAD-NEXT:    s_endpgm
+;
+; GFX90a-PRELOAD-LABEL: preload_block_count_x:
+; GFX90a-PRELOAD:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX90a-PRELOAD-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX90a-PRELOAD-NEXT:  ; %bb.0:
+; GFX90a-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX90a-PRELOAD-NEXT:    v_mov_b32_e32 v1, s8
+; GFX90a-PRELOAD-NEXT:    global_store_dword v0, v1, s[6:7]
+; GFX90a-PRELOAD-NEXT:    s_endpgm
+  %imp_arg_ptr = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %load = load i32, ptr addrspace(4) %imp_arg_ptr
+  store i32 %load, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @preload_block_count_y(ptr addrspace(1) %out) {
+; GFX940-NO-PRELOAD-LABEL: preload_block_count_y:
+; GFX940-NO-PRELOAD:       ; %bb.0:
+; GFX940-NO-PRELOAD-NEXT:    s_load_dword s4, s[0:1], 0xc
+; GFX940-NO-PRELOAD-NEXT:    s_load_dwordx2 s[2:3], s[0:1], 0x0
+; GFX940-NO-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX940-NO-PRELOAD-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX940-NO-PRELOAD-NEXT:    v_mov_b32_e32 v1, s4
+; GFX940-NO-PRELOAD-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
+; GFX940-NO-PRELOAD-NEXT:    s_endpgm
+;
+; GFX940-PRELOAD-LABEL: preload_block_count_y:
+; GFX940-PRELOAD:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX940-PRELOAD-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX940-PRELOAD-NEXT:  ; %bb.0:
+; GFX940-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX940-PRELOAD-NEXT:    v_mov_b32_e32 v1, s5
+; GFX940-PRELOAD-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
+; GFX940-PRELOAD-NEXT:    s_endpgm
+;
+; GFX90a-NO-PRELOAD-LABEL: preload_block_count_y:
+; GFX90a-NO-PRELOAD:       ; %bb.0:
+; GFX90a-NO-PRELOAD-NEXT:    s_load_dword s2, s[4:5], 0xc
+; GFX90a-NO-PRELOAD-NEXT:    s_load_dwordx2 s[0:1], s[4:5], 0x0
+; GFX90a-NO-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX90a-NO-PRELOAD-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX90a-NO-PRELOAD-NEXT:    v_mov_b32_e32 v1, s2
+; GFX90a-NO-PRELOAD-NEXT:    global_store_dword v0, v1, s[0:1]
+; GFX90a-NO-PRELOAD-NEXT:    s_endpgm
+;
+; GFX90a-PRELOAD-LABEL: preload_block_count_y:
+; GFX90a-PRELOAD:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX90a-PRELOAD-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX90a-PRELOAD-NEXT:  ; %bb.0:
+; GFX90a-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX90a-PRELOAD-NEXT:    v_mov_b32_e32 v1, s9
+; GFX90a-PRELOAD-NEXT:    global_store_dword v0, v1, s[6:7]
+; GFX90a-PRELOAD-NEXT:    s_endpgm
+  %imp_arg_ptr = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %gep = getelementptr i8, ptr addrspace(4) %imp_arg_ptr, i32 4
+  %load = load i32, ptr addrspace(4) %gep
+  store i32 %load, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @preload_block_count_z(ptr addrspace(1) %out) {
+; GFX940-NO-PRELOAD-LABEL: preload_block_count_z:
+; GFX940-NO-PRELOAD:       ; %bb.0:
+; GFX940-NO-PRELOAD-NEXT:    s_load_dword s4, s[0:1], 0x10
+; GFX940-NO-PRELOAD-NEXT:    s_load_dwordx2 s[2:3], s[0:1], 0x0
+; GFX940-NO-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX940-NO-PRELOAD-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX940-NO-PRELOAD-NEXT:    v_mov_b32_e32 v1, s4
+; GFX940-NO-PRELOAD-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
+; GFX940-NO-PRELOAD-NEXT:    s_endpgm
+;
+; GFX940-PRELOAD-LABEL: preload_block_count_z:
+; GFX940-PRELOAD:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX940-PRELOAD-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX940-PRELOAD-NEXT:  ; %bb.0:
+; GFX940-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX940-PRELOAD-NEXT:    v_mov_b32_e32 v1, s6
+; GFX940-PRELOAD-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
+; GFX940-PRELOAD-NEXT:    s_endpgm
+;
+; GFX90a-NO-PRELOAD-LABEL: preload_block_count_z:
+; GFX90a-NO-PRELOAD:       ; %bb.0:
+; GFX90a-NO-PRELOAD-NEXT:    s_load_dword s2, s[4:5], 0x10
+; GFX90a-NO-PRELOAD-NEXT:    s_load_dwordx2 s[0:1], s[4:5], 0x0
+; GFX90a-NO-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX90a-NO-PRELOAD-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX90a-NO-PRELOAD-NEXT:    v_mov_b32_e32 v1, s2
+; GFX90a-NO-PRELOAD-NEXT:    global_store_dword v0, v1, s[0:1]
+; GFX90a-NO-PRELOAD-NEXT:    s_endpgm
+;
+; GFX90a-PRELOAD-LABEL: preload_block_count_z:
+; GFX90a-PRELOAD:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX90a-PRELOAD-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX90a-PRELOAD-NEXT:  ; %bb.0:
+; GFX90a-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX90a-PRELOAD-NEXT:    v_mov_b32_e32 v1, s10
+; GFX90a-PRELOAD-NEXT:    global_store_dword v0, v1, s[6:7]
+; GFX90a-PRELOAD-NEXT:    s_endpgm
+  %imp_arg_ptr = call ptr addrspace(4) @llvm.amdgcn.implicitarg.ptr()
+  %gep = getelementptr i8, ptr addrspace(4) %imp_arg_ptr, i32 8
+  %load = load i32, ptr addrspace(4) %gep
+  store i32 %load, ptr addrspace(1) %out
+  ret void
+}
+
+define amdgpu_kernel void @preload_block_count_x_imparg_align_ptr_i8(ptr addrspace(1) %out, i8 %val) {
+; GFX940-NO-PRELOAD-LABEL: preload_block_count_x_imparg_align_ptr_i8:
+; GFX940-NO-PRELOAD:       ; %bb.0:
+; GFX940-NO-PRELOAD-NEXT:    s_load_dword s4, s[0:1], 0x8
+; GFX940-NO-PRELOAD-NEXT:    s_load_dword s5, s[0:1], 0x10
+; GFX940-NO-PRELOAD-NEXT:    s_load_dwordx2 s[2:3], s[0:1], 0x0
+; GFX940-NO-PRELOAD-NEXT:    v_mov_b32_e32 v0, 0
+; GFX940-NO-PRELOAD-NEXT:    s_waitcnt lgkmcnt(0)
+; GFX940-NO-PRELOAD-NEXT:    s_and_b32 s0, s4, 0xff
+; GFX940-NO-PRELOAD-NEXT:    s_add_i32 s0, s5, s0
+; GFX940-NO-PRELOAD-NEXT:    v_mov_b32_e32 v1, s0
+; GFX940-NO-PRELOAD-NEXT:    global_store_dword v0, v1, s[2:3] sc0 sc1
+; GFX940-NO-PRELOAD-NEXT:    s_endpgm
+;
+; GFX940-PRELOAD-LABEL: preload_block_count_x_imparg_align_ptr_i8:
+; GFX940-PRELOAD:         s_trap 2 ; Kernarg preload header. Trap with incompatible firmware that doesn't support preloading kernel arguments.
+; GFX940-PRELOAD-NEXT:    .fill 63, 4, 0xbf800000 ; s_nop 0
+; GFX940-PRELOAD-NEXT:  ; %bb.0:
+; GFX940-PRELOAD-NEXT:    s_and_b32 s0, s4, 0xff
+; GFX940-PRELOAD-NEX...
[truncated]

kerbowa · 2024-07-15T06:46:01Z

WIP since I still have some concerns about the function cloning and whether we can do it in a function pass properly, and also about having what are essentially reserved parameter name prefixes. I'm thinking about adding metadata so that a conflict wouldn't break anything - otherwise we would need to explicitly disallow conflicts with reserved names, but I'm not doing that currently.

github-actions · 2024-07-15T06:46:48Z

✅ With the latest revision this PR passed the C/C++ code formatter.

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

llvm/test/CodeGen/AMDGPU/preload-kernargs.ll

llvm/lib/Target/AMDGPU/AMDGPUHSAMetadataStreamer.cpp

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

arsenm · 2024-08-19T14:38:35Z

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

+        if (!Load || !Load->isSimple())
+          continue;


I can't think of a reason we wouldn't do this for volatile or atomic loads, but I guess can remove this later

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs-IR-lowering.ll

arsenm · 2024-08-19T14:49:16Z

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -2533,6 +2534,23 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
      unsigned NumAllocSGPRs =
          alignTo(ArgLoc.getLocVT().getFixedSizeInBits(), 32) / 32;

+      // Add padding SPGR to fix alignment for hidden arguments.


Missing the globalisel equvalent of this handling

Globalisel isn't supported with kernarg preloading. I'll work on adding support after this lands with #96030 being the first pre-req, it's a larger change.

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

arsenm · 2024-08-19T14:50:24Z

llvm/test/CodeGen/AMDGPU/preload-implicit-kernargs.ll

+  store <3 x i32> %ins.2, ptr addrspace(1) %out
+  ret void
+}
+


Are there tests somewhere where we have too many inreg arguments that cannot fit in user SGPRs during codegen?

There are many examples of this in preload-kernargs.ll with normal arguments. I added a test here where we don't have enough user SGPRs to preload an implicit_arg.

arsenm · 2024-09-25T10:30:24Z

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

+    StringRef Name = Intrinsic::getName(Intrinsic::amdgcn_implicitarg_ptr);
+    Function *ImplicitArgPtr = F.getParent()->getFunction(Name);
+    if (!ImplicitArgPtr)
+      return;


We really ought to have a version of getDeclaration that won't insert the declaration if it's not already there

arsenm · 2024-09-25T10:37:39Z

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

@@ -2534,6 +2535,20 @@ void SITargetLowering::allocatePreloadKernArgSGPRs(
      unsigned NumAllocSGPRs =
          alignTo(ArgLoc.getLocVT().getFixedSizeInBits(), 32) / 32;

+      // Add padding SPGR to fix alignment for hidden arguments.
+      if (!AlignedForImplictArgs && Arg.hasAttribute("amdgpu-work-group-id")) {


There is no such attribute as "amdgpu-work-group-id", all of this is dead?

Thanks for catching this, it must have become dead at some point in the review process and the tests autoupdated without noticing.

arsenm · 2024-10-01T10:47:26Z

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

+
+    AttrBuilder AB(Ctx);
+    AB.addAttribute(Attribute::InReg);
+    AB.addAttribute("amdgpu-hidden-argument");


Can you document this attribute in AMDGPUUsage?

arsenm · 2024-10-01T10:53:20Z

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

+    llvm_unreachable("Unexpected hidden argument.");
+  }
+
+  // Clones the function after adding implicit arguments to the argument list


Does this need to defend against arguments already marked as hidden arguments, in case the pass runs twice? i.e. running this pass a second time should be a no-op

arsenm · 2024-10-01T10:55:21Z

llvm/lib/Target/AMDGPU/SIISelLowering.cpp

+        if (OffsetBefore != LastExplicitArgOffset) {
+          unsigned PaddingSGPRs =
+              alignTo(LastExplicitArgOffset - OffsetBefore, 4) / 4;
+          Info.allocateUserSGPRs(*Subtarget, PaddingSGPRs);


Should check if this failed?

arsenm · 2024-10-04T13:37:36Z

llvm/docs/AMDGPUUsage.rst

+With code object V5 and later, hidden kernel arguments that are normally accessed
+through the Implicit Argument Ptr, may be preloaded into User SGPRs. These
+arguments are added to the kernel function signature and are marked with the
+attribute "amdgpu-hidden-argument". (See :ref:`amdgpu-llvm-ir-attributes-table`).


arsenm · 2024-10-04T13:37:49Z

llvm/include/llvm/IR/Function.h

@@ -433,6 +433,9 @@ class LLVM_ABI Function : public GlobalObject, public ilist_node<Function> {
  /// check if an attributes is in the list of attributes.
  bool hasParamAttribute(unsigned ArgNo, Attribute::AttrKind Kind) const;

+  /// check if an attributes is in the list of attributes.


Suggested change

/// check if an attributes is in the list of attributes.

/// Check if an attributes is in the list of attributes.

arsenm · 2024-10-04T13:38:35Z

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp

+
+    NF->copyAttributesFrom(&F);
+    NF->copyMetadata(&F, 0);
+    NF->setIsNewDbgInfoFormat(F.IsNewDbgInfoFormat);


This should probably be in copyAttributesFrom

Adds hidden kernel arguments to the function signature and marks them inreg if they should be preloaded into user SGPRs. The normal kernarg preloading logic then takes over with some additional checks for the correct implicitarg_ptr alignment. Special care is needed so that metadata for the hidden arguments is not added twice when generating the code object.

* commit 'FETCH_HEAD': [X86] getIntImmCostInst - pull out repeated Imm.getBitWidth() calls. NFC. [X86] Add test coverage for llvm#111323 [Driver] Use empty multilib file in another test (llvm#111352) [clang][OpenMP][test] Use x86_64-linux-gnu triple for test referencing avx512f feature (llvm#111337) [doc] Fix Kaleidoscope tutorial chapter 3 code snippet and full listing discrepancies (llvm#111289) [Flang][OpenMP] Improve entry block argument creation and binding (llvm#110267) [x86] combineMul - handle 0/-1 KnownBits cases before MUL_IMM logic (REAPPLIED) [llvm-dis] Fix non-deterministic disassembly across multiple inputs (llvm#110988) [lldb][test] TestDataFormatterLibcxxOptionalSimulator.py: change order of ifdefs [lldb][test] Add libcxx-simulators test for std::optional (llvm#111133) [x86] combineMul - use computeKnownBits directly to find MUL_IMM constant splat. (REAPPLIED) Reland "[lldb][test] TestDataFormatterLibcxxStringSimulator.py: add new padding layout" (llvm#111123) Revert "[x86] combineMul - use computeKnownBits directly to find MUL_IMM constant splat." update_test_checks: fix a simple regression (llvm#111347) [LegalizeVectorTypes] Always widen fabs (llvm#111298) [lsan] Make ReportUnsuspendedThreads return bool also for Fuchsia [mlir][vector] Add more tests for ConvertVectorToLLVM (6/n) (llvm#111121) [bazel] port 9144fed [SystemZ] Remove inlining threshold multiplier. (llvm#106058) [LegalizeVectorTypes] When widening don't check for libcalls if promoted (llvm#111297) [clang][Driver] Improve multilib custom error reporting (llvm#110804) [clang][Driver] Rename "FatalError" key to "Error" in multilib.yaml (llvm#110804) [LLVM][Maintainers] Update release managers (llvm#111164) [Clang][Driver] Add option to provide path for multilib's YAML config file (llvm#109640) [LoopVectorize] Remove redundant code in emitSCEVChecks (llvm#111132) [AMDGPU] Only emit SCOPE_SYS global_wb (llvm#110636) [ELF] Change Ctx::target to unique_ptr (llvm#111260) [ELF] Pass Ctx & to some free functions [RISCV] Only disassemble fcvtmod.w.d if the rounding mode is rtz. (llvm#111308) [Clang] Remove the special-casing for RequiresExprBodyDecl in BuildResolvedCallExpr() after fd87d76 (llvm#111277) [ELF] Pass Ctx & to InputFile [clang-format] Add AlignFunctionDeclarations to AlignConsecutiveDeclarations (llvm#108241) [AMDGPU] Support preloading hidden kernel arguments (llvm#98861) [ELF] Move static nextGroupId isInGroup to LinkerDriver [clangd] Add ArgumentLists config option under Completion (llvm#111322) [ELF] Pass Ctx & to SyntheticSections [ELF] Pass Ctx & to Symbols [ELF] Pass Ctx & to Symbols [ELF] getRelocTargetVA: pass Ctx and Relocation. NFC [clang-tidy] Avoid capturing a local variable in a static lambda in UseRangesCheck (llvm#111282) [VPlan] Use pointer to member 0 as VPInterleaveRecipe's pointer arg. (llvm#106431) [clangd] Simplify ternary expressions with std::optional::value_or (NFC) (llvm#111309) [libc++][format][2/3] Optimizes c-string arguments. (llvm#101805) [RISCV] Combine RVBUnary and RVKUnary into classes that are more similar to ALU(W)_r(r/i). NFC (llvm#111279) [ELF] Pass Ctx & to InputFiles [libc] GPU RPC interface: add return value to `rpc_host_call` (llvm#111288) Signed-off-by: kyvangka1610 <[email protected]>

Adds hidden kernel arguments to the function signature and marks them inreg if they should be preloaded into user SGPRs. The normal kernarg preloading logic then takes over with some additional checks for the correct implicitarg_ptr alignment. Special care is needed so that metadata for the hidden arguments is not added twice when generating the code object. Change-Id: Iea2929f1d56aecb9e15ed942bcc6b4a1831dd0e2

kerbowa requested a review from arsenm July 15, 2024 06:42

llvmbot added the backend:AMDGPU label Jul 15, 2024

arsenm reviewed Jul 23, 2024

View reviewed changes

kerbowa force-pushed the preload-hidden-kernargs branch from 550a6f3 to 40889af Compare July 29, 2024 05:20

arsenm reviewed Jul 29, 2024

View reviewed changes

llvm/lib/Target/AMDGPU/AMDGPULowerKernelArguments.cpp Outdated Show resolved Hide resolved

arsenm reviewed Aug 2, 2024

View reviewed changes

kerbowa force-pushed the preload-hidden-kernargs branch 2 times, most recently from 01d38d0 to e2f0d77 Compare August 5, 2024 06:38

arsenm reviewed Aug 19, 2024

View reviewed changes

kerbowa force-pushed the preload-hidden-kernargs branch from e2f0d77 to 92ae467 Compare September 16, 2024 01:31

llvmbot added the llvm:ir label Sep 16, 2024

arsenm reviewed Sep 25, 2024

View reviewed changes

kerbowa force-pushed the preload-hidden-kernargs branch from 92ae467 to 75de53e Compare September 30, 2024 17:44

arsenm reviewed Oct 1, 2024

View reviewed changes

kerbowa force-pushed the preload-hidden-kernargs branch from 75de53e to c621b1f Compare October 4, 2024 04:05

arsenm approved these changes Oct 4, 2024

View reviewed changes

kerbowa force-pushed the preload-hidden-kernargs branch from c621b1f to 1d543eb Compare October 6, 2024 23:33

kerbowa merged commit c4d8920 into llvm:main Oct 7, 2024
9 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[AMDGPU] Support preloading hidden kernel arguments #98861

[AMDGPU] Support preloading hidden kernel arguments #98861

kerbowa commented Jul 15, 2024

llvmbot commented Jul 15, 2024 •

edited

Loading

kerbowa commented Jul 15, 2024

github-actions bot commented Jul 15, 2024 •

edited

Loading

arsenm Aug 19, 2024

arsenm Aug 19, 2024

kerbowa Sep 16, 2024

arsenm Aug 19, 2024

kerbowa Sep 16, 2024

arsenm Sep 25, 2024

arsenm Sep 25, 2024

kerbowa Sep 30, 2024

arsenm Oct 1, 2024

arsenm Oct 1, 2024

arsenm Oct 1, 2024

arsenm Oct 4, 2024

arsenm Oct 4, 2024

arsenm Oct 4, 2024

	/// check if an attributes is in the list of attributes.
	/// Check if an attributes is in the list of attributes.

[AMDGPU] Support preloading hidden kernel arguments #98861

[AMDGPU] Support preloading hidden kernel arguments #98861

Conversation

kerbowa commented Jul 15, 2024

llvmbot commented Jul 15, 2024 • edited Loading

kerbowa commented Jul 15, 2024

github-actions bot commented Jul 15, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

llvmbot commented Jul 15, 2024 •

edited

Loading

github-actions bot commented Jul 15, 2024 •

edited

Loading