Fix over-copy of packed sub-byte tensors in OrtApi::GetValue by neilmsft · Pull Request #29157 · microsoft/onnxruntime

neilmsft · 2026-06-18T18:52:13Z

Summary

OrtApi::GetValue on a sequence of tensors copied elements using element_type->Size() * element_count bytes. For packed sub-byte types (INT4/UINT4, two elements per byte) this is ~2× the real storage size, causing a heap over-read of the source and overflow of the destination.

Fix

In PopulateTensorWithData (onnxruntime/core/session/onnxruntime_c_api.cc), copy the tensor's actual packed size via Tensor::SizeInBytes() instead of element_type->Size() * num_elems, and drop the now-unused elem_size parameter. SizeInBytes() is packing-aware: identical for ≥1-byte types (no behavior change) and correct for sub-byte types.

Testing

Added CApiTest.CreateGetSeqSubByteTensors: sequences of INT4/UINT4 tensors with an odd length (7 elements → 4 packed bytes), verifying GetValue() round-trips correctly.
Full onnxruntime_shared_lib_test suite passes (185/185); other paths unchanged.
Confirmed under AddressSanitizer: the old formula triggers a heap-buffer-overflow, the new one is clean.

Copilot

Pull request overview

Fixes a memory-safety bug in OrtApi::GetValue when copying tensors with packed sub-byte element types (INT4/UINT4) out of tensor sequences, ensuring the copy size matches the tensor’s actual packed storage size.

Changes:

Update PopulateTensorWithData to copy tensor.SizeInBytes() (packing-aware) instead of element_type->Size() * num_elems.
Simplify PopulateTensorWithData by removing the unused elem_size parameter.
Add a shared-lib C API test covering sequence GetValue() round-trips for INT4/UINT4 tensors with odd element counts (packed byte edge case).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/session/onnxruntime_c_api.cc	Fix tensor byte-copy sizing in `GetValue()` sequence element extraction for packed sub-byte tensors.
onnxruntime/test/shared_lib/test_nontensor_types.cc	Add regression test for `GetValue()` on sequences of INT4/UINT4 tensors with odd lengths (packed storage).

Comments suppressed due to low confidence (1)

onnxruntime/core/session/onnxruntime_c_api.cc:2035

In the string-tensor path, std::copy(str_span.begin(), str_span.end(), dst) copies num_elems strings into a buffer sized for len elements. The precondition only checks num_elems < len, so num_elems > len would write past the destination. Copy only len elements (and ignore extras) to keep this helper safe/consistent with the non-string path.

    const std::string* strings = reinterpret_cast<const std::string*>(data_elem);
    auto str_span = gsl::make_span(strings, num_elems);
    auto* dst = tensor.MutableData<std::string>();
    std::copy(str_span.begin(), str_span.end(), dst);

yuslepukhin · 2026-06-23T02:05:46Z

Issues not addressed by PR #29157.

EstimateTensorSizeInBytes in constant_folding.cc:177-200 — computes num_elements * 1 instead of ceil(num_elements / sub_elems) for all 5 sub-byte types. Overestimates tensor size by 2-4×, potentially skipping constant-folding opportunities for small sub-byte tensors.
Doc comment for GetTensorSizeInBytes in onnxruntime_c_api.h — states the formula is sizeof(element_type) * total_element_count, which is incorrect for sub-byte packed types.
No TypeToTensorType specializations for sub-byte types in the C++ header API — users cannot use the typed Value::CreateTensor() templates and must fall back to the void* overload. This is a design/usability gap rather than a correctness bug.
No documentation on GetElementCount() vs GetTensorSizeInBytes() distinction for sub-byte types — users who iterate GetElementCount() bytes from GetTensorData<uint8_t>() read past the buffer (UB).

Can we address at least some of those here or should we open another PR?

neilmsft · 2026-07-02T06:47:54Z

Issues not addressed by PR #29157.

EstimateTensorSizeInBytes in constant_folding.cc:177-200 — computes num_elements * 1 instead of ceil(num_elements / sub_elems) for all 5 sub-byte types. Overestimates tensor size by 2-4×, potentially skipping constant-folding opportunities for small sub-byte tensors.

Doc comment for GetTensorSizeInBytes in onnxruntime_c_api.h — states the formula is sizeof(element_type) * total_element_count, which is incorrect for sub-byte packed types.

No TypeToTensorType specializations for sub-byte types in the C++ header API — users cannot use the typed Value::CreateTensor() templates and must fall back to the void* overload. This is a design/usability gap rather than a correctness bug.

No documentation on GetElementCount() vs GetTensorSizeInBytes() distinction for sub-byte types — users who iterate GetElementCount() bytes from GetTensorData<uint8_t>() read past the buffer (UB).

Can we address at least some of those here or should we open another PR?

For the first one: Fixed. Size estimators are now packing-aware ( ceil(num_elements / sub_elems_per_byte) * element_size ; 2 for int4/uint4/float, 4 for int2/uint2), applied to all three paths (generic, Unique , ConstantOfShape ). This matches the packing-aware post-exec Tensor::SizeInBytes() check. Added ConstantFoldingSubByteOutputSizeIsPacked covering fold/block cases.

For the second: Fixed. Clarified the sizeof * count formula only holds for >= 1-byte types, sub-byte types return the smaller packed size.

For the fourth: Fixed in onnxruntime_cxx_api.h - I noted element count ≠ byte size for sub-byte types.

For the third: Deferring to a follow-up. As you said, it's a usability gap, not a correctness bug, and needs new public wrapper types plus reworking CreateTensor 's count * sizeof(T) sizing.

edgchen1 · 2026-07-03T00:57:08Z

