Fix heap buffer overflow from reusing a packed sub-byte buffer for a full-byte tensor

[Wei Wang (wangw-1991)]

Description

The allocation planner's SameSize() decided buffer reuse by comparing the C++ carrier size of the element type (elt_type->Size()) plus the logical shape. Packed sub-byte types (int4/uint4) have the same 1-byte carrier size as int8/uint8, but one carrier stores 2 logical elements. As a result a uint4[N] tensor (physical ceil(N/2) bytes) and a uint8[N] tensor (physical N bytes) were treated as the same size, and the smaller uint4 buffer was reused for the uint8 output. Writing the uint8 tensor into that half-sized buffer overflows it.

Fix

• SameSize() now also requires the sub-element packing density (GetNumSubElems()) to match, so tensors are only considered the same size when their physical storage bytes are actually equal.
• ExecutionFrame::AllocateMLValueTensorPreAllocateBuffer() adds a defense-in-depth capacity check that rejects a reuse when the requested tensor needs more storage bytes than the buffer being reused (the previous check only compared logical element counts).
• Added AllocationPlannerTest.AvoidReuseOfPackedSubByteBufferForFullByteTensor, which builds X(float) → uint4 → float → uint8 → float so a dead uint4[1024] output would be reused by a uint8[1024] output. Without the fix the planner reuses the buffer (alloc_kind == kReuse) and the output is corrupted; with the fix the test passes and the model runs correctly

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

🟡 Changes recommended

The runtime safeguard lacks direct coverage, and the regression test can continue into an unsafe write after a nonfatal assertion.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Fixes unsafe memory reuse between packed sub-byte and full-byte tensors.

Changes:

• Accounts for sub-element packing density during allocation planning.
• Adds a runtime storage-capacity check.
• Adds a uint4-to-uint8 regression test.

File summaries

File	Description
onnxruntime/core/framework/allocation_planner.cc	Prevents reuse across incompatible packing densities.
onnxruntime/core/framework/execution_frame.cc	Rejects undersized reused buffers at runtime.
onnxruntime/test/framework/allocation_planner_test.cc	Tests the packed-buffer regression.

Review details

• Files reviewed: 3/3 changed files
• Comments generated: 2
• Review effort level: Balanced

---

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Wei Wang (wangw-1991)]

Dmitri Smirnov (@yuslepukhin) Can you help review this PR or help find the suitable reviewers? Thanks,

[Copilot]

🔵 Needs a closer look

The new execution-frame capacity-check branch lacks direct unit coverage.

Review details

Suppressed comments (1)

onnxruntime/core/framework/execution_frame.cc:701

• The new capacity-check/error path is not exercised by the added regression test: once SameSize() rejects the uint4/uint8 pair, execution never reaches this branch. Please add a focused ExecutionFrame test that directly reuses an undersized packed tensor and asserts this call fails, so this defense-in-depth check cannot regress independently of the planner fix.

    size_t required_storage_bytes = 0;
    ORT_RETURN_IF_ERROR(
        Tensor::CalculateTensorStorageSize(element_type, shape, /*alignment*/ 0, required_storage_bytes));
    const size_t buffer_storage_bytes = reuse_tensor->SizeInBytes();
    if (required_storage_bytes > buffer_storage_bytes) {

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Akshay Sonawane (apsonawane)]

Findings

Minor: The runtime capacity guard lacks direct coverage. The planner test prevents the invalid reuse, so execution_frame.cc rejection path never executes. Add a focused ExecutionFrameTest.

Minor: The test uses uint4[1024], covering only even packing. An odd dimension such as 1023 would verify ceiling division in allocation_planner_test.cc.

No blocking correctness issue found. Requiring equal carrier size, logical shape, and GetNumSubElems() correctly guarantees equal physical storage. Reuse chains, dynamic shapes, aliases, memory patterns, and offsets remain safe.

[Copilot]

🔵 Needs a closer look

The new execution-frame rejection path lacks direct test coverage.

Review details

Suppressed comments (1)

onnxruntime/core/framework/execution_frame.cc:706

• The new capacity-rejection path is not exercised by the added regression test: once SameSize() is fixed, the plan no longer reaches this branch, and the existing direct execution-frame test only reuses equal-size float tensors. Please add a direct AllocateMLValueTensorPreAllocateBuffer() test with a packed source buffer and a larger full-byte requested tensor, asserting that this status is returned; otherwise this defense-in-depth check could regress independently without detection.

    if (required_storage_bytes > buffer_storage_bytes) {
      return ORT_MAKE_STATUS(
          ONNXRUNTIME, FAIL, "Cannot re-use buffer: requested tensor needs ", required_storage_bytes,
          " bytes of storage but the buffer being reused only has ", buffer_storage_bytes,
          " bytes (buffer shape ", reuse_tensor->Shape(), ", requested shape ", shape,
          "). This can happen when a packed sub-byte tensor is reused for a full-byte tensor of the same shape.");

• Files reviewed: 3/3 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Wei Wang (wangw-1991)]

Findings

Minor: The runtime capacity guard lacks direct coverage. The planner test prevents the invalid reuse, so execution_frame.cc rejection path never executes. Add a focused ExecutionFrameTest.

Minor: The test uses uint4[1024], covering only even packing. An odd dimension such as 1023 would verify ceiling division in allocation_planner_test.cc.

No blocking correctness issue found. Requiring equal carrier size, logical shape, and GetNumSubElems() correctly guarantees equal physical storage. Reuse chains, dynamic shapes, aliases, memory patterns, and offsets remain safe.

Akshay Sonawane (@apsonawane) Thanks, I have fixed the two minor issues. Please take a look again. Please help merge this PR if no other issues.

[Copilot]

🟢 Approval recommended

The planner and runtime defenses are correct, targeted, and adequately covered by regression tests.

Review details

• Files reviewed: 4/4 changed files
• Comments generated: 0 new
• Review effort level: Balanced

[Akshay Sonawane (apsonawane)]

Wei Wang (@wangw-1991) can you fix the lint error

[Wei Wang (wangw-1991)]

Wei Wang (Wei Wang (@wangw-1991)) can you fix the lint error

Akshay Sonawane (@apsonawane) Thanks, done. Please take a look again.

[Wei Wang (wangw-1991)]

Akshay Sonawane (@apsonawane) The failed "Linux CUDA CI" seems not related to this PR. Can you help merge this PR if no other issues.