onnxruntime
6843e599 - Validate serialized prepacked buffers (#33499)

Commit
7 days ago
Validate serialized prepacked buffers (#33499) This pull request introduces a mechanism to verify the equivalence of pre-packed weights in ONNX Runtime, ensuring that pre-packed weights loaded from disk or shared across sessions exactly match those generated by the current kernel. This helps prevent subtle bugs caused by mismatched buffer contents or layouts. The main changes include implementing an `IsEquivalentTo` method for `PrePackedWeights`, adding validation logic to session state handling, and providing comprehensive unit tests. **PrePackedWeights Equivalence and Validation:** * Added a new method `IsEquivalentTo` to the `PrePackedWeights` class to compare buffer contents, sizes, and ownership flags for equivalence. (`onnxruntime/core/framework/prepacked_weights.h`, `onnxruntime/core/framework/prepacked_weights.cc`) [[1]](diffhunk://#diff-d88f6d6584c1873f2a7596d8194b2dac02cb6ac609243f6c2ad850116c81bc75R30-R31) [[2]](diffhunk://#diff-72fe9886d7a2247e53ed822e181677c8dbcbd0cb2765adc9ff145ae5d6987f6dR34-R56) * Introduced a helper function `ValidatePrePackedWeights` in `session_state.cc` that uses `IsEquivalentTo` to ensure loaded or shared pre-packed weights match expected values, returning an error if not. (`onnxruntime/core/framework/session_state.cc`) **Session State Integration:** * Integrated `ValidatePrePackedWeights` into the prepacking flow in `SessionState::PrepackConstantInitializedTensors` to validate pre-packed weights at all relevant points: when using shared, disk-loaded, or cached weights. (`onnxruntime/core/framework/session_state.cc`) [[1]](diffhunk://#diff-2126cbcd226e0e59f46b87efcb45eca3d118901a764fba01d62049b1ad0bad14R637-R638) [[2]](diffhunk://#diff-2126cbcd226e0e59f46b87efcb45eca3d118901a764fba01d62049b1ad0bad14R669-R670) [[3]](diffhunk://#diff-2126cbcd226e0e59f46b87efcb45eca3d118901a764fba01d62049b1ad0bad14R743-R745) **Testing:** * Added a comprehensive unit test `PrePackedWeightsTest.EquivalenceRequiresMatchingLayoutAndContents` to verify the correctness of the new equivalence logic under various scenarios, including different buffer contents, sizes, missing buffers, and ownership flags. (`onnxruntime/test/framework/session_state_test.cc`) **Build/Include Cleanups:** * Included `<cstring>` where necessary to support buffer comparison using `std::memcmp`. (`onnxruntime/core/framework/prepacked_weights.cc`, `onnxruntime/test/framework/session_state_test.cc`) [[1]](diffhunk://#diff-72fe9886d7a2247e53ed822e181677c8dbcbd0cb2765adc9ff145ae5d6987f6dR5-R7) [[2]](diffhunk://#diff-2e498c0ec894ae05b0ddfb15f7bc4b927df6626c0497c47de85a5e808317d3f1R7)
Author
Parents
Loading