| 23 | : ImageCacheBlob(cache_size, 0, stats_enabled) {} |
| 24 | |
| 25 | void ImageCacheLargest::Add(const ImageKey& image_key, |
| 26 | const uint8_t *data, |
| 27 | const ImageShape& data_shape, |
| 28 | cudaStream_t stream) { |
| 29 | const std::size_t data_size = volume(data_shape); |
| 30 | std::unique_lock<std::mutex> lock(mutex_); |
| 31 | // If we haven't started caching |
| 32 | if (!start_caching_) { |
| 33 | // if we've already seen this image, start caching |
| 34 | start_caching_ = (images_.find(image_key) != images_.end()); |
| 35 | // if we decided to start caching, prepare the data |
| 36 | if (start_caching_) { |
| 37 | // replace images_ with the biggest_images |
| 38 | // and clean unnecessary data structures |
| 39 | total_seen_images_ = images_.size(); |
| 40 | images_.clear(); |
| 41 | while (!biggest_images_.empty()) { |
| 42 | images_.insert(biggest_images_.top().second); |
| 43 | biggest_images_.pop(); |
| 44 | } |
| 45 | } else { |
| 46 | // mark the image as seen |
| 47 | images_.insert(image_key); |
| 48 | |
| 49 | const bool data_fits = (biggest_images_total_ + data_size <= cache_size_); |
| 50 | is_full = is_full || !data_fits; |
| 51 | // if there is enough space, store the image as one of biggest |
| 52 | if (data_fits) { |
| 53 | biggest_images_.push({data_size, image_key}); |
| 54 | biggest_images_total_ += data_size; |
| 55 | } else if (data_size <= cache_size_) { |
| 56 | // If full, check whether the current image has higher priority |
| 57 | std::stack<QueueElement> to_be_discarded; |
| 58 | while (!biggest_images_.empty() |
| 59 | && biggest_images_total_ + data_size > cache_size_ |
| 60 | && biggest_images_.top().first < data_size) { |
| 61 | biggest_images_total_ -= biggest_images_.top().first; |
| 62 | to_be_discarded.push(biggest_images_.top()); |
| 63 | biggest_images_.pop(); |
| 64 | } |
| 65 | |
| 66 | // If we have enough space now, push the new image |
| 67 | if (biggest_images_total_ + data_size <= cache_size_) { |
| 68 | biggest_images_.push({data_size, image_key}); |
| 69 | biggest_images_total_ += data_size; |
| 70 | } |
| 71 | |
| 72 | // If there is extra space, push back the images we took out |
| 73 | while (!to_be_discarded.empty()) { |
| 74 | if (biggest_images_total_ + to_be_discarded.top().first <= cache_size_) { |
| 75 | biggest_images_total_ += to_be_discarded.top().first; |
| 76 | biggest_images_.push(std::move(to_be_discarded.top())); |
| 77 | } |
| 78 | to_be_discarded.pop(); |
| 79 | } |
| 80 | } |
| 81 | } |
| 82 | } |