Skip to content

Do scratch-buffer size math in size_t to avoid 32-bit overflow - #220

Open
EylonKrause wants to merge 1 commit into
google:masterfrom
EylonKrause:fix/scratch-size-overflow
Open

Do scratch-buffer size math in size_t to avoid 32-bit overflow#220
EylonKrause wants to merge 1 commit into
google:masterfrom
EylonKrause:fix/scratch-size-overflow

Conversation

@EylonKrause

Copy link
Copy Markdown

Problem

The scratch-buffer sizing for the public EightBitIntGemm float entry point does
its byte math in int:

scratch->AssureSize(n * ldc * sizeof(std::int32_t));      // (or m*ldc for transpose_c)

and Scratch::AssureSize both takes and stores the size as std::int32_t:

void AssureSize(std::int32_t required_size) { ... size_ = required_size; }
std::int32_t size_;

For large-but-legal shapes the product n * ldc * 4 overflows 32-bit int
(and is further truncated on the std::int32_t parameter). AssureSize then
under-allocates the scratch buffer, and the subsequent int32 result writes run
past it — a heap out-of-bounds write. Example: m=1, k=1, n=50000, ldc=50000
needs 50000500004 = 10 GB but n*ldc*4 overflows to a small (even negative)
int.

The internal packing paths have the same latent int * int overflow before the
size reaches Allocator::Reserve(std::size_t):

  • internal/pack.h: Reserve<std::uint8_t>(params_.l2_width * params_.l2_depth)
  • internal/unpack.h: Reserve<std::int32_t>(block_params_.l2_rows * block_params_.l2_cols)

Fix

Do the dimension arithmetic in 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.

Testing

test/test.cc (full GEMM correctness suite, built with
test/test.cc test/test_data.cc eight_bit_int_gemm/eight_bit_int_gemm.cc)
passes: All tests passed. The change is confined to size computations, so
correctness for in-range shapes is unaffected.

Disclosure: this contribution was authored with an AI coding assistant (Claude) and reviewed before submission.

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.
@EylonKrause
EylonKrause force-pushed the fix/scratch-size-overflow branch from f2556ab to 06cbfb0 Compare August 26, 2026 21:27
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant