| 53 | } |
| 54 | |
| 55 | void GPUTensorPoolAllocator::Init() { |
| 56 | bool tmp = false; |
| 57 | if (initing_.compare_exchange_strong(tmp, true)) { |
| 58 | auto lifetime_policy = mem_planner_->BestLifetimePolicy(); |
| 59 | |
| 60 | alignment_ = lifetime_policy->Alignment(); |
| 61 | alignment_offset_ = lifetime_policy->AlignmentOffset(); |
| 62 | |
| 63 | big_bytes_ = 0; |
| 64 | std::map<size_t, size_t> bin_to_offset; |
| 65 | |
| 66 | auto policy_bins = lifetime_policy->GetBins(); |
| 67 | large_bin_index_ = policy_bins.size(); |
| 68 | lifetime_bins_.resize(large_bin_index_); |
| 69 | |
| 70 | size_t max_alignment = 0; |
| 71 | |
| 72 | for (auto it = policy_bins.begin(); it != policy_bins.end(); |
| 73 | ++it) { |
| 74 | if ((*it)->BlockSize() > 0) { |
| 75 | // add padding between two bins |
| 76 | big_bytes_ = RoundedBytes(big_bytes_, (*it)->Alignment()); |
| 77 | bin_to_offset[(*it)->BinIndex()] = big_bytes_; |
| 78 | big_bytes_ += (*it)->TotalMem(); |
| 79 | max_alignment = std::max<size_t>(max_alignment, (*it)->Alignment()); |
| 80 | } |
| 81 | } |
| 82 | |
| 83 | auto policy_large_bins = lifetime_policy->GetLargeBins(); |
| 84 | for (auto it = policy_large_bins.begin(); |
| 85 | it != policy_large_bins.end(); ++it) { |
| 86 | auto bin_info = it->second; |
| 87 | if (bin_info->BlockSize() > 0) { |
| 88 | // add padding between two bins |
| 89 | big_bytes_ = RoundedBytes(big_bytes_, bin_info->Alignment()); |
| 90 | bin_to_offset[bin_info->BinIndex()] = big_bytes_; |
| 91 | big_bytes_ += bin_info->TotalMem(); |
| 92 | max_alignment = std::max<size_t>(max_alignment, bin_info->Alignment()); |
| 93 | } |
| 94 | } |
| 95 | |
| 96 | big_mem_begin_ = sub_allocator_->Alloc(max_alignment, big_bytes_); |
| 97 | if (big_bytes_ > 0 && big_mem_begin_ == nullptr) { |
| 98 | LOG(FATAL) << "OOM!!! Try to alloc(" |
| 99 | << max_alignment << ", " << big_bytes_ << ")"; |
| 100 | } |
| 101 | if (big_bytes_ > 0) { |
| 102 | big_mem_end_ = static_cast<char*>(big_mem_begin_) + big_bytes_; |
| 103 | } else { |
| 104 | big_mem_end_ = nullptr; |
| 105 | } |
| 106 | |
| 107 | // create bigger bin first |
| 108 | for (auto rit = policy_large_bins.rbegin(); |
| 109 | rit != policy_large_bins.rend(); ++rit) { |
| 110 | auto bin_info = rit->second; |
| 111 | Bin* bin = nullptr; |
| 112 | if (bin_info->BlockSize() > 0) { |
nothing calls this directly
no test coverage detected