Skip to content

Commit 2944894

Browse files
yuslepukhinyuslepukhin
committed
Add more tests
1 parent 883d509 commit 2944894

File tree

2 files changed

+99
-80
lines changed

2 files changed

+99
-80
lines changed

‎onnxruntime/core/framework/tensorprotoutils.cc‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1763,7 +1763,7 @@ static Status CopySparseData(size_t n_sparse_elements,
17631763
ORT_RETURN_IF_ERROR(UnpackInitializerData(indices, model_path, unpack_buffer));
17641764
indices_data = ReinterpretAsSpan<const int64_t>(gsl::make_span(unpack_buffer));
17651765
} else {
1766-
ORT_RETURN_IF_NOT(indices.int64_data_size() == static_cast<int64_t>(nnz_elements),
1766+
ORT_RETURN_IF_NOT(indices.int64_data_size() == nnz_elements,
17671767
"Sparse indices int64 data size does not match expected");
17681768
indices_data = gsl::make_span(indices.int64_data().data(), nnz_elements);
17691769
}
@@ -1778,7 +1778,7 @@ static Status CopySparseData(size_t n_sparse_elements,
17781778
unpack_buffer.clear();
17791779
unpack_buffer.shrink_to_fit();
17801780
} else {
1781-
ORT_RETURN_IF_NOT(indices.int32_data_size() == static_cast<int64_t>(nnz_elements),
1781+
ORT_RETURN_IF_NOT(indices.int32_data_size() == nnz_elements,
17821782
"Sparse indices int32 data size does not match expected");
17831783
indices_values.insert(indices_values.cend(), indices.int32_data().cbegin(), indices.int32_data().cend());
17841784
}
@@ -1792,13 +1792,14 @@ static Status CopySparseData(size_t n_sparse_elements,
17921792
ORT_RETURN_IF_ERROR(UnpackInitializerData(indices, model_path, unpack_buffer));
17931793
auto int16_span = ReinterpretAsSpan<const int16_t>(gsl::make_span(unpack_buffer));
17941794
indices_values.insert(indices_values.cend(), int16_span.begin(), int16_span.end());
1795-
indices_data = gsl::make_span(indices_values);
17961795
unpack_buffer.clear();
17971796
unpack_buffer.shrink_to_fit();
17981797
} else {
1799-
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH,
1800-
"Invalid SparseTensor indices. INT16 indices must be in the raw data of indices tensor");
1798+
ORT_RETURN_IF_NOT(indices.int32_data_size() == nnz_elements,
1799+
"Sparse indices int16 data size does not match expected");
1800+
indices_values.insert(indices_values.cend(), indices.int32_data().cbegin(), indices.int32_data().cend());
18011801
}
1802+
indices_data = gsl::make_span(indices_values);
18021803
break;
18031804
}
18041805
case ONNX_NAMESPACE::TensorProto_DataType_INT8: {
@@ -1808,13 +1809,14 @@ static Status CopySparseData(size_t n_sparse_elements,
18081809
ORT_RETURN_IF_ERROR(UnpackInitializerData(indices, model_path, unpack_buffer));
18091810
auto int8_span = ReinterpretAsSpan<const int8_t>(gsl::make_span(unpack_buffer));
18101811
indices_values.insert(indices_values.cend(), int8_span.begin(), int8_span.end());
1811-
indices_data = gsl::make_span(indices_values);
18121812
unpack_buffer.clear();
18131813
unpack_buffer.shrink_to_fit();
18141814
} else {
1815-
return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_GRAPH,
1816-
"Invalid SparseTensor indices. INT8 indices must be in the raw data of indices tensor");
1815+
ORT_RETURN_IF_NOT(indices.int32_data_size() == nnz_elements,
1816+
"Sparse indices int8 data size does not match expected");
1817+
indices_values.insert(indices_values.cend(), indices.int32_data().cbegin(), indices.int32_data().cend());
18171818
}
1819+
indices_data = gsl::make_span(indices_values);
18181820
break;
18191821
}
18201822
default:

‎onnxruntime/test/framework/tensorutils_test.cc‎

Lines changed: 89 additions & 72 deletions
Original file line numberDiff line numberDiff line change
@@ -586,11 +586,12 @@ TEST_F(PathValidationTest, ValidateExternalDataPathWithSymlinkOutside) {
586586
ASSERT_FALSE(utils::ValidateExternalDataPath(base_dir_, "outside_link.bin").IsOK());
587587
}
588588

