saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix bug in GatherElements (#2130)

Browse source 
* Fix bug in GatherElements

* Uncomment some tests

* Updates

* Nits

* Nits

* Nits
This commit is contained in:
Hariharan Seshadri 2019-10-15 07:54:42 -07:00 committed by GitHub
parent 7ef02f14d2
commit 6857bb8aba
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 109 additions and 51 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

51
onnxruntime/core/providers/cpu/tensor/gather_elements.cc
View file

@ -9,9 +9,10 @@ namespace onnxruntime {
ONNX_CPU_OPERATOR_KERNEL(
GatherElements,
11,
KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllTensorTypes())
.TypeConstraint("Tind", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
KernelDefBuilder()
.TypeConstraint("T", DataTypeImpl::AllTensorTypes())
.TypeConstraint("Tind", std::vector<MLDataType>{DataTypeImpl::GetTensorType<int32_t>(),
DataTypeImpl::GetTensorType<int64_t>()}),
GatherElements);
// Some helpers needed for GatherElements op -
@ -57,7 +58,7 @@ static inline void increment_over_inner_dim(std::vector<int64_t>& current_dims,
int64_t rank = static_cast<int64_t>(current_dims.size());
// 'reset' innermost dimension value
current_dims[0] = 0;
current_dims[rank - 1] = 0;
// nothing to increment over
if (rank == 1) {
@ -77,7 +78,7 @@ static inline void increment_over_inner_dim(std::vector<int64_t>& current_dims,
// parse indices_tensor and along the way validate its shape and contents
static std::vector<int64_t> parse_and_validate_indices_tensor(const Tensor* indices_tensor,
int64_t axis, const TensorShape& input_shape) {
int64_t axis, const TensorShape& input_shape) {
const auto& indices_shape = indices_tensor->Shape().GetDims();
int64_t indices_rank = static_cast<int64_t>(indices_shape.size());
for (int64_t i = 0; i < indices_rank; ++i) {
@ -86,8 +87,10 @@ static std::vector<int64_t> parse_and_validate_indices_tensor(const Tensor* indi
// value if within bounds of the corresponding 'data' shape
if (i != axis) {
if (indices_shape[i] < 0 || indices_shape[i] > input_shape[i])
ORT_THROW("GatherElements op: 'indices' shape should have values within bounds of 'data' shape. "
"Invalid value in indices shape is: ", indices_shape[i]);
ORT_THROW(
"GatherElements op: 'indices' shape should have values within bounds of 'data' shape. "
"Invalid value in indices shape is: ",
indices_shape[i]);
}
}
@ -131,21 +134,20 @@ static std::vector<int64_t> parse_and_validate_indices_tensor(const Tensor* indi
#pragma GCC diagnostic ignored "-Wclass-memaccess"
#endif
#endif
template<bool is_string, typename T>
template <bool is_string, typename T>
static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor,
Tensor* output_tensor, int64_t axis) {
// get pointer to input data
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
Tensor* output_tensor, int64_t axis) {
// get pointer to input data
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
const T* input_data = nullptr;
if (is_string) {
input_data = input_tensor->Data<T>();
input_data = input_tensor->Data<T>();
} else {
input_data = reinterpret_cast<const T*>(input_tensor->DataRaw());
input_data = reinterpret_cast<const T*>(input_tensor->DataRaw());
}
// get pointer to output data
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
T* output_data = nullptr;
if (is_string) {
output_data = output_tensor->MutableData<T>();
@ -176,18 +178,19 @@ static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor,
// process 1 chunk of 'inner dimension' length
for (int64_t i = 0; i < inner_dim_size; ++i) {
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
if (is_string) {
output_data[++output_counter] = input_data[base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i];
output_data[++output_counter] = input_data[base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i];
} else {
memcpy(output_data, input_data + (base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i) * element_size, element_size);
output_data += element_size;
memcpy(output_data,
input_data + (base_offset + (indices_data[++indices_counter] * input_shape_pitches[axis]) + i) * element_size, element_size);
output_data += element_size;
}
}
increment_over_inner_dim(process_dims, indices_shape);
}
}
}
// we special-case inner dim as we can weed-out some unnecessary computations in element offset calculations
else {
while (num_inner_dim-- != 0) {
@ -196,12 +199,12 @@ static void core_impl(const Tensor* input_tensor, const Tensor* indices_tensor,
// process 1 chunk of 'inner dimension' length
for (int64_t i = 0; i < inner_dim_size; ++i) {
// for innermost axis, input_shape_pitches[axis] = 1 (so no need to multiply)
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
if (is_string) {
output_data[++output_counter] = input_data[base_offset + indices_data[++indices_counter]];
// optimizer will remove the redundant if/else block based on 'is_string' template parameter
if (is_string) {
output_data[++output_counter] = input_data[base_offset + indices_data[++indices_counter]];
} else {
memcpy(output_data, input_data + (base_offset + indices_data[++indices_counter]) * element_size, element_size);
output_data += element_size;
output_data += element_size;
}
}

109
onnxruntime/test/providers/cpu/tensor/gather_elements_op_test.cc
View file

@ -8,46 +8,44 @@ namespace onnxruntime {
namespace test {
template <typename T>
void RunTypedTest()
{
void RunTypedTest() {
// int32_t indices - axis 0
OpTester test1("GatherElements", 11);
test1.AddAttribute<int64_t>("axis", 0LL);
test1.AddInput<T>("data", {2, 3},
{0, 1, 2, 3, 4, 5});
{0, 1, 2, 3, 4, 5});
test1.AddInput<int32_t>("indices", {1, 2}, {0, 1});
test1.AddOutput<T>("output", {1, 2},
{0, 4});
{0, 4});
test1.Run();
// int32_t indices - axis 1
OpTester test2("GatherElements", 11);
test2.AddAttribute<int64_t>("axis", 1LL);
test2.AddInput<T>("data", {2, 2},
{1, 2,
3, 4});
{1, 2,
3, 4});
test2.AddInput<int32_t>("indices", {2, 2},
{0, 0,
1, 0});
{0, 0,
1, 0});
test2.AddOutput<T>("output", {2, 2},
{1, 1,
4, 3});
{1, 1,
4, 3});
test2.Run();
// int64_t indices - axis 1
OpTester test3("GatherElements", 11);
test3.AddAttribute<int64_t>("axis", 1LL);
test3.AddInput<T>("data", {2, 2},
{1, 2,
3, 4});
{1, 2,
3, 4});
test3.AddInput<int64_t>("indices", {2, 2},
{0, 0,
1, 0});
{0, 0,
1, 0});
test3.AddOutput<T>("output", {2, 2},
{1, 1,
4, 3});
{1, 1,
4, 3});
test3.Run();
// negative indices - axis 1
@ -106,36 +104,69 @@ void RunTypedTest()
{0, 1});
test7.AddOutput<T>("output", {1, 2, 1}, {1, 4});
test7.Run();
// 2D input - axis 1
OpTester test8("GatherElements", 11);
test8.AddAttribute<int64_t>("axis", 1LL);
test8.AddInput<T>("data", {3, 3},
{1, 2, 3,
4, 5, 6,
7, 8, 9});
test8.AddInput<int64_t>("indices", {3, 2},
{1, 0, 0, 1, 0, 1});
test8.AddOutput<T>("output", {3, 2}, {2, 1, 4, 5, 7, 8});
test8.Run();
// 2D input - axis 1
OpTester test9("GatherElements", 11);
test9.AddAttribute<int64_t>("axis", 0LL);
test9.AddInput<T>("data", {3, 3},
{1, 2, 3,
4, 5, 6,
7, 8, 9});
test9.AddInput<int64_t>("indices", {3, 2},
{1, 0, 0, 1, 0, 1});
test9.AddOutput<T>("output", {3, 2}, {4, 2, 1, 5, 1, 5});
test9.Run();
// 1D input - axis 0
OpTester test10("GatherElements", 11);
test10.AddAttribute<int64_t>("axis", 0LL);
test10.AddInput<T>("data", {3},
{1, 2, 3});
test10.AddInput<int64_t>("indices", {2},
{1, 0});
test10.AddOutput<T>("output", {2}, {2, 1});
test10.Run();
}
template <>
void RunTypedTest<std::string>() {
// int32_t indices - axis 0
OpTester test1("GatherElements", 11);
test1.AddAttribute<int64_t>("axis", 0LL);
test1.AddInput<std::string>("data", {2, 3},
{"a", "b", "c", "d", "e", "f"});
{"a", "b", "c", "d", "e", "f"});
test1.AddInput<int32_t>("indices", {1, 2}, {0, 1});
test1.AddOutput<std::string>("output", {1, 2},
{"a", "e"});
{"a", "e"});
test1.Run();
// int32_t indices - axis 1
OpTester test2("GatherElements", 11);
test2.AddAttribute<int64_t>("axis", 1LL);
test2.AddInput<std::string>("data", {2, 2},
{"a", "b",
"c", "d"});
{"a", "b",
"c", "d"});
test2.AddInput<int32_t>("indices", {2, 2},
{0, 0,
1, 0});
{0, 0,
1, 0});
test2.AddOutput<std::string>("output", {2, 2},
{"a", "a",
"d", "c"});
{"a", "a",
"d", "c"});
test2.Run();
// negative indices - axis 1
// negative indices - axis 1
OpTester test3("GatherElements", 11);
test3.AddAttribute<int64_t>("axis", 1LL);
test3.AddInput<std::string>("data", {2, 2},
@ -193,6 +224,30 @@ void RunTypedTest<std::string>() {
test6.AddOutput<std::string>("output", {1, 2, 1},
{"a", "d"});
test6.Run();
// 2D input - axis 1
OpTester test7("GatherElements", 11);
test7.AddAttribute<int64_t>("axis", 1LL);
test7.AddInput<std::string>("data", {3, 3},
{"a", "b", "c",
"d", "e", "f",
"g", "h", "i"});
test7.AddInput<int64_t>("indices", {3, 2},
{1, 0, 0, 1, 0, 1});
test7.AddOutput<std::string>("output", {3, 2}, {"b", "a", "d", "e", "g", "h"});
test7.Run();
// 2D input - axis 2
OpTester test8("GatherElements", 11);
test8.AddAttribute<int64_t>("axis", 0LL);
test8.AddInput<std::string>("data", {3, 3},
{"a", "b", "c",
"d", "e", "f",
"g", "h", "i"});
test8.AddInput<int64_t>("indices", {3, 2},
{1, 0, 0, 1, 0, 1});
test8.AddOutput<std::string>("output", {3, 2}, {"d", "b", "a", "e", "a", "e"});
test8.Run();
}
TEST(GatherElementsOpTest, int8_t) {