From 5565b53d15ee3aa7e8c954067377c26ba8554993 Mon Sep 17 00:00:00 2001
From: David Belanger <david.belanger@amd.com>
Date: Wed, 24 Jan 2024 12:32:31 -0500
Subject: [PATCH] Updated amd_aql_queue for GFX12

Added GFX12 implementation for InitScratchSRD and for compute_tmpring.
Implementation for compute_tmpring could be combined with GFX11 with some
refactoring as a possible future improvement.

Signed-off-by: David Belanger <david.belanger@amd.com>
Change-Id: I8013cbe4438786bf41bbfd03f6a5d3b9ef51e7bf
Signed-off-by: Chris Freehill <cfreehil@amd.com>


[ROCm/ROCR-Runtime commit: def4a6c326380de40d8c96abf80f56c3a7a83067]
---
 .../hsa-runtime/core/inc/amd_aql_queue.h      |  2 +
 .../core/runtime/amd_aql_queue.cpp            | 61 +++++++++++++++++++
 2 files changed, 63 insertions(+)

diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_aql_queue.h b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_aql_queue.h
index 559415f23d..09f14f9410 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_aql_queue.h
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/inc/amd_aql_queue.h
@@ -252,9 +252,11 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
   void FillBufRsrcWord3();
   void FillBufRsrcWord3_Gfx10();
   void FillBufRsrcWord3_Gfx11();
+  void FillBufRsrcWord3_Gfx12();
   void FillComputeTmpRingSize();
   void FillAltComputeTmpRingSize();
   void FillComputeTmpRingSize_Gfx11();
+  void FillComputeTmpRingSize_Gfx12();
 
   void FreeMainScratchSpace();
   void FreeAltScratchSpace();
diff --git a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp
index 3dad89cf69..4720a6b32b 100644
--- a/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp
+++ b/projects/rocr-runtime/runtime/hsa-runtime/core/runtime/amd_aql_queue.cpp
@@ -1674,6 +1674,26 @@ void AqlQueue::FillBufRsrcWord3_Gfx11() {
   amd_queue_.scratch_resource_descriptor[3] = srd3.u32All;
 }
 
+void AqlQueue::FillBufRsrcWord3_Gfx12() {
+  SQ_BUF_RSRC_WORD3_GFX12 srd3;
+
+  srd3.bits.DST_SEL_X = SQ_SEL_X;
+  srd3.bits.DST_SEL_Y = SQ_SEL_Y;
+  srd3.bits.DST_SEL_Z = SQ_SEL_Z;
+  srd3.bits.DST_SEL_W = SQ_SEL_W;
+  srd3.bits.FORMAT = BUF_FORMAT_32_UINT;
+  srd3.bits.RESERVED1 = 0;
+  srd3.bits.INDEX_STRIDE = 0;  // filled in by CP
+  srd3.bits.ADD_TID_ENABLE = 1;
+  srd3.bits.WRITE_COMPRESS_ENABLE = 0;
+  srd3.bits.COMPRESSION_EN = 0;
+  srd3.bits.COMPRESSION_ACCESS_MODE = 0;
+  srd3.bits.OOB_SELECT = 2;  // no bounds check in swizzle mode
+  srd3.bits.TYPE = SQ_RSRC_BUF;
+
+  amd_queue_.scratch_resource_descriptor[3] = srd3.u32All;
+}
+
 // Set concurrent wavefront limits only when scratch is being used.
 void AqlQueue::FillComputeTmpRingSize() {
   COMPUTE_TMPRING_SIZE tmpring_size = {};
@@ -1772,11 +1792,52 @@ void AqlQueue::FillComputeTmpRingSize_Gfx11() {
   amd_queue_.compute_tmpring_size = tmpring_size.u32All;
 }
 
+// Set concurrent wavefront limits only when scratch is being used.
+void AqlQueue::FillComputeTmpRingSize_Gfx12() {
+  // For GFX12, struct field size changes.
+  // Consider refactoring code for GFX11/GFX12 if no other changes.
+  COMPUTE_TMPRING_SIZE_GFX12 tmpring_size = {};
+  if (queue_scratch_.main_size == 0) {
+    amd_queue_.compute_tmpring_size = tmpring_size.u32All;
+    return;
+  }
+
+  const auto& agent_props = agent_->properties();
+  const uint32_t num_xcc = agent_props.NumXcc;
+
+  // Determine the maximum number of waves device can support
+  uint32_t num_cus = agent_props.NumFComputeCores / (agent_props.NumSIMDPerCU * num_xcc);
+  uint32_t max_scratch_waves = num_cus * agent_props.MaxSlotsScratchCU;
+
+  // Scratch is allocated program COMPUTE_TMPRING_SIZE register
+  // Scratch Size per Wave is specified in terms of kilobytes
+  uint32_t wave_scratch = (((queue_scratch_.main_lanes_per_wave * queue_scratch_.main_size_per_thread) +
+                            queue_scratch_.mem_alignment_size - 1) /
+                           queue_scratch_.mem_alignment_size);
+
+  tmpring_size.bits.WAVESIZE = wave_scratch;
+  assert(wave_scratch == tmpring_size.bits.WAVESIZE && "WAVESIZE Overflow.");
+
+  uint32_t num_waves =
+      queue_scratch_.main_size / (tmpring_size.bits.WAVESIZE * queue_scratch_.mem_alignment_size);
+
+  // For GFX11 we specify number of waves per engine instead of total
+  num_waves /= agent_->properties().NumShaderBanks;
+  tmpring_size.bits.WAVES = std::min(num_waves, max_scratch_waves);
+  amd_queue_.compute_tmpring_size = tmpring_size.u32All;
+}
+
 // @brief Define the Scratch Buffer Descriptor and related parameters
 // that enable kernel access scratch memory
 void AqlQueue::InitScratchSRD() {
   switch (agent_->isa()->GetMajorVersion()) {
     case 12:
+      FillBufRsrcWord0();
+      FillBufRsrcWord1_Gfx11();
+      FillBufRsrcWord2();
+      FillBufRsrcWord3_Gfx12();
+      FillComputeTmpRingSize_Gfx12();
+      break;
     case 11:
       FillBufRsrcWord0();
       FillBufRsrcWord1_Gfx11();