Use std::ptrdiff_t instead of int when calculating memory size to avoid int overflow.

PiperOrigin-RevId: 360298662

3 files changed, 13 insertions, 13 deletions

diff --git a/ruy/mat.h b/ruy/mat.h
index 587b208..c2254f9 100644
--- a/ruy/mat.h
+++ b/ruy/mat.h
@@ -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;
 };