From 06cbfb07781a61c4ad5dc001a3f5f08f1148a9a1 Mon Sep 17 00:00:00 2001
From: Eylon Krause <eylon1909@gmail.com>
Date: Thu, 27 Aug 2026 00:22:59 +0300
Subject: [PATCH] Do scratch-buffer size math in size_t to avoid 32-bit
 overflow

The scratch sizing for the public EightBitIntGemm float entry point computes
`n * ldc * sizeof(std::int32_t)` (and the transpose_c case `m * ldc * ...`) in
`int`, and Scratch::AssureSize takes an `std::int32_t required_size` and stores
it in an `std::int32_t size_`. For large-but-legal shapes (e.g. m=1, k=1,
n=50000, ldc=50000) the product overflows 32 bits and/or is truncated when
passed to AssureSize, so the scratch buffer is under-allocated and the
subsequent int32 result writes run past it -- a heap out-of-bounds write.

The internal packing paths have the same latent overflow: PackedSideBlock
reserves `params_.l2_width * params_.l2_depth` and PackedResult reserves
`block_params_.l2_rows * block_params_.l2_cols`, both `int * int`, before the
result is handed to Allocator::Reserve(std::size_t).

Widen the dimension arithmetic to size_t at each site and make
Scratch::AssureSize / Scratch::size_ size_t so the computed requirement is
never truncated. No behavior change for in-range sizes.
---
 eight_bit_int_gemm/eight_bit_int_gemm.cc | 10 ++++++----
 internal/pack.h                          |  3 ++-
 internal/unpack.h                        |  5 +++--
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/eight_bit_int_gemm/eight_bit_int_gemm.cc b/eight_bit_int_gemm/eight_bit_int_gemm.cc
index a8d9b43..ebbd96c 100644
--- a/eight_bit_int_gemm/eight_bit_int_gemm.cc
+++ b/eight_bit_int_gemm/eight_bit_int_gemm.cc
@@ -144,7 +144,7 @@ class Scratch {
  public:
   Scratch() : buffer_(), buffer_32_(nullptr), size_(0) {}
 
-  void AssureSize(std::int32_t required_size) {
+  void AssureSize(std::size_t required_size) {
     if (size_ >= required_size) {
       return;
     }
@@ -167,7 +167,7 @@ class Scratch {
  private:
   std::unique_ptr<std::uint8_t[]> buffer_;
   std::uint8_t* buffer_32_;
-  std::int32_t size_;
+  std::size_t size_;
 };
 
 Scratch* global_scratch = nullptr;
@@ -356,9 +356,11 @@ void EightBitIntGemm(bool transpose_a, bool transpose_b, bool transpose_c,
   // TODO(maciekc): implement a float output stage, get rid of scratch memory.
   Scratch* scratch = GetOrCreateGlobalScratch();
   if (transpose_c) {
-    scratch->AssureSize(m * ldc * sizeof(std::int32_t));
+    scratch->AssureSize(static_cast<std::size_t>(m) * ldc *
+                        sizeof(std::int32_t));
   } else {
-    scratch->AssureSize(n * ldc * sizeof(std::int32_t));
+    scratch->AssureSize(static_cast<std::size_t>(n) * ldc *
+                        sizeof(std::int32_t));
   }
   std::int32_t* temp_c = reinterpret_cast<std::int32_t*>(scratch->buffer());
 
diff --git a/internal/pack.h b/internal/pack.h
index a9eb396..8ac7aa4 100644
--- a/internal/pack.h
+++ b/internal/pack.h
@@ -53,7 +53,8 @@ class PackedSideBlock {
       : allocator_(allocator), pos_(0) {
     GetSideBlockParams(side, &params_, block_params);
     data_handle_ =
-        allocator_->Reserve<std::uint8_t>(params_.l2_width * params_.l2_depth);
+        allocator_->Reserve<std::uint8_t>(
+            static_cast<std::size_t>(params_.l2_width) * params_.l2_depth);
     sums_of_each_slice_handle_ =
         allocator_->Reserve<std::int32_t>(params_.l2_width);
   }
diff --git a/internal/unpack.h b/internal/unpack.h
index 021f4aa..67ae5a1 100644
--- a/internal/unpack.h
+++ b/internal/unpack.h
@@ -31,8 +31,9 @@ class PackedResult {
  public:
   PackedResult(Allocator* _allocator, const BlockParams& _block_params)
       : allocator_(_allocator), block_params_(_block_params) {
-    matrix_handle_ = allocator_->Reserve<std::int32_t>(block_params_.l2_rows *
-                                                       block_params_.l2_cols);
+    matrix_handle_ = allocator_->Reserve<std::int32_t>(
+        static_cast<std::size_t>(block_params_.l2_rows) *
+        block_params_.l2_cols);
   }
 
   ~PackedResult() {}