589-
template <typename T>
589+
template <typename T, typename I>
590590
void TestSparseToDenseConversion(gsl::span<const int64_t> dense_shape,
591591
const std::vector<T>& values,
592-
gsl::span<const int64_t> indices,
592+
gsl::span<const I> indices,
593593
gsl::span<const int64_t> indices_shape,
594+
bool raw_data_indices,
594595
const std::vector<T>& expected_dense_data) {
595596
ONNX_NAMESPACE::SparseTensorProto sparse_proto;
596597
for (auto dim : dense_shape) {
@@ -615,8 +616,42 @@ void TestSparseToDenseConversion(gsl::span<const int64_t> dense_shape,
615616
auto* indices_tensor = sparse_proto.mutable_indices();
616617
indices_tensor->set_name("indices");
617618
indices_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
618-
for (int64_t idx : indices) {
619-
indices_tensor->add_int64_data(idx);
619+
if constexpr (std::is_same_v<I, int8_t>) {
620+
indices_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT8);
621+
if (raw_data_indices) {
622+
indices_tensor->set_raw_data(indices.data(), indices.size() * sizeof(I));
623+
} else {
624+
for (auto idx : indices) {
625+
indices_tensor->add_int32_data(static_cast<int32_t>(idx)); // indices are stored in int32_data for types < int32
626+
}
627+
}
628+
} else if constexpr (std::is_same_v<I, int16_t>) {
629+
indices_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT16);
630+
if (raw_data_indices) {
631+
indices_tensor->set_raw_data(indices.data(), indices.size() * sizeof(I));
632+
} else {
633+
for (auto idx : indices) {
634+
indices_tensor->add_int32_data(static_cast<int32_t>(idx));
635+
}
636+
}
637+
} else if constexpr (std::is_same_v<I, int32_t>) {
638+
indices_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT32);
639+
if (raw_data_indices) {
640+
indices_tensor->set_raw_data(indices.data(), indices.size() * sizeof(I));
641+
} else {
642+
for (auto idx : indices) {
643+
indices_tensor->add_int32_data(idx);
644+
}
645+
}
646+
} else if constexpr (std::is_same_v<I, int64_t>) {
647+
indices_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64);
648+
if (raw_data_indices) {
649+
indices_tensor->set_raw_data(indices.data(), indices.size() * sizeof(I));
650+
} else {
651+
for (auto idx : indices) {
652+
indices_tensor->add_int64_data(idx);
653+
}
654+
}
620655
}
621656
for (auto dim : indices_shape) {
622657
indices_tensor->add_dims(dim);
@@ -638,7 +673,7 @@ void TestSparseToDenseConversion(gsl::span<const int64_t> dense_shape,
638673
EXPECT_EQ(unpacked_data, expected_dense_data);
639674
}
640675

641-
TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Rank1Indices) {
676+
TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Rank1Indices64) {
642677
// Dense Shape: [2, 2] -> 4 elements
643678
// Indices: [0, 3] (linear)
644679
// Values: [1.0, 2.0]
@@ -649,7 +684,53 @@ TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Rank1Indices) {
649684
std::vector<int64_t> indices_shape = {2}; // [NNZ]
650685
std::vector<float> expected = {1.0f, 0.0f, 0.0f, 2.0f};
651686

652-
TestSparseToDenseConversion<float>(dense_shape, values, indices, indices_shape, expected);
687+
TestSparseToDenseConversion<float, int64_t>(dense_shape, values, indices, indices_shape, false, expected);
688+
TestSparseToDenseConversion<float, int64_t>(dense_shape, values, indices, indices_shape, true, expected);
689+
}
690+
691+
TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Rank1Indices32) {
692+
// Dense Shape: [2, 2] -> 4 elements
693+
// Indices: [0, 3] (linear)
694+
// Values: [1.0, 2.0]
695+
// Expected: [1.0, 0.0, 0.0, 2.0]
696+
std::vector<int64_t> dense_shape = {2, 2};
697+
std::vector<float> values = {1.0f, 2.0f};
698+
std::vector<int32_t> indices = {0, 3};
699+
std::vector<int64_t> indices_shape = {2}; // [NNZ]
700+
std::vector<float> expected = {1.0f, 0.0f, 0.0f, 2.0f};
701+
702+
TestSparseToDenseConversion<float, int32_t>(dense_shape, values, indices, indices_shape, false, expected);
703+
TestSparseToDenseConversion<float, int32_t>(dense_shape, values, indices, indices_shape, true, expected);
704+
}
705+
706+
TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Rank1Indices16) {
707+
// Dense Shape: [2, 2] -> 4 elements
708+
// Indices: [0, 3] (linear)
709+
// Values: [1.0, 2.0]
710+
// Expected: [1.0, 0.0, 0.0, 2.0]
711+
std::vector<int64_t> dense_shape = {2, 2};
712+
std::vector<float> values = {1.0f, 2.0f};
713+
std::vector<int16_t> indices = {0, 3};
714+
std::vector<int64_t> indices_shape = {2}; // [NNZ]
715+
std::vector<float> expected = {1.0f, 0.0f, 0.0f, 2.0f};
716+
717+
TestSparseToDenseConversion<float, int16_t>(dense_shape, values, indices, indices_shape, false, expected);
718+
TestSparseToDenseConversion<float, int16_t>(dense_shape, values, indices, indices_shape, true, expected);
719+
}
720+
721+
TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Rank1Indices8) {
722+
// Dense Shape: [2, 2] -> 4 elements
723+
// Indices: [0, 3] (linear)
724+
// Values: [1.0, 2.0]
725+
// Expected: [1.0, 0.0, 0.0, 2.0]
726+
std::vector<int64_t> dense_shape = {2, 2};
727+
std::vector<float> values = {1.0f, 2.0f};
728+
std::vector<int8_t> indices = {0, 3};
729+
std::vector<int64_t> indices_shape = {2}; // [NNZ]
730+
std::vector<float> expected = {1.0f, 0.0f, 0.0f, 2.0f};
731+
732+
TestSparseToDenseConversion<float, int8_t>(dense_shape, values, indices, indices_shape, false, expected);
733+
TestSparseToDenseConversion<float, int8_t>(dense_shape, values, indices, indices_shape, true, expected);
653734
}
654735

