Add a session config to share pre-packed weights of all initializers across sessions

include/onnxruntime/core/session/onnxruntime_session_options_config_keys.h

@@ -383,6 +383,19 @@ static const char* const kOrtSessionOptionsModelExternalInitializersFileFolderPa
 static const char* const kOrtSessionOptionsSavePrePackedConstantInitializers =
     "session.save_external_prepacked_constant_initializers";
 
+// Enrolls every constant initializer handled by the CPU EP in the shared pre-packed weights
+// container attached to the session, not only the initializers registered via
+// OrtApi::AddInitializer. Sessions created over the same model with the same container then
+// reuse one pre-packed copy per weight. Has no effect unless a prepacked weights container
+// is attached. Kernels that pre-pack without producing shareable buffers keep their
+// kernel-owned result and do not participate in sharing.
+//
+// - "0": Default. Only AddInitializer-registered (or transformer-tagged) initializers share.
+// - "1": All constant initializers on the CPU EP enroll in the container.
+// Sample usage: sess_options.add_session_config_entry(kOrtSessionOptionsSharePrepackedWeightsForAllInitializers, "1")
+static const char* const kOrtSessionOptionsSharePrepackedWeightsForAllInitializers =
+    "session.share_prepacked_weights_for_all_initializers";
+
 // Use this config when you want to collect memory stats for each node in the graph.
 // The file format is a CSV file with the following columns:
 // The file will be created if it does not exist, and will be overwritten if it does.

onnxruntime/core/framework/session_state.cc

