Added two cases in strided_slice

Signed-off-by: Chen Xin <jack.chen@verisilicon.com>
This commit is contained in:
Chen Xin 2022-09-29 14:25:08 +08:00 committed by Sven
parent 535c9da867
commit 20db77ee61
1 changed files with 78 additions and 0 deletions

View File

@ -109,3 +109,81 @@ TEST(StridedSlice, shape_) {
EXPECT_TRUE(ret) << "Failed at execute";
}
TEST(StridedSlice, shrinkmask_1) {
auto ctx = tim::vx::Context::Create();
auto graph = ctx->CreateGraph();
tim::vx::ShapeType input_shape({3, 2, 1});
tim::vx::ShapeType output_shape({3, 2});
tim::vx::TensorSpec input_spec(tim::vx::DataType::FLOAT32, input_shape,
tim::vx::TensorAttribute::INPUT);
tim::vx::TensorSpec output_spec(tim::vx::DataType::FLOAT32, output_shape,
tim::vx::TensorAttribute::OUTPUT);
auto input_tensor = graph->CreateTensor(input_spec);
auto output_tensor = graph->CreateTensor(output_spec);
std::vector<float> in_data = {1, 1, 2, 2, 3, 3};
std::vector<float> golden = {1, 1, 2, 2, 3, 3};
std::vector<int> begin = {0, 0, 0};
std::vector<int> end = {0, 0, 0};
std::vector<int> strides = {1, 1, 1};
// The ith bits in MASK_SHRINK will mask input_shape[i].
uint32_t MASK_BEGIN = 0, MASK_END = 0, MASK_SHRINK = 0b100;
auto op = graph->CreateOperation<tim::vx::ops::StridedSlice>(
begin, end, strides, MASK_BEGIN, MASK_END, MASK_SHRINK);
(*op).BindInputs({input_tensor}).BindOutputs({output_tensor});
input_tensor->CopyDataToTensor(in_data.data(),
in_data.size() * sizeof(float));
EXPECT_TRUE(graph->Compile());
EXPECT_TRUE(graph->Run());
std::vector<float> output(golden.size());
EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
EXPECT_EQ(golden, output);
}
TEST(StridedSlice, endmask_1) {
auto ctx = tim::vx::Context::Create();
auto graph = ctx->CreateGraph();
tim::vx::ShapeType input_shape({3, 2, 1});
tim::vx::ShapeType output_shape({3, 2, 1});
tim::vx::TensorSpec input_spec(tim::vx::DataType::FLOAT32, input_shape,
tim::vx::TensorAttribute::INPUT);
tim::vx::TensorSpec output_spec(tim::vx::DataType::FLOAT32, output_shape,
tim::vx::TensorAttribute::OUTPUT);
auto input_tensor = graph->CreateTensor(input_spec);
auto output_tensor = graph->CreateTensor(output_spec);
std::vector<float> in_data = {1, 1, 2, 2, 3, 3};
std::vector<float> golden = {1, 1, 2, 2, 3, 3};
std::vector<int> begin = {2, 0, 0};
std::vector<int> end = {3, 2, 1};
std::vector<int> strides = {1, 1, 1};
// The ith bits in MASK_BEGIN will mask begin[i]. MASK_END is similar.
uint32_t MASK_BEGIN = 0b001, MASK_END = 0, MASK_SHRINK = 0;
auto op = graph->CreateOperation<tim::vx::ops::StridedSlice>(
begin, end, strides, MASK_BEGIN, MASK_END, MASK_SHRINK);
(*op).BindInputs({input_tensor}).BindOutputs({output_tensor});
input_tensor->CopyDataToTensor(in_data.data(),
in_data.size() * sizeof(float));
EXPECT_TRUE(graph->Compile());
EXPECT_TRUE(graph->Run());
std::vector<float> output(golden.size());
EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
EXPECT_EQ(golden, output);
}