[onert] Templatize memory planners in train backend (#13892)

This PR templatize memory planners(Bump, FirstFit, WIC) in train backend. This memory planners currently only used for DisposableTensor, but it will be also used for LayerScopeTensor. ONE-DCO-1.0-Signed-off-by: seunghui youn <[email protected]>
Samsung · Sep 4, 2024 · 433edf9 · 433edf9
1 parent 8a22e18
commit 433edf9
Show file tree

Hide file tree

Showing 5 changed files with 58 additions and 33 deletions.
diff --git a/runtime/onert/backend/train/DisposableTensorIndex.h b/runtime/onert/backend/train/DisposableTensorIndex.h
@@ -88,8 +88,6 @@ class DisposableTensorIndex
   ir::OperandIndex _operand_index;
 };
 
-template <typename T> using DisposableTensorIndexMap = std::unordered_map<DisposableTensorIndex, T>;
-
 inline std::ostream &operator<<(std::ostream &o, const DisposableTensorIndex &i)
 {
   return operator<<(o, i.operand_index());

diff --git a/runtime/onert/backend/train/MemoryPlanner.cc b/runtime/onert/backend/train/MemoryPlanner.cc
@@ -16,6 +16,9 @@
 
 #include "MemoryPlanner.h"
 
+#include "DisposableTensorIndex.h"
+#include "LayerScopeTensorIndex.h"
+
 #include <util/logging.h>
 
 #include <cassert>
@@ -27,7 +30,7 @@ namespace backend
 namespace train
 {
 
-void BumpPlanner::claim(const DisposableTensorIndex &ind, size_t size)
+template <typename Index> void BumpPlanner<Index>::claim(const Index &ind, size_t size)
 {
   basic::Block blk{_capacity, size};
   _mem_plans[ind] = blk;
@@ -36,7 +39,7 @@ void BumpPlanner::claim(const DisposableTensorIndex &ind, size_t size)
   VERBOSE(BP_PLANNER) << "CLAIM(" << ind << "): " << blk.offset << ", " << blk.size << std::endl;
 }
 
-void BumpPlanner::release(const DisposableTensorIndex &ind)
+template <typename Index> void BumpPlanner<Index>::release(const Index &ind)
 {
   VERBOSE(BP_PLANNER) << "RELEASE(" << ind << "): "
                       << "NOTHING does" << std::endl;
@@ -56,7 +59,7 @@ void BumpPlanner::release(const DisposableTensorIndex &ind)
 //       point in time, it means the place at the offset can be claimed.
 // 2. In the loop for _claim_table, we can assume the current claim_base_offset value is bigger than
 //    the previous claim_base_offset.
-void FirstFitPlanner::claim(const DisposableTensorIndex &ind, size_t size)
+template <typename Index> void FirstFitPlanner<Index>::claim(const Index &ind, size_t size)
 {
   // Find the right position for claiming
   uint32_t next_offset = 0;
@@ -88,7 +91,7 @@ void FirstFitPlanner::claim(const DisposableTensorIndex &ind, size_t size)
   }
 }
 
-void FirstFitPlanner::release(const DisposableTensorIndex &ind)
+template <typename Index> void FirstFitPlanner<Index>::release(const Index &ind)
 {
   for (auto it = _claim_table.cbegin(); it != _claim_table.cend(); ++it)
   {
@@ -107,14 +110,15 @@ void FirstFitPlanner::release(const DisposableTensorIndex &ind)
   assert(!"Cannot release for given index. It has been not claimed or released already.");
 }
 
-WICPlanner::WICPlanner()
+template <typename Index>
+WICPlanner<Index>::WICPlanner()
   : _initialized(false), _capacity(0), _mem_plans(), _live_indices(), _interference_graph(),
     _indices()
 {
   // DO NOTHING
 }
 
-void WICPlanner::claim(const DisposableTensorIndex &ind, size_t size)
+template <typename Index> void WICPlanner<Index>::claim(const Index &ind, size_t size)
 {
   _indices.emplace(size, ind);
   _interference_graph[ind].insert(_interference_graph[ind].end(), _live_indices.cbegin(),
@@ -128,7 +132,7 @@ void WICPlanner::claim(const DisposableTensorIndex &ind, size_t size)
   VERBOSE(WIC_PLANNER) << "claim(" << ind << "): [" << size << "sz]" << std::endl;
 }
 
-void WICPlanner::release(const DisposableTensorIndex &ind)
+template <typename Index> void WICPlanner<Index>::release(const Index &ind)
 {
   _live_indices.erase(ind);
   VERBOSE(WIC_PLANNER) << "release(" << ind << ")" << std::endl;
@@ -143,7 +147,8 @@ void WICPlanner::release(const DisposableTensorIndex &ind)
  * 3. Allocate memory block for sorted operands
  *   - Find free memory block which does not overlap with interfered operands
  */
-void WICPlanner::buildMemoryPlans()
+
+template <typename Index> void WICPlanner<Index>::buildMemoryPlans()
 {
   for (const auto &[size, ind] : _indices)
   {
@@ -194,13 +199,22 @@ void WICPlanner::buildMemoryPlans()
   _indices.clear();
 }
 
-std::unordered_map<DisposableTensorIndex, basic::Block> &WICPlanner::memory_plans()
+template <typename Index> typename WICPlanner<Index>::MemoryPlans &WICPlanner<Index>::memory_plans()
 {
   if (!_initialized)
     buildMemoryPlans();
   return _mem_plans;
 }
 
+template class BumpPlanner<DisposableTensorIndex>;
+template class BumpPlanner<LayerScopeTensorIndex>;
+
+template class FirstFitPlanner<DisposableTensorIndex>;
+template class FirstFitPlanner<LayerScopeTensorIndex>;
+
+template class WICPlanner<DisposableTensorIndex>;
+template class WICPlanner<LayerScopeTensorIndex>;
+
 } // namespace train
 } // namespace backend
 } // namespace onert
diff --git a/runtime/onert/backend/train/MemoryPlanner.h b/runtime/onert/backend/train/MemoryPlanner.h
@@ -41,20 +41,23 @@ namespace train
 /**
  * @brief Class to plan memory by bump way
  */
-class BumpPlanner : public basic::IMemoryPlanner<DisposableTensorIndex>
+template <typename Index> class BumpPlanner : public basic::IMemoryPlanner<Index>
 {
+private:
+  using MemoryPlans = typename basic::IMemoryPlanner<Index>::MemoryPlans;
+
 public:
   /**
    * @brief Claim memory for tensor by bump way
    * @param[in] index The tensor index
    * @param[in] size The size of the memory
    */
-  void claim(const DisposableTensorIndex &, size_t) override;
+  void claim(const Index &, size_t) override;
   /**
    * @brief Release memory for tensor by bump way
    * @param[in] index The tensor index
    */
-  void release(const DisposableTensorIndex &) override;
+  void release(const Index &) override;
   /**
    * @brief Get capacity for memory planning
    * @return The value of capacity
@@ -74,20 +77,23 @@ class BumpPlanner : public basic::IMemoryPlanner<DisposableTensorIndex>
 /**
  * @brief Class to plan memory by firstfit way
  */
-class FirstFitPlanner : public basic::IMemoryPlanner<DisposableTensorIndex>
+template <typename Index> class FirstFitPlanner : public basic::IMemoryPlanner<Index>
 {
+private:
+  using MemoryPlans = typename basic::IMemoryPlanner<Index>::MemoryPlans;
+
 public:
   /**
    * @brief Claim memory for tensor by firstfit way
    * @param[in] index The tensor index
    * @param[in] size The size of the memory
    */
-  void claim(const DisposableTensorIndex &, size_t) override;
+  void claim(const Index &, size_t) override;
   /**
    * @brief Release memory for tensor by firstfit way
    * @param[in] index The tensor index
    */
-  void release(const DisposableTensorIndex &) override;
+  void release(const Index &) override;
   /**
    * @brief Get capacity for memory planning
    * @return The value of capacity
@@ -103,14 +109,17 @@ class FirstFitPlanner : public basic::IMemoryPlanner<DisposableTensorIndex>
   uint32_t _capacity = 0;
   MemoryPlans _mem_plans;
   // Use std::map because claim() assumes that _claim_table is sorted by uint32_t(base_offset)
-  std::map<uint32_t, DisposableTensorIndex> _claim_table;
+  std::map<uint32_t, Index> _claim_table;
 };
 
 /**
  * @brief Class to plan memory by Weighted Interval Color algorithm
  */
-class WICPlanner : public basic::IMemoryPlanner<DisposableTensorIndex>
+template <typename Index> class WICPlanner : public basic::IMemoryPlanner<Index>
 {
+private:
+  using MemoryPlans = typename basic::IMemoryPlanner<Index>::MemoryPlans;
+
 public:
   WICPlanner();
 
@@ -119,12 +128,12 @@ class WICPlanner : public basic::IMemoryPlanner<DisposableTensorIndex>
    * @param[in] index The tensor index
    * @param[in] size The size of the memory
    */
-  void claim(const DisposableTensorIndex &, size_t) override;
+  void claim(const Index &, size_t) override;
   /**
    * @brief Release memory for tensor by WIC algorithm
    * @param[in] index The tensor index
    */
-  void release(const DisposableTensorIndex &) override;
+  void release(const Index &) override;
   /**
    * @brief Get capacity for memory planning
    * @return The value of capacity
@@ -147,10 +156,10 @@ class WICPlanner : public basic::IMemoryPlanner<DisposableTensorIndex>
   bool _initialized;
   uint32_t _capacity;
   MemoryPlans _mem_plans;
-  std::unordered_set<DisposableTensorIndex> _live_indices;
-  DisposableTensorIndexMap<std::vector<DisposableTensorIndex>> _interference_graph;
+  std::unordered_set<Index> _live_indices;
+  std::unordered_map<Index, std::vector<Index>> _interference_graph;
   // Sort tensors by descending order of size
-  std::multimap<uint32_t, DisposableTensorIndex, std::greater<uint32_t>> _indices;
+  std::multimap<uint32_t, Index, std::greater<uint32_t>> _indices;
 };
 
 } // namespace train

diff --git a/runtime/onert/backend/train/MemoryPlanner.test.cc b/runtime/onert/backend/train/MemoryPlanner.test.cc
@@ -16,16 +16,19 @@
 
 #include <gtest/gtest.h>
 
+#include "DisposableTensorIndex.h"
 #include "MemoryPlanner.h"
 #include "ir/Index.h"
 
 using namespace onert::backend::train;
 using onert::ir::OperandIndex;
 using onert::ir::OperationIndex;
 
+// TODO: Add test testcase for {Bump, FirstFit, WIC}Planner<LayerScopeTensor>
+
 TEST(BumpPlanner, claim_test)
 {
-  BumpPlanner planner;
+  BumpPlanner<DisposableTensorIndex> planner;
 
   auto claim = [&planner](uint32_t op_index, uint32_t operand_index, size_t size,
                           uint32_t expected_offset) {
@@ -55,7 +58,7 @@ TEST(BumpPlanner, claim_test)
 
 TEST(FirstFitPlanner, claim_release_test)
 {
-  FirstFitPlanner planner;
+  FirstFitPlanner<DisposableTensorIndex> planner;
 
   auto claim = [&planner](uint32_t op_index, uint32_t operand_index, size_t size,
                           uint32_t expected_offset) {
@@ -148,7 +151,7 @@ TEST(FirstFitPlanner, claim_release_test)
 
 TEST(FirstFitPlanner, neg_release_non_existing_index)
 {
-  FirstFitPlanner planner;
+  FirstFitPlanner<DisposableTensorIndex> planner;
 
   auto claim = [&planner](uint32_t op_index, uint32_t operand_index, size_t size,
                           uint32_t expected_offset) {
@@ -184,7 +187,7 @@ TEST(FirstFitPlanner, neg_release_non_existing_index)
 
 TEST(FirstFitPlanner, neg_release_twice)
 {
-  FirstFitPlanner planner;
+  FirstFitPlanner<DisposableTensorIndex> planner;
 
   auto claim = [&planner](uint32_t op_index, uint32_t operand_index, size_t size,
                           uint32_t expected_offset) {
@@ -223,7 +226,7 @@ TEST(FirstFitPlanner, neg_release_twice)
 
 TEST(WICPlanner, claim_release_test)
 {
-  WICPlanner planner;
+  WICPlanner<DisposableTensorIndex> planner;
 
   auto claim = [&planner](uint32_t op_index, uint32_t operand_index, size_t size) {
     DisposableTensorIndex mem_idx{OperationIndex{op_index}, OperandIndex{operand_index}};

diff --git a/runtime/onert/backend/train/MemoryPlannerFactory.cc b/runtime/onert/backend/train/MemoryPlannerFactory.cc
@@ -29,21 +29,22 @@ MemoryPlannerFactory &MemoryPlannerFactory::get()
   return instance;
 }
 
+// TODO: Update to use template varialbe instead of DisposableTensorIndex
 basic::IMemoryPlanner<DisposableTensorIndex> *MemoryPlannerFactory::create(const std::string &key)
 {
   if (key == "FirstFit")
   {
-    return new FirstFitPlanner;
+    return new FirstFitPlanner<DisposableTensorIndex>();
   }
   else if (key == "Bump")
   {
-    return new BumpPlanner;
+    return new BumpPlanner<DisposableTensorIndex>();
   }
   else if (key == "WIC")
   {
-    return new WICPlanner;
+    return new WICPlanner<DisposableTensorIndex>();
   }
-  return new FirstFitPlanner; // Default Planner
+  return new FirstFitPlanner<DisposableTensorIndex>(); // Default Planner
 }
 
 } // namespace train