Cache scratch allocations.

Avoids calling to KFD to map/unmap scratch allocations for
every large scratch using dispatch.

Change-Id: I9fab5705251ec82b03e4f2f2ca6da7cdccabefb9

runtime/hsa-runtime/core/inc/agent.h

@@ -61,8 +61,6 @@ class Signal;
 typedef void (*HsaEventCallback)(hsa_status_t status, hsa_queue_t* source,
                                  void* data);
 
-class MemoryRegion;
-
 // Agent is intended to be an pure interface class and may be wrapped or
 // replaced by tools libraries. All funtions other than Convert, node_id,
 // device_type, and public_handle must be virtual.
@@ -260,6 +258,10 @@ class Agent : public Checked<0xF6BC25EB17E6F917> {
     return stat;
   }
 
+  virtual void Trim() {
+    for (auto region : regions()) region->Trim();
+  }
+
  protected:
   // Intention here is to have a polymorphic update procedure for public_handle_
   // which is callable on any Agent* but only from some class dervied from

runtime/hsa-runtime/core/inc/amd_gpu_agent.h

@@ -55,6 +55,7 @@
 #include "core/inc/blit.h"
 #include "core/inc/signal.h"
 #include "core/inc/cache.h"
+#include "core/inc/scratch_cache.h"
 #include "core/util/small_heap.h"
 #include "core/util/locks.h"
 #include "core/util/lazy_ptr.h"
