1 /** 2 * Copyright 2021 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 #include "src/train/opt_allocator.h" 17 #include <limits> 18 #include "nnacl/op_base.h" 19 20 namespace mindspore { FindFree(size_t size)21size_t OptAllocator::FindFree(size_t size) { 22 size_t min_size = std::numeric_limits<size_t>::max(); 23 size_t min_addr = std::numeric_limits<size_t>::max(); 24 for (auto const &itr : arena_) { 25 // best fit 26 if (itr.second >= size) { 27 if (min_size > itr.second) { 28 min_size = itr.second; 29 min_addr = itr.first; 30 } 31 } 32 } 33 return min_addr; 34 } 35 Reorder(size_t addr)36void OptAllocator::Reorder(size_t addr) { 37 size_t length = arena_[addr]; 38 size_t post = addr + length; 39 // connect to upper block 40 auto it = arena_.find(post); 41 if (it != arena_.end()) { 42 size_t post_size = it->second; 43 arena_[addr] = length + post_size; 44 arena_.erase(post); 45 } 46 // connect to lower block 47 auto itr = arena_.lower_bound(addr); 48 if (itr != arena_.begin()) { 49 itr--; 50 size_t last = itr->first; 51 if ((last + arena_[last]) == addr) { 52 arena_[last] = arena_[last] + arena_[addr]; 53 arena_.erase(addr); 54 } 55 } 56 } 57 Malloc(size_t size)58size_t OptAllocator::Malloc(size_t size) { 59 size = UP_DIV(size, align_size_) * align_size_; 60 size_t addr = FindFree(size); 61 // free block not found 62 if (addr == std::numeric_limits<size_t>::max()) { 63 if (!arena_.empty()) { 64 addr = arena_.rbegin()->first; 65 if (addr + arena_[addr] < heap_) { 66 addr = heap_; 67 } else { 68 arena_.erase(addr); 69 } 70 } else { 71 addr = heap_; 72 } 73 heap_ = addr + size; 74 } else { 75 if (arena_[addr] > size) { 76 arena_[addr + size] = arena_[addr] - size; 77 } 78 arena_.erase(addr); 79 } 80 alloc_[addr] = size; 81 return addr; 82 } 83 Free(size_t addr)84void OptAllocator::Free(size_t addr) { 85 arena_[addr] = alloc_[addr]; 86 alloc_.erase(addr); 87 Reorder(addr); 88 } 89 } // namespace mindspore 90