655736
TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Rank2Indices_COO) {
@@ -666,7 +747,8 @@ TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Rank2Indices_COO) {
666747
0, 20, 0,
667748
0, 0, 30};
668749

669-
TestSparseToDenseConversion<int32_t>(dense_shape, values, indices, indices_shape, expected);
750+
TestSparseToDenseConversion<int32_t, int64_t>(dense_shape, values, indices, indices_shape, false, expected);
751+
TestSparseToDenseConversion<int32_t, int64_t>(dense_shape, values, indices, indices_shape, true, expected);
670752
}
671753

672754
TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_InvalidIndicesRank) {
@@ -795,70 +877,5 @@ TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_NegativeDenseShape) {
795877
EXPECT_THAT(status.ErrorMessage(), testing::HasSubstr("Sparse tensor dense dims expected to be positive"));
796878
}
797879

798-
template <typename T>
799-
void TestSparseToDenseConversionWithInt8Indices(const std::vector<int64_t>& dense_shape,
800-
const std::vector<T>& values,
801-
const std::vector<int8_t>& indices,
802-
const std::vector<int64_t>& indices_shape,
803-
const std::vector<T>& expected_dense_data) {
804-
ONNX_NAMESPACE::SparseTensorProto sparse_proto;
805-
for (auto dim : dense_shape) {
806-
sparse_proto.add_dims(dim);
807-
}
808-
809-
// Create values tensor
810-
auto* values_tensor = sparse_proto.mutable_values();
811-
values_tensor->set_name("values");
812-
if constexpr (std::is_same_v<T, float>) {
813-
values_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT);
814-
for (float v : values) values_tensor->add_float_data(v);
815-
} else if constexpr (std::is_same_v<T, int32_t>) {
816-
values_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT32);
817-
for (int32_t v : values) values_tensor->add_int32_data(v);
818-
}
819-
values_tensor->add_dims(values.size());
820-
821-
// Create indices tensor with INT8 data in raw_data
822-
auto* indices_tensor = sparse_proto.mutable_indices();
823-
indices_tensor->set_name("indices");
824-
indices_tensor->set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT8);
825-
826-
std::string raw_indices_data;
827-
raw_indices_data.resize(indices.size() * sizeof(int8_t));
828-
std::memcpy(raw_indices_data.data(), indices.data(), raw_indices_data.size());
829-
indices_tensor->set_raw_data(raw_indices_data);
830-
831-
for (auto dim : indices_shape) {
832-
indices_tensor->add_dims(dim);
833-
}
834-
835-
ONNX_NAMESPACE::TensorProto dense_proto;
836-
std::filesystem::path model_path;
837-
ASSERT_STATUS_OK(utils::SparseTensorProtoToDenseTensorProto(sparse_proto, model_path, dense_proto));
838-
839-
// Verify dense proto
840-
ASSERT_EQ(dense_proto.dims_size(), dense_shape.size());
841-
for (size_t i = 0; i < (size_t)dense_shape.size(); ++i) {
842-
ASSERT_EQ(dense_proto.dims(static_cast<int>(i)), dense_shape[i]);
843-
}
844-
845-
std::vector<T> unpacked_data(expected_dense_data.size());
846-
ASSERT_STATUS_OK(utils::UnpackTensor(dense_proto, model_path, unpacked_data.data(), unpacked_data.size()));
847-
848-
EXPECT_EQ(unpacked_data, expected_dense_data);
849-
}
850-
851-
TEST(TensorProtoUtilsTest, SparseTensorProtoToDense_Int8Indices) {
852-
// Dense Shape: [5], Indices: [1, 4], Values: [10, 20]
853-
// Expected: [0, 10, 0, 0, 20]
854-
std::vector<int64_t> dense_shape = {5};
855-
std::vector<int32_t> values = {10, 20};
856-
std::vector<int8_t> indices = {1, 4};
857-
std::vector<int64_t> indices_shape = {2};
858-
std::vector<int32_t> expected = {0, 10, 0, 0, 20};
859-
860-
TestSparseToDenseConversionWithInt8Indices(dense_shape, values, indices, indices_shape, expected);
861-
}
862-
863880
} // namespace test
864881
} // namespace onnxruntime

0 commit comments

Comments
 (0)