@@ -63,18 +64,7 @@ namespace rocr {
 namespace AMD {
 class MemoryRegion;
 
-// @brief Contains scratch memory information.
-struct ScratchInfo {
-  void* queue_base;
-  size_t size;
-  size_t size_per_thread;
-  uint32_t lanes_per_wave;
-  ptrdiff_t queue_process_offset;
-  bool large;
-  bool retry;
-  hsa_signal_t queue_retry;
-  uint64_t wanted_slots;
-};
+typedef ScratchCache::ScratchInfo ScratchInfo;
 
 // @brief Interface to represent a GPU agent.
 class GpuAgentInt : public core::Agent {
@@ -331,6 +321,8 @@ class GpuAgent : public GpuAgentInt {
     return memory_max_frequency_;
   }
 
+  void Trim() override;
+
  protected:
   static const uint32_t minAqlSize_ = 0x1000;   // 4KB min
   static const uint32_t maxAqlSize_ = 0x20000;  // 8MB max
@@ -494,6 +486,10 @@ class GpuAgent : public GpuAgentInt {
   // @brief Deregister scratch notification signals.
   void ClearScratchNotifiers() { scratch_notifiers_.clear(); }
 
+  // @brief Releases scratch back to the driver.
+  // caller must hold scratch_lock_.
+  void ReleaseScratch(void* base, size_t size, bool large);
+
   // Bind index of peer device that is connected via xGMI links
   lazy_ptr<core::Blit>& GetXgmiBlit(const core::Agent& peer_agent);
 
@@ -517,6 +513,8 @@ class GpuAgent : public GpuAgentInt {
     KernelMutex lock_;
   } gws_queue_;
 
+  ScratchCache scratch_cache_;
+
   DISALLOW_COPY_AND_ASSIGN(GpuAgent);
 };
 

runtime/hsa-runtime/core/inc/amd_memory_region.h

@@ -136,6 +136,8 @@ class MemoryRegion : public core::MemoryRegion {
   hsa_status_t AssignAgent(void* ptr, size_t size, const core::Agent& agent,
                            hsa_access_permission_t access) const;
 
+  void Trim() const;
+
   __forceinline bool IsLocalMemory() const {
     return ((mem_props_.HeapType == HSA_HEAPTYPE_FRAME_BUFFER_PRIVATE) ||
             (mem_props_.HeapType == HSA_HEAPTYPE_FRAME_BUFFER_PUBLIC));

runtime/hsa-runtime/core/inc/memory_region.h

@@ -47,8 +47,9 @@
 
 #include <vector>
 
-#include "core/inc/agent.h"
+#include "core/inc/hsa_internal.h"
 #include "core/inc/checked.h"
+#include "core/util/utils.h"
 
 namespace rocr {
 namespace core {
@@ -106,6 +107,9 @@ class MemoryRegion : public Checked<0x9C961F19EE175BB3> {
   virtual hsa_status_t AssignAgent(void* ptr, size_t size, const Agent& agent,
                                    hsa_access_permission_t access) const = 0;
 
+  // Releases any cached memory that may be held within the allocator.
+  virtual void Trim() const {}
+
   __forceinline bool fine_grain() const { return fine_grain_; }
 
   __forceinline bool full_profile() const { return full_profile_; }

runtime/hsa-runtime/core/inc/scratch_cache.h

@@ -0,0 +1,191 @@
+////////////////////////////////////////////////////////////////////////////////
+//
+// The University of Illinois/NCSA
+// Open Source License (NCSA)
+//
+// Copyright (c) 2020-2020, Advanced Micro Devices, Inc. All rights reserved.
+//
+// Developed by:
+//
+//                 AMD Research and AMD HSA Software Development
+//
+//                 Advanced Micro Devices, Inc.
+//
+//                 www.amd.com
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to
+// deal with the Software without restriction, including without limitation
+// the rights to use, copy, modify, merge, publish, distribute, sublicense,
+// and/or sell copies of the Software, and to permit persons to whom the
+// Software is furnished to do so, subject to the following conditions:
+//
+//  - Redistributions of source code must retain the above copyright notice,
+//    this list of conditions and the following disclaimers.
+//  - Redistributions in binary form must reproduce the above copyright
+//    notice, this list of conditions and the following disclaimers in
+//    the documentation and/or other materials provided with the distribution.
+//  - Neither the names of Advanced Micro Devices, Inc,
+//    nor the names of its contributors may be used to endorse or promote
+//    products derived from this Software without specific prior written
+//    permission.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
+// THE CONTRIBUTORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR
+// OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE,
+// ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
+// DEALINGS WITH THE SOFTWARE.
+//
+////////////////////////////////////////////////////////////////////////////////
+
+#ifndef HSA_RUNTIME_CORE_INC_SCRATCH_CACHE_H_
+#define HSA_RUNTIME_CORE_INC_SCRATCH_CACHE_H_
+
+#include "core/inc/amd_gpu_agent.h"
+#include "core/util/locks.h"
+#include "core/util/utils.h"
+
+#include <map>
+#include <functional>
+
+namespace rocr {
+namespace AMD {
+
+class ScratchCache {
+ public:
+  struct node {
+    enum STATE { FREE = 0, ALLOC = 1, TRIM = 2, STEAL = 4 };
+    void* base;
+    bool large;
+    uint32_t state;
+
+    node() : base(nullptr), state(FREE) {}
+
+    bool isFree() const { return state == FREE; }
+    bool trimPending() const { return state == (ALLOC | TRIM); }
+
+    void trim() {
+      assert(!isFree() && "Trim of free scratch node.");
+      state |= TRIM;
+    }
+    void free() {
+      assert(!isFree() && "Free of free scratch node.");
+      state = FREE;
+    }
+    void alloc() {
+      assert(isFree() && "Alloc of non-free scratch node.");
+      state = ALLOC;
+    }
+  };
+
+  typedef ::std::multimap<size_t, node> map_t;
+  typedef map_t::iterator ref_t;
+  typedef ::std::function<void(void*, size_t, bool)> deallocator_t;
+
+  // @brief Contains scratch memory information.
+  struct ScratchInfo {
+    void* queue_base;
+    // Size to fill the machine with size_per_thread
+    size_t size;
+    // Size to satisfy the present dispatch without throttling.
+    size_t dispatch_size;
+    size_t size_per_thread;
+    uint32_t lanes_per_wave;
+    ptrdiff_t queue_process_offset;
+    bool large;
+    bool retry;
+    hsa_signal_t queue_retry;
+    uint64_t wanted_slots;
+    ScratchCache::ref_t scratch_node;
+  };
+
+  ScratchCache(const ScratchCache& rhs) = delete;
+  ScratchCache(ScratchCache&& rhs) = delete;
+  ScratchCache& operator=(const ScratchCache& rhs) = delete;
+  ScratchCache& operator=(ScratchCache&& rhs) = delete;
+
+  ScratchCache(deallocator_t deallocator) : dealloc(deallocator) {}
+
+  ~ScratchCache() { assert(map.empty() && "ScratchCache not empty at shutdown."); }
+
+  bool alloc(ScratchInfo& info) {
+    ref_t it = map.upper_bound(info.size - 1);
+    if (it == map.end()) return false;
+
+    // Small requests must have an exact size match and be small.
+    if (!info.large) {
+      while ((it != map.end()) && (it->first == info.size)) {
+        if (it->second.isFree() && (!it->second.large)) {
+          it->second.alloc();
+          info.queue_base = it->second.base;
+          info.scratch_node = it;
+          return true;
+        }
+        it++;
+      }
+      return false;
+    }
+
+    // Large requests may use a small allocation and do not require an exact size match.
+    while (it != map.end()) {
+      if (it->second.isFree()) {
+        it->second.alloc();
+        info.queue_base = it->second.base;
+        info.size = it->first;
+        info.scratch_node = it;
+        return true;
+      }
+      it++;
+    }
+    return false;
+  }
+
+  void free(ScratchInfo& info) {
+    assert(!info.scratch_node->second.isFree() && "free called on free scratch node.");
+    auto it = info.scratch_node;
+    if (it->second.trimPending()) {
+      dealloc(it->second.base, it->first, it->second.large);
+      map.erase(it);
+      return;
+    }
+    it->second.free();
+  }
+
+  bool trim(bool trim_nodes_in_use) {
+    bool ret = !map.empty();
+    auto it = map.begin();
+    while (it != map.end()) {
+      if (it->second.isFree()) {
+        dealloc(it->second.base, it->first, it->second.large);
+        auto temp = it;
+        it++;
+        map.erase(temp);
+      } else {
+        if (trim_nodes_in_use) it->second.trim();
+        it++;
+      }
+    }
+    return ret;
+  }
+
+  void insert(ScratchInfo& info) {
+    node n;
+    n.base = info.queue_base;
+    n.large = info.large;
+    n.alloc();
+
+    auto it = map.insert(std::make_pair(info.size, n));
+    info.scratch_node = it;
+  }
+
+ private:
+  map_t map;
+  deallocator_t dealloc;
+};
+
+}  // namespace AMD
+}  // namespace rocr
+
+#endif  // header guard