Skip to content

Commit 4dbb05f

Browse files
authored
[core] allow using initializer allocator for prepack (#26617)
### Description This PR makes ORT to prefer initializer allocator when calling `OpKernel::PrePack`. If an EP does not register an initializer allocator (currently only WebGPU does this), the behavior is kept unchanged. ### Motivation and Context Helps to improve the memory usage when doing prepack.
1 parent 607d5e4 commit 4dbb05f

File tree

1 file changed

+2
-2
lines changed

1 file changed

+2
-2
lines changed

onnxruntime/core/framework/session_state.cc

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -588,15 +588,15 @@ Status SessionState::PrepackConstantInitializedTensors(
588588
// within this session. Or if the weight is not present on disk,
589589
// we store the newly minted pre-packed data.
590590

591-
AllocatorPtr session_cpu_alloc = GetAllocator(kernel->Info().GetDevice(OrtMemType::OrtMemTypeDefault));
591+
AllocatorPtr session_initializer_alloc = GetInitializerAllocator(kernel->Info().GetDevice(OrtMemType::OrtMemTypeDefault));
592592
PrePackedWeights weights_to_be_filled_in;
593593
// The reason we invoke PrePack() before looking into the container for any pre-packed weight
594594
// cached by another instance of the same op_type (for the same constant initializer) is because
595595
// to truly know if we can use a cached pre-packed weight, we would have to compare the cached
596596
// pre-packed weight with the pre-packed weight generated by this instance of the same op_type because
597597
// other static properties of the node like node attributes could play a role in the pre-packed
598598
// weights' contents.
599-
ORT_RETURN_IF_ERROR(kernel->PrePack(const_initialized_tensor, input_idx, session_cpu_alloc,
599+
ORT_RETURN_IF_ERROR(kernel->PrePack(const_initialized_tensor, input_idx, session_initializer_alloc,
600600
is_packed,
601601
&weights_to_be_filled_in));
602602

0 commit comments

Comments
 (0)