[AMDGPU] Add DAG ISel support for preloaded kernel arguments

This patch adds the DAG isel changes for kernel argument preloading.
These changes are not usable with older firmware but subsequent patches
in the series will make the codegen backwards compatible. This patch
should only be submitted alongside that subsequent patch.

Preloading here begins from the start of the kernel arguments until the
amount of arguments indicated by the CL flag
amdgpu-kernarg-preload-count.

Aggregates and arguments passed by-ref are not supported.

Special care for the alignment of the kernarg segment is needed as well
as consideration of the alignment of addressable SGPR tuples when we
cannot directly use misaligned large tuples that the arguments are
loaded to.

Reviewed By: bcahoon

Differential Revision: https://reviews.llvm.org/D158579

Author: kerbowa

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.cpp

@@ -60,6 +60,7 @@ bool AMDGPUArgumentUsageInfo::doFinalization(Module &M) {
   return false;
 }
 
+// TODO: Print preload kernargs?
 void AMDGPUArgumentUsageInfo::print(raw_ostream &OS, const Module *M) const {
   for (const auto &FI : ArgInfoMap) {
     OS << "Arguments for " << FI.first->getName() << '\n'
@@ -148,7 +149,7 @@ AMDGPUFunctionArgInfo::getPreloadedValue(
   llvm_unreachable("unexpected preloaded value type");
 }
 
-constexpr AMDGPUFunctionArgInfo AMDGPUFunctionArgInfo::fixedABILayout() {
+AMDGPUFunctionArgInfo AMDGPUFunctionArgInfo::fixedABILayout() {
   AMDGPUFunctionArgInfo AI;
   AI.PrivateSegmentBuffer
     = ArgDescriptor::createRegister(AMDGPU::SGPR0_SGPR1_SGPR2_SGPR3);

llvm/lib/Target/AMDGPU/AMDGPUArgumentUsageInfo.h

@@ -9,6 +9,7 @@
 #ifndef LLVM_LIB_TARGET_AMDGPU_AMDGPUARGUMENTUSAGEINFO_H
 #define LLVM_LIB_TARGET_AMDGPU_AMDGPUARGUMENTUSAGEINFO_H
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/CodeGen/Register.h"
 #include "llvm/Pass.h"
 
@@ -37,22 +38,19 @@ struct ArgDescriptor {
   bool IsSet : 1;
 
 public:
-  constexpr ArgDescriptor(unsigned Val = 0, unsigned Mask = ~0u,
-                bool IsStack = false, bool IsSet = false)
-    : Reg(Val), Mask(Mask), IsStack(IsStack), IsSet(IsSet) {}
+  ArgDescriptor(unsigned Val = 0, unsigned Mask = ~0u, bool IsStack = false,
+                bool IsSet = false)
+      : Reg(Val), Mask(Mask), IsStack(IsStack), IsSet(IsSet) {}
 
-  static constexpr ArgDescriptor createRegister(Register Reg,
-                                                unsigned Mask = ~0u) {
+  static ArgDescriptor createRegister(Register Reg, unsigned Mask = ~0u) {
     return ArgDescriptor(Reg, Mask, false, true);
   }
 
-  static constexpr ArgDescriptor createStack(unsigned Offset,
-                                             unsigned Mask = ~0u) {
+  static ArgDescriptor createStack(unsigned Offset, unsigned Mask = ~0u) {
     return ArgDescriptor(Offset, Mask, true, true);
   }
 
-  static constexpr ArgDescriptor createArg(const ArgDescriptor &Arg,
-                                           unsigned Mask) {
+  static ArgDescriptor createArg(const ArgDescriptor &Arg, unsigned Mask) {
     return ArgDescriptor(Arg.Reg, Mask, Arg.IsStack, Arg.IsSet);
   }
 
@@ -94,6 +92,11 @@ inline raw_ostream &operator<<(raw_ostream &OS, const ArgDescriptor &Arg) {
   return OS;
 }
 
+struct KernArgPreloadDescriptor : public ArgDescriptor {
+  KernArgPreloadDescriptor() {}
+  SmallVector<MCRegister> Regs;
+};
+
 struct AMDGPUFunctionArgInfo {
   enum PreloadedValue {
     // SGPRS:
@@ -151,10 +154,13 @@ struct AMDGPUFunctionArgInfo {
   ArgDescriptor WorkItemIDY;
   ArgDescriptor WorkItemIDZ;
 
+  // Map the index of preloaded kernel arguments to its descriptor.
+  SmallDenseMap<int, KernArgPreloadDescriptor> PreloadKernArgs{};
+
   std::tuple<const ArgDescriptor *, const TargetRegisterClass *, LLT>
   getPreloadedValue(PreloadedValue Value) const;
 
-  static constexpr AMDGPUFunctionArgInfo fixedABILayout();
+  static AMDGPUFunctionArgInfo fixedABILayout();
 };
 
 class AMDGPUArgumentUsageInfo : public ImmutablePass {

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

@@ -205,6 +205,11 @@ void AMDGPUAsmPrinter::emitFunctionBodyStart() {
 
   if (STM.isAmdHsaOS())
     HSAMetadataStream->emitKernel(*MF, CurrentProgramInfo);
+
+  if (MFI.getNumKernargPreloadedSGPRs() > 0) {
+    assert(AMDGPU::hasKernargPreload(STM));
+    getTargetStreamer()->EmitKernargPreloadHeader(*getGlobalSTI());
+  }
 }
 
 void AMDGPUAsmPrinter::emitFunctionBodyEnd() {
@@ -417,6 +422,7 @@ amdhsa::kernel_descriptor_t AMDGPUAsmPrinter::getAmdhsaKernelDescriptor(
     const SIProgramInfo &PI) const {
   const GCNSubtarget &STM = MF.getSubtarget<GCNSubtarget>();
   const Function &F = MF.getFunction();
+  const SIMachineFunctionInfo *Info = MF.getInfo<SIMachineFunctionInfo>();
 
   amdhsa::kernel_descriptor_t KernelDescriptor;
   memset(&KernelDescriptor, 0x0, sizeof(KernelDescriptor));
@@ -440,6 +446,10 @@ amdhsa::kernel_descriptor_t AMDGPUAsmPrinter::getAmdhsaKernelDescriptor(
     KernelDescriptor.compute_pgm_rsrc3 =
       CurrentProgramInfo.ComputePGMRSrc3GFX90A;
 
+  if (AMDGPU::hasKernargPreload(STM))
+    KernelDescriptor.kernarg_preload =
+        static_cast<uint16_t>(Info->getNumKernargPreloadedSGPRs());
+
   return KernelDescriptor;
 }
 

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

@@ -1027,7 +1027,8 @@ const AMDGPUSubtarget &AMDGPUSubtarget::get(const TargetMachine &TM, const Funct
 }
 
 GCNUserSGPRUsageInfo::GCNUserSGPRUsageInfo(const Function &F,
-                                           const GCNSubtarget &ST) {
+                                           const GCNSubtarget &ST)
+    : ST(ST) {
   const CallingConv::ID CC = F.getCallingConv();
   const bool IsKernel =
       CC == CallingConv::AMDGPU_KERNEL || CC == CallingConv::SPIR_KERNEL;
@@ -1068,30 +1069,35 @@ GCNUserSGPRUsageInfo::GCNUserSGPRUsageInfo(const Function &F,
       !ST.flatScratchIsArchitected()) {
     FlatScratchInit = true;
   }
-}
 
-unsigned GCNUserSGPRUsageInfo::getNumUsedUserSGPRs() const {
-  unsigned NumUserSGPRs = 0;
   if (hasImplicitBufferPtr())
-    NumUserSGPRs += getNumUserSGPRForField(ImplicitBufferPtrID);
+    NumUsedUserSGPRs += getNumUserSGPRForField(ImplicitBufferPtrID);
 
   if (hasPrivateSegmentBuffer())
-    NumUserSGPRs += getNumUserSGPRForField(PrivateSegmentBufferID);
+    NumUsedUserSGPRs += getNumUserSGPRForField(PrivateSegmentBufferID);
 
   if (hasDispatchPtr())
-    NumUserSGPRs += getNumUserSGPRForField(DispatchPtrID);
+    NumUsedUserSGPRs += getNumUserSGPRForField(DispatchPtrID);
 
   if (hasQueuePtr())
-    NumUserSGPRs += getNumUserSGPRForField(QueuePtrID);
+    NumUsedUserSGPRs += getNumUserSGPRForField(QueuePtrID);
 
   if (hasKernargSegmentPtr())
-    NumUserSGPRs += getNumUserSGPRForField(KernargSegmentPtrID);
+    NumUsedUserSGPRs += getNumUserSGPRForField(KernargSegmentPtrID);
 
   if (hasDispatchID())
-    NumUserSGPRs += getNumUserSGPRForField(DispatchIdID);
+    NumUsedUserSGPRs += getNumUserSGPRForField(DispatchIdID);
 
   if (hasFlatScratchInit())
-    NumUserSGPRs += getNumUserSGPRForField(FlatScratchInitID);
+    NumUsedUserSGPRs += getNumUserSGPRForField(FlatScratchInitID);
+}
+
+void GCNUserSGPRUsageInfo::allocKernargPreloadSGPRs(unsigned NumSGPRs) {
+  assert(NumKernargPreloadSGPRs + NumSGPRs <= AMDGPU::getMaxNumUserSGPRs(ST));
+  NumKernargPreloadSGPRs += NumSGPRs;
+  NumUsedUserSGPRs += NumSGPRs;
+}
 
-  return NumUserSGPRs;
+unsigned GCNUserSGPRUsageInfo::getNumFreeUserSGPRs() {
+  return AMDGPU::getMaxNumUserSGPRs(ST) - NumUsedUserSGPRs;
 }

llvm/lib/Target/AMDGPU/GCNSubtarget.h

@@ -1394,8 +1394,6 @@ class GCNSubtarget final : public AMDGPUGenSubtargetInfo,
 
 class GCNUserSGPRUsageInfo {
 public:
-  unsigned getNumUsedUserSGPRs() const;
-
   bool hasImplicitBufferPtr() const { return ImplicitBufferPtr; }
 
   bool hasPrivateSegmentBuffer() const { return PrivateSegmentBuffer; }
@@ -1410,6 +1408,14 @@ class GCNUserSGPRUsageInfo {
 
   bool hasFlatScratchInit() const { return FlatScratchInit; }
 
+  unsigned getNumKernargPreloadSGPRs() const { return NumKernargPreloadSGPRs; }
+
+  unsigned getNumUsedUserSGPRs() const { return NumUsedUserSGPRs; }
+
+  unsigned getNumFreeUserSGPRs();
+
+  void allocKernargPreloadSGPRs(unsigned NumSGPRs);
+
   enum UserSGPRID : unsigned {
     ImplicitBufferPtrID = 0,
     PrivateSegmentBufferID = 1,
@@ -1447,6 +1453,8 @@ class GCNUserSGPRUsageInfo {
   GCNUserSGPRUsageInfo(const Function &F, const GCNSubtarget &ST);
 
 private:
+  const GCNSubtarget &ST;
+
   // Private memory buffer
   // Compute directly in sgpr[0:1]
   // Other shaders indirect 64-bits at sgpr[0:1]
@@ -1463,6 +1471,10 @@ class GCNUserSGPRUsageInfo {
   bool DispatchID = false;
 
   bool FlatScratchInit = false;
+
+  unsigned NumKernargPreloadSGPRs = 0;
+
+  unsigned NumUsedUserSGPRs = 0;
 };
 
 } // end namespace llvm

llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUTargetStreamer.cpp

@@ -830,6 +830,24 @@ bool AMDGPUTargetELFStreamer::EmitHSAMetadata(
   return true;
 }
 
+bool AMDGPUTargetAsmStreamer::EmitKernargPreloadHeader(
+    const MCSubtargetInfo &STI) {
+  for (int i = 0; i < 64; ++i) {
+    OS << "\ts_nop 0\n";
+  }
+  return true;
+}
+
+bool AMDGPUTargetELFStreamer::EmitKernargPreloadHeader(
+    const MCSubtargetInfo &STI) {
+  const uint32_t Encoded_s_nop = 0xbf800000;
+  MCStreamer &OS = getStreamer();
+  for (int i = 0; i < 64; ++i) {
+    OS.emitInt32(Encoded_s_nop);
+  }
+  return true;
+}
+
 bool AMDGPUTargetELFStreamer::EmitCodeEnd(const MCSubtargetInfo &STI) {
   const uint32_t Encoded_s_code_end = 0xbf9f0000;
   const uint32_t Encoded_s_nop = 0xbf800000;

llvm/lib/Target/AMDGPU/MCTargetDesc/AMDGPUTargetStreamer.h

@@ -90,6 +90,11 @@ class AMDGPUTargetStreamer : public MCTargetStreamer {
   /// \returns True on success, false on failure.
   virtual bool EmitCodeEnd(const MCSubtargetInfo &STI) { return true; }
 
+  /// \returns True on success, false on failure.
+  virtual bool EmitKernargPreloadHeader(const MCSubtargetInfo &STI) {
+    return true;
+  }
+
   virtual void EmitAmdhsaKernelDescriptor(
       const MCSubtargetInfo &STI, StringRef KernelName,
       const amdhsa::kernel_descriptor_t &KernelDescriptor, uint64_t NextVGPR,
@@ -154,6 +159,9 @@ class AMDGPUTargetAsmStreamer final : public AMDGPUTargetStreamer {
   /// \returns True on success, false on failure.
   bool EmitCodeEnd(const MCSubtargetInfo &STI) override;
 
+  /// \returns True on success, false on failure.
+  bool EmitKernargPreloadHeader(const MCSubtargetInfo &STI) override;
+
   void EmitAmdhsaKernelDescriptor(
       const MCSubtargetInfo &STI, StringRef KernelName,
       const amdhsa::kernel_descriptor_t &KernelDescriptor, uint64_t NextVGPR,
@@ -215,6 +223,9 @@ class AMDGPUTargetELFStreamer final : public AMDGPUTargetStreamer {
   /// \returns True on success, false on failure.
   bool EmitCodeEnd(const MCSubtargetInfo &STI) override;
 
+  /// \returns True on success, false on failure.
+  bool EmitKernargPreloadHeader(const MCSubtargetInfo &STI) override;
+
   void EmitAmdhsaKernelDescriptor(
       const MCSubtargetInfo &STI, StringRef KernelName,
       const amdhsa::kernel_descriptor_t &KernelDescriptor, uint64_t NextVGPR,