[AMDGPU] Add support for architected flat scratch

[rampitec]

Summary

Add support for the readonly flat Scratch register initialized
by the SPI.

[t-tye]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

496–504

Does the architected flat scratch register pair also need to be made live in when scratch is enabled and architected flat scratch is supported?

[rampitec]

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

496–504

No, it does not. It is unallocatable anyway.

[t-tye]

The documentation parts LGTM.

[arsenm]

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

727–728

Presumably this is a reserved register which we wouldn't need to add to liveins?

[rampitec]

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

727–728

It's not a flat scratch, that is SGPR to init it. We normally add the livein in the emitEntryFunctionFlatScratchInit().

[kzhuravl]

LGTM, but please run clang-format, thanks!

[rampitec]

clang-formatted.

llvm/docs/AMDGPUUsage.rst

458:448 7 bits  *See separate bits below.*      Enable the setup of the
 
                                                      The total number of SGPR
                                                      user data registers
                                                      requested must not exceed
                                                      16 and match value in
                                                      ``compute_pgm_rsrc2.user_sgpr.user_sgpr_count``.
                                                      Any requests beyond 16
                                                      will be ignored.
-     >448    1 bit   ENABLE_SGPR_PRIVATE_SEGMENT
-                     _BUFFER
+     >448    1 bit   ENABLE_SGPR_PRIVATE_SEGMENT     If the *Target Properties*
+                     _BUFFER                         column of
+                                                     :ref:`amdgpu-processor-table`
+                                                     specifies *Architected flat
+                                                     scratch* then not supported
+                                                     and must be 0,
      >449    1 bit   ENABLE_SGPR_DISPATCH_PTR
      >450    1 bit   ENABLE_SGPR_QUEUE_PTR
      >451    1 bit   ENABLE_SGPR_KERNARG_SEGMENT_PTR
      >452    1 bit   ENABLE_SGPR_DISPATCH_ID
-     >453    1 bit   ENABLE_SGPR_FLAT_SCRATCH_INIT
-
+     >453    1 bit   ENABLE_SGPR_FLAT_SCRATCH_INIT   If the *Target Properties*
+                                                     column of
+                                                     :ref:`amdgpu-processor-table`
+                                                     specifies *Architected flat
+                                                     scratch* then not supported
+                                                     and must be 0,
      >454    1 bit   ENABLE_SGPR_PRIVATE_SEGMENT
                      _SIZE
      457:455 3 bits                                  Reserved, must be 0.
      458     1 bit   ENABLE_WAVEFRONT_SIZE32         GFX6-GFX9
                                                        Reserved, must be 0.
                                                      GFX10
                                                        - If 0 execute in
                                                          wavefront size 64 mode.

 ..
 
   .. table:: compute_pgm_rsrc2 for GFX6-GFX10
      :name: amdgpu-amdhsa-compute_pgm_rsrc2-gfx6-gfx10-table
 
      ======= ======= =============================== ===========================================================================
      Bits    Size    Field Name                      Description
      ======= ======= =============================== ===========================================================================
-     0       1 bit   ENABLE_PRIVATE_SEGMENT          Enable the setup of the
+     0       1 bit   ENABLE_PRIVATE_SEGMENT          * Enable the setup of the
                                                        private segment.
-
-                                                     In addition, enable the
+                                                     * If the *Target Properties*
+                                                       column of
+                                                       :ref:`amdgpu-processor-table`
+                                                       does not specify
+                                                       *Architected flat
+                                                       scratch* then enable the
                                                        setup of the SGPR
                                                        wavefront scratch offset
                                                        system register (see
                                                        :ref:`amdgpu-amdhsa-initial-kernel-execution-state`).
+                                                     * If the *Target Properties*
+                                                       column of
+                                                       :ref:`amdgpu-processor-table`
+                                                       specifies *Architected
+                                                       flat scratch* then enable
+                                                       the setup of the
+                                                       FLAT_SCRATCH register
+                                                       pair (see
+                                                       :ref:`amdgpu-amdhsa-initial-kernel-execution-state`).
 
                                                      Used by CP to set up
                                                      ``COMPUTE_PGM_RSRC2.SCRATCH_EN``.
      5:1     5 bits  USER_SGPR_COUNT                 The total number of SGPR
                                                      user data registers
                                                      requested. This number must
                                                      match the number of user
                                                      data registers enabled.
* If the *Target Properties* column of :ref:`amdgpu-processor-table`
   and move the result as a 64-bit value to the FLAT_SCRATCH SGPR register pair
   which is SGPRn-6 and SGPRn-5. It is used as the FLAT SCRATCH BASE in flat
   memory instructions.
 
   The Scratch Wavefront Offset must also be used as an offset with Private
   segment address when using the Scratch Segment Buffer (see
   :ref:`amdgpu-amdhsa-kernel-prolog-private-segment-buffer`).
 
+* If the *Target Properties* column of :ref:`amdgpu-processor-table`
+  specifies *Architected flat scratch*:
+
+  If ENABLE_PRIVATE_SEGMENT is enabled in
+  :ref:`amdgpu-amdhsa-compute_pgm_rsrc2-gfx6-gfx10-table` then the FLAT_SCRATCH
+  register pair will be initialized to the 64-bit address of the base of scratch
+  backing memory being managed by SPI for the queue executing the kernel
+  dispatch plus the value of the wave's Scratch Wavefront Offset for use as the
+  flat scratch base in flat memory instructions.
+
 .. _amdgpu-amdhsa-kernel-prolog-private-segment-buffer:
 
 Private Segment Buffer
 ++++++++++++++++++++++
 
