[AMDGPU] Clarify scratch initialization

[t-tye]

Summary

• Clarify documentation on initializing scratch.
• Rename compute_pgm_rsrc2 field for enabling scratch from ENABLE_SGPR_PRIVATE_SEGMENT_WAVEFRONT_OFFSET to ENABLE_PRIVATE_SEGMENT to match hardware definition.

[t-tye]

Fix missing blank line.

[rampitec]

LGTM, thanks!

llvm/docs/AMDGPUUsage.rst

415:384 4 bytes COMPUTE_PGM_RSRC1               Compute Shader (CS)
                                                      :ref:`amdgpu-amdhsa-compute_pgm_rsrc1-gfx6-gfx10-table`.
      447:416 4 bytes COMPUTE_PGM_RSRC2               Compute Shader (CS)
                                                      program settings used by
                                                      CP to set up
                                                      ``COMPUTE_PGM_RSRC2``
                                                      configuration
                                                      register. See
                                                      :ref:`amdgpu-amdhsa-compute_pgm_rsrc2-gfx6-gfx10-table`.
-     448     1 bit   ENABLE_SGPR_PRIVATE_SEGMENT     Enable the setup of the
-                     _BUFFER                         SGPR user data registers
+     458:448 7 bits  *See separate bits below.*      Enable the setup of the
+                                                     SGPR user data registers
                                                      (see
                                                      :ref:`amdgpu-amdhsa-initial-kernel-execution-state`).
 
                                                      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.
-     449     1 bit   ENABLE_SGPR_DISPATCH_PTR        *see above*
-     450     1 bit   ENABLE_SGPR_QUEUE_PTR           *see above*
-     451     1 bit   ENABLE_SGPR_KERNARG_SEGMENT_PTR *see above*
-     452     1 bit   ENABLE_SGPR_DISPATCH_ID         *see above*
-     453     1 bit   ENABLE_SGPR_FLAT_SCRATCH_INIT   *see above*
-     454     1 bit   ENABLE_SGPR_PRIVATE_SEGMENT     *see above*
+     >448    1 bit   ENABLE_SGPR_PRIVATE_SEGMENT
+                     _BUFFER
+     >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
+
+     >454    1 bit   ENABLE_SGPR_PRIVATE_SEGMENT
                      _SIZE
      457:455 3 bits                                  Reserved, must be 0.
      458     1 bit   ENABLE_WAVEFRONT_SIZE32         GFX6-9
                                                        Reserved, must be 0.
                                                      GFX10
                                                        - If 0 execute in
                                                          wavefront size 64 mode.
                                                        - If 1 execute in

 ..
 
   .. 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_SGPR_PRIVATE_SEGMENT     Enable the setup of the
