onnxruntime
df44eabd - Introduce checks to prevent buffer overflow, add tests (#28713)

Commit
70 days ago
Introduce checks to prevent buffer overflow, add tests (#28713) 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]](diffhunk://#diff-c62fec6d9ac7d24c7d6befe4e18317d2690385051c01a6412651f01e190de1beR2247-R2288) [[2]](diffhunk://#diff-c62fec6d9ac7d24c7d6befe4e18317d2690385051c01a6412651f01e190de1beR2332-R2346) **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.
Author
Parents
Loading