-Private Segment Buffer SGPR register is used to initialize 4 SGPRs
+If the *Target Properties* column of :ref:`amdgpu-processor-table` specifies
+*Architected flat scratch* then a Private Segment Buffer is not supported.
+Instead the flat SCRATCH instructions are used.
+
+Otherwise, Private Segment Buffer SGPR register is used to initialize 4 SGPRs
 that are used as a V# to access scratch. CP uses the value provided by the
 runtime. It is used, together with Scratch Wavefront Offset as an offset, to
 access the private memory space using a segment address. See
 :ref:`amdgpu-amdhsa-initial-kernel-execution-state`.
 
 The scratch V# is a four-aligned SGPR and always selected for the kernel as
 follows:
 

llvm/lib/Target/AMDGPU/AMDGPU.td

 >;
 
 def FeaturePackedTID : SubtargetFeature<"packed-tid",
   "HasPackedTID",
   "true",
   "Workitem IDs are packed into v0 at kernel launch"
 >;
 
+def FeatureArchitectedFlatScratch : SubtargetFeature<"architected-flat-scratch",
+  "HasArchitectedFlatScratch",
+  "true",
+  "Flat Scratch register is a readonly SPI initialized architected register"
+>;
+
 // Dummy feature used to disable assembler instructions.
 def FeatureDisable : SubtargetFeature<"",
   "FeatureDisable","true",
   "Dummy feature to disable assembler instructions"
 >;
 
 class GCNSubtargetFeatureGeneration <string Value,
                                      string FeatureName,

llvm/lib/Target/AMDGPU/AMDGPUAsmPrinter.cpp

