[AMDGPU] Update code object metadata for kernarg preload

Tracks the registers that explicit and hidden arguments are preloaded to
with new code object metadata.

IR arguments may be split across multiple parts by isel, and SGPR tuple
alignment means that an argument may be spread across multiple
registers.

To support this, some of the utilities for hidden kernel arguments are
moved to `AMDGPUArgumentUsageInfo.h`. Additional bookkeeping is also
needed for tracking purposes.

Author: kerbowa

llvm/include/llvm/Support/AMDGPUMetadata.h

@@ -47,7 +47,7 @@ constexpr uint32_t VersionMinorV5 = 2;
 /// HSA metadata major version for code object V6.
 constexpr uint32_t VersionMajorV6 = 1;
 /// HSA metadata minor version for code object V6.
-constexpr uint32_t VersionMinorV6 = 2;
+constexpr uint32_t VersionMinorV6 = 3;
 
 /// Old HSA metadata beginning assembler directive for V2. This is only used for
 /// diagnostics now.

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp

@@ -16,12 +16,15 @@
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
+using namespace llvm::KernArgPreload;
 
 #define DEBUG_TYPE "amdgpu-argument-reg-usage-info"
 
 INITIALIZE_PASS(AMDGPUArgumentUsageInfo, DEBUG_TYPE,
                 "Argument Register Usage Information Storage", false, true)
 
+constexpr HiddenArgInfo HiddenArgUtils::HiddenArgs[END_HIDDEN_ARGS];
+
 void ArgDescriptor::print(raw_ostream &OS,
                           const TargetRegisterInfo *TRI) const {
   if (!isSet()) {
@@ -176,6 +179,37 @@ AMDGPUFunctionArgInfo AMDGPUFunctionArgInfo::fixedABILayout() {
   return AI;
 }
 
+SmallVector<const KernArgPreloadDescriptor *, 4>
+AMDGPUFunctionArgInfo::getPreloadDescriptorsForArgIdx(unsigned ArgIdx) const {
+  SmallVector<const KernArgPreloadDescriptor *, 4> Results;
+  for (const auto &KV : PreloadKernArgs) {
+    if (KV.second.OrigArgIdx == ArgIdx)
+      Results.push_back(&KV.second);
+  }
+
+  llvm::stable_sort(Results, [](const KernArgPreloadDescriptor *A,
+                                const KernArgPreloadDescriptor *B) {
+    return A->PartIdx < B->PartIdx;
+  });
+
+  return Results;
+}
+
+std::optional<const KernArgPreloadDescriptor *>
+AMDGPUFunctionArgInfo::getHiddenArgPreloadDescriptor(HiddenArg HA) const {
+  assert(HA < END_HIDDEN_ARGS);
+
+  auto HiddenArgIt = PreloadHiddenArgsIndexMap.find(HA);
+  if (HiddenArgIt == PreloadHiddenArgsIndexMap.end())
+    return std::nullopt;
+
+  auto KernArgIt = PreloadKernArgs.find(HiddenArgIt->second);
+  if (KernArgIt == PreloadKernArgs.end())
+    return std::nullopt;
+
+  return &KernArgIt->second;
+}
+
 const AMDGPUFunctionArgInfo &
 AMDGPUArgumentUsageInfo::lookupFuncArgInfo(const Function &F) const {
   auto I = ArgInfoMap.find(&F);

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h

@@ -11,7 +11,10 @@
 
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/Analysis/ValueTracking.h"
 #include "llvm/CodeGen/Register.h"
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Type.h"
 #include "llvm/Pass.h"
 
 namespace llvm {
@@ -95,11 +98,78 @@ inline raw_ostream &operator<<(raw_ostream &OS, const ArgDescriptor &Arg) {
   return OS;
 }
 
-struct KernArgPreloadDescriptor : public ArgDescriptor {
-  KernArgPreloadDescriptor() {}
-  SmallVector<MCRegister> Regs;
+namespace KernArgPreload {
+
+enum HiddenArg {
+  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
 };
 
+// Stores information about a specific hidden argument.
+struct HiddenArgInfo {
+  // Offset in bytes from the location in the kernearg segment pointed to by
+  // the implicitarg pointer.
+  uint8_t Offset;
+  // The size of the hidden argument in bytes.
+  uint8_t Size;
+  // The name of the hidden argument in the kernel signature.
+  const char *Name;
+};
+
+struct HiddenArgUtils {
+  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 getHiddenArgFromOffset(unsigned Offset) {
+    for (unsigned I = 0; I < END_HIDDEN_ARGS; ++I)
+      if (HiddenArgs[I].Offset == Offset)
+        return static_cast<HiddenArg>(I);
+
+    return END_HIDDEN_ARGS;
+  }
+
+  static Type *getHiddenArgType(LLVMContext &Ctx, HiddenArg HA) {
+    if (HA < END_HIDDEN_ARGS)
+      return static_cast<Type *>(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.");
+  }
+};
+
+struct KernArgPreloadDescriptor {
+  // Id of the original argument in the IR kernel function argument list.
+  unsigned OrigArgIdx = 0;
+
+  // If this IR argument was split into multiple parts, this is the index of the
+  // part in the original argument.
+  unsigned PartIdx = 0;
+
+  // The registers that the argument is preloaded into. The argument may be
+  // split accross multilpe registers.
+  SmallVector<MCRegister, 2> Regs;
+};
+
+} // namespace KernArgPreload
+
 struct AMDGPUFunctionArgInfo {
   // clang-format off
   enum PreloadedValue {
@@ -161,14 +231,27 @@ struct AMDGPUFunctionArgInfo {
   ArgDescriptor WorkItemIDZ;
 
   // Map the index of preloaded kernel arguments to its descriptor.
-  SmallDenseMap<int, KernArgPreloadDescriptor> PreloadKernArgs{};
+  SmallDenseMap<int, KernArgPreload::KernArgPreloadDescriptor>
+      PreloadKernArgs{};
+  // Map hidden argument to the index of it's descriptor.
+  SmallDenseMap<KernArgPreload::HiddenArg, int> PreloadHiddenArgsIndexMap{};
   // The first user SGPR allocated for kernarg preloading.
   Register FirstKernArgPreloadReg;
 
   std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>
   getPreloadedValue(PreloadedValue Value) const;
 
   static AMDGPUFunctionArgInfo fixedABILayout();
+
+  // Returns preload argument descriptors for an IR argument index. Isel may
+  // split IR arguments into multiple parts, the return vector holds all parts
+  // associated with an IR argument in the kernel signature.
+  SmallVector<const KernArgPreload::KernArgPreloadDescriptor *, 4>
+  getPreloadDescriptorsForArgIdx(unsigned ArgIdx) const;
+
+  // Returns the hidden arguments `KernArgPreloadDescriptor` if it is preloaded.
+  std::optional<const KernArgPreload::KernArgPreloadDescriptor *>
+  getHiddenArgPreloadDescriptor(KernArgPreload::HiddenArg HA) const;
 };
 
 class AMDGPUArgumentUsageInfo : public ImmutablePass {