Fix INT64 initializer data propagation

[GopalakrishnanN]

Description

Fixes the in-memory INT64 initializer materialization path in Graph::SaveShapeValuesFromDataPropagation.

The INT64 branch now validates the tensor element count and resizes the destination TensorShapeVector before copying, matching the behavior of the other initializer materialization branches.

Testing

• Added a regression model that exercises the in-memory-external INT64 initializer path during graph re-resolution.
• Built onnxruntime_test_all with AddressSanitizer enabled on Windows Debug.
• Verified the regression fails on unfixed origin/main under ASan.
• Verified the same regression passes with this fix under ASan:

[  PASSED  ] 1 test.

Focused test:

ShapeInferenceV2Test.PartialDataPropagationTest

[Copilot]

Pull request overview

Fixes a buffer overrun in Graph::SaveShapeValuesFromDataPropagation when materializing in-memory-external INT64 initializers during shape/data propagation, and adds a regression model + session-load test to cover the scenario.

Changes:

• Fix INT64 initializer materialization by validating element count, resizing the destination TensorShapeVector, and then copying.
• Add a new regression model to exercise the in-memory-external INT64 initializer path during graph re-resolution.
• Extend ShapeInferenceV2Test.PartialDataPropagationTest to load the new model under BASIC optimizations.

Reviewed changes

Copilot reviewed 3 out of 4 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/core/graph/graph.cc	Fixes INT64 in-memory-external initializer copying by validating size and resizing before memcpy.
onnxruntime/test/testdata/test_shape_data_propagation_unsqueeze_inmemory_int64.py	Adds a generator script for the regression ONNX model that triggers the affected initializer path.
onnxruntime/test/framework/shape_inference_test.cc	Adds a regression test case that loads the new model to ensure the session can be created without error.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 4 changed files in this pull request and generated no new comments.

[hariharans29]

Review — PR #28778 (Fix INT64 initializer data propagation)

