Skip to content

Fix external tensor copy validation, dense-to-sparse bounds, and ORT format initializer state #33643

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
m00dy (professor-moody) wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Fix external tensor copy validation, dense-to-sparse bounds, and ORT format initializer state #33643

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4b7828f
Fix external tensor copy validation, dense-to-sparse bounds, and ORT …
professor-moody Oct 27, 2026
81d4e45
Add regression tests for external copy capacity, sub-byte dense-to-sp…
professor-moody Oct 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 23 additions & 11 deletions onnxruntime/core/framework/tensorprotoutils.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<Tensor>();
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!");
}
Expand All @@ -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");
Comment on lines +2077 to +2080

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<Tensor>();
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
Expand Down Expand Up @@ -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());
Expand All @@ -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");
Comment on lines +2967 to +2968
size_t element_size = ml_data->Size();

std::vector<uint8_t> 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();
Expand Down
2 changes: 2 additions & 0 deletions onnxruntime/core/graph/graph.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.";
Expand Down Expand Up @@ -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.";
Expand Down
145 changes: 145 additions & 0 deletions onnxruntime/test/framework/ort_model_only_test.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<uint8_t> BuildOrtModelWithDuplicateInitializer(DuplicateInitializerReplacement replacement) {
return BuildOrtModelBuffer([replacement](flatbuffers::FlatBufferBuilder& builder) {
const int64_t arg_dim = replacement == DuplicateInitializerReplacement::kDense ? 4 : 40;
std::vector<flatbuffers::Offset<fbs::ValueInfo>> node_args{
fbs::CreateValueInfoDirect(builder, "w", "", CreateFloatTensorTypeInfo(builder, arg_dim)),
fbs::CreateValueInfoDirect(builder, "y", "", CreateFloatTensorTypeInfo(builder, arg_dim))};

std::vector<int64_t> large_dims{40};
std::vector<float> large_values(40, 1.0f);
std::vector<uint8_t> large_raw(reinterpret_cast<const uint8_t*>(large_values.data()),
reinterpret_cast<const uint8_t*>(large_values.data() + large_values.size()));
std::vector<flatbuffers::Offset<fbs::Tensor>> initializers{
fbs::CreateTensorDirect(builder, "w", "", &large_dims, fbs::TensorDataType::FLOAT, &large_raw)};

std::vector<int64_t> small_dims{4};
const std::vector<float> small_values{2.0f, 3.0f, 4.0f, 5.0f};
std::vector<uint8_t> small_raw(reinterpret_cast<const uint8_t*>(small_values.data()),
reinterpret_cast<const uint8_t*>(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<int64_t> sparse_nnz_dims{2};
const std::vector<float> sparse_values{5.0f, -6.0f};
std::vector<uint8_t> sparse_values_raw(
reinterpret_cast<const uint8_t*>(sparse_values.data()),
reinterpret_cast<const uint8_t*>(sparse_values.data() + sparse_values.size()));
const std::vector<int64_t> sparse_indices{3, 17};
std::vector<uint8_t> sparse_indices_raw(
reinterpret_cast<const uint8_t*>(sparse_indices.data()),
reinterpret_cast<const uint8_t*>(sparse_indices.data() + sparse_indices.size()));
std::vector<int64_t> sparse_dense_dims{40};
std::vector<flatbuffers::Offset<fbs::SparseTensor>> 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<flatbuffers::Offset<flatbuffers::String>> inputs{builder.CreateSharedString("w")};
std::vector<flatbuffers::Offset<flatbuffers::String>> outputs{builder.CreateSharedString("y")};
std::vector<flatbuffers::Offset<flatbuffers::String>> implicit_inputs;
std::vector<flatbuffers::Offset<flatbuffers::String>> graph_inputs;
std::vector<int32_t> input_arg_counts{1};
std::vector<flatbuffers::Offset<fbs::Node>> 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<float> UnpackFloatInitializer(const TensorProto& initializer) {
const TensorShape shape = utils::GetTensorShapeFromTensorProto(initializer);
std::vector<float> values(static_cast<size_t>(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<int>(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<Tensor>().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<int>(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<Tensor>().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<int>(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<float> 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<Tensor>().Shape();
}
#endif // !defined(DISABLE_SPARSE_TENSORS)
#endif // !defined(ORT_MINIMAL_BUILD)

#if !defined(ORT_MINIMAL_BUILD)
Expand Down
Loading