Implement alternate scratch

The alternate scratch memory is used for dispatches that have a low number of waves but relatively large wave size. This allows us to keep the tmpring_size.bits.WAVES field of the main scratch to full occupancy. Change-Id: I32d240fac4b7d38200d1eebc1b0fdc8a823920d3 [ROCm/ROCR-Runtime commit: a7a3358067]
2023-11-15 16:44:11 +00:00
@@ -215,6 +215,9 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
  /// @brief Async reclaim main scratch memory
  void AsyncReclaimMainScratch();

+  /// @brief Async reclaim alternate scratch memory
+  void AsyncReclaimAltScratch();
+
 protected:
  bool _IsA(Queue::rtti_t id) const override { return id == &rtti_id_; }

@@ -241,9 +244,11 @@ class AqlQueue : public core::Queue, private core::LocalSignal, public core::Doo
  void FillBufRsrcWord3_Gfx10();
  void FillBufRsrcWord3_Gfx11();
  void FillComputeTmpRingSize();
+  void FillAltComputeTmpRingSize();
  void FillComputeTmpRingSize_Gfx11();

  void FreeMainScratchSpace();
+  void FreeAltScratchSpace();

  /// @brief Halt the queue without destroying it or fencing memory.
  void Suspend();
@@ -104,12 +104,24 @@ class GpuAgentInt : public core::Agent {
  // information.
  virtual void AcquireQueueMainScratch(ScratchInfo& scratch) = 0;

+  // @brief Carve scratch memory for alt from scratch pool.
+  //
+  // @param [in/out] scratch Structure to be populated with the carved memory
+  // information.
+  virtual void AcquireQueueAltScratch(ScratchInfo& scratch) = 0;
+
  // @brief Release scratch memory from main back to scratch pool.
  //
  // @param [in/out] scratch Scratch memory previously acquired with call to
  // ::AcquireQueueMainScratch.
  virtual void ReleaseQueueMainScratch(ScratchInfo& base) = 0;

+  // @brief Release scratch memory back from alternate to scratch pool.
+  //
+  // @param [in/out] scratch Scratch memory  previously acquired with call to
+  // ::AcquireQueueAltcratch.
+  virtual void ReleaseQueueAltScratch(ScratchInfo& base) = 0;
+
  // @brief Translate the kernel start and end dispatch timestamp from agent
  // domain to host domain.
  //
@@ -278,6 +290,9 @@ class GpuAgent : public GpuAgentInt {
  void AcquireQueueMainScratch(ScratchInfo& scratch) override;
  void ReleaseQueueMainScratch(ScratchInfo& scratch) override;

+  void AcquireQueueAltScratch(ScratchInfo& scratch) override;
+  void ReleaseQueueAltScratch(ScratchInfo& scratch) override;
+
  // @brief Override from AMD::GpuAgentInt.
  void TranslateTime(core::Signal* signal, hsa_amd_profiling_dispatch_time_t& time) override;

@@ -92,6 +92,7 @@ class ScratchCache {

    bool large;
    size_t use_once_limit;
+    size_t use_alt_limit;
    bool async_reclaim;  // This version of CP FW supports async_reclaim
    bool retry;
    uint32_t mem_alignment_size;  // Populated into SRD
@@ -106,6 +107,18 @@ class ScratchCache {
    void* main_queue_base;
    ptrdiff_t main_queue_process_offset;
    ScratchCache::ref_t main_scratch_node;
+
+    size_t alt_size;
+    size_t alt_size_per_thread;    // Populated into SRD
+    uint32_t alt_lanes_per_wave;   // Populated into SRD
+    uint32_t alt_waves_per_group;  // Used during waves reduction
+
+    uint64_t alt_dispatch_limit_x;
+    uint64_t alt_dispatch_limit_y;
+    uint64_t alt_dispatch_limit_z;
+    void* alt_queue_base;
+    ptrdiff_t alt_queue_process_offset;
+    ScratchCache::ref_t alt_scratch_node;
  };

  ScratchCache(const ScratchCache& rhs) = delete;
@@ -198,6 +211,46 @@ class ScratchCache {
    return ret;
  }

+  bool allocAlt(ScratchInfo& info) {
+    ref_t it = map.upper_bound(info.alt_size - 1);
+    if (it == map.end()) return false;
+
+    // Alt requests should have exact size
+    while ((it != map.end()) && (it->first == info.alt_size)) {
+      if (it->second.isFree() && (!it->second.large)) {
+        it->second.alloc();
+        info.alt_queue_base = it->second.base;
+        info.alt_scratch_node = it;
+        available_bytes_ -= it->first;
+        return true;
+      }
+      it++;
+    }
+    return false;
+  }
+
+  void freeAlt(ScratchInfo& info) {
+    assert(!info.alt_scratch_node->second.isFree() && "free called on free scratch node.");
+    auto it = info.alt_scratch_node;
+    if (it->second.trimPending()) {
+      dealloc(it->second.base, it->first, it->second.large);
+      map.erase(it);
+      return;
+    }
+    it->second.free();
+    available_bytes_ += it->first;
+  }
+
+  void insertAlt(ScratchInfo& info) {
+    node n;
+    n.base = info.alt_queue_base;
+    n.large = false;
+    n.alloc();
+
+    auto it = map.insert(std::make_pair(info.alt_size, n));
+    info.alt_scratch_node = it;
+  }
+
  size_t free_bytes() const { return available_bytes_; }
  size_t reserved_bytes() const { return reserved_.first; }

@@ -217,9 +217,15 @@ AqlQueue::AqlQueue(GpuAgent* agent, size_t req_size_pkts, HSAuint32 node_id, Scr
    queue_scratch_.mem_alignment_size = 1024;

  queue_scratch_.use_once_limit = core::Runtime::runtime_singleton_->flag().scratch_single_limit();
+  queue_scratch_.use_alt_limit = 0;
+
  queue_scratch_.async_reclaim = agent_->AsyncScratchReclaimEnabled();
-  if (queue_scratch_.async_reclaim)
+  if (queue_scratch_.async_reclaim) {
    queue_scratch_.use_once_limit = agent_->ScratchSingleLimitAsyncThreshold();
+    queue_scratch_.use_alt_limit = core::Runtime::runtime_singleton_->flag().enable_scratch_alt()
+        ? (queue_scratch_.use_once_limit / 4)
+        : 0;
+  }

  MAKE_NAMED_SCOPE_GUARD(EventGuard, [&]() {
    ScopedAcquire<KernelMutex> _lock(&queue_lock_);
@@ -289,6 +295,7 @@ AqlQueue::AqlQueue(GpuAgent* agent, size_t req_size_pkts, HSAuint32 node_id, Scr
  MAKE_NAMED_SCOPE_GUARD(QueueGuard, [&]() { hsaKmtDestroyQueue(queue_id_); });

  amd_queue_.scratch_last_used_index = UINT64_MAX;
+  amd_queue_.alt_scratch_last_used_index = UINT64_MAX;

  // On the first queue creation, reserve some scratch memory on this agent.
  agent_->ReserveScratch();
@@ -356,6 +363,8 @@ AqlQueue::~AqlQueue() {

  Inactivate();
  agent_->ReleaseQueueMainScratch(queue_scratch_);
+  agent_->ReleaseQueueAltScratch(queue_scratch_);
+
  FreeRegisteredRingBuffer();
  exception_signal_->DestroySignal();
  HSA::hsa_signal_destroy(amd_queue_.queue_inactive_signal);
@@ -790,9 +799,14 @@ void AqlQueue::CheckScratchLimits() {
  if (!scratch.async_reclaim) return;

  scratch.use_once_limit = agent_->ScratchSingleLimitAsyncThreshold();
+  scratch.use_alt_limit = core::Runtime::runtime_singleton_->flag().enable_scratch_alt()
+      ? (scratch.use_once_limit / 4)
+      : 0;

  if (scratch.main_size > scratch.use_once_limit) AsyncReclaimMainScratch();

+  if (scratch.alt_size > scratch.use_alt_limit) AsyncReclaimAltScratch();
+
  return;
 }

@@ -828,6 +842,38 @@ void AqlQueue::AsyncReclaimMainScratch() {
  }
 }

+void AqlQueue::FreeAltScratchSpace() {
+  auto& scratch = queue_scratch_;
+  agent_->ReleaseQueueAltScratch(scratch);
+  scratch.alt_size = 0;
+  scratch.alt_size_per_thread = 0;
+  scratch.alt_queue_process_offset = 0;
+  InitScratchSRD();
+
+  HSA::hsa_signal_store_relaxed(amd_queue_.queue_inactive_signal, 0);
+}
+
+void AqlQueue::AsyncReclaimAltScratch() {
+  auto& scratch = queue_scratch_;
+  if (!scratch.async_reclaim || !scratch.alt_size) return;
+
+  // Notify CP that we are trying to reclaim scratch. CP will assume scratch is reclaimed on next
+  // dispatch
+  amd_queue_.alt_scratch_wave64_lane_byte_size = 0;
+  uint64_t last_used = atomic::Exchange(
+      &amd_queue_.alt_scratch_last_used_index, UINT64_MAX,
+      std::memory_order_relaxed);  // TODO: Confirm this is the correct memory order
+  // Wait for scratch to be idle.
+  while (true) {
+    uint64_t last = amd_queue_.alt_scratch_last_used_index;
+
+    if (std::min(last, last_used) < amd_queue_.read_dispatch_id) {
+      FreeAltScratchSpace();
+      return;
+    }
+  }
+}
+
 void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
                                         hsa_signal_value_t& waitVal, bool& changeWait) {
  // Insufficient scratch - recoverable, don't process dynamic scratch if errors are present.
@@ -844,6 +890,25 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
   * size_t use_once_limit = 128 MB      // When async reclaim not supported
   *                       = 1GB per-XCC // When async reclaim is supported
   *
+   * size_t use_alt_limit  = 256 MB per-XCC
+   *
+   * if (async-scratch-reclaim-supported
+   *     && dispatch_slots < max_scratch_slots
+   *     && dispatch_size < use_alt_limit) {
+   *   // This dispatch wants less waves than number of slots, use alternate scratch
+   *   // alt_tmpring_size will have limited waves
+   *  use_alt()
+   * } else if (all_slots_size <= use_once_limit) {
+   *  use_main()
+   *
+   *  //If we failed to allocate memory to fill all slots, scratch.use_once will be set
+   *  if (scratch.use_once) {
+   *    use_once
+   *  } else if (all_slots_size > scratch.alt_size) {
+   *    //Primary scratch is large enough to handle needs of alt-scratch
+   *    free_alt()
+   *  }
+   * }
   *
   *******************************************************************************************/

@@ -925,6 +990,35 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
  const uint64_t device_size = size_per_thread * lanes_per_wave * device_slots;
  const uint64_t dispatch_size = size_per_thread * lanes_per_wave * dispatch_slots;

+  // scratch.use_alt_limit will be 0 if alt scratch is not supported or disabled
+  if (dispatch_size < scratch.use_alt_limit && dispatch_slots < device_slots) {
+    // Try to use ALT scratch
+    agent_->ReleaseQueueAltScratch(scratch);
+
+    scratch.alt_size = dispatch_size;
+    scratch.alt_size_per_thread = size_per_thread;
+    scratch.alt_lanes_per_wave = lanes_per_wave;
+    scratch.alt_waves_per_group = waves_per_group;
+
+    agent_->AcquireQueueAltScratch(scratch);
+    if (scratch.alt_queue_base) {
+      scratch.alt_dispatch_limit_x = pkt.dispatch.grid_size_x;
+      scratch.alt_dispatch_limit_y = pkt.dispatch.grid_size_y;
+      scratch.alt_dispatch_limit_z = pkt.dispatch.grid_size_z;
+
+      // Update queue SRD
+      InitScratchSRD();
+      // Restart the queue.
+      HSA::hsa_signal_store_screlease(amd_queue_.queue_inactive_signal, 0);
+
+      return;
+    }
+    // Could not allocate enough memory for alternate scratch fallback to primary scratch
+    scratch.alt_size = 0;
+    scratch.alt_size_per_thread = 0;
+  }
+
+  // Use PRIMARY scratch
  agent_->ReleaseQueueMainScratch(scratch);
  scratch.main_size = device_size;
  scratch.main_size_per_thread = size_per_thread;
@@ -954,6 +1048,10 @@ void AqlQueue::HandleInsufficientScratch(hsa_signal_value_t& error_code,
                               << HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE);
      pkt.dispatch.header |= (HSA_FENCE_SCOPE_SYSTEM << HSA_PACKET_HEADER_SCRELEASE_FENCE_SCOPE);
    }
+  } else if (scratch.alt_size && scratch.main_size > scratch.alt_size) {
+    // Not using use-scratch-once, and dispatches that would fit in alt-scratch would also fit in
+    // main scratch. No need for alt-scratch.
+    AsyncReclaimAltScratch();
  }

  // Reset scratch memory related entities for the queue
@@ -1004,7 +1102,8 @@ bool AqlQueue::DynamicQueueEventsHandler(hsa_signal_value_t error_code, void* ar
      queue->HandleInsufficientScratch(error_code, waitVal, changeWait);

      // Out of scratch - promote error
-      if (queue->queue_scratch_.main_queue_base == nullptr)
+      if (queue->queue_scratch_.main_queue_base == nullptr &&
+          queue->queue_scratch_.alt_queue_base == nullptr)
        errorCode = HSA_STATUS_ERROR_OUT_OF_RESOURCES;


@@ -1514,6 +1613,38 @@ void AqlQueue::FillComputeTmpRingSize() {
         "Invalid scratch wave count.  Must be divisible by #SEs.");
 }

+// Set concurrent wavefront limits only when scratch is being used.
+void AqlQueue::FillAltComputeTmpRingSize() {
+  COMPUTE_TMPRING_SIZE tmpring_size = {};
+  if (queue_scratch_.alt_size == 0) {
+    amd_queue_.alt_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;
+  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_.alt_lanes_per_wave * queue_scratch_.alt_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_.alt_size / num_xcc) /
+      (tmpring_size.bits.WAVESIZE * queue_scratch_.mem_alignment_size);
+
+  tmpring_size.bits.WAVES = std::min(num_waves, max_scratch_waves);
+  amd_queue_.alt_compute_tmpring_size = tmpring_size.u32All;
+  assert((tmpring_size.bits.WAVES % (agent_props.NumShaderBanks / num_xcc) == 0) &&
+         "Invalid scratch wave count.  Must be divisible by #SEs.");
+}
+
 // Set concurrent wavefront limits only when scratch is being used.
 void AqlQueue::FillComputeTmpRingSize_Gfx11() {
  COMPUTE_TMPRING_SIZE_GFX11 tmpring_size = {};
@@ -1572,16 +1703,19 @@ void AqlQueue::InitScratchSRD() {
      FillBufRsrcWord2();
      FillBufRsrcWord3();
      FillComputeTmpRingSize();
+      FillAltComputeTmpRingSize();
      break;
  }

  // Populate flat scratch parameters in amd_queue_.
  amd_queue_.scratch_backing_memory_location = queue_scratch_.main_queue_process_offset;
+  amd_queue_.alt_scratch_backing_memory_location = queue_scratch_.alt_queue_process_offset;

  const auto& agent_props = agent_->properties();
  const uint32_t num_xcc = agent_props.NumXcc;
  // report size per XCC
  amd_queue_.scratch_backing_memory_byte_size = queue_scratch_.main_size / num_xcc;
+  amd_queue_.alt_scratch_backing_memory_byte_size = queue_scratch_.alt_size / num_xcc;

  // For backwards compatibility this field records the per-lane scratch
  // for a 64 lane wavefront. If scratch was allocated for 32 lane waves
@@ -1589,6 +1723,13 @@ void AqlQueue::InitScratchSRD() {
  amd_queue_.scratch_wave64_lane_byte_size =
      uint32_t((queue_scratch_.main_size_per_thread * queue_scratch_.main_lanes_per_wave) / 64);

+  amd_queue_.alt_scratch_wave64_lane_byte_size =
+      uint32_t((queue_scratch_.alt_size_per_thread * queue_scratch_.alt_lanes_per_wave) / 64);
+
+  amd_queue_.alt_scratch_dispatch_limit_x = queue_scratch_.alt_dispatch_limit_x;
+  amd_queue_.alt_scratch_dispatch_limit_y = queue_scratch_.alt_dispatch_limit_y;
+  amd_queue_.alt_scratch_dispatch_limit_z = queue_scratch_.alt_dispatch_limit_z;
+
  return;
 }

@@ -1808,6 +1808,81 @@ void GpuAgent::ReleaseQueueMainScratch(ScratchInfo& scratch) {
  scratch.main_queue_base = nullptr;
 }

+void GpuAgent::AcquireQueueAltScratch(ScratchInfo& scratch) {
+  assert(scratch.async_reclaim && "Acquire Alt Scratch when FW does not support it");
+  assert(scratch.alt_queue_base == nullptr &&
+         "AcquireQueueAltScratch called while holding alt scratch.");
+
+  // Fail scratch allocation if per wave limits are exceeded.
+  uint64_t size_per_wave = AlignUp(scratch.alt_size_per_thread * properties_.WaveFrontSize, 1024);
+  if (size_per_wave > MAX_WAVE_SCRATCH) return;
+
+  ScopedAcquire<KernelMutex> lock(&scratch_lock_);
+
+  // Ensure mapping will be in whole pages.
+  scratch.alt_size = AlignUp(scratch.alt_size, 4096);
+
+  /*
+  Sequence of attempts is:
+    check cache
+    attempt a new allocation
+    trim unused blocks from cache
+    attempt a new allocation
+    check cache for sufficient used block, steal and wait (not implemented)
+    trim used blocks from cache, evaluate retry
+  */
+
+  // Lambda called in place.
+  // Used to allow exit from nested loops.
+  [&]() {
+    // Check scratch cache
+    if (scratch_cache_.allocAlt(scratch)) return;
+
+    // Attempt new allocation.
+    for (int i = 0; i < 2; i++) {
+      scratch.alt_queue_base = scratch_pool_.alloc(scratch.alt_size);
+      if (scratch.alt_queue_base != nullptr) {
+        HSAuint64 alternate_va;
+        if ((profile_ == HSA_PROFILE_FULL) ||
+            (hsaKmtMapMemoryToGPU(scratch.alt_queue_base, scratch.alt_size, &alternate_va) ==
+             HSAKMT_STATUS_SUCCESS)) {
+          scratch_cache_.insertAlt(scratch);
+          return;
+        }
+      }
+
+      // Scratch request failed allocation or mapping.
+      scratch_pool_.free(scratch.alt_queue_base);
+      scratch.alt_queue_base = nullptr;
+
+      // Release cached scratch and retry.
+      // First iteration trims unused blocks, second trims all. 3rd uses reserved memory
+      switch (i) {
+        case 0:
+          scratch_cache_.trim(false);
+          break;
+        case 1:
+          scratch_cache_.trim(true);
+          break;
+      }
+    }
+
+    if (core::Runtime::runtime_singleton_->flag().enable_queue_fault_message())
+      debug_print("  Could not allocate alt scratch.\n");
+    return;
+  }();
+
+  scratch.alt_queue_process_offset = uintptr_t(scratch.alt_queue_base);
+}
+
+void GpuAgent::ReleaseQueueAltScratch(ScratchInfo& scratch) {
+  ScopedAcquire<KernelMutex> lock(&scratch_lock_);
+  if (scratch.alt_queue_base == nullptr) return;
+
+  scratch_cache_.freeAlt(scratch);
+  scratch.alt_queue_base = nullptr;
+}
+
 void GpuAgent::ReleaseScratch(void* base, size_t size, bool large) {
  if (profile_ == HSA_PROFILE_BASE) {
    if (HSAKMT_STATUS_SUCCESS != hsaKmtUnmapMemoryToGPU(base)) {
@@ -1830,6 +1905,7 @@ void GpuAgent::AsyncReclaimScratchQueues() {
  for (auto iter : aql_queues_) {
    auto aqlQueue = static_cast<AqlQueue*>(iter);
    aqlQueue->AsyncReclaimMainScratch();
+    aqlQueue->AsyncReclaimAltScratch();
  }
 }

@@ -138,9 +138,14 @@ class Flag {
    }

    // On GPUs that support asynchronous scratch reclaim this can be used to disable this feature.
+    // Disabling asynchronous scratch reclaim also disables use of alternate scratch
+    // HSA_ENABLE_SCRATCH_ALT
    var = os::GetEnvVar("HSA_ENABLE_SCRATCH_ASYNC_RECLAIM");
    enable_scratch_async_reclaim_ = (var == "0") ? false : true;

+    var = os::GetEnvVar("HSA_ENABLE_SCRATCH_ALT");
+    enable_scratch_alt_ = (var == "0") || !enable_scratch_async_reclaim_ ? false : true;
+
    tools_lib_names_ = os::GetEnvVar("HSA_TOOLS_LIB");

    var = os::GetEnvVar("HSA_TOOLS_REPORT_LOAD_FAILURE");
@@ -270,6 +275,8 @@ class Flag {

  bool enable_scratch_async_reclaim() const { return enable_scratch_async_reclaim_; }

+  bool enable_scratch_alt() const { return enable_scratch_alt_; }
+
  size_t scratch_single_limit_async() const { return scratch_single_limit_async_; }

  std::string tools_lib_names() const { return tools_lib_names_; }
@@ -352,6 +359,7 @@ class Flag {
  size_t scratch_single_limit_;
  size_t scratch_single_limit_async_;
  bool enable_scratch_async_reclaim_;
+  bool enable_scratch_alt_;

  std::string tools_lib_names_;
  std::string svm_profile_;