Introduce checks to prevent buffer overflow, add tests

[yuslepukhin]

This pull request improves the correctness and robustness of the InPlaceAccumulator and InPlaceAccumulatorV2 gradient accumulation kernels by adding new unit tests and enforcing stricter shape validation in both CPU and CUDA implementations. The changes ensure that shape mismatches are caught early, and that optional outputs are properly handled in all execution providers.

Test coverage improvements:

• Added a test to verify that InPlaceAccumulator correctly passes through the old_sum unchanged and does not consume the value input when the optional update_signal is false.
• Added tests for InPlaceAccumulatorV2 to check that shape mismatches between accumulation_buffer and value are detected and handled as errors, covering both overwrite and accumulate branches.
• Added tests to verify that InPlaceAccumulatorV2 correctly handles the case where the optional accumulation_buffer_out output is omitted, for both CPU and CUDA providers. [1] [2]

Kernel validation improvements:

• Added explicit shape validation to the CPU implementation of InPlaceAccumulatorV2, ensuring that the shapes of the accumulation buffer and the value tensor match.
• Added the same shape validation to the CUDA implementation of InPlaceAccumulatorV2, preventing out-of-bounds memory accesses.

[Copilot]

Pull request overview

Adds explicit shape validation to the InPlaceAccumulatorV2 CPU and CUDA kernels to prevent out-of-bounds memory access on mismatched accumulation_buffer/value shapes, and adds unit tests covering the new validation, the no-update path of InPlaceAccumulator, and the optional accumulation_buffer_out output path.

Changes:

• Add ORT_RETURN_IF_NOT shape equality check at the entry of both CPU and CUDA InPlaceAccumulatorV2 implementations.
• Add tests verifying shape mismatch is rejected (overwrite and accumulate branches).
• Add tests for InPlaceAccumulator no-update pass-through and for InPlaceAccumulatorV2 with the optional accumulation output omitted (CPU and CUDA).

Reviewed changes

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

File	Description
orttraining/orttraining/training_ops/cpu/optimizer/gradient_control.cc	Adds shape equality check between accumulation_buffer and new_value before performing in-place accumulation.
orttraining/orttraining/training_ops/cuda/optimizer/gradient_control.cc	Adds the same shape equality check on the CUDA side before the memcpy/cast/accumulate paths.
orttraining/orttraining/test/gradient/gradient_ops_test.cc	New tests for the no-update path, V2 shape-mismatch (overwrite and accumulate), and omitted optional accumulation output on CPU and CUDA.

---

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

[tianleiwu]

Summary

Clean, focused fix for a real buffer overflow vulnerability. The shape validation is correctly placed before any memcpy/cudaMemcpyAsync in both CPU and CUDA kernels, uses the right error-handling macro (ORT_RETURN_IF_NOT), and includes both shapes in the diagnostic message for easy debugging.

Kernel Validation

Positive:

• Check placement is correct in both implementations — runs before any data copy that uses the value tensor's size.
• Error messages are consistent between CPU and CUDA, include both shapes.
• ORT_RETURN_IF_NOT is the correct pattern for runtime input validation (recoverable error from user data, not a programming error).

Test Coverage

Positive:

• InPlaceAccumulatorFloat32_NoUpdate covers the previously untested update_signal=false pass-through path.
• Shape mismatch tests exercise both overwrite and accumulate branches.
• Optional output tests verify both CPU and GPU handle omitted accumulation_buffer_out gracefully.

Suggestion:

• The shape mismatch test only runs with DefaultCpuExecutionProvider(). Since the same validation was added to the CUDA kernel, a CUDA-side shape mismatch test (guarded by #if defined(USE_CUDA)) would confirm that path as well. Minor gap since the code is identical.