Fix prepacked weight reference lifetime

[Akshay Sonawane (apsonawane)]

This pull request introduces a new method to the PrepackedWeightsForGraph class to improve how prepacked weights are managed, particularly when saving models. The main change is the addition of a DiscardAndReplaceWithReferenceIfSaving method, which either replaces a prepacked weight with a reference or removes it, depending on whether the save mode is enabled. Corresponding updates are made to usage sites and comprehensive unit tests are added to verify the new behavior.

Enhancements to prepacked weights management:

• Added the DiscardAndReplaceWithReferenceIfSaving method to the PrepackedWeightsForGraph class. This method replaces an existing entry with a reference to another weight when in save mode, or removes it from the container otherwise. [1] [2]
• Updated the call site in SessionState::PrepackConstantInitializedTensors to use the new method, ensuring correct handling of prepacked weights during model save operations.

Testing improvements:

• Added new unit tests in tensorutils_test.cc to verify the behavior of DiscardAndReplaceWithReferenceIfSaving in both save mode and non-save mode, ensuring correct reference replacement and removal of weights as appropriate.

[Copilot]

Pull request overview

This PR fixes a lifetime issue in prepacked-weight handling by introducing a new PrepackedWeightsForGraph API that can safely discard an existing per-graph entry and (when saving) replace it with a non-owning reference to a stable, shared prepacked weight.

Changes:

• Added PrepackedWeightsForGraph::DiscardAndReplaceWithReferenceIfSaving to discard existing entries and, in save mode, store a referring copy to a provided weight.
• Updated SessionState::PrepackConstantInitializedTensors to use the new API when a shared prepacked weight already exists.
• Added unit tests for the new discard/replace behavior.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/framework/prepacked_weights_container.h	Declares the new discard/replace API for save-mode reference handling.
onnxruntime/core/framework/prepacked_weights_container.cc	Implements discard + replace-with-reference (save mode) or erase (non-save mode).
onnxruntime/core/framework/session_state.cc	Switches a call site to the new API to avoid creating references to soon-to-be-freed weights.
onnxruntime/test/framework/tensorutils_test.cc	Adds unit tests validating discard + reference replacement and non-save erasure behavior.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.