diff options
author | Chao Mei <chaomei@google.com> | 2021-03-01 16:02:51 -0800 |
---|---|---|
committer | Copybara-Service <copybara-worker@google.com> | 2021-03-01 16:03:11 -0800 |
commit | 54774a7a2cf85963777289193629d4bd42de4a59 (patch) | |
tree | 597d4f6c62019c7f2505c8ad256d148e2aa0647b | |
parent | be760b63149d8205dfb3ca66d78a049dc1ab7772 (diff) | |
download | ruy-54774a7a2cf85963777289193629d4bd42de4a59.tar.gz |
Use std::ptrdiff_t instead of int when calculating memory size to avoid int overflow.
PiperOrigin-RevId: 360298662
-rw-r--r-- | ruy/mat.h | 8 | ||||
-rw-r--r-- | ruy/prepacked_cache.cc | 10 | ||||
-rw-r--r-- | ruy/prepacked_cache.h | 8 |
3 files changed, 13 insertions, 13 deletions
@@ -327,7 +327,7 @@ inline bool IsColMajor(const MatLayout& layout) { return layout.order == Order::kColMajor; } -inline int FlatSize(const MatLayout& layout) { +inline std::ptrdiff_t FlatSize(const MatLayout& layout) { const int outerdim = layout.order == Order::kColMajor ? layout.cols : layout.rows; return layout.stride * outerdim; @@ -349,7 +349,7 @@ inline bool IsColMajor(const PMatLayout& layout) { return layout.order == Order::kColMajor; } -inline int FlatSize(const PMatLayout& layout) { +inline std::ptrdiff_t FlatSize(const PMatLayout& layout) { const int outerdim = layout.order == Order::kColMajor ? layout.cols : layout.rows; return layout.stride * outerdim; @@ -429,11 +429,11 @@ Scalar Element(const PMat<Scalar>& mat, int row, int col) { // Helpers for PEMat. -inline int DataBytes(const PEMat& packed) { +inline std::ptrdiff_t DataBytes(const PEMat& packed) { return FlatSize(packed.layout) * packed.data_type.size; } -inline int SumsBytes(const PEMat& packed) { +inline std::ptrdiff_t SumsBytes(const PEMat& packed) { // Packed matrices are only relevant for Ruy's TrMul implementations. For // TrMul, the number of sums is always equal to the number of columns. return packed.layout.cols * packed.sums_type.size; diff --git a/ruy/prepacked_cache.cc b/ruy/prepacked_cache.cc index ee891cb..5080ca9 100644 --- a/ruy/prepacked_cache.cc +++ b/ruy/prepacked_cache.cc @@ -26,10 +26,10 @@ namespace { // Allocates the `data` and `sums` buffers, and sets the corresponding // pointer fields, in a PEMat whose other fields, particularly `layout` // and the runtime data types, are already populated. -int AllocateBuffers(PEMat* packed_matrix) { - const int data_bytes = DataBytes(*packed_matrix); +std::ptrdiff_t AllocateBuffers(PEMat* packed_matrix) { + const std::ptrdiff_t data_bytes = DataBytes(*packed_matrix); packed_matrix->data = detail::SystemAlignedAlloc(data_bytes); - int sums_bytes = 0; + std::ptrdiff_t sums_bytes = 0; if (!packed_matrix->sums_type.is_floating_point) { // Integer quantized matrices also need the `sums` buffer. sums_bytes = SumsBytes(*packed_matrix); @@ -93,7 +93,7 @@ PrepackedCache::Action PrepackedCache::Get(const void* src_data, } // No existing entry found. Allocate new buffers now and insert in the cache. - const int new_bytes = AllocateBuffers(packed_matrix); + const std::ptrdiff_t new_bytes = AllocateBuffers(packed_matrix); EjectUntilRoomFor(new_bytes); Entry entry{*packed_matrix, timestamp_++}; cache_.emplace(key, entry); @@ -101,7 +101,7 @@ PrepackedCache::Action PrepackedCache::Get(const void* src_data, return Action::kInsertedNewEntry; } -void PrepackedCache::EjectUntilRoomFor(int new_bytes) { +void PrepackedCache::EjectUntilRoomFor(std::ptrdiff_t new_bytes) { profiler::ScopeLabel label("PrepackedCacheEjection"); // While we are above the threshold of ejection, eject the LRU entry. while (!cache_.empty() && buffers_bytes_ + new_bytes > max_buffers_bytes_) { diff --git a/ruy/prepacked_cache.h b/ruy/prepacked_cache.h index cb3a113..c58593e 100644 --- a/ruy/prepacked_cache.h +++ b/ruy/prepacked_cache.h @@ -101,7 +101,7 @@ class PrepackedCache final { ~PrepackedCache(); // Returns the total size in bytes of buffers held in this cache. - int BuffersBytes() const { return buffers_bytes_; } + std::ptrdiff_t BuffersBytes() const { return buffers_bytes_; } // Returns the number of packed matrices held in this cache. int MatrixCount() const { return cache_.size(); } @@ -128,11 +128,11 @@ class PrepackedCache final { private: void EjectOne(); - void EjectUntilRoomFor(int new_bytes); + void EjectUntilRoomFor(std::ptrdiff_t new_bytes); std::unordered_map<Key, Entry, KeyHash> cache_; - const int max_buffers_bytes_; - int buffers_bytes_ = 0; + const std::ptrdiff_t max_buffers_bytes_; + std::ptrdiff_t buffers_bytes_ = 0; Timestamp timestamp_ = 0; }; |