• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
1 /**
2  * Copyright 2020 Huawei Technologies Co., Ltd
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  * http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #ifndef MINDSPORE_LITE_SRC_RUNTIME_INNER_ALLOCATOR_H_
18 #define MINDSPORE_LITE_SRC_RUNTIME_INNER_ALLOCATOR_H_
19 
20 #include <memory>
21 #include <string>
22 #include <vector>
23 #include <mutex>
24 #include <map>
25 #include <unordered_map>
26 #include <unordered_set>
27 #include <atomic>
28 #include "include/api/allocator.h"
29 
30 namespace mindspore {
31 struct AllocatorContext {
32   int shiftFactor;
33   bool lockFlag;
34 };
35 
36 class DefaultAllocator : public Allocator {
37  public:
38   explicit DefaultAllocator(size_t aligned_size = 32);
39   ~DefaultAllocator() override;
40   void SetContext(const AllocatorContext &ctx);
41   void *Malloc(size_t size) override;
42   void Free(void *ptr) override;
43   int RefCount(void *ptr) override;
44   int SetRefCount(void *ptr, int ref_count) override;
45   int DecRefCount(void *ptr, int ref_count) override;
46   int IncRefCount(void *ptr, int ref_count) override;
total_size()47   size_t total_size() const { return this->total_size_; }
48   void Clear();
49 
50  private:
51   void Lock();
52   void UnLock();
53   bool ReuseMemory(size_t free_size, size_t size) const;
54   struct MemBuf {
55     std::atomic_int ref_count_ = {0};
56     size_t size = 0;
57     void *buf = nullptr;
58   };
59 
60   std::mutex lock_;
61   size_t total_size_ = 0;
62   // <membuf->buf, membuf>
63   std::unordered_map<void *, MemBuf *> allocatedList_;
64   std::multimap<size_t, MemBuf *> freeList_;
65   // 6 is empirical value
66   unsigned shiftFactor_ = 6;
67   bool lockFlag_ = true;
68   size_t max_malloc_size_ = 0;
69 };
70 
71 constexpr int64_t MAX_MALLOC_SIZE = static_cast<size_t>(2000) * 1024 * 1024;
72 constexpr int64_t MAX_THREAD_POOL_SIZE = static_cast<size_t>(3000) * 1024 * 1024;
73 
74 }  // namespace mindspore
75 
76 #endif  // MINDSPORE_LITE_SRC_RUNTIME_INNER_ALLOCATOR_H_
77