AMDGPUAsmPrinter::SIFunctionResourceInfo AMDGPUAsmPrinter::analyzeResourceUsage(
   const SIMachineFunctionInfo *MFI = MF.getInfo<SIMachineFunctionInfo>();
   const GCNSubtarget &ST = MF.getSubtarget<GCNSubtarget>();
   const MachineFrameInfo &FrameInfo = MF.getFrameInfo();
   const MachineRegisterInfo &MRI = MF.getRegInfo();
   const SIInstrInfo *TII = ST.getInstrInfo();
   const SIRegisterInfo &TRI = TII->getRegisterInfo();
 
   Info.UsesFlatScratch = MRI.isPhysRegUsed(AMDGPU::FLAT_SCR_LO) ||
-                         MRI.isPhysRegUsed(AMDGPU::FLAT_SCR_HI);
+                         MRI.isPhysRegUsed(AMDGPU::FLAT_SCR_HI) ||
+                         MRI.isLiveIn(MFI->getPreloadedReg(
+                             AMDGPUFunctionArgInfo::FLAT_SCRATCH_INIT));
 
   // Even if FLAT_SCRATCH is implicitly used, it has no effect if flat
   // instructions aren't used to access the scratch buffer. Inline assembly may
   // need it though.
   //
   // If we only have implicit uses of flat_scr on flat instructions, it is not
   // really needed.
   if (Info.UsesFlatScratch && !MFI->hasFlatScratchInit() &&

llvm/lib/Target/AMDGPU/AMDGPUSubtarget.cpp

GCNSubtarget::GCNSubtarget(const Triple &TT, StringRef GPU, StringRef FS,
     HasRegisterBanking(false),
     HasVOP3Literal(false),
     HasNoDataDepHazard(false),
     FlatAddressSpace(false),
     FlatInstOffsets(false),
     FlatGlobalInsts(false),
     FlatScratchInsts(false),
     ScalarFlatScratchInsts(false),
+    HasArchitectedFlatScratch(false),
     AddNoCarryInsts(false),
     HasUnpackedD16VMem(false),
     LDSMisalignedBug(false),
     HasMFMAInlineLiteralBug(false),
     UnalignedBufferAccess(false),
     UnalignedDSAccess(false),
     HasPackedTID(false),
 
GCNSubtarget::GCNSubtarget(const Triple &TT, StringRef GPU, StringRef FS,
   InlineAsmLoweringInfo.reset(new InlineAsmLowering(getTargetLowering()));
   Legalizer.reset(new AMDGPULegalizerInfo(*this, TM));
   RegBankInfo.reset(new AMDGPURegisterBankInfo(*this));
   InstSelector.reset(new AMDGPUInstructionSelector(
   *this, *static_cast<AMDGPURegisterBankInfo *>(RegBankInfo.get()), TM));
 }
 
 bool GCNSubtarget::enableFlatScratch() const {
-  return EnableFlatScratch && hasFlatScratchInsts();
+  return flatScratchIsArchitected() ||
+         (EnableFlatScratch && hasFlatScratchInsts());
 }
 
 unsigned GCNSubtarget::getConstantBusLimit(unsigned Opcode) const {
   if (getGeneration() < GFX10)
     return 1;
 
   switch (Opcode) {
   case AMDGPU::V_LSHLREV_B64_e64:

llvm/lib/Target/AMDGPU/AsmParser/AMDGPUAsmParser.cpp

public:
   bool hasInv2PiInlineImm() const {
     return getFeatureBits()[AMDGPU::FeatureInv2PiInlineImm];
   }
 
   bool hasFlatOffsets() const {
     return getFeatureBits()[AMDGPU::FeatureFlatInstOffsets];
   }
 
+  bool hasArchitectedFlatScratch() const {
+    return getFeatureBits()[AMDGPU::FeatureArchitectedFlatScratch];
+  }
+
   bool hasSGPR102_SGPR103() const {
     return !isVI() && !isGFX9();
   }
 
   bool hasSGPR104_SGPR105() const { return isGFX10Plus(); }
 
   bool hasIntClamp() const {
     return getFeatureBits()[AMDGPU::FeatureIntClamp];
if (ID == ".amdhsa_group_segment_fixed_size") {
       if (!isUInt<sizeof(KD.private_segment_fixed_size) * CHAR_BIT>(Val))
         return OutOfRangeError(ValRange);
       KD.private_segment_fixed_size = Val;
     } else if (ID == ".amdhsa_kernarg_size") {
       if (!isUInt<sizeof(KD.kernarg_size) * CHAR_BIT>(Val))
         return OutOfRangeError(ValRange);
       KD.kernarg_size = Val;
     } else if (ID == ".amdhsa_user_sgpr_private_segment_buffer") {
+      if (hasArchitectedFlatScratch())
+        return Error(IDRange.Start,
+                     "directive is not supported with architected flat scratch",
+                     IDRange);
       PARSE_BITS_ENTRY(KD.kernel_code_properties,
                        KERNEL_CODE_PROPERTY_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER,
                        Val, ValRange);
       if (Val)
         UserSGPRCount += 4;
     } else if (ID == ".amdhsa_user_sgpr_dispatch_ptr") {
       PARSE_BITS_ENTRY(KD.kernel_code_properties,
                        KERNEL_CODE_PROPERTY_ENABLE_SGPR_DISPATCH_PTR, Val,
if (ID == ".amdhsa_group_segment_fixed_size") {
         UserSGPRCount += 2;
     } else if (ID == ".amdhsa_user_sgpr_dispatch_id") {
       PARSE_BITS_ENTRY(KD.kernel_code_properties,
                        KERNEL_CODE_PROPERTY_ENABLE_SGPR_DISPATCH_ID, Val,
                        ValRange);
       if (Val)
         UserSGPRCount += 2;
     } else if (ID == ".amdhsa_user_sgpr_flat_scratch_init") {
+      if (hasArchitectedFlatScratch())
+        return Error(IDRange.Start,
+                     "directive is not supported with architected flat scratch",
+                     IDRange);
       PARSE_BITS_ENTRY(KD.kernel_code_properties,
                        KERNEL_CODE_PROPERTY_ENABLE_SGPR_FLAT_SCRATCH_INIT, Val,
                        ValRange);
       if (Val)
         UserSGPRCount += 2;
     } else if (ID == ".amdhsa_user_sgpr_private_segment_size") {
       PARSE_BITS_ENTRY(KD.kernel_code_properties,
                        KERNEL_CODE_PROPERTY_ENABLE_SGPR_PRIVATE_SEGMENT_SIZE,
                        Val, ValRange);
       if (Val)
         UserSGPRCount += 1;
     } else if (ID == ".amdhsa_wavefront_size32") {
       if (IVersion.Major < 10)
         return Error(IDRange.Start, "directive requires gfx10+", IDRange);
       EnableWavefrontSize32 = Val;
       PARSE_BITS_ENTRY(KD.kernel_code_properties,
                        KERNEL_CODE_PROPERTY_ENABLE_WAVEFRONT_SIZE32,
                        Val, ValRange);
     } else if (ID == ".amdhsa_system_sgpr_private_segment_wavefront_offset") {
-      PARSE_BITS_ENTRY(
-          KD.compute_pgm_rsrc2,
-          COMPUTE_PGM_RSRC2_ENABLE_PRIVATE_SEGMENT, Val,
-          ValRange);
+      if (hasArchitectedFlatScratch())
+        return Error(IDRange.Start,
+                     "directive is not supported with architected flat scratch",
+                     IDRange);
+      PARSE_BITS_ENTRY(KD.compute_pgm_rsrc2,
+                       COMPUTE_PGM_RSRC2_ENABLE_PRIVATE_SEGMENT, Val, ValRange);
+    } else if (ID == ".amdhsa_enable_private_segment") {
+      if (!hasArchitectedFlatScratch())
+        return Error(
+            IDRange.Start,
+            "directive is not supported without architected flat scratch",
+            IDRange);
+      PARSE_BITS_ENTRY(KD.compute_pgm_rsrc2,
+                       COMPUTE_PGM_RSRC2_ENABLE_PRIVATE_SEGMENT, Val, ValRange);
     } else if (ID == ".amdhsa_system_sgpr_workgroup_id_x") {
       PARSE_BITS_ENTRY(KD.compute_pgm_rsrc2,
                        COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_X, Val,
                        ValRange);
     } else if (ID == ".amdhsa_system_sgpr_workgroup_id_y") {
       PARSE_BITS_ENTRY(KD.compute_pgm_rsrc2,
                        COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Y, Val,
                        ValRange);
if (ID == ".amdhsa_group_segment_fixed_size") {
       AccumOffset = Val;
     } else if (ID == ".amdhsa_reserve_vcc") {
       if (!isUInt<1>(Val))
         return OutOfRangeError(ValRange);
       ReserveVCC = Val;
     } else if (ID == ".amdhsa_reserve_flat_scratch") {
       if (IVersion.Major < 7)
         return Error(IDRange.Start, "directive requires gfx7+", IDRange);
+      if (hasArchitectedFlatScratch())
+        return Error(IDRange.Start,
+                     "directive is not supported with architected flat scratch",
+                     IDRange);
       if (!isUInt<1>(Val))
         return OutOfRangeError(ValRange);
       ReserveFlatScr = Val;
     } else if (ID == ".amdhsa_reserve_xnack_mask") {
       if (IVersion.Major < 8)
         return Error(IDRange.Start, "directive requires gfx8+", IDRange);
       if (!isUInt<1>(Val))
         return OutOfRangeError(ValRange);

llvm/lib/Target/AMDGPU/Disassembler/AMDGPUDisassembler.h

public:
   const MCInstrInfo *getMCII() const { return MCII.get(); }
 
   bool isVI() const;
   bool isGFX9() const;
   bool isGFX90A() const;
   bool isGFX9Plus() const;
   bool isGFX10() const;
   bool isGFX10Plus() const;
+
+  bool hasArchitectedFlatScratch() const;
 };
 
 //===----------------------------------------------------------------------===//
 // AMDGPUSymbolizer
 //===----------------------------------------------------------------------===//
 
 class AMDGPUSymbolizer : public MCSymbolizer {
 private:

llvm/lib/Target/AMDGPU/Disassembler/AMDGPUDisassembler.cpp

 bool AMDGPUDisassembler::isGFX9Plus() const { return AMDGPU::isGFX9Plus(STI); }
 
 bool AMDGPUDisassembler::isGFX10() const { return AMDGPU::isGFX10(STI); }
 
 bool AMDGPUDisassembler::isGFX10Plus() const {
   return AMDGPU::isGFX10Plus(STI);
 }
 
+bool AMDGPUDisassembler::hasArchitectedFlatScratch() const {
+  return STI.getFeatureBits()[AMDGPU::FeatureArchitectedFlatScratch];
+}
+
 //===----------------------------------------------------------------------===//
 // AMDGPU specific symbol handling
 //===----------------------------------------------------------------------===//
 #define PRINT_DIRECTIVE(DIRECTIVE, MASK)                                       \
   do {                                                                         \
     KdStream << Indent << DIRECTIVE " "                                        \
              << ((FourByteBuffer & MASK) >> (MASK##_SHIFT)) << '\n';           \
   } while (0)
MCDisassembler::DecodeStatus AMDGPUDisassembler::decodeCOMPUTE_PGM_RSRC1(
 
   if (isGFX10Plus() && GranulatedWavefrontSGPRCount)
     return MCDisassembler::Fail;
 
   uint32_t NextFreeSGPR = (GranulatedWavefrontSGPRCount + 1) *
                           AMDGPU::IsaInfo::getSGPREncodingGranule(&STI);
 
   KdStream << Indent << ".amdhsa_reserve_vcc " << 0 << '\n';
+  if (!hasArchitectedFlatScratch())
     KdStream << Indent << ".amdhsa_reserve_flat_scratch " << 0 << '\n';
   KdStream << Indent << ".amdhsa_reserve_xnack_mask " << 0 << '\n';
   KdStream << Indent << ".amdhsa_next_free_sgpr " << NextFreeSGPR << "\n";
 
   if (FourByteBuffer & COMPUTE_PGM_RSRC1_PRIORITY)
     return MCDisassembler::Fail;
 
   PRINT_DIRECTIVE(".amdhsa_float_round_mode_32",
                   COMPUTE_PGM_RSRC1_FLOAT_ROUND_MODE_32);
MCDisassembler::DecodeStatus AMDGPUDisassembler::decodeCOMPUTE_PGM_RSRC1(
   return MCDisassembler::Success;
 }
 
 // NOLINTNEXTLINE(readability-identifier-naming)
 MCDisassembler::DecodeStatus AMDGPUDisassembler::decodeCOMPUTE_PGM_RSRC2(
     uint32_t FourByteBuffer, raw_string_ostream &KdStream) const {
   using namespace amdhsa;
   StringRef Indent = "\t";
-  PRINT_DIRECTIVE(
-      ".amdhsa_system_sgpr_private_segment_wavefront_offset",
+  if (hasArchitectedFlatScratch())
+    PRINT_DIRECTIVE(".amdhsa_enable_private_segment",
+                    COMPUTE_PGM_RSRC2_ENABLE_PRIVATE_SEGMENT);
+  else
+    PRINT_DIRECTIVE(".amdhsa_system_sgpr_private_segment_wavefront_offset",
                     COMPUTE_PGM_RSRC2_ENABLE_PRIVATE_SEGMENT);
   PRINT_DIRECTIVE(".amdhsa_system_sgpr_workgroup_id_x",
                   COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_X);
   PRINT_DIRECTIVE(".amdhsa_system_sgpr_workgroup_id_y",
                   COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Y);
   PRINT_DIRECTIVE(".amdhsa_system_sgpr_workgroup_id_z",
                   COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Z);
   PRINT_DIRECTIVE(".amdhsa_system_sgpr_workgroup_info",
                   COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_INFO);
if (decodeCOMPUTE_PGM_RSRC2(FourByteBuffer, KdStream) ==
       return MCDisassembler::Fail;
     }
     return MCDisassembler::Success;
 
   case amdhsa::KERNEL_CODE_PROPERTIES_OFFSET:
     using namespace amdhsa;
     TwoByteBuffer = DE.getU16(Cursor);
 
+    if (!hasArchitectedFlatScratch())
       PRINT_DIRECTIVE(".amdhsa_user_sgpr_private_segment_buffer",
                       KERNEL_CODE_PROPERTY_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER);
     PRINT_DIRECTIVE(".amdhsa_user_sgpr_dispatch_ptr",
                     KERNEL_CODE_PROPERTY_ENABLE_SGPR_DISPATCH_PTR);
     PRINT_DIRECTIVE(".amdhsa_user_sgpr_queue_ptr",
                     KERNEL_CODE_PROPERTY_ENABLE_SGPR_QUEUE_PTR);
     PRINT_DIRECTIVE(".amdhsa_user_sgpr_kernarg_segment_ptr",
                     KERNEL_CODE_PROPERTY_ENABLE_SGPR_KERNARG_SEGMENT_PTR);
     PRINT_DIRECTIVE(".amdhsa_user_sgpr_dispatch_id",
                     KERNEL_CODE_PROPERTY_ENABLE_SGPR_DISPATCH_ID);
+    if (!hasArchitectedFlatScratch())
       PRINT_DIRECTIVE(".amdhsa_user_sgpr_flat_scratch_init",
                       KERNEL_CODE_PROPERTY_ENABLE_SGPR_FLAT_SCRATCH_INIT);
     PRINT_DIRECTIVE(".amdhsa_user_sgpr_private_segment_size",
                     KERNEL_CODE_PROPERTY_ENABLE_SGPR_PRIVATE_SEGMENT_SIZE);
 
     if (TwoByteBuffer & KERNEL_CODE_PROPERTY_RESERVED0)
       return MCDisassembler::Fail;
 
     // Reserved for GFX9
     if (isGFX9() &&

llvm/lib/Target/AMDGPU/GCNSubtarget.h

protected:
   bool HasRegisterBanking;
   bool HasVOP3Literal;
   bool HasNoDataDepHazard;
   bool FlatAddressSpace;
   bool FlatInstOffsets;
   bool FlatGlobalInsts;
   bool FlatScratchInsts;
   bool ScalarFlatScratchInsts;
+  bool HasArchitectedFlatScratch;
   bool AddNoCarryInsts;
   bool HasUnpackedD16VMem;
   bool R600ALUInst;
   bool CaymanISA;
   bool CFALUBug;
   bool LDSMisalignedBug;
   bool HasMFMAInlineLiteralBug;
   bool HasVertexCache;
unsigned computeOccupancy(const Function &F, unsigned LDSSize = 0,
                             unsigned NumSGPRs = 0, unsigned NumVGPRs = 0) const;
 
   /// \returns true if the flat_scratch register should be initialized with the
   /// pointer to the wave's scratch memory rather than a size and offset.
   bool flatScratchIsPointer() const {
     return getGeneration() >= AMDGPUSubtarget::GFX9;
   }
 
+  /// \returns true if the flat_scratch register is initialized by the HW.
+  /// In this case it is readonly.
+  bool flatScratchIsArchitected() const { return HasArchitectedFlatScratch; }
+
   /// \returns true if the machine has merged shaders in which s0-s7 are
   /// reserved by the hardware and user SGPRs start at s8
   bool hasMergedShaders() const {
     return getGeneration() >= GFX9;
   }
 
   /// \returns SGPR allocation granularity supported by the subtarget.
   unsigned getSGPRAllocGranule() const {

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

STREAM << "\t\t" << DIRECTIVE << " "                                         \
          << AMDHSA_BITS_GET(KERNEL_DESC.MEMBER_NAME, FIELD_NAME) << '\n';
 
   OS << "\t\t.amdhsa_group_segment_fixed_size " << KD.group_segment_fixed_size
      << '\n';
   OS << "\t\t.amdhsa_private_segment_fixed_size "
      << KD.private_segment_fixed_size << '\n';
   OS << "\t\t.amdhsa_kernarg_size " << KD.kernarg_size << '\n';
 
-  PRINT_FIELD(OS, ".amdhsa_user_sgpr_private_segment_buffer", KD,
-              kernel_code_properties,
-              amdhsa::KERNEL_CODE_PROPERTY_ENABLE_SGPR_PRIVATE_SEGMENT_BUFFER);
+  if (!hasArchitectedFlatScratch(STI))
+    PRINT_FIELD(
+        OS, ".amdhsa_user_sgpr_private_segment_buffer", KD,
   PRINT_FIELD(OS, ".amdhsa_user_sgpr_dispatch_ptr", KD,
               kernel_code_properties,
               amdhsa::KERNEL_CODE_PROPERTY_ENABLE_SGPR_DISPATCH_PTR);
   PRINT_FIELD(OS, ".amdhsa_user_sgpr_queue_ptr", KD,
               kernel_code_properties,
               amdhsa::KERNEL_CODE_PROPERTY_ENABLE_SGPR_QUEUE_PTR);
   PRINT_FIELD(OS, ".amdhsa_user_sgpr_kernarg_segment_ptr", KD,
               kernel_code_properties,
               amdhsa::KERNEL_CODE_PROPERTY_ENABLE_SGPR_KERNARG_SEGMENT_PTR);
   PRINT_FIELD(OS, ".amdhsa_user_sgpr_dispatch_id", KD,
               kernel_code_properties,
               amdhsa::KERNEL_CODE_PROPERTY_ENABLE_SGPR_DISPATCH_ID);
+  if (!hasArchitectedFlatScratch(STI))
     PRINT_FIELD(OS, ".amdhsa_user_sgpr_flat_scratch_init", KD,
                 kernel_code_properties,
                 amdhsa::KERNEL_CODE_PROPERTY_ENABLE_SGPR_FLAT_SCRATCH_INIT);
   PRINT_FIELD(OS, ".amdhsa_user_sgpr_private_segment_size", KD,
               kernel_code_properties,
               amdhsa::KERNEL_CODE_PROPERTY_ENABLE_SGPR_PRIVATE_SEGMENT_SIZE);
   if (IVersion.Major >= 10)
     PRINT_FIELD(OS, ".amdhsa_wavefront_size32", KD,
                 kernel_code_properties,
                 amdhsa::KERNEL_CODE_PROPERTY_ENABLE_WAVEFRONT_SIZE32);
-  PRINT_FIELD(
-      OS, ".amdhsa_system_sgpr_private_segment_wavefront_offset", KD,
-      compute_pgm_rsrc2,
+  PRINT_FIELD(OS,
+              (hasArchitectedFlatScratch(STI)
+                   ? ".amdhsa_enable_private_segment"
+                   : ".amdhsa_system_sgpr_private_segment_wavefront_offset"),
+              KD, compute_pgm_rsrc2,
               amdhsa::COMPUTE_PGM_RSRC2_ENABLE_PRIVATE_SEGMENT);
   PRINT_FIELD(OS, ".amdhsa_system_sgpr_workgroup_id_x", KD,
               compute_pgm_rsrc2,
               amdhsa::COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_X);
   PRINT_FIELD(OS, ".amdhsa_system_sgpr_workgroup_id_y", KD,
               compute_pgm_rsrc2,
               amdhsa::COMPUTE_PGM_RSRC2_ENABLE_SGPR_WORKGROUP_ID_Y);
   PRINT_FIELD(OS, ".amdhsa_system_sgpr_workgroup_id_z", KD,
               compute_pgm_rsrc2,
#define PRINT_FIELD(STREAM, DIRECTIVE, KERNEL_DESC, MEMBER_NAME, FIELD_NAME)   \
   if (AMDGPU::isGFX90A(STI))
     OS << "\t\t.amdhsa_accum_offset " <<
       (AMDHSA_BITS_GET(KD.compute_pgm_rsrc3,
                        amdhsa::COMPUTE_PGM_RSRC3_GFX90A_ACCUM_OFFSET) + 1) * 4
       << '\n';
 
   if (!ReserveVCC)
     OS << "\t\t.amdhsa_reserve_vcc " << ReserveVCC << '\n';
-  if (IVersion.Major >= 7 && !ReserveFlatScr)
+  if (IVersion.Major >= 7 && !ReserveFlatScr && !hasArchitectedFlatScratch(STI))
     OS << "\t\t.amdhsa_reserve_flat_scratch " << ReserveFlatScr << '\n';
 
   if (Optional<uint8_t> HsaAbiVer = getHsaAbiVersion(&STI)) {
     switch (*HsaAbiVer) {
     default:
       break;
     case ELF::ELFABIVERSION_AMDGPU_HSA_V2:
       break;

llvm/lib/Target/AMDGPU/SIFrameLowering.cpp

void SIFrameLowering::emitEntryFunctionPrologue(MachineFunction &MF,
   }
 
   if (hasFP(MF)) {
     Register FPReg = MFI->getFrameOffsetReg();
     assert(FPReg != AMDGPU::FP_REG);
     BuildMI(MBB, I, DL, TII->get(AMDGPU::S_MOV_B32), FPReg).addImm(0);
   }
 
-  if (MFI->hasFlatScratchInit() || ScratchRsrcReg) {
+  if ((MFI->hasFlatScratchInit() || ScratchRsrcReg) &&
+      !ST.flatScratchIsArchitected()) {
     MRI.addLiveIn(PreloadedScratchWaveOffsetReg);
     MBB.addLiveIn(PreloadedScratchWaveOffsetReg);
   }
 
   if (MFI->hasFlatScratchInit()) {
     emitEntryFunctionFlatScratchInit(MF, MBB, I, DL, ScratchWaveOffsetReg);
   }
 
   if (ScratchRsrcReg) {
     emitEntryFunctionScratchRsrcRegSetup(MF, MBB, I, DL,
                                          PreloadedScratchRsrcReg,
                                          ScratchRsrcReg, ScratchWaveOffsetReg);
   }
 }
 

llvm/lib/Target/AMDGPU/SIMachineFunctionInfo.cpp

SIMachineFunctionInfo::SIMachineFunctionInfo(const MachineFunction &MF)
 
   bool HasStackObjects = F.hasFnAttribute("amdgpu-stack-objects");
   if (isEntryFunction()) {
     // X, XY, and XYZ are the only supported combinations, so make sure Y is
     // enabled if Z is.
     if (WorkItemIDZ)
       WorkItemIDY = true;
 
+    if (!ST.flatScratchIsArchitected()) {
       PrivateSegmentWaveByteOffset = true;
 
       // HS and GS always have the scratch wave offset in SGPR5 on GFX9.
       if (ST.getGeneration() >= AMDGPUSubtarget::GFX9 &&
           (CC == CallingConv::AMDGPU_HS || CC == CallingConv::AMDGPU_GS))
         ArgInfo.PrivateSegmentWaveByteOffset =
             ArgDescriptor::createRegister(AMDGPU::SGPR5);
     }
+  }
 
   bool isAmdHsaOrMesa = ST.isAmdHsaOrMesa(F);
   if (isAmdHsaOrMesa) {
     if (!ST.enableFlatScratch())
       PrivateSegmentBuffer = true;
 
     if (UseFixedABI) {
       DispatchPtr = true;
SIMachineFunctionInfo::SIMachineFunctionInfo(const MachineFunction &MF)
   } else if (ST.isMesaGfxShader(F)) {
     ImplicitBufferPtr = true;
   }
 
   if (UseFixedABI || F.hasFnAttribute("amdgpu-kernarg-segment-ptr"))
     KernargSegmentPtr = true;
 
   if (ST.hasFlatAddressSpace() && isEntryFunction() &&
-      (isAmdHsaOrMesa || ST.enableFlatScratch())) {
+      (isAmdHsaOrMesa || ST.enableFlatScratch()) &&
+      !ST.flatScratchIsArchitected()) {
     // TODO: This could be refined a lot. The attribute is a poor way of
     // detecting calls or stack objects that may require it before argument
     // lowering.
     if (HasCalls || HasStackObjects || ST.enableFlatScratch())
       FlatScratchInit = true;
   }
 
   Attribute A = F.getFnAttribute("amdgpu-git-ptr-high");

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.h

 bool isGFX9(const MCSubtargetInfo &STI);
 bool isGFX9Plus(const MCSubtargetInfo &STI);
 bool isGFX10(const MCSubtargetInfo &STI);
 bool isGFX10Plus(const MCSubtargetInfo &STI);
 bool isGCN3Encoding(const MCSubtargetInfo &STI);
 bool isGFX10_BEncoding(const MCSubtargetInfo &STI);
 bool hasGFX10_3Insts(const MCSubtargetInfo &STI);
 bool isGFX90A(const MCSubtargetInfo &STI);
+bool hasArchitectedFlatScratch(const MCSubtargetInfo &STI);
 
 /// Is Reg - scalar register
 bool isSGPR(unsigned Reg, const MCRegisterInfo* TRI);
 
 /// Is there any intersection between registers
 bool isRegIntersect(unsigned Reg0, unsigned Reg1, const MCRegisterInfo* TRI);
 
 /// If \p Reg is a pseudo reg, return the correct hardware register given

llvm/lib/Target/AMDGPU/Utils/AMDGPUBaseInfo.cpp

 bool hasGFX10_3Insts(const MCSubtargetInfo &STI) {
   return STI.getFeatureBits()[AMDGPU::FeatureGFX10_3Insts];
 }
 
 bool isGFX90A(const MCSubtargetInfo &STI) {
   return STI.getFeatureBits()[AMDGPU::FeatureGFX90AInsts];
 }
 
+bool hasArchitectedFlatScratch(const MCSubtargetInfo &STI) {
+  return STI.getFeatureBits()[AMDGPU::FeatureArchitectedFlatScratch];
+}
+
 bool isSGPR(unsigned Reg, const MCRegisterInfo* TRI) {
   const MCRegisterClass SGPRClass = TRI->getRegClass(AMDGPU::SReg_32RegClassID);
   const unsigned FirstSubReg = TRI->getSubReg(Reg, AMDGPU::sub0);
   return SGPRClass.contains(FirstSubReg != 0 ? FirstSubReg : Reg) ||
     Reg == AMDGPU::SCC;
 }
 
 bool isRegIntersect(unsigned Reg0, unsigned Reg1, const MCRegisterInfo* TRI) {

llvm/test/CodeGen/AMDGPU/GlobalISel/flat-scratch-init.ll

-; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefix=GCN %s
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 < %s | FileCheck -check-prefixes=GCN,RW-FLAT %s
+; RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -mattr=+architected-flat-scratch < %s | FileCheck -check-prefixes=GCN,RO-FLAT %s
 
 ; Make sure flat_scratch_init is set
 
 ; GCN-LABEL: {{^}}stack_object_addrspacecast_in_kernel_no_calls:
-; GCN:     s_add_u32 flat_scratch_lo, s4, s7
-; GCN:     s_addc_u32 flat_scratch_hi, s5, 0
-; GCN:     flat_store_dword
-; GCN:     .amdhsa_user_sgpr_flat_scratch_init 1
-; GCN:     .amdhsa_system_sgpr_private_segment_wavefront_offset
-; GCN-NOT: .amdhsa_reserve_flat_scratch
-; GCN:     COMPUTE_PGM_RSRC2:SCRATCH_EN: 1
-; GCN:     COMPUTE_PGM_RSRC2:USER_SGPR: 6
+; RW-FLAT:     s_add_u32 flat_scratch_lo, s4, s7
+; RW-FLAT:     s_addc_u32 flat_scratch_hi, s5, 0
+; RO-FLAT-NOT: flat_scratch
+; GCN:         flat_store_dword
+; RO-FLAT-NOT: .amdhsa_user_sgpr_private_segment_buffer
+; RW-FLAT:     .amdhsa_user_sgpr_flat_scratch_init 1
+; RO-FLAT-NOT: .amdhsa_user_sgpr_flat_scratch_init
+; RW-FLAT:     .amdhsa_system_sgpr_private_segment_wavefront_offset
+; RW-FLAT-NOT: .amdhsa_enable_private_segment
+; RO-FLAT-NOT: .amdhsa_system_sgpr_private_segment_wavefront_offset
+; RO-FLAT:     .amdhsa_enable_private_segment 1
+; GCN-NOT:     .amdhsa_reserve_flat_scratch
+; GCN:         COMPUTE_PGM_RSRC2:SCRATCH_EN: 1
+; RW-FLAT:     COMPUTE_PGM_RSRC2:USER_SGPR: 6
+; RO-FLAT:     COMPUTE_PGM_RSRC2:USER_SGPR: 0
 define amdgpu_kernel void @stack_object_addrspacecast_in_kernel_no_calls() {
   %alloca = alloca i32, addrspace(5)
   %cast = addrspacecast i32 addrspace(5)* %alloca to i32*
   store volatile i32 0, i32* %cast
   ret void
 }
 
 ; TODO: Could optimize out in this case
 ; GCN-LABEL: {{^}}stack_object_in_kernel_no_calls:
-; GCN:     s_add_u32 flat_scratch_lo, s4, s7
-; GCN:     s_addc_u32 flat_scratch_hi, s5, 0
-; GCN:     buffer_store_dword
-; GCN:     .amdhsa_user_sgpr_private_segment_buffer 1
-; GCN:     .amdhsa_user_sgpr_flat_scratch_init 1
-; GCN:     .amdhsa_system_sgpr_private_segment_wavefront_offset 1
-; GCN-NOT: .amdhsa_reserve_flat_scratch
-; GCN:     COMPUTE_PGM_RSRC2:SCRATCH_EN: 1
-; GCN:     COMPUTE_PGM_RSRC2:USER_SGPR: 6
+; RW-FLAT:     s_add_u32 flat_scratch_lo, s4, s7
+; RW-FLAT:     s_addc_u32 flat_scratch_hi, s5, 0
+; RO-FLAT-NOT: flat_scratch
+; RW-FLAT:     buffer_store_dword
+; RO-FLAT:     scratch_store_dword
+; RW-FLAT:     .amdhsa_user_sgpr_private_segment_buffer 1
+; RO-FLAT-NOT: .amdhsa_user_sgpr_private_segment_buffer
+; RW-FLAT:     .amdhsa_user_sgpr_flat_scratch_init 1
+; RO-FLAT-NOT: .amdhsa_user_sgpr_flat_scratch_init
+; RW-FLAT:     .amdhsa_system_sgpr_private_segment_wavefront_offset 1
+; RW-FLAT-NOT: .amdhsa_enable_private_segment
+; RO-FLAT-NOT: .amdhsa_system_sgpr_private_segment_wavefront_offset
+; RO-FLAT:     .amdhsa_enable_private_segment 1
+; GCN-NOT:     .amdhsa_reserve_flat_scratch
+; GCN:         COMPUTE_PGM_RSRC2:SCRATCH_EN: 1
+; RW-FLAT:     COMPUTE_PGM_RSRC2:USER_SGPR: 6
+; RO-FLAT:     COMPUTE_PGM_RSRC2:USER_SGPR: 0
 define amdgpu_kernel void @stack_object_in_kernel_no_calls() {
   %alloca = alloca i32, addrspace(5)
   store volatile i32 0, i32 addrspace(5)* %alloca
   ret void
 }
 
 ; GCN-LABEL: {{^}}kernel_no_calls_no_stack:
-; GCN-NOT: flat_scratch
-; GCN:     .amdhsa_user_sgpr_private_segment_buffer 1
-; GCN:     .amdhsa_user_sgpr_flat_scratch_init 0
-; GCN:     .amdhsa_system_sgpr_private_segment_wavefront_offset 0
-; GCN:     .amdhsa_reserve_flat_scratch 0
-; GCN:     COMPUTE_PGM_RSRC2:SCRATCH_EN: 0
-; GCN:     COMPUTE_PGM_RSRC2:USER_SGPR: 4
+; GCN-NOT:    flat_scratch
+; RW-FLAT:     .amdhsa_user_sgpr_private_segment_buffer 1
+; RO-FLAT-NOT: .amdhsa_user_sgpr_private_segment_buffer
+; RW-FLAT:     .amdhsa_user_sgpr_flat_scratch_init 0
+; RO-FLAT-NOT: .amdhsa_user_sgpr_flat_scratch_init
+; RW-FLAT:     .amdhsa_system_sgpr_private_segment_wavefront_offset 0
+; RW-FLAT-NOT: .amdhsa_enable_private_segment
+; RO-FLAT-NOT: .amdhsa_system_sgpr_private_segment_wavefront_offset
+; RO-FLAT:     .amdhsa_enable_private_segment 0
+; RW-FLAT:     .amdhsa_reserve_flat_scratch 0
+; RO-FLAT-NOT: .amdhsa_reserve_flat_scratch 0
+; GCN:         COMPUTE_PGM_RSRC2:SCRATCH_EN: 0
+; RW-FLAT:     COMPUTE_PGM_RSRC2:USER_SGPR: 4
+; RO-FLAT:     COMPUTE_PGM_RSRC2:USER_SGPR: 0
 define amdgpu_kernel void @kernel_no_calls_no_stack() {
   ret void
 }