[AMDGPU] Add backend support for new PAL ELF Metadata 3.0

PAL Metadata 3.0 introduces an explicit structure in metadata for the
programmable registers written out by the compiler backend.
Rather than using opaque registers which can change between different
architectures and requires encoding the bitfield information in the backend,
which may change between versions.

This is the initial minimal implementation that enables the use of PAL Metadata
3.0.

The change itself should be NFC for non-PAL, although the way RSRC2 register is
handled has been changed slightly.

The test is fairly minimal, but checks that the metadata format looks as
expected and verifies a couple of special cases such as tgid_[xyz]_en handling
and PsInputAddr/Ena which also change to explicit fields.

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

Author: dstutt

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

@@ -441,7 +441,7 @@ amdhsa::kernel_descriptor_t AMDGPUAsmPrinter::getAmdhsaKernelDescriptor(
 
   assert(isUInt<32>(PI.ScratchSize));
   assert(isUInt<32>(PI.getComputePGMRSrc1()));
-  assert(isUInt<32>(PI.ComputePGMRSrc2));
+  assert(isUInt<32>(PI.getComputePGMRSrc2()));
 
   KernelDescriptor.group_segment_fixed_size = PI.LDSSize;
   KernelDescriptor.private_segment_fixed_size = PI.ScratchSize;
@@ -450,7 +450,7 @@ amdhsa::kernel_descriptor_t AMDGPUAsmPrinter::getAmdhsaKernelDescriptor(
   KernelDescriptor.kernarg_size = STM.getKernArgSegmentSize(F, MaxKernArgAlign);
 
   KernelDescriptor.compute_pgm_rsrc1 = PI.getComputePGMRSrc1();
-  KernelDescriptor.compute_pgm_rsrc2 = PI.ComputePGMRSrc2;
+  KernelDescriptor.compute_pgm_rsrc2 = PI.getComputePGMRSrc2();
   KernelDescriptor.kernel_code_properties = getAmdhsaKernelCodeProperties(MF);
 
   assert(STM.hasGFX90AInsts() || CurrentProgramInfo.ComputePGMRSrc3GFX90A == 0);
@@ -579,28 +579,27 @@ bool AMDGPUAsmPrinter::runOnMachineFunction(MachineFunction &MF) {
     OutStreamer->emitRawComment(
       " WaveLimiterHint : " + Twine(MFI->needsWaveLimiter()), false);
 
-    OutStreamer->emitRawComment(
-      " COMPUTE_PGM_RSRC2:SCRATCH_EN: " +
-      Twine(G_00B84C_SCRATCH_EN(CurrentProgramInfo.ComputePGMRSrc2)), false);
-    OutStreamer->emitRawComment(
-      " COMPUTE_PGM_RSRC2:USER_SGPR: " +
-      Twine(G_00B84C_USER_SGPR(CurrentProgramInfo.ComputePGMRSrc2)), false);
-    OutStreamer->emitRawComment(
-      " COMPUTE_PGM_RSRC2:TRAP_HANDLER: " +
-      Twine(G_00B84C_TRAP_HANDLER(CurrentProgramInfo.ComputePGMRSrc2)), false);
-    OutStreamer->emitRawComment(
-      " COMPUTE_PGM_RSRC2:TGID_X_EN: " +
-      Twine(G_00B84C_TGID_X_EN(CurrentProgramInfo.ComputePGMRSrc2)), false);
-    OutStreamer->emitRawComment(
-      " COMPUTE_PGM_RSRC2:TGID_Y_EN: " +
-      Twine(G_00B84C_TGID_Y_EN(CurrentProgramInfo.ComputePGMRSrc2)), false);
-    OutStreamer->emitRawComment(
-      " COMPUTE_PGM_RSRC2:TGID_Z_EN: " +
-      Twine(G_00B84C_TGID_Z_EN(CurrentProgramInfo.ComputePGMRSrc2)), false);
-    OutStreamer->emitRawComment(
-      " COMPUTE_PGM_RSRC2:TIDIG_COMP_CNT: " +
-      Twine(G_00B84C_TIDIG_COMP_CNT(CurrentProgramInfo.ComputePGMRSrc2)),
-      false);
+    OutStreamer->emitRawComment(" COMPUTE_PGM_RSRC2:SCRATCH_EN: " +
+                                    Twine(CurrentProgramInfo.ScratchEnable),
+                                false);
+    OutStreamer->emitRawComment(" COMPUTE_PGM_RSRC2:USER_SGPR: " +
+                                    Twine(CurrentProgramInfo.UserSGPR),
+                                false);
+    OutStreamer->emitRawComment(" COMPUTE_PGM_RSRC2:TRAP_HANDLER: " +
+                                    Twine(CurrentProgramInfo.TrapHandlerEnable),
+                                false);
+    OutStreamer->emitRawComment(" COMPUTE_PGM_RSRC2:TGID_X_EN: " +
+                                    Twine(CurrentProgramInfo.TGIdXEnable),
+                                false);
+    OutStreamer->emitRawComment(" COMPUTE_PGM_RSRC2:TGID_Y_EN: " +
+                                    Twine(CurrentProgramInfo.TGIdYEnable),
+                                false);
+    OutStreamer->emitRawComment(" COMPUTE_PGM_RSRC2:TGID_Z_EN: " +
+                                    Twine(CurrentProgramInfo.TGIdZEnable),
+                                false);
+    OutStreamer->emitRawComment(" COMPUTE_PGM_RSRC2:TIDIG_COMP_CNT: " +
+                                    Twine(CurrentProgramInfo.TIdIGCompCount),
+                                false);
 
     assert(STM.hasGFX90AInsts() ||
            CurrentProgramInfo.ComputePGMRSrc3GFX90A == 0);
@@ -922,22 +921,21 @@ void AMDGPUAsmPrinter::getSIProgramInfo(SIProgramInfo &ProgInfo,
   // anything to disable it if we know the stack isn't used here. We may still
   // have emitted code reading it to initialize scratch, but if that's unused
   // reading garbage should be OK.
-  const bool EnablePrivateSegment =
+  ProgInfo.ScratchEnable =
       ProgInfo.ScratchBlocks > 0 || ProgInfo.DynamicCallStack;
-  ProgInfo.ComputePGMRSrc2 =
-      S_00B84C_SCRATCH_EN(EnablePrivateSegment) |
-      S_00B84C_USER_SGPR(MFI->getNumUserSGPRs()) |
-      // For AMDHSA, TRAP_HANDLER must be zero, as it is populated by the CP.
-      S_00B84C_TRAP_HANDLER(STM.isAmdHsaOS() ? 0 : STM.isTrapHandlerEnabled()) |
-      S_00B84C_TGID_X_EN(MFI->hasWorkGroupIDX()) |
-      S_00B84C_TGID_Y_EN(MFI->hasWorkGroupIDY()) |
-      S_00B84C_TGID_Z_EN(MFI->hasWorkGroupIDZ()) |
-      S_00B84C_TG_SIZE_EN(MFI->hasWorkGroupInfo()) |
-      S_00B84C_TIDIG_COMP_CNT(TIDIGCompCnt) |
-      S_00B84C_EXCP_EN_MSB(0) |
-      // For AMDHSA, LDS_SIZE must be zero, as it is populated by the CP.
-      S_00B84C_LDS_SIZE(STM.isAmdHsaOS() ? 0 : ProgInfo.LDSBlocks) |
-      S_00B84C_EXCP_EN(0);
+  ProgInfo.UserSGPR = MFI->getNumUserSGPRs();
+  // For AMDHSA, TRAP_HANDLER must be zero, as it is populated by the CP.
+  ProgInfo.TrapHandlerEnable =
+      STM.isAmdHsaOS() ? 0 : STM.isTrapHandlerEnabled();
+  ProgInfo.TGIdXEnable = MFI->hasWorkGroupIDX();
+  ProgInfo.TGIdYEnable = MFI->hasWorkGroupIDY();
+  ProgInfo.TGIdZEnable = MFI->hasWorkGroupIDZ();
+  ProgInfo.TGSizeEnable = MFI->hasWorkGroupInfo();
+  ProgInfo.TIdIGCompCount = TIDIGCompCnt;
+  ProgInfo.EXCPEnMSB = 0;
+  // For AMDHSA, LDS_SIZE must be zero, as it is populated by the CP.
+  ProgInfo.LdsSize = STM.isAmdHsaOS() ? 0 : ProgInfo.LDSBlocks;
+  ProgInfo.EXCPEnable = 0;
 
   if (STM.hasGFX90AInsts()) {
     AMDHSA_BITS_SET(ProgInfo.ComputePGMRSrc3GFX90A,
@@ -978,7 +976,7 @@ void AMDGPUAsmPrinter::EmitProgramInfoSI(const MachineFunction &MF,
     OutStreamer->emitInt32(CurrentProgramInfo.getComputePGMRSrc1());
 
     OutStreamer->emitInt32(R_00B84C_COMPUTE_PGM_RSRC2);
-    OutStreamer->emitInt32(CurrentProgramInfo.ComputePGMRSrc2);
+    OutStreamer->emitInt32(CurrentProgramInfo.getComputePGMRSrc2());
 
     OutStreamer->emitInt32(R_00B860_COMPUTE_TMPRING_SIZE);
     OutStreamer->emitInt32(
@@ -1038,25 +1036,87 @@ void AMDGPUAsmPrinter::EmitPALMetadata(const MachineFunction &MF,
   }
 
   MD->setNumUsedSgprs(CC, CurrentProgramInfo.NumSGPRsForWavesPerEU);
-  MD->setRsrc1(CC, CurrentProgramInfo.getPGMRSrc1(CC));
-  if (AMDGPU::isCompute(CC)) {
-    MD->setRsrc2(CC, CurrentProgramInfo.ComputePGMRSrc2);
+  if (MD->getPALMajorVersion() < 3) {
+    MD->setRsrc1(CC, CurrentProgramInfo.getPGMRSrc1(CC));
+    if (AMDGPU::isCompute(CC)) {
+      MD->setRsrc2(CC, CurrentProgramInfo.getComputePGMRSrc2());
+    } else {
+      if (CurrentProgramInfo.ScratchBlocks > 0)
+        MD->setRsrc2(CC, S_00B84C_SCRATCH_EN(1));
+    }
   } else {
-    if (CurrentProgramInfo.ScratchBlocks > 0)
-      MD->setRsrc2(CC, S_00B84C_SCRATCH_EN(1));
+    // Priority?
+    MD->setHwStage(CC, ".float_mode", CurrentProgramInfo.FloatMode);
+    // Priv?
+    // DX10Clamp?
+    MD->setHwStage(CC, ".debug_mode", (bool)CurrentProgramInfo.DebugMode);
+    MD->setHwStage(CC, ".ieee_mode", (bool)CurrentProgramInfo.IEEEMode);
+    MD->setHwStage(CC, ".wgp_mode", (bool)CurrentProgramInfo.WgpMode);
+    MD->setHwStage(CC, ".mem_ordered", (bool)CurrentProgramInfo.MemOrdered);
+
+    if (AMDGPU::isCompute(CC)) {
+      MD->setHwStage(CC, ".scratch_en", (bool)CurrentProgramInfo.ScratchEnable);
+      MD->setHwStage(CC, ".trap_present",
+                     (bool)CurrentProgramInfo.TrapHandlerEnable);
+
+      // Compute registers
+      // If the front-end has set tgid_x/y/z_en - assert that the
+      // CurrentProgramInfo is consistent (usually set with function attributes
+      // amdgpu-no-workgroup-id-x etc.).
+      assert(MD->checkComputeRegisters(".tgid_x_en",
+                                       (bool)CurrentProgramInfo.TGIdXEnable));
+      assert(MD->checkComputeRegisters(".tgid_y_en",
+                                       (bool)CurrentProgramInfo.TGIdYEnable));
+      assert(MD->checkComputeRegisters(".tgid_z_en",
+                                       (bool)CurrentProgramInfo.TGIdZEnable));
+
+      // EXCPEnMSB?
+      const unsigned LdsDwGranularity = 128;
+      MD->setHwStage(CC, ".lds_size",
+                     (unsigned)(CurrentProgramInfo.LdsSize * LdsDwGranularity *
+                                sizeof(uint32_t)));
+      MD->setHwStage(CC, ".excp_en", CurrentProgramInfo.EXCPEnable);
+    } else {
+      MD->setHwStage(CC, ".scratch_en", (bool)CurrentProgramInfo.ScratchEnable);
+    }
   }
+
   // ScratchSize is in bytes, 16 aligned.
   MD->setScratchSize(CC, alignTo(CurrentProgramInfo.ScratchSize, 16));
   if (MF.getFunction().getCallingConv() == CallingConv::AMDGPU_PS) {
     unsigned ExtraLDSSize = STM.getGeneration() >= AMDGPUSubtarget::GFX11
                                 ? divideCeil(CurrentProgramInfo.LDSBlocks, 2)
                                 : CurrentProgramInfo.LDSBlocks;
-    MD->setRsrc2(CC, S_00B02C_EXTRA_LDS_SIZE(ExtraLDSSize));
-    MD->setSpiPsInputEna(MFI->getPSInputEnable());
-    MD->setSpiPsInputAddr(MFI->getPSInputAddr());
+    if (MD->getPALMajorVersion() < 3) {
+      MD->setRsrc2(CC, S_00B02C_EXTRA_LDS_SIZE(ExtraLDSSize));
+      MD->setSpiPsInputEna(MFI->getPSInputEnable());
+      MD->setSpiPsInputAddr(MFI->getPSInputAddr());
+    } else {
+      // Graphics registers
+      MD->setGraphicsRegisters(".ps_extra_lds_size", ExtraLDSSize);
+      // Set PsInputEna and PsInputAddr .spi_ps_input_ena and .spi_ps_input_addr
+      static StringLiteral const PsInputFields[] = {
+          ".persp_sample_ena",    ".persp_center_ena",
+          ".persp_centroid_ena",  ".persp_pull_model_ena",
+          ".linear_sample_ena",   ".linear_center_ena",
+          ".linear_centroid_ena", ".line_stipple_tex_ena",
+          ".pos_x_float_ena",     ".pos_y_float_ena",
+          ".pos_z_float_ena",     ".pos_w_float_ena",
+          ".front_face_ena",      ".ancillary_ena",
+          ".sample_coverage_ena", ".pos_fixed_pt_ena"};
+      unsigned PSInputEna = MFI->getPSInputEnable();
+      unsigned PSInputAddr = MFI->getPSInputAddr();
+      for (auto [Idx, Field] : enumerate(PsInputFields)) {
+        MD->setGraphicsRegisters(".spi_ps_input_ena", Field,
+                                 (bool)((PSInputEna >> Idx) & 1));
+        MD->setGraphicsRegisters(".spi_ps_input_addr", Field,
+                                 (bool)((PSInputAddr >> Idx) & 1));
+      }
+    }
   }
 
-  if (STM.isWave32())
+  // For version 3 and above the wave front size is already set in the metadata
+  if (MD->getPALMajorVersion() < 3 && STM.isWave32())
     MD->setWave32(MF.getFunction().getCallingConv());
 }
 
@@ -1068,7 +1128,7 @@ void AMDGPUAsmPrinter::emitPALFunctionMetadata(const MachineFunction &MF) {
   // Set compute registers
   MD->setRsrc1(CallingConv::AMDGPU_CS,
                CurrentProgramInfo.getPGMRSrc1(CallingConv::AMDGPU_CS));
-  MD->setRsrc2(CallingConv::AMDGPU_CS, CurrentProgramInfo.ComputePGMRSrc2);
+  MD->setRsrc2(CallingConv::AMDGPU_CS, CurrentProgramInfo.getComputePGMRSrc2());
 
   // Set optional info
   MD->setFunctionLdsSize(MF, CurrentProgramInfo.LDSSize);
@@ -1104,7 +1164,7 @@ void AMDGPUAsmPrinter::getAmdKernelCode(amd_kernel_code_t &Out,
 
   Out.compute_pgm_resource_registers =
       CurrentProgramInfo.getComputePGMRSrc1() |
-      (CurrentProgramInfo.ComputePGMRSrc2 << 32);
+      (CurrentProgramInfo.getComputePGMRSrc2() << 32);
   Out.code_properties |= AMD_CODE_PROPERTY_IS_PTR64;
 
   if (CurrentProgramInfo.DynamicCallStack)

llvm/lib/Target/AMDGPU/SIProgramInfo.cpp

@@ -54,3 +54,23 @@ uint64_t SIProgramInfo::getPGMRSrc1(CallingConv::ID CC) const {
   }
   return Reg;
 }
+
+uint64_t SIProgramInfo::getComputePGMRSrc2() const {
+  uint64_t Reg =
+      S_00B84C_SCRATCH_EN(ScratchEnable) | S_00B84C_USER_SGPR(UserSGPR) |
+      S_00B84C_TRAP_HANDLER(TrapHandlerEnable) |
+      S_00B84C_TGID_X_EN(TGIdXEnable) | S_00B84C_TGID_Y_EN(TGIdYEnable) |
+      S_00B84C_TGID_Z_EN(TGIdZEnable) | S_00B84C_TG_SIZE_EN(TGSizeEnable) |
+      S_00B84C_TIDIG_COMP_CNT(TIdIGCompCount) |
+      S_00B84C_EXCP_EN_MSB(EXCPEnMSB) | S_00B84C_LDS_SIZE(LdsSize) |
+      S_00B84C_EXCP_EN(EXCPEnable);
+
+  return Reg;
+}
+
+uint64_t SIProgramInfo::getPGMRSrc2(CallingConv::ID CC) const {
+  if (AMDGPU::isCompute(CC))
+    return getComputePGMRSrc2();
+
+  return 0;
+}

llvm/lib/Target/AMDGPU/SIProgramInfo.h

@@ -36,11 +36,23 @@ struct SIProgramInfo {
     uint32_t MemOrdered = 0; // GFX10+
     uint64_t ScratchSize = 0;
 
-    // Fields set in PGM_RSRC2 pm4 packet.
+    // State used to calculate fields set in PGM_RSRC2 pm4 packet.
     uint32_t LDSBlocks = 0;
     uint32_t ScratchBlocks = 0;
 
-    uint64_t ComputePGMRSrc2 = 0;
+    // Fields set in PGM_RSRC2 pm4 packet
+    uint32_t ScratchEnable = 0;
+    uint32_t UserSGPR = 0;
+    uint32_t TrapHandlerEnable = 0;
+    uint32_t TGIdXEnable = 0;
+    uint32_t TGIdYEnable = 0;
+    uint32_t TGIdZEnable = 0;
+    uint32_t TGSizeEnable = 0;
+    uint32_t TIdIGCompCount = 0;
+    uint32_t EXCPEnMSB = 0;
+    uint32_t LdsSize = 0;
+    uint32_t EXCPEnable = 0;
+
     uint64_t ComputePGMRSrc3GFX90A = 0;
 
     uint32_t NumVGPR = 0;
@@ -75,6 +87,10 @@ struct SIProgramInfo {
     /// Compute the value of the ComputePGMRsrc1 register.
     uint64_t getComputePGMRSrc1() const;
     uint64_t getPGMRSrc1(CallingConv::ID CC) const;
+
+    /// Compute the value of the ComputePGMRsrc2 register.
+    uint64_t getComputePGMRSrc2() const;
+    uint64_t getPGMRSrc2(CallingConv::ID CC) const;
 };
 
 } // namespace llvm