diff --git a/onnxruntime/core/framework/tensorprotoutils.cc b/onnxruntime/core/framework/tensorprotoutils.cc index 3916a320980ba..2660b64398f21 100644 --- a/onnxruntime/core/framework/tensorprotoutils.cc +++ b/onnxruntime/core/framework/tensorprotoutils.cc @@ -2063,17 +2063,6 @@ Status TensorProtoToTensor(const Env& env, const std::filesystem::path& model_pa return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "tensor can't contain negative dims"); } - if (HasExternalData(tensor_proto)) { - OrtValue ort_value; - ORT_RETURN_IF_ERROR(GetExtDataFromTensorProto(env, model_path, tensor_proto, ort_value)); - const auto& ext_tensor = ort_value.Get(); - MakeCpuTensorCopy(ext_tensor, tensor); - // MakeCpuTensorCopy memcpy's external bytes verbatim. Bool external initializers may carry - // bytes outside the canonical {0, 1} set, so normalize them here as well (see NormalizeBoolBytes). - NormalizeBoolTensorIfNeeded(tensor); - return Status::OK(); - } - if (tensor_shape != tensor.Shape()) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "TensorProtoToTensor() tensor shape mismatch!"); } @@ -2085,6 +2074,22 @@ Status TensorProtoToTensor(const Env& env, const std::filesystem::path& model_pa " can not be written into Tensor type ", DataTypeImpl::ToString(tensor.DataType())); } + size_t source_bytes; + ORT_RETURN_IF_ERROR(Tensor::CalculateTensorStorageSize(source_type, tensor_shape, 0, source_bytes)); + ORT_RETURN_IF(source_bytes > tensor.SizeInBytes(), + "TensorProto data exceeds destination tensor capacity"); + + if (HasExternalData(tensor_proto)) { + OrtValue ort_value; + ORT_RETURN_IF_ERROR(GetExtDataFromTensorProto(env, model_path, tensor_proto, ort_value)); + const auto& ext_tensor = ort_value.Get(); + MakeCpuTensorCopy(ext_tensor, tensor); + // MakeCpuTensorCopy memcpy's external bytes verbatim. Bool external initializers may carry + // bytes outside the canonical {0, 1} set, so normalize them here as well (see NormalizeBoolBytes). + NormalizeBoolTensorIfNeeded(tensor); + return Status::OK(); + } + // Below we handle the case where TensorProto contains data in itself // find raw data in proto buf @@ -2943,6 +2948,9 @@ common::Status DenseTensorToSparseTensorProto(const ONNX_NAMESPACE::TensorProto& } const auto data_type = dense_proto.data_type(); + if (!HasExternalData(dense_proto)) { + ORT_RETURN_IF_ERROR(ValidateEmbeddedTensorProtoDataSizeAndShape(dense_proto)); + } SparseTensorProto sparse_proto; auto& values = *sparse_proto.mutable_values(); values.set_name(dense_proto.name()); @@ -2956,10 +2964,14 @@ common::Status DenseTensorToSparseTensorProto(const ONNX_NAMESPACE::TensorProto& } auto ml_data = DataTypeImpl::TensorTypeFromONNXEnum(data_type)->GetElementType(); + ORT_RETURN_IF(ml_data->AsPrimitiveDataType()->HasSubElems(), + "Dense-to-sparse conversion does not support sub-byte tensor data types"); size_t element_size = ml_data->Size(); std::vector dense_raw_data; ORT_RETURN_IF_ERROR(UnpackInitializerData(dense_proto, model_path, dense_raw_data)); + ORT_RETURN_IF_NOT(dense_raw_data.size() == n_dense_elements * element_size, + "Dense tensor data size does not match its shape"); size_t nnz = 0; void* dense_data = dense_raw_data.data(); diff --git a/onnxruntime/core/graph/graph.cc b/onnxruntime/core/graph/graph.cc index ac48926d34392..771023a0a05c3 100644 --- a/onnxruntime/core/graph/graph.cc +++ b/onnxruntime/core/graph/graph.cc @@ -7022,6 +7022,7 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph ORT_RETURN_IF_ERROR(fbs::utils::LoadInitializerOrtFormat(*fbs_tensor, *initializer, load_options)); auto p = name_to_initial_tensor_.emplace(initializer->name(), initializer); if (!p.second) { + ORT_IGNORE_RETURN_VALUE(ortvalue_initializers_.erase(initializer->name())); LOGS(logger_, WARNING) << "Duplicate initializer (dense or ConstantNode): '" << initializer->name() << "' the model will use the latest encountered initializer" << ". Please, fix your model."; @@ -7061,6 +7062,7 @@ common::Status Graph::LoadFromOrtFormat(const onnxruntime::fbs::Graph& fbs_graph ORT_RETURN_IF_ERROR(utils::SparseTensorProtoToDenseTensorProto(sparse_initializer, model_path, initializer)); auto p = name_to_initial_tensor_.emplace(initializer.name(), &initializer); if (!p.second) { + ORT_IGNORE_RETURN_VALUE(ortvalue_initializers_.erase(initializer.name())); LOGS(logger_, WARNING) << "Duplicate initializer (dense, sparse or ConstantNode): '" << initializer.name() << "' the model will use the latest encountered initializer" << ". Please, fix your model."; diff --git a/onnxruntime/test/framework/ort_model_only_test.cc b/onnxruntime/test/framework/ort_model_only_test.cc index 2192ea56db714..105d7b49080f5 100644 --- a/onnxruntime/test/framework/ort_model_only_test.cc +++ b/onnxruntime/test/framework/ort_model_only_test.cc @@ -476,6 +476,151 @@ TEST(OrtModelTest, NormalizesValidVariadicInputArgCounts) { EXPECT_THAT(node->InputArgCount(), testing::ElementsAre(2)); } + +namespace { + +enum class DuplicateInitializerReplacement { + kNone, + kDense, + kSparse, +}; + +// Builds a graph `y = Identity(w)` whose first `w` initializer holds 40 floats (160 bytes). Loading it with +// can_use_flatbuffer_for_initializers enabled caches an OrtValue over the flatbuffer bytes. The optional second +// `w` initializer replaces it: a 4-float (16 byte) dense initializer, or a sparse initializer with dense shape [40] +// holding 5.0f at index 3 and -6.0f at index 17. Neither replacement is cached. +std::vector BuildOrtModelWithDuplicateInitializer(DuplicateInitializerReplacement replacement) { + return BuildOrtModelBuffer([replacement](flatbuffers::FlatBufferBuilder& builder) { + const int64_t arg_dim = replacement == DuplicateInitializerReplacement::kDense ? 4 : 40; + std::vector> node_args{ + fbs::CreateValueInfoDirect(builder, "w", "", CreateFloatTensorTypeInfo(builder, arg_dim)), + fbs::CreateValueInfoDirect(builder, "y", "", CreateFloatTensorTypeInfo(builder, arg_dim))}; + + std::vector large_dims{40}; + std::vector large_values(40, 1.0f); + std::vector large_raw(reinterpret_cast(large_values.data()), + reinterpret_cast(large_values.data() + large_values.size())); + std::vector> initializers{ + fbs::CreateTensorDirect(builder, "w", "", &large_dims, fbs::TensorDataType::FLOAT, &large_raw)}; + + std::vector small_dims{4}; + const std::vector small_values{2.0f, 3.0f, 4.0f, 5.0f}; + std::vector small_raw(reinterpret_cast(small_values.data()), + reinterpret_cast(small_values.data() + small_values.size())); + if (replacement == DuplicateInitializerReplacement::kDense) { + initializers.push_back( + fbs::CreateTensorDirect(builder, "w", "", &small_dims, fbs::TensorDataType::FLOAT, &small_raw)); + } + + std::vector sparse_nnz_dims{2}; + const std::vector sparse_values{5.0f, -6.0f}; + std::vector sparse_values_raw( + reinterpret_cast(sparse_values.data()), + reinterpret_cast(sparse_values.data() + sparse_values.size())); + const std::vector sparse_indices{3, 17}; + std::vector sparse_indices_raw( + reinterpret_cast(sparse_indices.data()), + reinterpret_cast(sparse_indices.data() + sparse_indices.size())); + std::vector sparse_dense_dims{40}; + std::vector> sparse_initializers; + if (replacement == DuplicateInitializerReplacement::kSparse) { + const auto values = fbs::CreateTensorDirect(builder, "w", "", &sparse_nnz_dims, fbs::TensorDataType::FLOAT, + &sparse_values_raw); + const auto indices = fbs::CreateTensorDirect(builder, "w_indices", "", &sparse_nnz_dims, + fbs::TensorDataType::INT64, &sparse_indices_raw); + sparse_initializers.push_back(fbs::CreateSparseTensorDirect(builder, values, indices, &sparse_dense_dims)); + } + + std::vector> inputs{builder.CreateSharedString("w")}; + std::vector> outputs{builder.CreateSharedString("y")}; + std::vector> implicit_inputs; + std::vector> graph_inputs; + std::vector input_arg_counts{1}; + std::vector> nodes{ + fbs::CreateNodeDirect(builder, "identity", "", "", 16, 0, "Identity", + fbs::NodeType::Primitive, nullptr, + &inputs, &outputs, nullptr, + &input_arg_counts, &implicit_inputs)}; + + return fbs::CreateGraphDirect(builder, &initializers, &node_args, &nodes, 1, nullptr, &graph_inputs, &outputs, + sparse_initializers.empty() ? nullptr : &sparse_initializers); + }); +} + +SessionOptions SessionOptionsUsingOrtBytesForInitializers() { + SessionOptions so; + ORT_THROW_IF_ERROR(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigLoadModelFormat, "ORT")); + ORT_THROW_IF_ERROR(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigUseORTModelBytesDirectly, "1")); + ORT_THROW_IF_ERROR(so.config_options.AddConfigEntry(kOrtSessionOptionsConfigUseORTModelBytesForInitializers, "1")); + return so; +} + +std::vector UnpackFloatInitializer(const TensorProto& initializer) { + const TensorShape shape = utils::GetTensorShapeFromTensorProto(initializer); + std::vector values(static_cast(shape.Size())); + const void* raw_data = utils::HasRawData(initializer) ? initializer.raw_data().data() : nullptr; + const size_t raw_data_len = utils::HasRawData(initializer) ? initializer.raw_data().size() : 0; + ORT_THROW_IF_ERROR(utils::UnpackTensor(initializer, raw_data, raw_data_len, values.data(), values.size())); + return values; +} + +} // namespace + +// Control: without a duplicate, an initializer larger than 127 bytes is cached as an OrtValue over the ORT +// format bytes, so the replacement tests below start from a populated cache. +TEST(OrtModelTest, OrtFormatLargeInitializerIsCachedWhenUsingBytesForInitializers) { + const auto buffer = BuildOrtModelWithDuplicateInitializer(DuplicateInitializerReplacement::kNone); + InferenceSessionWrapper session_object{SessionOptionsUsingOrtBytesForInitializers(), GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(buffer.data(), static_cast(buffer.size()))); + + const auto& graph = session_object.GetGraph(); + const TensorProto* initializer = nullptr; + ASSERT_TRUE(graph.GetInitializedTensor("w", initializer)); + ASSERT_NE(initializer, nullptr); + EXPECT_TRUE(utils::HasExternalDataInMemory(*initializer)); + OrtValue cached; + ASSERT_TRUE(graph.GetOrtValueInitializer("w", cached)); + EXPECT_EQ(cached.Get().Shape(), TensorShape({40})); +} + +TEST(OrtModelTest, OrtFormatDuplicateDenseInitializerReplacesCachedValue) { + const auto buffer = BuildOrtModelWithDuplicateInitializer(DuplicateInitializerReplacement::kDense); + InferenceSessionWrapper session_object{SessionOptionsUsingOrtBytesForInitializers(), GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(buffer.data(), static_cast(buffer.size()))); + + const auto& graph = session_object.GetGraph(); + const TensorProto* initializer = nullptr; + ASSERT_TRUE(graph.GetInitializedTensor("w", initializer)); + ASSERT_NE(initializer, nullptr); + EXPECT_FALSE(utils::HasExternalData(*initializer)); + EXPECT_THAT(UnpackFloatInitializer(*initializer), testing::ElementsAre(2.0f, 3.0f, 4.0f, 5.0f)); + + OrtValue cached; + EXPECT_FALSE(graph.GetOrtValueInitializer("w", cached)) + << "stale cached OrtValue with shape " << cached.Get().Shape(); +} + +#if !defined(DISABLE_SPARSE_TENSORS) +TEST(OrtModelTest, OrtFormatDuplicateSparseInitializerReplacesCachedValue) { + const auto buffer = BuildOrtModelWithDuplicateInitializer(DuplicateInitializerReplacement::kSparse); + InferenceSessionWrapper session_object{SessionOptionsUsingOrtBytesForInitializers(), GetEnvironment()}; + ASSERT_STATUS_OK(session_object.Load(buffer.data(), static_cast(buffer.size()))); + + const auto& graph = session_object.GetGraph(); + const TensorProto* initializer = nullptr; + ASSERT_TRUE(graph.GetInitializedTensor("w", initializer)); + ASSERT_NE(initializer, nullptr); + EXPECT_FALSE(utils::HasExternalData(*initializer)); + std::vector expected(40, 0.0f); + expected[3] = 5.0f; + expected[17] = -6.0f; + EXPECT_THAT(UnpackFloatInitializer(*initializer), testing::ElementsAreArray(expected)); + + OrtValue cached; + EXPECT_FALSE(graph.GetOrtValueInitializer("w", cached)) + << "stale cached OrtValue with shape " << cached.Get().Shape(); +} +#endif // !defined(DISABLE_SPARSE_TENSORS) #endif // !defined(ORT_MINIMAL_BUILD) #if !defined(ORT_MINIMAL_BUILD) diff --git a/onnxruntime/test/framework/tensorutils_test.cc b/onnxruntime/test/framework/tensorutils_test.cc index 779c6cfff5366..c7e920a862b2e 100644 --- a/onnxruntime/test/framework/tensorutils_test.cc +++ b/onnxruntime/test/framework/tensorutils_test.cc @@ -2212,5 +2212,149 @@ TEST(TensorProtoDataSizeShapeValidationTest, UnpackInitializerDataRejectsInlineR } #endif // !defined(__wasm__) +#if !defined(ORT_MINIMAL_BUILD) && !defined(__wasm__) +// Positive control: valid matching external tensor copy succeeds. +TEST(TensorProtoUtilsTest, ExternalTensorCopyPreservesMatchingDestination) { + std::basic_string filename(ORT_TSTR("matching_tensor_XXXXXX")); + const std::vector expected{1.0f, 0.0f, -2.0f, 3.0f}; + TensorProto proto; + CreateTensorWithExternalData(TensorProto_DataType_FLOAT, expected, filename, proto); + std::unique_ptr file_deleter( + const_cast(filename.c_str()), DeleteFileFromDisk); + Tensor destination(DataTypeImpl::GetType(), TensorShape({4}), CPUAllocator::DefaultInstance()); + ASSERT_STATUS_OK(utils::TensorProtoToTensor(Env::Default(), {}, proto, destination)); + EXPECT_THAT(gsl::make_span(destination.Data(), expected.size()), + ::testing::ElementsAreArray(expected)); +} + +// UINT8 and UINT4 both report an element size of 1 and the shapes match, so only the byte-capacity check +// separates them: four UINT8 elements need four bytes, while four packed UINT4 elements occupy two. +TEST(TensorProtoUtilsTest, ExternalTensorCopyRejectsInsufficientDestinationCapacity) { + std::basic_string filename(ORT_TSTR("capacity_tensor_XXXXXX")); + const std::vector external_values{1, 2, 3, 4}; + TensorProto proto; + CreateTensorWithExternalData(TensorProto_DataType_UINT8, external_values, filename, proto); + std::unique_ptr file_deleter( + const_cast(filename.c_str()), DeleteFileFromDisk); + + Tensor destination(DataTypeImpl::GetType(), TensorShape({4}), CPUAllocator::DefaultInstance()); + ASSERT_EQ(DataTypeImpl::GetType()->Size(), destination.DataType()->Size()); + ASSERT_EQ(destination.SizeInBytes(), size_t{2}); + auto* destination_bytes = static_cast(destination.MutableDataRaw()); + destination_bytes[0] = 0xA5; + destination_bytes[1] = 0x5A; + + Status status = utils::TensorProtoToTensor(Env::Default(), {}, proto, destination); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("exceeds destination tensor capacity")); + EXPECT_EQ(destination_bytes[0], 0xA5); + EXPECT_EQ(destination_bytes[1], 0x5A); +} + +// The element-size compatibility check must also run for external data: FLOAT (4 bytes per element) cannot be +// written into a UINT8 destination of the same shape. +TEST(TensorProtoUtilsTest, ExternalTensorCopyRejectsWiderSourceElementType) { + std::basic_string filename(ORT_TSTR("wider_tensor_XXXXXX")); + const std::vector external_values{1.0f, 0.0f, -2.0f, 3.0f}; + TensorProto proto; + CreateTensorWithExternalData(TensorProto_DataType_FLOAT, external_values, filename, proto); + std::unique_ptr file_deleter( + const_cast(filename.c_str()), DeleteFileFromDisk); + + Tensor destination(DataTypeImpl::GetType(), TensorShape({4}), CPUAllocator::DefaultInstance()); + ASSERT_EQ(destination.SizeInBytes(), size_t{4}); + auto destination_bytes = gsl::make_span(static_cast(destination.MutableDataRaw()), size_t{4}); + std::fill(destination_bytes.begin(), destination_bytes.end(), uint8_t{0xA5}); + + Status status = utils::TensorProtoToTensor(Env::Default(), {}, proto, destination); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), ::testing::HasSubstr("can not be written into Tensor type")); + EXPECT_THAT(destination_bytes, ::testing::Each(uint8_t{0xA5})); +} + +#if !defined(DISABLE_SPARSE_TENSORS) +TEST(TensorProtoUtilsTest, DenseSparseConversionPreservesTypedFloatValues) { + TensorProto proto; + proto.set_name("weights"); + proto.set_data_type(TensorProto_DataType_FLOAT); + proto.add_dims(2); + proto.add_dims(2); + for (float value : {0.0f, 1.0f, -2.0f, 0.0f}) { + proto.add_float_data(value); + } + SparseTensorProto sparse; + ASSERT_STATUS_OK(utils::DenseTensorToSparseTensorProto(proto, {}, sparse)); + EXPECT_EQ(sparse.values().name(), "weights"); + ASSERT_EQ(sparse.values().dims_size(), 1); + EXPECT_EQ(sparse.values().dims(0), 2); + TensorProto restored; + ASSERT_STATUS_OK(utils::SparseTensorProtoToDenseTensorProto(sparse, {}, restored)); + std::vector actual(4); + ASSERT_STATUS_OK(utils::UnpackTensor(restored, restored.raw_data().data(), restored.raw_data().size(), + actual.data(), actual.size())); + EXPECT_THAT(actual, ::testing::ElementsAre(0.0f, 1.0f, -2.0f, 0.0f)); +} + +TEST(TensorProtoUtilsTest, RejectsMismatchedExternalTensorShapeAndOversizedDenseToSparseInput) { + std::basic_string filename(ORT_TSTR("mismatched_tensor_XXXXXX")); + const std::vector external_values{1.0f, 0.0f, -2.0f, 3.0f}; + TensorProto external_proto; + CreateTensorWithExternalData(TensorProto_DataType_FLOAT, external_values, filename, external_proto); + std::unique_ptr file_deleter( + const_cast(filename.c_str()), DeleteFileFromDisk); + Tensor mismatched_destination(DataTypeImpl::GetType(), TensorShape({2, 2}), CPUAllocator::DefaultInstance()); + Status external_status = utils::TensorProtoToTensor(Env::Default(), {}, external_proto, mismatched_destination); + EXPECT_FALSE(external_status.IsOK()); + + TensorProto oversized_dense_proto; + oversized_dense_proto.set_name("oversized_weights"); + oversized_dense_proto.set_data_type(TensorProto_DataType_FLOAT); + oversized_dense_proto.add_dims(2); + for (float value : {0.0f, 1.0f, -2.0f, 0.0f}) { + oversized_dense_proto.add_float_data(value); + } + SparseTensorProto sparse; + Status sparse_status = utils::DenseTensorToSparseTensorProto(oversized_dense_proto, {}, sparse); + EXPECT_FALSE(sparse_status.IsOK()); +} + +// Well-formed packed sub-byte tensors are rejected as unsupported. The one-element cases have a one-byte buffer +// that also satisfies the later data-size check, so they depend on the explicit sub-byte rejection alone. +TEST(TensorProtoUtilsTest, DenseToSparseRejectsPackedSubByteTypes) { + struct PackedCase { + TensorProto_DataType data_type; + int64_t num_elements; + std::vector packed; + }; + const std::vector cases{ + {TensorProto_DataType_INT4, 1, {0x07}}, + {TensorProto_DataType_INT4, 4, {0x21, 0x43}}, + {TensorProto_DataType_UINT4, 1, {0x0F}}, + {TensorProto_DataType_UINT4, 4, {0x21, 0x43}}, + {TensorProto_DataType_INT2, 1, {0x01}}, + {TensorProto_DataType_INT2, 8, {0x1B, 0xE4}}, + {TensorProto_DataType_UINT2, 1, {0x03}}, + {TensorProto_DataType_UINT2, 8, {0x1B, 0xE4}}, + }; + for (const auto& test_case : cases) { + SCOPED_TRACE(::testing::Message() << "data_type=" << test_case.data_type + << " num_elements=" << test_case.num_elements); + TensorProto proto; + proto.set_name("packed_weights"); + proto.set_data_type(test_case.data_type); + proto.add_dims(test_case.num_elements); + proto.set_raw_data(test_case.packed.data(), test_case.packed.size()); + ASSERT_STATUS_OK(utils::ValidateEmbeddedTensorProtoDataSizeAndShape(proto)); + + SparseTensorProto sparse; + Status status = utils::DenseTensorToSparseTensorProto(proto, {}, sparse); + EXPECT_FALSE(status.IsOK()); + EXPECT_THAT(status.ErrorMessage(), + ::testing::HasSubstr("does not support sub-byte tensor data types")); + } +} +#endif // !defined(DISABLE_SPARSE_TENSORS) +#endif // !defined(ORT_MINIMAL_BUILD) && !defined(__wasm__) + } // namespace test } // namespace onnxruntime