mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
[NNAPI EP] Add Gather op support (#11824)
* initial gather support nnapi * update * minor update * address pr comments * add int32 indices test case for nnapi * remove nnapi ep limitation for added UT * add link for memcpy type punning usage
This commit is contained in:
Rachel Guo
GitHub
parent
02457ec30a
commit
1a1c360a80
5 changed files with 364 additions and 107 deletions
|
|
@ -2436,6 +2436,82 @@ Status FlattenOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, cons
|
|||
|
||||
#pragma endregion
|
||||
|
||||
#pragma region op_gather
|
||||
|
||||
class GatherOpBuilder : public BaseOpBuilder {
|
||||
public:
|
||||
void AddInitializersToSkip(ModelBuilder& model_builder, const NodeUnit& node_unit) const override;
|
||||
|
||||
private:
|
||||
Status AddToModelBuilderImpl(ModelBuilder& model_builder, const NodeUnit& node_unit) const override;
|
||||
};
|
||||
|
||||
void GatherOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const NodeUnit& node_unit) const {
|
||||
// Skip the second input `indices` for Gather
|
||||
const auto& inputs = node_unit.Inputs();
|
||||
model_builder.AddInitializerToSkip(inputs[1].node_arg.Name()); // indices
|
||||
}
|
||||
|
||||
Status GatherOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const NodeUnit& node_unit) const {
|
||||
auto& shaper(model_builder.GetShaper());
|
||||
const auto& operand_indices(model_builder.GetOperandIndices());
|
||||
const auto& operand_types(model_builder.GetOperandTypes());
|
||||
const auto& initializers(model_builder.GetInitializerTensors());
|
||||
const auto& input1 = node_unit.Inputs()[0].node_arg.Name();
|
||||
const auto& input2 = node_unit.Inputs()[1].node_arg.Name(); // "indices"
|
||||
const auto& output = node_unit.Outputs()[0].node_arg.Name();
|
||||
|
||||
NodeAttrHelper helper(node_unit);
|
||||
int32_t rank = static_cast<int32_t>(shaper[input1].size());
|
||||
int32_t axis = static_cast<int32_t>(HandleNegativeAxis(helper.Get("axis", 0), rank));
|
||||
|
||||
std::vector<uint32_t> input_indices;
|
||||
input_indices.push_back(operand_indices.at(input1));
|
||||
ADD_SCALAR_OPERAND(model_builder, input_indices, axis);
|
||||
|
||||
// Add indices operand into nnapi
|
||||
const auto& indices_tensor = *initializers.at(input2);
|
||||
std::vector<uint8_t> unpacked_tensor;
|
||||
ORT_RETURN_IF_ERROR(onnxruntime::utils::UnpackInitializerData(indices_tensor, unpacked_tensor));
|
||||
|
||||
const auto data_type = indices_tensor.data_type();
|
||||
const auto indices_shape = indices_tensor.dims();
|
||||
uint32_t size = 1;
|
||||
Shape indices_dimen;
|
||||
indices_dimen.reserve(indices_tensor.dims_size());
|
||||
for (auto i = 0; i < indices_tensor.dims_size(); i++) {
|
||||
size *= indices_shape[i];
|
||||
indices_dimen.push_back(static_cast<uint32_t>(indices_shape[i]));
|
||||
}
|
||||
|
||||
std::vector<int32_t> indices(size);
|
||||
// see https://gist.github.com/shafik/848ae25ee209f698763cffee272a58f8#type-punning-arrays for the usage of memcpy here
|
||||
if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT64) {
|
||||
for (uint32_t i = 0; i < size; i++) {
|
||||
int64_t index_i64;
|
||||
memcpy(&index_i64, unpacked_tensor.data() + i * sizeof(int64_t), sizeof(int64_t));
|
||||
indices[i] = SafeInt<int32_t>(index_i64);
|
||||
}
|
||||
} else if (data_type == ONNX_NAMESPACE::TensorProto_DataType_INT32) {
|
||||
for (uint32_t i = 0; i < size; i++) {
|
||||
int32_t index;
|
||||
memcpy(&index, unpacked_tensor.data() + i * sizeof(int32_t), sizeof(int32_t));
|
||||
indices[i] = SafeInt<int32_t>(index);
|
||||
}
|
||||
}
|
||||
|
||||
OperandType indices_operand_type(Type::TENSOR_INT32, indices_dimen);
|
||||
ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(input2, indices.data(), indices_operand_type));
|
||||
input_indices.push_back(operand_indices.at(input2));
|
||||
ORT_RETURN_IF_ERROR(shaper.Gather(input1, input2, axis, output));
|
||||
const OperandType output_operand_type(operand_types.at(input1).type, shaper[output]);
|
||||
|
||||
return model_builder.AddOperation(ANEURALNETWORKS_GATHER, input_indices,
|
||||
{output}, {output_operand_type});
|
||||
}
|
||||
|
||||
#pragma endregion
|
||||
|
||||
#pragma region op_minmax
|
||||
|
||||
class MinMaxOpBuilder : public BaseOpBuilder {
|
||||
|
|
@ -2721,6 +2797,7 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() {
|
|||
NNAPI_EP_ADD_SINGLE_OP_BUILDER("DequantizeLinear", DequantizeLinearOpBuilder);
|
||||
NNAPI_EP_ADD_SINGLE_OP_BUILDER("Elu", EluOpBuilder);
|
||||
NNAPI_EP_ADD_SINGLE_OP_BUILDER("Flatten", FlattenOpBuilder);
|
||||
NNAPI_EP_ADD_SINGLE_OP_BUILDER("Gather", GatherOpBuilder);
|
||||
NNAPI_EP_ADD_SINGLE_OP_BUILDER("Identity", IdentityOpBuilder);
|
||||
NNAPI_EP_ADD_SINGLE_OP_BUILDER("LRN", LRNOpBuilder);
|
||||
NNAPI_EP_ADD_SINGLE_OP_BUILDER("QuantizeLinear", QuantizeLinearOpBuilder);
|
||||
|
|
|
|||
|
|
@ -2067,6 +2067,50 @@ bool FlattenOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& /* i
|
|||
|
||||
#pragma endregion
|
||||
|
||||
#pragma region op_gather
|
||||
|
||||
class GatherOpSupportChecker : public BaseOpSupportChecker {
|
||||
private:
|
||||
int32_t GetMinSupportedNNAPIFeatureLevel(const NodeUnit& /* node_unit */,
|
||||
const OpSupportCheckParams& /* params */) const override {
|
||||
return ANEURALNETWORKS_FEATURE_LEVEL_3;
|
||||
}
|
||||
|
||||
bool IsOpSupportedImpl(const InitializedTensorSet& initializers, const NodeUnit& node_unit,
|
||||
const OpSupportCheckParams& params) const override;
|
||||
};
|
||||
|
||||
bool GatherOpSupportChecker::IsOpSupportedImpl(const InitializedTensorSet& initializers, const NodeUnit& node_unit,
|
||||
const OpSupportCheckParams& /* params */) const {
|
||||
const auto& inputs = node_unit.Inputs();
|
||||
Shape input_shape;
|
||||
|
||||
if (!GetShape(inputs[0].node_arg, input_shape)) {
|
||||
return false;
|
||||
}
|
||||
|
||||
if (input_shape.size() > 4 || input_shape.empty()) {
|
||||
LOGS_DEFAULT(VERBOSE) << "Gather only supports up to 1-4d shape, input is "
|
||||
<< input_shape.size() << "d shape";
|
||||
return false;
|
||||
}
|
||||
|
||||
if (std::any_of(input_shape.cbegin(), input_shape.cend(), [](int32_t i) { return i == 0; })) {
|
||||
LOGS_DEFAULT(VERBOSE) << "Gather doesn't support dynamic input shape";
|
||||
return false;
|
||||
}
|
||||
|
||||
const auto& indices_name = inputs[1].node_arg.Name();
|
||||
if (!Contains(initializers, indices_name)) {
|
||||
LOGS_DEFAULT(VERBOSE) << "Indices of Gather must be known";
|
||||
return false;
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
#pragma endregion
|
||||
|
||||
#pragma region op_minmax
|
||||
|
||||
class MinMaxOpSupportChecker : public BaseOpSupportChecker {
|
||||
|
|
@ -2210,6 +2254,7 @@ static OpSupportCheckerRegistrations CreateOpSupportCheckerRegistrations() {
|
|||
NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("DequantizeLinear", DequantizeLinearOpSupportChecker);
|
||||
NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("Elu", EluOpSupportChecker);
|
||||
NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("Flatten", FlattenOpSupportChecker);
|
||||
NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("Gather", GatherOpSupportChecker);
|
||||
NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("LRN", LRNOpSupportChecker);
|
||||
NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("QuantizeLinear", QuantizeLinearOpSupportChecker);
|
||||
NNAPI_EP_ADD_SINGLE_OP_SUPPORT_CHECKER("Reshape", ReshapeOpSupportChecker);
|
||||
|
|
|
|||
|
|
@ -136,6 +136,13 @@ Status Shaper::DepthToSpace(const std::string& input_name,
|
|||
SHAPER_FUNC(DepthToSpace, input_name, blocksize, nchw, output_name);
|
||||
}
|
||||
|
||||
Status Shaper::Gather(const std::string& input_name1,
|
||||
const std::string& input_name2,
|
||||
const int32_t axis,
|
||||
const std::string& output_name) {
|
||||
SHAPER_FUNC(Gather, input_name1, input_name2, axis, output_name);
|
||||
}
|
||||
|
||||
Status Shaper::ResizeUsingScales(const std::string& input_name,
|
||||
const float scale_h, const float scale_w,
|
||||
bool nchw,
|
||||
|
|
@ -428,6 +435,30 @@ Status Shaper::DepthToSpaceImpl(const std::string& input_name,
|
|||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Shaper::GatherImpl(const std::string& input_name1,
|
||||
const std::string& input_name2,
|
||||
const int32_t axis,
|
||||
const std::string& output_name) {
|
||||
const Shape& input_dimen = shape_map_.at(input_name1);
|
||||
const Shape& indices_dimen = shape_map_.at(input_name2);
|
||||
|
||||
std::vector<uint32_t> output_dimen;
|
||||
output_dimen.reserve(indices_dimen.size() + input_dimen.size() - 1);
|
||||
|
||||
// Calculate the output dim
|
||||
for (int32_t i = 0; i < axis; ++i)
|
||||
output_dimen.push_back(input_dimen[i]);
|
||||
|
||||
for (const auto dim : indices_dimen)
|
||||
output_dimen.push_back(dim);
|
||||
|
||||
for (size_t i = axis + 1; i < input_dimen.size(); ++i)
|
||||
output_dimen.push_back(input_dimen[i]);
|
||||
|
||||
shape_map_[output_name] = output_dimen;
|
||||
return Status::OK();
|
||||
}
|
||||
|
||||
Status Shaper::ResizeUsingScalesImpl(const std::string& input_name,
|
||||
const float scale_h, const float scale_w,
|
||||
bool nchw,
|
||||
|
|
|
|||
|
|
@ -64,6 +64,11 @@ class Shaper {
|
|||
bool nchw,
|
||||
const std::string& output_name);
|
||||
|
||||
common::Status Gather(const std::string& input_name1,
|
||||
const std::string& input_name2,
|
||||
const int32_t axis,
|
||||
const std::string& output_name);
|
||||
|
||||
common::Status ResizeUsingScales(const std::string& input_name,
|
||||
const float scale_h, const float scale_w,
|
||||
bool nchw,
|
||||
|
|
@ -117,6 +122,10 @@ class Shaper {
|
|||
const int32_t blocksize,
|
||||
bool nchw,
|
||||
const std::string& output_name);
|
||||
common::Status GatherImpl(const std::string& input_name1,
|
||||
const std::string& input_name2,
|
||||
const int32_t axis,
|
||||
const std::string& output_name);
|
||||
common::Status ResizeUsingScalesImpl(const std::string& input_name,
|
||||
const float scale_h, const float scale_w,
|
||||
bool nchw,
|
||||
|
|
|
|||
|
|
@ -13,58 +13,76 @@ namespace test {
|
|||
// Those tests will fallback to other EPs
|
||||
|
||||
TEST(GatherOpTest, Gather_axis0) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 0LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {1}, {1LL});
|
||||
test.AddOutput<float>("output", {1, 3, 4},
|
||||
{10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.Run();
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 0LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {1}, {1LL}, indices_is_initializer);
|
||||
test.AddOutput<float>("output", {1, 3, 4},
|
||||
{10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.Run();
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_negative_axis) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", -3LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {1}, {1LL});
|
||||
test.AddOutput<float>("output", {1, 3, 4},
|
||||
{10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.Run();
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", -3LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {1}, {1LL}, indices_is_initializer);
|
||||
test.AddOutput<float>("output", {1, 3, 4},
|
||||
{10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.Run();
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_invalid_axis) {
|
||||
OpTester test("Gather");
|
||||
// Invalid axis not in range [-r, r-1]
|
||||
test.AddAttribute<int64_t>("axis", -10LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {1}, {1LL});
|
||||
test.AddOutput<float>("output", {1, 3, 4},
|
||||
{10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.Run(OpTester::ExpectResult::kExpectFailure, "axis must be in [-r, r-1]");
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
// Invalid axis not in range [-r, r-1]
|
||||
test.AddAttribute<int64_t>("axis", -10LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {1}, {1LL}, indices_is_initializer);
|
||||
test.AddOutput<float>("output", {1, 3, 4},
|
||||
{10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.Run(OpTester::ExpectResult::kExpectFailure, "axis must be in [-r, r-1]");
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_invalid_index_cpu) {
|
||||
|
|
@ -107,76 +125,153 @@ TEST(GatherOpTest, Gather_invalid_index_gpu) {
|
|||
#endif
|
||||
|
||||
TEST(GatherOpTest, Gather_axis1) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 1LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {2}, {2LL, 0LL});
|
||||
test.AddOutput<float>("output", {2, 2, 4},
|
||||
{2.0f, 2.1f, 2.2f, 2.3f,
|
||||
0.0f, 0.1f, 0.2f, 0.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f});
|
||||
test.Run();
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 1LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {2}, {2LL, 0LL}, indices_is_initializer);
|
||||
test.AddOutput<float>("output", {2, 2, 4},
|
||||
{2.0f, 2.1f, 2.2f, 2.3f,
|
||||
0.0f, 0.1f, 0.2f, 0.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f});
|
||||
test.Run();
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_axis2) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 2LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {3}, {1LL, 0LL, 2LL});
|
||||
test.AddOutput<float>("output", {2, 3, 3},
|
||||
{0.1f, 0.0f, 0.2f,
|
||||
1.1f, 1.0f, 1.2f,
|
||||
2.1f, 2.0f, 2.2f,
|
||||
10.1f, 10.0f, 10.2f,
|
||||
11.1f, 11.0f, 11.2f,
|
||||
12.1f, 12.0f, 12.2f});
|
||||
test.Run();
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 2LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int64_t>("indices", {3}, {1LL, 0LL, 2LL}, indices_is_initializer);
|
||||
test.AddOutput<float>("output", {2, 3, 3},
|
||||
{0.1f, 0.0f, 0.2f,
|
||||
1.1f, 1.0f, 1.2f,
|
||||
2.1f, 2.0f, 2.2f,
|
||||
10.1f, 10.0f, 10.2f,
|
||||
11.1f, 11.0f, 11.2f,
|
||||
12.1f, 12.0f, 12.2f});
|
||||
test.Run();
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_axis0_indices2d) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 0LL);
|
||||
test.AddInput<float>("data", {3, 3},
|
||||
{0.0f, 0.1f, 0.2f,
|
||||
1.0f, 1.1f, 1.2f,
|
||||
2.0f, 2.1f, 2.2f});
|
||||
test.AddInput<int64_t>("indices", {2LL, 2LL},
|
||||
{1LL, 0LL,
|
||||
2LL, 1LL});
|
||||
test.AddOutput<float>("output", {2, 2, 3},
|
||||
{1.0f, 1.1f, 1.2f, 0.0f, 0.1f, 0.2f,
|
||||
2.0f, 2.1f, 2.2f, 1.0f, 1.1f, 1.2f});
|
||||
test.Run();
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 0LL);
|
||||
test.AddInput<float>("data", {3, 3},
|
||||
{0.0f, 0.1f, 0.2f,
|
||||
1.0f, 1.1f, 1.2f,
|
||||
2.0f, 2.1f, 2.2f});
|
||||
test.AddInput<int64_t>("indices", {2LL, 2LL},
|
||||
{1LL, 0LL,
|
||||
2LL, 1LL},
|
||||
indices_is_initializer);
|
||||
test.AddOutput<float>("output", {2, 2, 3},
|
||||
{1.0f, 1.1f, 1.2f, 0.0f, 0.1f, 0.2f,
|
||||
2.0f, 2.1f, 2.2f, 1.0f, 1.1f, 1.2f});
|
||||
test.Run();
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_axis1_indices2d) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 1LL);
|
||||
test.AddInput<float>("data", {3, 3},
|
||||
{0.0f, 0.1f, 0.2f,
|
||||
1.0f, 1.1f, 1.2f,
|
||||
2.0f, 2.1f, 2.2f});
|
||||
test.AddInput<int64_t>("indices", {2LL, 2LL},
|
||||
{1LL, 0LL,
|
||||
2LL, 1LL});
|
||||
test.AddOutput<float>("output", {3, 2, 2},
|
||||
{0.1f, 0.0f, 0.2f, 0.1f,
|
||||
1.1f, 1.0f, 1.2f, 1.1f,
|
||||
2.1f, 2.0f, 2.2f, 2.1f});
|
||||
test.Run();
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 1LL);
|
||||
test.AddInput<float>("data", {3, 3},
|
||||
{0.0f, 0.1f, 0.2f,
|
||||
1.0f, 1.1f, 1.2f,
|
||||
2.0f, 2.1f, 2.2f});
|
||||
test.AddInput<int64_t>("indices", {2LL, 2LL},
|
||||
{1LL, 0LL,
|
||||
2LL, 1LL},
|
||||
indices_is_initializer);
|
||||
test.AddOutput<float>("output", {3, 2, 2},
|
||||
{0.1f, 0.0f, 0.2f, 0.1f,
|
||||
1.1f, 1.0f, 1.2f, 1.1f,
|
||||
2.1f, 2.0f, 2.2f, 2.1f});
|
||||
test.Run();
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_axis0_indicesInt32) {
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
// NNAPI EP only supports float input data for now,
|
||||
// the following two test cases cover int32_t indices with float input other than int64_t type for Nnapi
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 0LL);
|
||||
test.AddInput<float>("data", {2, 3, 4},
|
||||
{0.0f, 0.1f, 0.2f, 0.3f,
|
||||
1.0f, 1.1f, 1.2f, 1.3f,
|
||||
2.0f, 2.1f, 2.2f, 2.3f,
|
||||
10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
test.AddInput<int32_t>("indices", {1}, {1}, indices_is_initializer);
|
||||
test.AddOutput<float>("output", {1, 3, 4},
|
||||
{10.0f, 10.1f, 10.2f, 10.3f,
|
||||
11.0f, 11.1f, 11.2f, 11.3f,
|
||||
12.0f, 12.1f, 12.2f, 12.3f});
|
||||
|
||||
test.Run();
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_axis0_indices2dInt32) {
|
||||
// To test for NNAPI EP, we need the indices to be initializers
|
||||
auto run_test = [](bool indices_is_initializer) {
|
||||
OpTester test("Gather");
|
||||
test.AddAttribute<int64_t>("axis", 0LL);
|
||||
test.AddInput<float>("data", {3, 3},
|
||||
{0.0f, 0.1f, 0.2f,
|
||||
1.0f, 1.1f, 1.2f,
|
||||
2.0f, 2.1f, 2.2f});
|
||||
test.AddInput<int32_t>("indices", {2, 2},
|
||||
{1, 0,
|
||||
2, 1},
|
||||
indices_is_initializer);
|
||||
test.AddOutput<float>("output", {2, 2, 3},
|
||||
{1.0f, 1.1f, 1.2f, 0.0f, 0.1f, 0.2f,
|
||||
2.0f, 2.1f, 2.2f, 1.0f, 1.1f, 1.2f});
|
||||
|
||||
test.Run();
|
||||
};
|
||||
|
||||
run_test(false);
|
||||
run_test(true);
|
||||
}
|
||||
|
||||
TEST(GatherOpTest, Gather_axis1_indices2d_int32) {
|
||||
|
|
|
|||
Loading…
Reference in a new issue