Verdict: approve with one small ask. Real bug, correct fix, scoped to the actual defect. The INT64 branch was missing both the element-count enforcement and the destination resize() that the INT32 branch above it had, so the memcpy could overrun (or in the lucky case, write into uninitialized capacity that input_values.data() of a default-constructed TensorShapeVector doesn't own). ASan caught it.

What it fixes

In Graph::SaveShapeValuesFromDataPropagation (onnxruntime/core/graph/graph.cc), the two initializer-materialization branches were inconsistent:

• INT32 branch (pre-existing, correct): validates data_span.size() == element_cnt, then input_values.resize(element_cnt), then element-wise copy.
• INT64 branch (bug): grabbed tensor.Data<int64_t>() directly and did memcpy(input_values.data(), src, element_cnt * sizeof(int64_t)) against a vector that was never sized.

The INT64 branch now mirrors INT32: size-check via the new enforce_tensor_element_count_matches lambda, resize(element_cnt), then memcpy.

The trigger model is well-chosen: Shape → Identity → Unsqueeze with 16 axes. Under ORT_ENABLE_BASIC the Identity gets folded, graph resolution re-runs data propagation, and by that point the unsq_axes initializer has been promoted to in-memory external storage — which is exactly the GetOrtValueInitializer(..., true) path that holds the buggy branch.

What's right

• Fix is at the right layer. No band-aid on the call site; the materialization function itself now upholds its postcondition.
• Refactored the enforce into a lambda so INT32 and INT64 share one message and one check. Tidy and reduces drift — if a third dtype gets added, the precedent is in front of the next contributor.
• DataAsSpan<int64_t>() instead of raw Data<int64_t>() gives .size() for the check at zero extra cost. Right primitive.
• Regression test loads the model under ORT_ENABLE_BASIC, which is the only level that triggers the re-resolution path. Loading under default optimizations wouldn't have caught this, so the explicit level setting is load-bearing — the comment in the test calls that out. Good.
• Generator script (.py) committed alongside the .onnx so the binary asset is reproducible. The repo's convention.
• memcpy is fine here despite the TensorShapeVector element type being a signed 8-byte int — same width, same alignment, trivially copyable. No need to switch to std::copy_n.

Comments

1. The ORT_ENFORCE should be ORT_RETURN_IF_NOT (or Status). This is the one thing I'd push back on. Graph::SaveShapeValuesFromDataPropagation already returns Status, and the function is reachable from session initialization on user-supplied models. A mismatch between utils::GetTensorShapeFromTensorProto() and the actual tensor element count is a model/serialization invariant violation, not a programmer bug — that's exactly the case ORT_RETURN_IF_NOT is for. ORT_ENFORCE will throw (or abort() in no-exceptions builds), and you'll see it as a crash on customer telemetry instead of a clean failed-Load. The INT32 branch above made the same mistake; ideally fix both in the same PR.

ORT_RETURN_IF_NOT(tensor_element_count == element_cnt,
                  "The element count from Tensor for initializer '", input_name,
                  "' should match utils::GetTensorShapeFromTensorProto(). Tensor: ",
                  tensor_element_count, ", TensorProto: ", element_cnt);

If you'd rather keep this PR surgical, file a follow-up — but please don't propagate the ORT_ENFORCE pattern further.

2. ORT_ENFORCE(session.GetOutputCount() == 1) in the test. Use ASSERT_EQ(session.GetOutputCount(), 1U). This is a gtest; using ORT_ENFORCE throws instead of recording a gtest failure, which gives a worse diagnostic on CI. The other Model #N blocks in the same function presumably already use gtest assertions — be consistent with those.

3. Branch name unsqueeze-datapropagation-stack-overflow says "stack overflow" but the diff and ASan output suggest a heap overrun (the destination is a TensorShapeVector, which is InlinedVector<int64_t> — heap-backed once it spills past the SBO threshold of 16 elements, which is exactly the test's axis_count = 16). Not a code issue, just noting that the original symptom description in the branch name and the fix description don't fully line up. Worth a line in the PR body about what the original repro actually crashed on, for the bug-DB record.

4. The test loads the model and checks output count, but doesn't run inference. That's enough to exercise the buggy path (the bug fires during Session(...) construction), but a single Run with a 1×1×1×…×1 input wouldn't hurt and would lock down that the Unsqueeze still produces the right shape. Optional.

5. The lambda captures [&] — fine in this scope, but it only actually needs input_name and element_cnt. Explicit capture ([&input_name, element_cnt]) is the local style in places, but the file isn't consistent about it. Leave as-is.

6. Consider extracting axis_count = 16 to a named constant. The number 16 is exactly the InlinedVector SBO threshold for int64_t (16 × 8 = 128 bytes, which is the default), and that's almost certainly why this model triggers the bug and a 4-axis model didn't. A # 16 == TensorShapeVector inline-buffer threshold; spills to heap and exposes the missing resize() comment in the generator script would save the next person an hour of bisecting.

What's not covered (and probably fine)

• The fix only touches the INT64 branch. If a third integer dtype gets propagated through this path in the future (UINT64? INT16?), the new contributor needs to follow the same pattern — but since both branches now look identical structurally, that's easy to get right. No action needed.
• Data propagation through Shape → Squeeze (mirror case) likely had the same shape-of-bug. Worth a separate quick audit, but not blocking this PR.

Bottom line

Ship it after the ORT_ENFORCE → ORT_RETURN_IF_NOT change (and the matching gtest assertion in the test). The fix is correct, the test reproduces the bug deterministically on an unfixed build, and the lambda refactor is the right shape. The two cleanup asks are the difference between "fixes one symptom" and "leaves the function in a better state than it found it." No reason to block; can be a follow-up commit.

[hariharans29]

Review — PR #28778 (Fix INT64 initializer data propagation)

Verdict: approve with one small ask. Real bug, correct fix, scoped to the actual defect. The INT64 branch was missing both the element-count enforcement and the destination resize() that the INT32 branch above it had, so the memcpy could overrun (or in the lucky case, write into uninitialized capacity that input_values.data() of a default-constructed TensorShapeVector doesn't own). ASan caught it.

What it fixes

In Graph::SaveShapeValuesFromDataPropagation (onnxruntime/core/graph/graph.cc), the two initializer-materialization branches were inconsistent:

• INT32 branch (pre-existing, correct): validates data_span.size() == element_cnt, then input_values.resize(element_cnt), then element-wise copy.
• INT64 branch (bug): grabbed tensor.Data<int64_t>() directly and did memcpy(input_values.data(), src, element_cnt * sizeof(int64_t)) against a vector that was never sized.

The INT64 branch now mirrors INT32: size-check via the new enforce_tensor_element_count_matches lambda, resize(element_cnt), then memcpy.

The trigger model is well-chosen: Shape → Identity → Unsqueeze with 16 axes. Under ORT_ENABLE_BASIC the Identity gets folded, graph resolution re-runs data propagation, and by that point the unsq_axes initializer has been promoted to in-memory external storage — which is exactly the GetOrtValueInitializer(..., true) path that holds the buggy branch.

What's right

• Fix is at the right layer. No band-aid on the call site; the materialization function itself now upholds its postcondition.
• Refactored the enforce into a lambda so INT32 and INT64 share one message and one check. Tidy and reduces drift — if a third dtype gets added, the precedent is in front of the next contributor.
• DataAsSpan<int64_t>() instead of raw Data<int64_t>() gives .size() for the check at zero extra cost. Right primitive.
• Regression test loads the model under ORT_ENABLE_BASIC, which is the only level that triggers the re-resolution path. Loading under default optimizations wouldn't have caught this, so the explicit level setting is load-bearing — the comment in the test calls that out. Good.
• Generator script (.py) committed alongside the .onnx so the binary asset is reproducible. The repo's convention.
• memcpy is fine here despite the TensorShapeVector element type being a signed 8-byte int — same width, same alignment, trivially copyable. No need to switch to std::copy_n.

Comments

1. The ORT_ENFORCE should be ORT_RETURN_IF_NOT (or Status). This is the one thing I'd push back on. Graph::SaveShapeValuesFromDataPropagation already returns Status, and the function is reachable from session initialization on user-supplied models. A mismatch between utils::GetTensorShapeFromTensorProto() and the actual tensor element count is a model/serialization invariant violation, not a programmer bug — that's exactly the case ORT_RETURN_IF_NOT is for. ORT_ENFORCE will throw (or abort() in no-exceptions builds), and you'll see it as a crash on customer telemetry instead of a clean failed-Load. The INT32 branch above made the same mistake; ideally fix both in the same PR.

ORT_RETURN_IF_NOT(tensor_element_count == element_cnt,
                  "The element count from Tensor for initializer '", input_name,
                  "' should match utils::GetTensorShapeFromTensorProto(). Tensor: ",
                  tensor_element_count, ", TensorProto: ", element_cnt);

If you'd rather keep this PR surgical, file a follow-up — but please don't propagate the ORT_ENFORCE pattern further.

2. ORT_ENFORCE(session.GetOutputCount() == 1) in the test. Use ASSERT_EQ(session.GetOutputCount(), 1U). This is a gtest; using ORT_ENFORCE throws instead of recording a gtest failure, which gives a worse diagnostic on CI. The other Model #N blocks in the same function presumably already use gtest assertions — be consistent with those.

3. Branch name unsqueeze-datapropagation-stack-overflow says "stack overflow" but the diff and ASan output suggest a heap overrun (the destination is a TensorShapeVector, which is InlinedVector<int64_t> — heap-backed once it spills past the SBO threshold of 16 elements, which is exactly the test's axis_count = 16). Not a code issue, just noting that the original symptom description in the branch name and the fix description don't fully line up. Worth a line in the PR body about what the original repro actually crashed on, for the bug-DB record.

4. The test loads the model and checks output count, but doesn't run inference. That's enough to exercise the buggy path (the bug fires during Session(...) construction), but a single Run with a 1×1×1×…×1 input wouldn't hurt and would lock down that the Unsqueeze still produces the right shape. Optional.

5. The lambda captures [&] — fine in this scope, but it only actually needs input_name and element_cnt. Explicit capture ([&input_name, element_cnt]) is the local style in places, but the file isn't consistent about it. Leave as-is.

6. Consider extracting axis_count = 16 to a named constant. The number 16 is exactly the InlinedVector SBO threshold for int64_t (16 × 8 = 128 bytes, which is the default), and that's almost certainly why this model triggers the bug and a 4-axis model didn't. A # 16 == TensorShapeVector inline-buffer threshold; spills to heap and exposes the missing resize() comment in the generator script would save the next person an hour of bisecting.

What's not covered (and probably fine)

• The fix only touches the INT64 branch. If a third integer dtype gets propagated through this path in the future (UINT64? INT16?), the new contributor needs to follow the same pattern — but since both branches now look identical structurally, that's easy to get right. No action needed.
• Data propagation through Shape → Squeeze (mirror case) likely had the same shape-of-bug. Worth a separate quick audit, but not blocking this PR.

Bottom line

Ship it after the ORT_ENFORCE → ORT_RETURN_IF_NOT change (and the matching gtest assertion in the test). The fix is correct, the test reproduces the bug deterministically on an unfixed build, and the lambda refactor is the right shape. The two cleanup asks are the difference between "fixes one symptom" and "leaves the function in a better state than it found it." No reason to block; can be a follow-up commit.

AI suggests using ORT_RETURN_IF_NOT() instead of ORT_ENFORCE(), but that is changing away from the previous behavior I think. So it might be to keep as is.

[tianleiwu]

Review Summary

Correct, well-scoped fix for a heap buffer overflow in Graph::SaveShapeValuesFromDataPropagation on the in-memory-external INT64 initializer path. Approving.

Root cause: The previous INT64 branch copied into input_values without sizing it first:

const int64_t* src = tensor.Data<int64_t>();
memcpy(input_values.data(), src, element_cnt * sizeof(int64_t));

Since input_values (a TensorShapeVector) was never resize()d, the memcpy wrote past its small-buffer/empty storage whenever the axes initializer was materialized as in-memory external data — a real out-of-bounds write surfaced by ASan.

What's good:

• Overflow fixed: input_values.resize(element_cnt) is now done before the memcpy, matching the INT32 branch and both non-in-memory branches. The copy length element_cnt * sizeof(int64_t) is bounded by the validated element count, so source and destination lengths agree.
• Checked element count: Switching to tensor.DataAsSpan<int64_t>() provides a checked size for validation, consistent with the INT32 path.
• Clearer diagnostics: The shared enforce_tensor_element_count_matches lambda replaces the duplicated, garbled ORT_ENFORCE message ("...samefrom...") and now reports the initializer name plus observed vs. expected counts.
• Targeted regression: The new Shape -> Identity -> Unsqueeze model reproduces the exact path where ORT_ENABLE_BASIC drops the Identity, re-resolution re-runs data propagation after large initializers become in-memory external OrtValues, and the INT64 axes initializer hits the fixed branch. axis_count = 16 keeps the checked-in model small while still exceeding the small-buffer threshold.

Prior threads: All earlier review threads are already resolved on this head (Copilot's ORT_ENFORCE message and oversized axis_count, skottmckay's request for a comment explaining the broken path, and the CodeQL duplicate-import note).

No blocking concerns and no new findings.