Avoid overflow in CPU TensorScatter indices

[Akshay Sonawane (apsonawane)]

This pull request addresses edge cases in the TensorScatter operator, particularly around handling large or potentially overflowing write_indices values. It also improves test coverage for these scenarios, ensuring robust and correct behavior in both linear and circular modes.

Bug fixes and safety improvements

• Fixed a potential overflow bug in linear mode by updating the bounds check to prevent write_indices values that could cause overflow when added to sequence_length. (onnxruntime/core/providers/cpu/llm/tensorscatter.cc)
• Refined the circular mode logic to correctly handle very large write_indices values, ensuring correct wraparound behavior without overflow. (onnxruntime/core/providers/cpu/llm/tensorscatter.cc)

Test coverage enhancements

• Added a test for linear mode that verifies the operator fails gracefully when write_indices addition would overflow, ensuring the new bounds check is enforced. (onnxruntime/test/providers/cpu/llm/tensorscatter_op_test.cc)
• Added a test for circular mode to confirm that very large write_indices values wrap correctly without overflow, verifying correct wraparound logic. (onnxruntime/test/providers/cpu/llm/tensorscatter_op_test.cc)
• Included <limits> header to support the use of std::numeric_limits<int64_t>::max() in tests. (onnxruntime/test/providers/cpu/llm/tensorscatter_op_test.cc)

[Ti-Tai Wang (titaiwangms)]

Major review finding

onnxruntime/core/providers/cpu/llm/tensorscatter.cc:146

Circular mode now evaluates wi % max_sequence_length before entering the copy loop. A shape with a zero-length cache axis and zero-length update is accepted by the existing sequence_length <= max_sequence_length check. With, for example, cache/update shapes [1, 0, 1], prefix_count is still 1, so this executes 0 % 0, which is undefined behavior.

Previously the zero-length loop performed no modulo operation and returned successfully.

Please return as a no-op when sequence_length == 0, or explicitly reject a zero cache axis before computing wi_mod.

[Ti-Tai Wang (titaiwangms)]

Major review finding

onnxruntime/core/providers/cpu/llm/tensorscatter.cc:93

The zero-sequence-length fix now loops over every batch entry even when write_indices is omitted and there is nothing to validate. Zero-volume tensors can have arbitrarily large nonzero dimensions without allocating corresponding data. For example, past_cache and update shaped [INT64_MAX, 0, 1] have zero elements and pass shape validation, but the new early-return path performs INT64_MAX iterations with wi hardcoded to zero.

This turns a zero-byte inference into an effectively unbounded CPU loop. When write_indices == nullptr, please return immediately for sequence_length == 0; retain the validation loop only when caller-provided indices actually need checking.