+  /// For packed sub-byte types (e.g. int4/uint4) this is the element count, which is NOT the same
+  /// as the storage size in bytes: multiple elements share a storage byte, so the raw buffer is
+  /// smaller than the element count. Use Ort::Value::GetTensorSizeInBytes() when sizing or bounds-


"which is NOT the same as the storage size in bytes"

nit: this is true for element types that are larger than one byte too. this sentence feels overly specific in this context - perhaps we can drop it. I think the recommendation in the next sentence to use GetTensorSizeInBytes() is still good to have.

Done. dropped the sub-byte-specific sentence, kept the GetTensorSizeInBytes()

edgchen1 · 2026-07-03T00:59:54Z

+                                                          int64_t num_elements, size_t element_size) {
+  const int64_t sub_elems_per_byte = GetNumSubElemsPerByteForConstantFolding(elem_type);
+  const int64_t num_storage_units = (num_elements + (sub_elems_per_byte - 1)) / sub_elems_per_byte;
+  return SafeInt<int64_t>(num_storage_units) * static_cast<int64_t>(element_size);


we don't need to cast the operand of a SafeInt multiply, SafeInt will handle checking it already.

https://github.com/dcleblanc/SafeInt/blob/4cafc9196c4da9c817992b20f5253ef967685bf8/SafeInt.hpp#L5836-L5837

Suggested change

return SafeInt<int64_t>(num_storage_units) * static_cast<int64_t>(element_size);

return SafeInt<int64_t>(num_storage_units) * element_size;

Done. The multiply no longer casts the operand.

edgchen1 · 2026-07-03T01:10:34Z

+// Number of sub-byte elements packed into a single storage byte for packed sub-byte element
+// types (e.g. int4/uint4 pack 2 elements per byte, int2/uint2 pack 4). Returns 1 for all other
+// (>= 1 byte) element types. Keep in sync with the sub-byte types registered in data_types.cc.
+static int64_t GetNumSubElemsPerByteForConstantFolding(ONNX_NAMESPACE::TensorProto_DataType elem_type) {
+  switch (elem_type) {
+    case ONNX_NAMESPACE::TensorProto_DataType_INT4:
+    case ONNX_NAMESPACE::TensorProto_DataType_UINT4:
+    case ONNX_NAMESPACE::TensorProto_DataType_FLOAT4E2M1:
+      return 2;
+    case ONNX_NAMESPACE::TensorProto_DataType_INT2:
+    case ONNX_NAMESPACE::TensorProto_DataType_UINT2:
+      return 4;
+    default:
+      return 1;
+  }
+}


is it possible to reuse the Tensor class member functions, like Tensor::CalculateTensorStorageSize()?

onnxruntime/onnxruntime/core/framework/tensor.cc

Line 74 in 3eb4a35

size_t Tensor::CalculateTensorStorageSize(MLDataType elt_type, const TensorShape& shape) {

that implementation already accounts for the sub-byte type sizes so we don't have to keep this in sync here.

edgchen1 · 2026-07-03T01:11:36Z

+                                                               input_element_size);
+    } else {
+      // The remaining Unique outputs are int64 index/count tensors.
+      total_size += SafeInt<int64_t>(input_num_elements) * static_cast<int64_t>(sizeof(int64_t));


Suggested change

total_size += SafeInt<int64_t>(input_num_elements) * static_cast<int64_t>(sizeof(int64_t));

total_size += SafeInt<int64_t>(input_num_elements) * sizeof(int64_t);

Fix memcpy over-copy of packed sub-byte tensors

d49467d

neilmsft requested review from devang-ml and edgchen1 June 18, 2026 19:00

update comment

e05b6d9

tianleiwu requested a review from Copilot June 18, 2026 22:07

Copilot started reviewing on behalf of tianleiwu June 18, 2026 22:07 View session

Copilot AI reviewed Jun 18, 2026

View reviewed changes

Comment thread onnxruntime/test/shared_lib/test_nontensor_types.cc Outdated

edgchen1 reviewed Jun 19, 2026

View reviewed changes

Comment thread onnxruntime/core/session/onnxruntime_c_api.cc Outdated

Comment thread onnxruntime/test/shared_lib/test_nontensor_types.cc Outdated

Comment thread onnxruntime/test/shared_lib/test_nontensor_types.cc Outdated

kibae mentioned this pull request Jun 20, 2026

CPU EP mis-loads packed UINT2 Constant initializer (treats UInt2x4 as unpacked storage; INT2 unaffected) #29172

Open

address const and other update

38461d0

edgchen1 reviewed Jun 22, 2026

View reviewed changes

Comment thread onnxruntime/test/shared_lib/test_nontensor_types.cc

yuslepukhin mentioned this pull request Jun 23, 2026

[Mobile] QNN EP yields Segmentation Fault on PadNodeGroupSelector Check if dq_nodes are empty #28717

Closed

neilmsft added 2 commits July 1, 2026 23:48

Address follow up comments

6aa364f

Update header

44a147b

edgchen1 reviewed Jul 3, 2026

View reviewed changes

address comments

2854d77

	return SafeInt<int64_t>(num_storage_units) * static_cast<int64_t>(element_size);
	return SafeInt<int64_t>(num_storage_units) * element_size;

	total_size += SafeInt<int64_t>(input_num_elements) * static_cast<int64_t>(sizeof(int64_t));
	total_size += SafeInt<int64_t>(input_num_elements) * sizeof(int64_t);

Uh oh!

Conversation

neilmsft commented Jun 18, 2026

Summary

Fix

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yuslepukhin commented Jun 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

neilmsft commented Jul 2, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yuslepukhin commented Jun 23, 2026 •

edited

Loading