@@ -478,7 +478,14 @@ static std::string GenerateKeyForPrepackedWeightsMap(const std::string& op_type,
 Status SessionState::PrepackConstantInitializedTensors(
     InlinedHashMap<std::string, size_t>& constant_initializers_use_count,
     const std::unordered_map<std::string, const OrtValue*>& initializers_to_share_map) {
-  auto prepacked_constant_weights = [this, &constant_initializers_use_count, &initializers_to_share_map](
+  // Enrolls every constant initializer in the shared pre-packed weights container.
+  // See kOrtSessionOptionsSharePrepackedWeightsForAllInitializers.
+  const bool share_all_initializers =
+      sess_options_.config_options.GetConfigOrDefault(
+          kOrtSessionOptionsSharePrepackedWeightsForAllInitializers, "0") == "1";
+
+  auto prepacked_constant_weights = [this, &constant_initializers_use_count, &initializers_to_share_map,
+                                     share_all_initializers](
                                         bool should_cache_prepacked_weights_for_shared_initializers) -> Status {
     for (auto& node : GetGraphViewer().Nodes()) {
       if (sess_options_.IsLoadCancellationFlagSet()) {
@@ -513,7 +520,7 @@ Status SessionState::PrepackConstantInitializedTensors(
                 // hash, never the tag value (see the rationale at the key computation).
                 const bool enroll_tagged_initializer =
                     (st->graph_.GetSharedPrepackInitializerId(input_name) != nullptr);
-                if ((is_shared_initializer || enroll_tagged_initializer) &&
+                if ((is_shared_initializer || enroll_tagged_initializer || share_all_initializers) &&
                     should_cache_prepacked_weights_for_shared_initializers &&
                     node.GetExecutionProviderType() == kCpuExecutionProvider) {
                   // caching of pre-packed weights' turned ON
@@ -535,13 +542,18 @@ Status SessionState::PrepackConstantInitializedTensors(
                   if (is_packed) {
                     // BUG CHECK: Ensure that a kernel either filled in the pre-packed weights
                     // to be cached, or explicitly marked the packed weights as kernel-owned.
+                    // Initializers that are only enrolled by the share-all option are exempt:
+                    // a kernel that keeps its packed data to itself without declaring so
+                    // (e.g. fp16 LayerNormalization) does not participate in sharing.
                     ORT_RETURN_IF_NOT(!weights_to_be_filled_in.buffers_.empty() ||
-                                          weights_to_be_filled_in.has_kernel_owned_packed_weights_,
+                                          weights_to_be_filled_in.has_kernel_owned_packed_weights_ ||
+                                          !(is_shared_initializer || enroll_tagged_initializer),
                                       "The kernel corresponding to the node ", node.Name(),
                                       " doesn't have an implementation that can cache computed pre-packed weights");
                   }
 
-                  if (is_packed && !weights_to_be_filled_in.has_kernel_owned_packed_weights_) {
+                  if (is_packed && !weights_to_be_filled_in.has_kernel_owned_packed_weights_ &&
+                      !weights_to_be_filled_in.buffers_.empty()) {
                     const auto& op_type = node.OpType();
 
                     // Sanity check

onnxruntime/test/framework/session_state_test.cc

@@ -1313,6 +1313,122 @@ TEST_F(SessionStateTestSharedInitalizersWithPrePacking, BrokenKernelWithoutCache
       "doesn't have an implementation that can cache computed pre-packed weights");
 }
 
+// Pre-packing enabled + no shared initializers + pre-packed weights container +
+// share_prepacked_weights_for_all_initializers = "1" =
+// the model's own constant initializers enroll in the container, so their
+// pre-packed weights are shared across sessions.
+TEST_F(SessionStateTestSharedInitalizersWithPrePacking, ShareAllInitializersSharesAcrossSessions) {
+  SessionOptions sess_options;
+  sess_options.enable_mem_pattern = true;
+  sess_options.execution_mode = ExecutionMode::ORT_SEQUENTIAL;
+  sess_options.use_deterministic_compute = false;
+  sess_options.enable_mem_reuse = true;
+  // Enable pre-packing
+  sess_options.config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "0";
+  // Enroll all constant initializers in the shared pre-packed weights container
+  sess_options.config_options.configurations[kOrtSessionOptionsSharePrepackedWeightsForAllInitializers] = "1";
+
+  // Enable pre-packed weights container
+  PrepackedWeightsContainer prepacked_weights_container;
+
+  // First session/model
+  Model model_1("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
+                domain_to_version, std::vector<ONNX_NAMESPACE::FunctionProto>(),
+                DefaultLoggingManager().DefaultLogger());
+
+  CreateSimpleGraph(model_1.MainGraph());
+  PlaceAllNodesToCPUEP(model_1.MainGraph());
+  SessionState session_state_1(model_1.MainGraph(),
+                               execution_providers,
+                               tp.get(),
+                               nullptr, /*inter_op_thread_pool*/
+                               dtm,
+                               edlm,
+                               DefaultLoggingManager().DefaultLogger(),
+                               profiler,
+                               sess_options,
+                               &prepacked_weights_container);
+
+  ASSERT_STATUS_OK(session_state_1.FinalizeSessionState(std::basic_string<PATH_CHAR_TYPE>(),
+                                                        kernel_registry_manager));
+
+  const auto* kernel_1 = reinterpret_cast<const PrePackingTestOpKernel*>(session_state_1.GetKernel(0));
+  ASSERT_EQ(session_state_1.GetNumberOfPrepacksCounter(), static_cast<size_t>(1));
+  ASSERT_EQ(kernel_1->prepack_calls_count, 1);
+  // The kernel was handed a container-owned buffer
+  ASSERT_EQ(kernel_1->store_pre_packed_weight_calls_count, 1);
+  // The first session computed the pre-pack; nothing was cached yet
+  ASSERT_EQ(session_state_1.GetUsedSharedPrePackedWeightCounter(), static_cast<size_t>(0));
+
+  // Second session/model
+  Model model_2("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
+                domain_to_version, std::vector<ONNX_NAMESPACE::FunctionProto>(),
+                DefaultLoggingManager().DefaultLogger());
+
+  CreateSimpleGraph(model_2.MainGraph());
+  PlaceAllNodesToCPUEP(model_2.MainGraph());
+  SessionState session_state_2(model_2.MainGraph(),
+                               execution_providers,
+                               tp.get(),
+                               nullptr, /*inter_op_thread_pool*/
+                               dtm,
+                               edlm,
+                               DefaultLoggingManager().DefaultLogger(),
+                               profiler,
+                               sess_options,
+                               &prepacked_weights_container);
+
+  ASSERT_STATUS_OK(session_state_2.FinalizeSessionState(std::basic_string<PATH_CHAR_TYPE>(),
+                                                        kernel_registry_manager));
+
+  const auto* kernel_2 = reinterpret_cast<const PrePackingTestOpKernel*>(session_state_2.GetKernel(0));
+  ASSERT_EQ(session_state_2.GetNumberOfPrepacksCounter(), static_cast<size_t>(1));
+  ASSERT_EQ(kernel_2->prepack_calls_count, 1);
+  ASSERT_EQ(kernel_2->store_pre_packed_weight_calls_count, 1);
+  // The second session reuses the pre-packed weight cached by the first one
+  ASSERT_EQ(session_state_2.GetUsedSharedPrePackedWeightCounter(), static_cast<size_t>(1));
+}
+
+// A kernel that pre-packs without producing cacheable buffers does not participate
+// in sharing under share_prepacked_weights_for_all_initializers; finalization
+// succeeds. An initializer registered via AddInitializer keeps the strict check
+// (see BrokenKernelWithoutCacheableBuffersFails).
+TEST_F(SessionStateTestSharedInitalizersWithPrePacking, ShareAllInitializersSkipsKernelOwnedPrepacks) {
+  SessionOptions sess_options;
+  sess_options.enable_mem_pattern = true;
+  sess_options.execution_mode = ExecutionMode::ORT_SEQUENTIAL;
+  sess_options.use_deterministic_compute = false;
+  sess_options.enable_mem_reuse = true;
+  sess_options.config_options.configurations[kOrtSessionOptionsConfigDisablePrepacking] = "0";
+  sess_options.config_options.configurations[kOrtSessionOptionsSharePrepackedWeightsForAllInitializers] = "1";
+
+  PrepackedWeightsContainer prepacked_weights_container;
+
+  Model model("graph_main", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(),
+              domain_to_version, std::vector<ONNX_NAMESPACE::FunctionProto>(),
+              DefaultLoggingManager().DefaultLogger());
+
+  CreateSimpleGraph(model.MainGraph(), "BrokenPrePackingTest");
+  PlaceAllNodesToCPUEP(model.MainGraph());
+  SessionState session_state(model.MainGraph(),
+                             execution_providers,
+                             tp.get(),
+                             nullptr, /*inter_op_thread_pool*/
+                             dtm,
+                             edlm,
+                             DefaultLoggingManager().DefaultLogger(),
+                             profiler,
+                             sess_options,
+                             &prepacked_weights_container);
+
+  ASSERT_STATUS_OK(session_state.FinalizeSessionState(std::basic_string<PATH_CHAR_TYPE>(),
+                                                      kernel_registry_manager));
+
+  // The pre-pack ran, but nothing was cached or reused
+  ASSERT_EQ(session_state.GetNumberOfPrepacksCounter(), static_cast<size_t>(1));
+  ASSERT_EQ(session_state.GetUsedSharedPrePackedWeightCounter(), static_cast<size_t>(0));
+}
+
 // Pre-packing enabled + shared initializers +
 // pre-packed weights container + subgraphs =
 // caching enabled in pre-packed weights used in subgraphs