[AMDGPU] Update code object metadata for kernarg preload

llvm/include/llvm/Support/AMDGPUMetadata.h

@@ -48,7 +48,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,31 @@ AMDGPUFunctionArgInfo AMDGPUFunctionArgInfo::fixedABILayout() {
   return AI;
 }
 
+SmallVector<const KernArgPreloadDescriptor *, 4>
+AMDGPUFunctionArgInfo::getPreloadDescriptorsForArgIdx(unsigned ArgIdx) const {
+  SmallVector<const KernArgPreloadDescriptor *, 4> Results;
+  for (unsigned PartIdx = 0; PartIdx < PreloadKernArgs.size(); ++PartIdx) {
+    const auto &Desc = PreloadKernArgs[PartIdx];
+    if (Desc.OrigArgIdx == ArgIdx)
+      Results.push_back(&Desc);
+  }
+
+  return Results;
+}
+
+const KernArgPreloadDescriptor *
+AMDGPUFunctionArgInfo::getHiddenArgPreloadDescriptor(HiddenArg HA) const {
+  assert(HA < END_HIDDEN_ARGS);
+
+  auto HiddenArgIt = PreloadHiddenArgsIndexMap.find(HA);
+  if (HiddenArgIt == PreloadHiddenArgsIndexMap.end())
+    return nullptr;
+
+  const KernArgPreloadDescriptor &Desc = PreloadKernArgs[HiddenArgIt->second];
+  assert(Desc.IsValid && "Hidden argument preload descriptor not valid.");
+  return &Desc;
+}
+
 const AMDGPUFunctionArgInfo &
 AMDGPUArgumentUsageInfo::lookupFuncArgInfo(const Function &F) const {
   auto I = ArgInfoMap.find(&F);

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h

@@ -11,7 +11,11 @@
 
 #include "MCTargetDesc/AMDGPUMCTargetDesc.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/IndexedMap.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 +99,81 @@ inline raw_ostream &operator<<(raw_ostream &OS, const ArgDescriptor &Arg) {
   return OS;
 }
 
-struct KernArgPreloadDescriptor : public ArgDescriptor {
-  KernArgPreloadDescriptor() {}
-  SmallVector<MCRegister> Regs;
+namespace KernArgPreload {
+
+enum HiddenArg : unsigned {
+  HIDDEN_BLOCK_COUNT_X = 0,
+  HIDDEN_BLOCK_COUNT_Y = 1,
+  HIDDEN_BLOCK_COUNT_Z = 2,
+  HIDDEN_GROUP_SIZE_X = 3,
+  HIDDEN_GROUP_SIZE_Y = 4,
+  HIDDEN_GROUP_SIZE_Z = 5,
+  HIDDEN_REMAINDER_X = 6,
+  HIDDEN_REMAINDER_Y = 7,
+  HIDDEN_REMAINDER_Z = 8,
+  END_HIDDEN_ARGS = HIDDEN_REMAINDER_Z + 1
 };
 
+// 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 = 0;
+  // The size of the hidden argument in bytes.
+  uint8_t Size = 0;
+  // The name of the hidden argument in the kernel signature.
+  const char *Name = nullptr;
+};
+
+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 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 across multiple registers.
+  SmallVector<MCRegister, 2> Regs;
+
+  bool IsValid = false;
+};
+
+} // namespace KernArgPreload
+
 struct AMDGPUFunctionArgInfo {
   // clang-format off
   enum PreloadedValue {
@@ -160,15 +234,34 @@ struct AMDGPUFunctionArgInfo {
   ArgDescriptor WorkItemIDY;
   ArgDescriptor WorkItemIDZ;
 
+  struct PreloadArgIndexFunctor {
+    using argument_type = unsigned;
+    unsigned operator()(unsigned Idx) const { return Idx; }
+  };
+
   // Map the index of preloaded kernel arguments to its descriptor.
-  SmallDenseMap<int, KernArgPreloadDescriptor> PreloadKernArgs{};
+  IndexedMap<KernArgPreload::KernArgPreloadDescriptor, PreloadArgIndexFunctor>
+      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.
+  const KernArgPreload::KernArgPreloadDescriptor *
+  getHiddenArgPreloadDescriptor(KernArgPreload::HiddenArg HA) const;
 };
 
 class AMDGPUArgumentUsageInfo : public ImmutablePass {