-                     _WAVEFRONT_OFFSET               SGPR wavefront scratch offset
+     0       1 bit   ENABLE_PRIVATE_SEGMENT          Enable the setup of the
+                                                     private segment.
+
+                                                     In addition, enable the
+                                                     setup of the SGPR
+                                                     wavefront scratch offset
                                                      system register (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
:ref:`amdgpu-amdhsa-sgpr-register-set-up-order-table`.
   .. table:: SGPR Register Set Up Order
      :name: amdgpu-amdhsa-sgpr-register-set-up-order-table
 
      ========== ========================== ====== ==============================
      SGPR Order Name                       Number Description
                 (kernel descriptor enable  of
                 field)                     SGPRs
      ========== ========================== ====== ==============================
-     First      Private Segment Buffer     4      V# that can be used, together
-                (enable_sgpr_private              with Scratch Wavefront Offset
-                _segment_buffer)                  as an offset, to access the
+     First      Private Segment Buffer     4      This is 4 SGPRs:
+                (enable_sgpr_private
+                _segment_buffer)
+                                                  V# that can be used,
+                                                  together with Scratch
+                                                  Wavefront Offset as an
+                                                  offset, to access the
                                                   private memory space using a
                                                   segment address.
-
-                                                  CP uses the value provided by
-                                                  the runtime.
+                                                  CP uses the value provided
+                                                  by the runtime.
      then       Dispatch Ptr               2      64-bit address of AQL dispatch
                 (enable_sgpr_dispatch_ptr)        packet for kernel dispatch
                                                   actually executing.
      then       Queue Ptr                  2      64-bit address of amd_queue_t
                 (enable_sgpr_queue_ptr)           object for AQL queue on which
                                                   the dispatch packet was
                                                   queued.
      then       Kernarg Segment Ptr        2      64-bit address of Kernarg
then       Work-Group Id Y            1      32-bit work-group id in Y
                 (enable_sgpr_workgroup_id         dimension of grid for
                 _Y)                               wavefront.
      then       Work-Group Id Z            1      32-bit work-group id in Z
                 (enable_sgpr_workgroup_id         dimension of grid for
                 _Z)                               wavefront.
      then       Work-Group Info            1      {first_wavefront, 14'b0000,
                 (enable_sgpr_workgroup            ordered_append_term[10:0],
                 _info)                            threadgroup_size_in_wavefronts[5:0]}
-     then       Scratch Wavefront Offset   1      32-bit byte offset from base
-                (enable_sgpr_private              of scratch base of queue
-                _segment_wavefront_offset)        executing the kernel
+     then       Scratch Wavefront Offset   1      This is 1 SGPR:
+                (enable_sgpr_private
+                _segment_wavefront_offset)
+                                                  32-bit byte offset from base
+                                                  of scratch base of queue
+                                                  executing the kernel
                                                   dispatch. Must be used as an
-                                                  offset with Private
-                                                  segment address when using
-                                                  Scratch Segment Buffer. It
-                                                  must be used to set up FLAT
-                                                  SCRATCH for flat addressing
-                                                  (see
+                                                  offset with Private segment
+                                                  address when using Scratch
+                                                  Segment Buffer. It must be
+                                                  used to set up FLAT SCRATCH
+                                                  for flat addressing (see
                                                   :ref:`amdgpu-amdhsa-kernel-prolog-flat-scratch`).
      ========== ========================== ====== ==============================
 
 The order of the VGPR registers is defined, but the compiler can specify which
 ones are actually setup in the kernel descriptor using the ``enable_vgpr*`` bit
 fields (see :ref:`amdgpu-amdhsa-kernel-descriptor`). The register numbers used
 for enabled registers are dense starting at VGPR0: the first enabled register is
 VGPR0, the next enabled register is VGPR1 etc.; disabled registers do not have a

 2. Work-group Id registers X, Y, Z are set by ADC which supports any
    combination including none.
 3. Scratch Wavefront Offset is set by SPI in a per wavefront basis which is why
    its value cannot be included with the flat scratch init value which is per
    queue.
 4. The VGPRs are set by SPI which only supports specifying either (X), (X, Y)
    or (X, Y, Z).
 
-Flat Scratch register pair are adjacent SGPRs so they can be moved as a 64-bit
-value to the hardware required SGPRn-3 and SGPRn-4 respectively.
+See :ref:`amdgpu-amdhsa-kernel-prolog-flat-scratch` for Flat Scratch register
+pair initialization.
 
 The global segment can be accessed either using buffer instructions (GFX6 which
 has V# 64-bit address support), flat instructions (GFX7-GFX10), or global
 instructions (GFX9-GFX10).
 
 If buffer operations are used, then the compiler can generate a V# with the
 following properties:
 

 kernel prolog. If a frame pointer is not required then all uses of the frame
 pointer are replaced with immediate ``0`` offsets.
 
 .. _amdgpu-amdhsa-kernel-prolog-flat-scratch:
 
 Flat Scratch
 ++++++++++++
 
+GFX6
+  Flat scratch is not supported.
+
+GFX7-GFX10
   If the kernel or any function it calls may use flat operations to access
   scratch memory, the prolog code must set up the FLAT_SCRATCH register pair
   (FLAT_SCRATCH_LO/FLAT_SCRATCH_HI which are in SGPRn-4/SGPRn-3). Initialization
   uses Flat Scratch Init and Scratch Wavefront Offset SGPR registers (see
   :ref:`amdgpu-amdhsa-initial-kernel-execution-state`):
 
-GFX6
-  Flat scratch is not supported.
-
   GFX7-GFX8
 
     1. The low word of Flat Scratch Init is 32-bit byte offset from
        ``SH_HIDDEN_PRIVATE_BASE_VIMID`` to the base of scratch backing memory
        being managed by SPI for the queue executing the kernel dispatch. This is
        the same value used in the Scratch Segment Buffer V# base address. The
        prolog must add the value of Scratch Wavefront Offset to get the
        wavefront's byte scratch backing memory offset from
        ``SH_HIDDEN_PRIVATE_BASE_VIMID``. Since FLAT_SCRATCH_LO is in units of 256
        bytes, the offset must be right shifted by 8 before moving into
        FLAT_SCRATCH_LO.
     2. The second word of Flat Scratch Init is 32-bit byte size of a single
        work-items scratch memory usage. This is directly loaded from the kernel
        dispatch packet Private Segment Byte Size and rounded up to a multiple of
        DWORD. Having CP load it once avoids loading it at the beginning of every
        wavefront. The prolog must move it to FLAT_SCRATCH_LO for use as FLAT
        SCRATCH SIZE.
 
   GFX9-GFX10
     The Flat Scratch Init is the 64-bit address of the base of scratch backing
     memory being managed by SPI for the queue executing the kernel dispatch. The
-  prolog must add the value of Scratch Wavefront Offset and moved to the
-  FLAT_SCRATCH pair for use as the flat scratch base in flat memory
-  instructions.
+    prolog must add the value of the wave's Scratch Wavefront Offset and moved
+    as a 64-bit value to the FLAT_SCRATCH pair for use as the flat scratch base
+    in flat memory instructions.
 
 .. _amdgpu-amdhsa-kernel-prolog-private-segment-buffer:
 
 Private Segment Buffer
 ++++++++++++++++++++++
 
 A set of four SGPRs beginning at a four-aligned SGPR index are always selected
 to serve as the scratch V# for the kernel as follows:
.. table:: AMDHSA Kernel Assembler Directives
      ``.amdhsa_user_sgpr_flat_scratch_init``                  0                   GFX6-GFX10   Controls ENABLE_SGPR_FLAT_SCRATCH_INIT in
                                                                                                :ref:`amdgpu-amdhsa-kernel-descriptor-v3-table`.
      ``.amdhsa_user_sgpr_private_segment_size``               0                   GFX6-GFX10   Controls ENABLE_SGPR_PRIVATE_SEGMENT_SIZE in
                                                                                                :ref:`amdgpu-amdhsa-kernel-descriptor-v3-table`.
      ``.amdhsa_wavefront_size32``                             Target              GFX10        Controls ENABLE_WAVEFRONT_SIZE32 in
                                                               Feature                          :ref:`amdgpu-amdhsa-kernel-descriptor-v3-table`.
                                                               Specific
                                                               (wavefrontsize64)
-     ``.amdhsa_system_sgpr_private_segment_wavefront_offset`` 0                   GFX6-GFX10   Controls ENABLE_SGPR_PRIVATE_SEGMENT_WAVEFRONT_OFFSET in
+     ``.amdhsa_system_sgpr_private_segment_wavefront_offset`` 0                   GFX6-GFX10   Controls ENABLE_PRIVATE_SEGMENT in
                                                                                                :ref:`amdgpu-amdhsa-compute_pgm_rsrc2-gfx6-gfx10-table`.
      ``.amdhsa_system_sgpr_workgroup_id_x``                   1                   GFX6-GFX10   Controls ENABLE_SGPR_WORKGROUP_ID_X in
                                                                                                :ref:`amdgpu-amdhsa-compute_pgm_rsrc2-gfx6-gfx10-table`.
      ``.amdhsa_system_sgpr_workgroup_id_y``                   0                   GFX6-GFX10   Controls ENABLE_SGPR_WORKGROUP_ID_Y in
                                                                                                :ref:`amdgpu-amdhsa-compute_pgm_rsrc2-gfx6-gfx10-table`.
      ``.amdhsa_system_sgpr_workgroup_id_z``                   0                   GFX6-GFX10   Controls ENABLE_SGPR_WORKGROUP_ID_Z in
                                                                                                :ref:`amdgpu-amdhsa-compute_pgm_rsrc2-gfx6-gfx10-table`.
      ``.amdhsa_system_sgpr_workgroup_info``                   0                   GFX6-GFX10   Controls ENABLE_SGPR_WORKGROUP_INFO in

llvm/include/llvm/Support/AMDHSAKernelDescriptor.h

enum : int32_t {
   COMPUTE_PGM_RSRC1(FWD_PROGRESS, 31, 1), // GFX10+
 };
 #undef COMPUTE_PGM_RSRC1
 
 // Compute program resource register 2. Must match hardware definition.
 #define COMPUTE_PGM_RSRC2(NAME, SHIFT, WIDTH) \
   AMDHSA_BITS_ENUM_ENTRY(COMPUTE_PGM_RSRC2_ ## NAME, SHIFT, WIDTH)
 enum : int32_t {
-  COMPUTE_PGM_RSRC2(ENABLE_SGPR_PRIVATE_SEGMENT_WAVEFRONT_OFFSET, 0, 1),
+  COMPUTE_PGM_RSRC2(ENABLE_PRIVATE_SEGMENT, 0, 1),
   COMPUTE_PGM_RSRC2(USER_SGPR_COUNT, 1, 5),
   COMPUTE_PGM_RSRC2(ENABLE_TRAP_HANDLER, 6, 1),
   COMPUTE_PGM_RSRC2(ENABLE_SGPR_WORKGROUP_ID_X, 7, 1),
   COMPUTE_PGM_RSRC2(ENABLE_SGPR_WORKGROUP_ID_Y, 8, 1),
   COMPUTE_PGM_RSRC2(ENABLE_SGPR_WORKGROUP_ID_Z, 9, 1),
   COMPUTE_PGM_RSRC2(ENABLE_SGPR_WORKGROUP_INFO, 10, 1),
   COMPUTE_PGM_RSRC2(ENABLE_VGPR_WORKITEM_ID, 11, 2),
   COMPUTE_PGM_RSRC2(ENABLE_EXCEPTION_ADDRESS_WATCH, 13, 1),

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

 //===- AMDGPUAsmParser.cpp - Parse SI asm to MCInst instructions ----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "AMDGPU.h"
if (ID == ".amdhsa_group_segment_fixed_size") {
                                  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_SGPR_PRIVATE_SEGMENT_WAVEFRONT_OFFSET, Val,
+          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,

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

 //===- AMDGPUDisassembler.cpp - Disassembler for AMDGPU ISA ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 //===----------------------------------------------------------------------===//

 
 // 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",
-      COMPUTE_PGM_RSRC2_ENABLE_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);

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

 //===-- AMDGPUTargetStreamer.cpp - Mips Target Streamer Methods -----------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file provides AMDGPU specific target streamer methods.
PRINT_FIELD(OS, ".amdhsa_user_sgpr_private_segment_size", KD,
               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,
-      amdhsa::COMPUTE_PGM_RSRC2_ENABLE_SGPR_PRIVATE_SEGMENT_WAVEFRONT_OFFSET);
+      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,