diff --git a/src/plugins/intel_npu/src/compiler_adapter/src/weightless_graph.cpp b/src/plugins/intel_npu/src/compiler_adapter/src/weightless_graph.cpp index 2fdef4e298f1e6..8d931006f6408f 100644 --- a/src/plugins/intel_npu/src/compiler_adapter/src/weightless_graph.cpp +++ b/src/plugins/intel_npu/src/compiler_adapter/src/weightless_graph.cpp @@ -459,6 +459,10 @@ WeightlessGraph::InputData WeightlessGraph::allocate_inputs( // Due to the large number of init inputs, allocating a single buffer for all of them is more efficient. "View // tensors" are used for separating them. + // TODO in case of cross-context weight sharing the current implementation has been allocating a single buffer for several weights coul be shared + // In this case we don't need to allocate the signle buffer for such weight, we can reuse it instead. + // But we following the current logic, we still have to allocate a single buffer for non-shared weights. + // Implication: we need to return multiple ZeroTensors instead of a single one. const std::shared_ptr initInputsAllocatedTensor = std::make_shared(_zeroInitStruct, ov::element::Type_t::u8, ov::Shape({initInputsByteSize}), true); diff --git a/src/plugins/intel_npu/src/plugin/include/shared_weights_assigner.hpp b/src/plugins/intel_npu/src/plugin/include/shared_weights_assigner.hpp new file mode 100644 index 00000000000000..e97ebeb4afcbec --- /dev/null +++ b/src/plugins/intel_npu/src/plugin/include/shared_weights_assigner.hpp @@ -0,0 +1,82 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#pragma once + +#include +#include +#include +#include + +#include "openvino/runtime/aligned_buffer.hpp" + +namespace ov { +class Model; +namespace op { +namespace v0 { +class Constant; +} // namespace v0 +} // namespace op +} // namespace ov + +namespace ov { +namespace intel_npu { + +class SharedWeightsAssigner { +public: + using SharedConstant = std::shared_ptr; + using PartitionedConstants = std::vector>; + using SharedSourceAndConstants = std::pair, std::vector>; + using SharedSourcesWithConstants = std::vector; + + struct Options { + std::vector shared_device_contexts; + size_t single_weight_shared_source_size_max = 0; + bool preserve_weightless_cache_attr = true; + std::function source_id_generator; + }; + + struct Statistic { + std::vector partition_constant_counts; + size_t collected_constants_count = 0; + size_t total_shared_constant_bytes = 0; + size_t total_non_shared_constant_bytes_released = 0; + + std::string to_string() const; + }; + + struct CollectResult { + Statistic statistic; + PartitionedConstants partitioned_constants; + }; + + explicit SharedWeightsAssigner(Options options); + CollectResult collect_and_partition(const std::shared_ptr& model); + SharedSourcesWithConstants mutate_model_with_constant_sharing(PartitionedConstants&& partitioned_constants); + +private: + bool constant_can_be_shared(const ov::op::v0::Constant& constant) const; + + std::vector collect_weights_to_share(const std::shared_ptr& model) const; + + PartitionedConstants partition_constants_by_size(std::vector&& constants) const; + + SharedSourcesWithConstants make_constant_shareable( + PartitionedConstants&& partitioned_constants) const; + + std::shared_ptr make_shared_source(const std::vector& partition) const; + + size_t get_constant_aligned_size(const ov::op::v0::Constant& constant) const; + + static size_t align_bytes(size_t bytes, size_t alignment); + + std::vector m_shared_device_contexts; + std::function m_source_id_generator; + bool m_preserve_weightless_cache_attr = true; + size_t m_min_relocate_bytes = 0; + size_t m_single_weight_shared_source_size_max = 0; +}; + +} // namespace intel_npu +} // namespace ov diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp index 47efb2987c72f6..b21a838daa9e9b 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.cpp @@ -1,6 +1,8 @@ // Copyright (C) 2018-2026 Intel Corporation // SPDX-License-Identifier: Apache-2.0 // + +#include #include "llm_compiled_model.hpp" #include "embedding/embedding_infer_request.hpp" @@ -42,12 +44,15 @@ #include "openvino/pass/pattern/op/wrap_type.hpp" #include "openvino/pass/stateful_to_stateless.hpp" #include "openvino/pass/validate.hpp" +#include "openvino/runtime/device_id_parser.hpp" #include "openvino/runtime/iasync_infer_request.hpp" +#include "openvino/runtime/internal_properties.hpp" #include "openvino/runtime/properties.hpp" #include "partitioning/patterns/fold_const.hpp" #include "partitioning/patterns/moe.hpp" #include "partitioning/patterns/pre_compute.hpp" #include "partitioning/patterns/sdpa.hpp" +#include "shared_weights_assigner.hpp" #include "serialization.hpp" #include "transformations/convert_precision.hpp" #include "util.hpp" @@ -684,6 +689,47 @@ void ov::npuw::LLMCompiledModel::compile_generate_model_variants( } } +void ov::npuw::LLMCompiledModel::assign_shared_weight_to_model_if_possible(const std::shared_ptr model, const std::shared_ptr& plugin, +const ov::AnyMap& properties) { + constexpr size_t single_weight_shared_source_size_max = static_cast(2ULL * 1024 * 1024 * 1024); + + NPUW_ASSERT(model && "Model for assigning shared weights must not be null"); + NPUW_ASSERT(plugin && "Plugin for assigning shared weights must not be null"); + auto shared_weight_property_it = properties.find("SHARED_WEIGHTS"); + if (shared_weight_property_it == properties.end()) { + return; + } + + auto shared_device_contexts = + ov::DeviceIDParser::get_hetero_devices(shared_weight_property_it->second.as()); + ::ov::intel_npu::SharedWeightsAssigner::Options shared_weights_assigner_options; + shared_weights_assigner_options.shared_device_contexts = std::move(shared_device_contexts); + shared_weights_assigner_options.single_weight_shared_source_size_max = single_weight_shared_source_size_max; + shared_weights_assigner_options.preserve_weightless_cache_attr = (std::getenv("NO_WEIGHTLESS_ATTR") == nullptr); + ::ov::intel_npu::SharedWeightsAssigner shared_weights_assigner(std::move(shared_weights_assigner_options)); + auto collect_result = shared_weights_assigner.collect_and_partition(model); + + LOG_INFO("[NPUW] SHARED_WEIGHTS: " << collect_result.statistic.to_string()); + for (size_t i = 0; i < collect_result.statistic.partition_constant_counts.size(); ++i) { + LOG_INFO("[NPUW] SHARED_WEIGHTS: partition " << i + 1 << "/" + << collect_result.statistic.partition_constant_counts.size() + << ", constants count: " << collect_result.statistic.partition_constant_counts[i]); + } + + auto shared_sources_with_constants = + shared_weights_assigner.mutate_model_with_constant_sharing(std::move(collect_result.partitioned_constants)); + + m_shared_weight_sources.clear(); + for (const auto& [shared_source, constants] : shared_sources_with_constants) { + LOG_INFO("[NPUW] SHARED_WEIGHTS: allocated shared source buffer: source_id: " << shared_source->get_descriptor()->get_id() + << ", ptr: " << static_cast(shared_source->get_ptr()) + << ", size: " << shared_source->size() + << ", holds shared weight count: " << constants.size()); + // Keep source buffers alive for the lifetime of this compiled model. + m_shared_weight_sources.push_back(shared_source); + } +} + ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& model, const std::shared_ptr& plugin, const ov::AnyMap& properties, @@ -820,6 +866,9 @@ ov::npuw::LLMCompiledModel::LLMCompiledModel(const std::shared_ptr& m LOG_INFO("Continuous prefill is enabled"); } + LOG_DEBUG("Assigning shared weights to model if possible."); + assign_shared_weight_to_model_if_possible(model, plugin, properties); + const uint32_t batch_dim = m_cfg.get<::intel_npu::NPUW_LLM_BATCH_DIM>(); const uint32_t seq_len_dim = m_cfg.get<::intel_npu::NPUW_LLM_SEQ_LEN_DIM>(); KVAxesPosition axes{batch_dim, seq_len_dim}; diff --git a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp index fffa92f8fbc382..60bda1ec367ac1 100644 --- a/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp +++ b/src/plugins/intel_npu/src/plugin/npuw/llm_compiled_model.hpp @@ -8,6 +8,7 @@ #include "compiled_model.hpp" #include "npuw_transformations/kv_axes_position.hpp" +#include "openvino/core/weight_sharing_util.hpp" namespace ov { namespace test { @@ -168,7 +169,14 @@ class LLMCompiledModel : public ov::npuw::ICompiledModel { const std::shared_ptr& plugin, const ov::AnyMap& generate_config); + // Relocate page-aligned constants into shared cross-device memory if SHARED_WEIGHTS is set. + void assign_shared_weight_to_model_if_possible(const std::shared_ptr model, + const std::shared_ptr& plugin, + const ov::AnyMap& properties); + bool m_is_eagle = false; + + std::vector> m_shared_weight_sources; }; } // namespace npuw diff --git a/src/plugins/intel_npu/src/plugin/sources.cmake b/src/plugins/intel_npu/src/plugin/sources.cmake index b7a7c3e4fd6b02..4936cb26b31b0a 100644 --- a/src/plugins/intel_npu/src/plugin/sources.cmake +++ b/src/plugins/intel_npu/src/plugin/sources.cmake @@ -14,6 +14,7 @@ set(SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/include/plugin_property_manager.hpp ${CMAKE_CURRENT_SOURCE_DIR}/include/property_registration.hpp ${CMAKE_CURRENT_SOURCE_DIR}/include/remote_context.hpp + ${CMAKE_CURRENT_SOURCE_DIR}/include/shared_weights_assigner.hpp ${CMAKE_CURRENT_SOURCE_DIR}/include/transformations.hpp ${CMAKE_CURRENT_SOURCE_DIR}/src/async_infer_request.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/backends_registry.cpp @@ -26,6 +27,7 @@ set(SOURCES ${CMAKE_CURRENT_SOURCE_DIR}/src/plugin.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/plugin_property_manager.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/remote_context.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/shared_weights_assigner.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/transformations.cpp ) diff --git a/src/plugins/intel_npu/src/plugin/src/shared_weights_assigner.cpp b/src/plugins/intel_npu/src/plugin/src/shared_weights_assigner.cpp new file mode 100644 index 00000000000000..48df7e12ce222a --- /dev/null +++ b/src/plugins/intel_npu/src/plugin/src/shared_weights_assigner.cpp @@ -0,0 +1,212 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include "shared_weights_assigner.hpp" + +#include +#include +#include + +#include "openvino/core/except.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/util/node_util.hpp" +#include "openvino/openvino.hpp" +#include "openvino/core/rt_info/weightless_caching_attributes.hpp" +#include "openvino/core/weight_sharing_util.hpp" +#include "openvino/runtime/shared_buffer.hpp" +#include "openvino/util/mmap_object.hpp" + +namespace ov { +namespace intel_npu { + +std::string SharedWeightsAssigner::Statistic::to_string() const { + std::ostringstream oss; + oss << "collected_constants_count=" << collected_constants_count + << ", total_shared_constant_bytes=" << total_shared_constant_bytes + << ", total_non_shared_constant_bytes_released=" << total_non_shared_constant_bytes_released + << ", partitions=" << partition_constant_counts.size() << " ["; + for (size_t i = 0; i < partition_constant_counts.size(); ++i) { + if (i != 0) { + oss << ","; + } + oss << partition_constant_counts[i]; + } + oss << "]"; + return oss.str(); +} + +SharedWeightsAssigner::SharedWeightsAssigner(Options options) + : m_shared_device_contexts(std::move(options.shared_device_contexts)), + m_source_id_generator(std::move(options.source_id_generator)), + m_preserve_weightless_cache_attr(options.preserve_weightless_cache_attr), + m_single_weight_shared_source_size_max(options.single_weight_shared_source_size_max) { + if (!m_source_id_generator) { + m_source_id_generator = []() { + static std::atomic source_id_counter{1}; + return source_id_counter.fetch_add(1, std::memory_order_relaxed); + }; + } + m_min_relocate_bytes = static_cast(::ov::util::get_system_page_size()); +} + +SharedWeightsAssigner::CollectResult SharedWeightsAssigner::collect_and_partition(const std::shared_ptr& model) { + OPENVINO_ASSERT(model && "Model for assigning shared weights must not be null"); + OPENVINO_ASSERT(m_single_weight_shared_source_size_max > 0, + "single_weight_shared_source_size_max must be greater than zero"); + + CollectResult result; + + auto constants_to_share = collect_weights_to_share(model); + result.statistic.collected_constants_count = constants_to_share.size(); + + result.partitioned_constants = partition_constants_by_size(std::move(constants_to_share)); + result.statistic.partition_constant_counts.reserve(result.partitioned_constants.size()); + for (const auto& partition : result.partitioned_constants) { + result.statistic.partition_constant_counts.push_back(partition.size()); + for (const auto& constant : partition) { + result.statistic.total_shared_constant_bytes += get_constant_aligned_size(*constant); + result.statistic.total_non_shared_constant_bytes_released += constant->get_byte_size(); + } + } + + return result; +} + +SharedWeightsAssigner::SharedSourcesWithConstants SharedWeightsAssigner::mutate_model_with_constant_sharing( + PartitionedConstants&& partitioned_constants) { + SharedSourcesWithConstants shared_sources_with_constants; + shared_sources_with_constants.clear(); + shared_sources_with_constants = make_constant_shareable(std::move(partitioned_constants)); + return shared_sources_with_constants; +} + +bool SharedWeightsAssigner::constant_can_be_shared(const ov::op::v0::Constant& constant) const { + if (constant.get_byte_size() < m_min_relocate_bytes || + constant.get_byte_size() > m_single_weight_shared_source_size_max) { + return false; + } + + bool needs_conversion = false; + // TODO check types of remote context + // The code below is written in assumption that we have NPU and GPU as the remote contexts + (void)m_shared_device_contexts; + if (ov::shape_size(constant.get_shape()) == 1 && constant.get_output_element_type(0) == ov::element::f64) { + // If a constant has element type f64 but contains no elements (empty tensor), + // GPU have to convert it to f32 because the GPU plugin only supports the f32 data type internally. + needs_conversion = true; + } else if (constant.get_output_element_type(0) == ov::element::u16 || + constant.get_output_element_type(0) == ov::element::i16) { + needs_conversion = true; + } + return !needs_conversion; +} + +std::vector SharedWeightsAssigner::collect_weights_to_share( + const std::shared_ptr& model) const { + std::vector constants_to_share; + for (const auto& op : model->get_ops()) { + auto shared_weight_candidate = std::dynamic_pointer_cast(op); + if (!shared_weight_candidate) { + continue; + } + + if (!constant_can_be_shared(*shared_weight_candidate)) { + continue; + } + constants_to_share.push_back(shared_weight_candidate); + } + return constants_to_share; +} + +SharedWeightsAssigner::PartitionedConstants SharedWeightsAssigner::partition_constants_by_size( + std::vector&& constants) const { + PartitionedConstants partitioned_constants; + std::vector current_partition; + size_t current_partition_size = 0; + for (auto&& constant : constants) { + size_t constant_size = get_constant_aligned_size(*constant); + if (current_partition_size + constant_size > m_single_weight_shared_source_size_max) { + if (!current_partition.empty()) { + partitioned_constants.emplace_back(std::move(current_partition)); + current_partition.clear(); + current_partition_size = 0; + } + } + current_partition.push_back(constant); + current_partition_size += constant_size; + } + if (!current_partition.empty()) { + partitioned_constants.emplace_back(std::move(current_partition)); + } + return partitioned_constants; +} + +SharedWeightsAssigner::SharedSourcesWithConstants SharedWeightsAssigner::make_constant_shareable( + PartitionedConstants&& partitioned_constants) const { + SharedSourcesWithConstants shared_sources; + for (const auto& partition : partitioned_constants) { + auto shared_source = make_shared_source(partition); + size_t constant_id = 0; // constants ID is a weight offset in the shared source buffer + std::vector shared_constants; + for (const auto& constant : partition) { + auto const_descriptor = + ::ov::create_base_descriptor(shared_source->get_descriptor()->get_id(), constant_id, shared_source); + auto constant_shared_buffer = std::make_shared<::ov::SharedBuffer>>( + shared_source->get_ptr() + constant_id, + constant->get_byte_size(), + shared_source, + const_descriptor); + constant_id += get_constant_aligned_size(*constant); + auto shared_constant = + std::make_shared(constant->get_element_type(), constant->get_shape(), constant_shared_buffer); + shared_constant->set_friendly_name(constant->get_friendly_name()); + ov::copy_runtime_info(constant, shared_constant); + std::memcpy(constant_shared_buffer->get_ptr(), constant->get_data_ptr(), constant->get_byte_size()); + + // Preserve the weightless-cache attribute: copy_runtime_info drops it (is_copyable()==false). + if (m_preserve_weightless_cache_attr) { + ov::copy_weightless_cache_attr(constant, shared_constant); + } + + ov::replace_node(constant, shared_constant); + ov::weight_sharing::Extension::hint_evict(*constant); + shared_constants.push_back(shared_constant); + } + shared_sources.emplace_back(std::move(shared_source), std::move(shared_constants)); + } + return shared_sources; +} + +std::shared_ptr SharedWeightsAssigner::make_shared_source( + const std::vector& partition) const { + size_t total_partition_size = 0; + for (const auto& constant : partition) { + total_partition_size += get_constant_aligned_size(*constant); + } + + // TODO not a unique ID in general: as it can clash with mmap weight source if generation. + // The uniqueness must meet the conditions: + // 1) persistent across different processes. + // 2) unique across different weight banks in the same process. + // 3) distinguishable from mmap sources for weightless cache. + const size_t source_id = m_source_id_generator(); + auto raw = std::make_shared(total_partition_size, m_min_relocate_bytes); + + return std::make_shared<::ov::SharedBuffer>>( + raw->get_ptr(), + raw->size(), + raw, + ::ov::create_base_descriptor(source_id, 0, raw)); +} + +size_t SharedWeightsAssigner::get_constant_aligned_size(const ov::op::v0::Constant& constant) const { + return align_bytes(constant.get_byte_size(), m_min_relocate_bytes); +} + +size_t SharedWeightsAssigner::align_bytes(size_t bytes, size_t alignment) { + return ((bytes + alignment - 1) / alignment) * alignment; +} + +} // namespace intel_npu +} // namespace ov diff --git a/src/plugins/intel_npu/tests/unit/CMakeLists.txt b/src/plugins/intel_npu/tests/unit/CMakeLists.txt index 1d3d6a8f7e35a1..b9fccdd0307edc 100644 --- a/src/plugins/intel_npu/tests/unit/CMakeLists.txt +++ b/src/plugins/intel_npu/tests/unit/CMakeLists.txt @@ -45,6 +45,7 @@ ov_add_test_target( ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/src/executor.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/src/metadata.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/src/blob_format_importers.cpp + ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/src/shared_weights_assigner.cpp # npuw core ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/compiled_model.cpp ${OpenVINO_SOURCE_DIR}/src/plugins/intel_npu/src/plugin/npuw/flux2_compiled_model.cpp @@ -184,5 +185,6 @@ target_sources(${TARGET_NAME} PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/npuw/pyramid_attention_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/sdpa_pattern_matcher_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/host_flash_attention_test.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/npuw/shared_weights_assigner_test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/npuw/moe_k_tag_test.cpp ) diff --git a/src/plugins/intel_npu/tests/unit/npuw/shared_weights_assigner_test.cpp b/src/plugins/intel_npu/tests/unit/npuw/shared_weights_assigner_test.cpp new file mode 100644 index 00000000000000..30238b3c8d8c22 --- /dev/null +++ b/src/plugins/intel_npu/tests/unit/npuw/shared_weights_assigner_test.cpp @@ -0,0 +1,137 @@ +// Copyright (C) 2018-2026 Intel Corporation +// SPDX-License-Identifier: Apache-2.0 +// + +#include + +#include +#include +#include +#include +#include + +#include "openvino/op/add.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/multiply.hpp" +#include "openvino/op/parameter.hpp" +#include "openvino/op/result.hpp" +#include "openvino/util/mmap_object.hpp" +#include "shared_weights_assigner.hpp" + +namespace { + +using SharedConstant = ov::intel_npu::SharedWeightsAssigner::SharedConstant; +using PartitionedConstants = ov::intel_npu::SharedWeightsAssigner::PartitionedConstants; +using SharedSourcesWithConstants = ov::intel_npu::SharedWeightsAssigner::SharedSourcesWithConstants; + +std::shared_ptr make_test_model(size_t element_count, + std::unordered_map>& by_name) { + auto input = std::make_shared(ov::element::u8, ov::Shape{element_count}); + + std::vector data_a(element_count, 1); + std::vector data_b(element_count, 2); + std::vector data_c(element_count, 3); + + auto c1 = ov::op::v0::Constant::create(ov::element::u8, ov::Shape{element_count}, data_a); + auto c2 = ov::op::v0::Constant::create(ov::element::u8, ov::Shape{element_count}, data_b); + auto c3 = ov::op::v0::Constant::create(ov::element::u8, ov::Shape{element_count}, data_c); + + c1->set_friendly_name("c1"); + c2->set_friendly_name("c2"); + c3->set_friendly_name("c3"); + + by_name["c1"] = c1; + by_name["c2"] = c2; + by_name["c3"] = c3; + + auto add1 = std::make_shared(input, c1); + auto add2 = std::make_shared(add1, c2); + auto mul = std::make_shared(add2, c3); + + auto result = std::make_shared(mul); + return std::make_shared(ov::ResultVector{result}, ov::ParameterVector{input}); +} + +std::unordered_map> collect_named_constants( + const std::shared_ptr& model) { + std::unordered_map> result; + for (const auto& op : model->get_ops()) { + auto constant = std::dynamic_pointer_cast(op); + if (!constant) { + continue; + } + result[constant->get_friendly_name()] = constant; + } + return result; +} + +ov::intel_npu::SharedWeightsAssigner::Options make_options(size_t max_source_size) { + ov::intel_npu::SharedWeightsAssigner::Options options; + options.single_weight_shared_source_size_max = max_source_size; + options.preserve_weightless_cache_attr = true; + options.source_id_generator = [id = static_cast(100)]() mutable { + return id++; + }; + return options; +} + +TEST(SharedWeightsAssignerTest, CollectAndPartitionProvidesStatsWithoutMutation) { + const size_t page_size = static_cast(ov::util::get_system_page_size()); + + std::unordered_map> original_by_name; + auto model = make_test_model(page_size, original_by_name); + + ov::intel_npu::SharedWeightsAssigner assigner(make_options(2 * page_size)); + auto collect_result = assigner.collect_and_partition(model); + + EXPECT_EQ(collect_result.statistic.collected_constants_count, 3u); + ASSERT_EQ(collect_result.statistic.partition_constant_counts.size(), 2u); + EXPECT_EQ(collect_result.statistic.partition_constant_counts[0], 2u); + EXPECT_EQ(collect_result.statistic.partition_constant_counts[1], 1u); + EXPECT_EQ(collect_result.statistic.total_shared_constant_bytes, 3 * page_size); + EXPECT_EQ(collect_result.statistic.total_non_shared_constant_bytes_released, 3 * page_size); + + auto current_by_name = collect_named_constants(model); + ASSERT_EQ(current_by_name.size(), 3u); + EXPECT_EQ(current_by_name.at("c1").get(), original_by_name.at("c1").get()); + EXPECT_EQ(current_by_name.at("c2").get(), original_by_name.at("c2").get()); + EXPECT_EQ(current_by_name.at("c3").get(), original_by_name.at("c3").get()); +} + +TEST(SharedWeightsAssignerTest, MutateModelWithConstantSharingReturnsExpectedBuffers) { + const size_t page_size = static_cast(ov::util::get_system_page_size()); + + std::unordered_map> original_by_name; + auto model = make_test_model(page_size, original_by_name); + + ov::intel_npu::SharedWeightsAssigner assigner(make_options(2 * page_size)); + auto collect_result = assigner.collect_and_partition(model); + auto shared_sources_with_constants = + assigner.mutate_model_with_constant_sharing(std::move(collect_result.partitioned_constants)); + + ASSERT_EQ(shared_sources_with_constants.size(), 2u); + EXPECT_EQ(shared_sources_with_constants[0].first->size(), 2 * page_size); + EXPECT_EQ(shared_sources_with_constants[1].first->size(), page_size); + EXPECT_EQ(shared_sources_with_constants[0].second.size(), 2u); + EXPECT_EQ(shared_sources_with_constants[1].second.size(), 1u); + + // All returned constants should point inside their partition buffer ranges. + for (const auto& [source, shared_constants] : shared_sources_with_constants) { + const auto source_begin = reinterpret_cast(source->get_ptr()); + const auto source_end = source_begin + source->size(); + for (const auto& shared_constant : shared_constants) { + const auto data_begin = reinterpret_cast(shared_constant->get_data_ptr()); + const auto data_end = data_begin + shared_constant->get_byte_size(); + EXPECT_GE(data_begin, source_begin); + EXPECT_LE(data_end, source_end); + } + } + + auto mutated_by_name = collect_named_constants(model); + ASSERT_EQ(mutated_by_name.size(), 3u); + EXPECT_NE(mutated_by_name.at("c1").get(), original_by_name.at("c1").get()); + EXPECT_NE(mutated_by_name.at("c2").get(), original_by_name.at("c2").get()); + EXPECT_NE(mutated_by_name.at("c3").get(), original_by_name.at("c3").get()); +} + +} // namespace