IntenseWebs/openvino
3
0 Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[GPU] gather nd shape agnostic kernel implementation (#17940)

Browse Source 
* gather nd shape agnostic kernel implementation

* add func test

* fix minor bugs

* minor bug fixes

* fix win build error
This commit is contained in:
Wilson Seok authored and GitHub committed 2023-06-10 00:28:00 -07:00
1 parent c413825845
commit cff083f83d
7 files changed
+426 -24
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

src/plugins/intel_gpu/src/graph/gather_nd.cpp
+2
View File
@@ -100,6 +100,8 @@ std::vector<layout> gather_nd_inst::calc_output_layouts(gather_nd_node const& /*
ov::op::v8::shape_infer(&op, input_shapes, output_shapes);
}
OPENVINO_ASSERT(!output_shapes[0].rank().is_dynamic(),
"[GPU] Doesn't support output dynamic rank in gather_nd");
format output_format = format::adjust_to_rank(input_layout.format, output_shapes[0].size());
return { layout{output_shapes[0], output_type, output_format} };
src/plugins/intel_gpu/src/graph/impls/ocl/gather_nd.cpp
+34 -11
View File
@@ -35,22 +35,45 @@ struct gather_nd_impl : typed_primitive_impl_ocl<gather_nd> {
params.inputs.push_back(convert_data_tensor(impl_param.get_input_layout(1)));
return {params, optional_params};
}
void update_dispatch_data(const kernel_impl_params& impl_param) override {
auto kernel_params = get_kernel_params(impl_param);
(_kernel_data.update_dispatch_data_func)(kernel_params.first, _kernel_data);
}
};
namespace detail {
attach_gather_nd_impl::attach_gather_nd_impl() {
implementation_map<gather_nd>::add(impl_types::ocl, typed_primitive_impl_ocl<gather_nd>::create<gather_nd_impl>, {
std::make_tuple(data_types::f32, format::bfyx),
std::make_tuple(data_types::f16, format::bfyx),
std::make_tuple(data_types::i32, format::bfyx),
std::make_tuple(data_types::f32, format::bfzyx),
std::make_tuple(data_types::f16, format::bfzyx),
std::make_tuple(data_types::i32, format::bfzyx),
std::make_tuple(data_types::f32, format::bfwzyx),
std::make_tuple(data_types::f16, format::bfwzyx),
std::make_tuple(data_types::i32, format::bfwzyx),
});
auto types = {
data_types::f32,
data_types::f16,
data_types::i32
};
auto static_formats = {
format::bfyx,
format::bfzyx,
format::bfwzyx
};
implementation_map<gather_nd>::add(impl_types::ocl,
shape_types::static_shape,
typed_primitive_impl_ocl<gather_nd>::create<gather_nd_impl>,
types,
static_formats);
auto dyn_formats = {
format::bfyx,
format::bfzyx,
format::bfwzyx
};
implementation_map<gather_nd>::add(impl_types::ocl,
shape_types::dynamic_shape,
typed_primitive_impl_ocl<gather_nd>::create<gather_nd_impl>,
types,
dyn_formats);
}
} // namespace detail
src/plugins/intel_gpu/src/kernel_selector/cl_kernels/gather_nd_ref.cl
+21 -5
View File
@@ -33,15 +33,31 @@
#define INDICES_MAX_DIM 6
KERNEL(gather_nd_ref)(const __global INPUT0_TYPE* data,
const __global INPUT1_TYPE* indices,
__global OUTPUT_TYPE* output
KERNEL(gather_nd_ref)(
OPTIONAL_SHAPE_INFO_ARG
const __global INPUT0_TYPE* data,
const __global INPUT1_TYPE* indices,
__global OUTPUT_TYPE* output
#if HAS_FUSED_OPS_DECLS
, FUSED_OPS_DECLS
, FUSED_OPS_DECLS
#endif
)
{
#if IS_DYNAMIC
uint wi_slice = 1;
#if INPUT0_DIMS == 4
uint input_dim[4] = {INPUT0_BATCH_NUM, INPUT0_FEATURE_NUM, INPUT0_SIZE_Y, INPUT0_SIZE_X};
#elif INPUT0_DIMS == 5
uint input_dim[5] = {INPUT0_BATCH_NUM, INPUT0_FEATURE_NUM, INPUT0_SIZE_Z, INPUT0_SIZE_Y, INPUT0_SIZE_X};
#else
uint input_dim[6] = {INPUT0_BATCH_NUM, INPUT0_FEATURE_NUM, INPUT0_SIZE_W, INPUT0_SIZE_Z, INPUT0_SIZE_Y, INPUT0_SIZE_X};
#endif
for (uint i = BATCH_DIMS + INPUT1_SIZE_X; i < INPUT0_DIMS; i++)
wi_slice *= input_dim[i];
#define WI_SLICE_SIZE wi_slice
#else
#define WI_SLICE_SIZE WI_SLICE_SIZE_STATIC
#endif
const uint dim0 = get_global_id(0);
const uint dim1 = get_global_id(1);
const uint dim2 = get_global_id(2);
src/plugins/intel_gpu/src/kernel_selector/kernels/gather/gather_nd_kernel_ref.cpp
+32 -7
View File
@@ -29,6 +29,7 @@ ParamsKey GatherNDKernelRef::GetSupportedKey() const {
k.EnableTensorPitches();
k.EnableBatching();
k.EnableDifferentTypes();
k.EnableDynamicShapesSupport();
return k;
}
@@ -45,7 +46,7 @@ static inline std::vector<std::string> GetDefaultOrder(size_t size) {
return default_order;
}
CommonDispatchData GatherNDKernelRef::SetDefault(const gather_nd_params& params, const optional_params&) const {
CommonDispatchData GatherNDKernelRef::SetDefault(const gather_nd_params& params) const {
CommonDispatchData dispatchData;
auto indices_dims = params.inputs[1].LogicalDims();
@@ -118,7 +119,7 @@ JitConstants GatherNDKernelRef::GetJitConstants(const gather_nd_params& params)
jit.AddConstant(MakeJitConstant("INDICES_RANK", params.indices_rank));
jit.AddConstant(MakeJitConstant("BATCH_DIMS", params.batch_dims));
jit.AddConstant(MakeJitConstant("BATCH_MERGED_OUTPUT", params.batch_merged_output));
jit.AddConstant(MakeJitConstant("WI_SLICE_SIZE", GetSliceSize(params)));
jit.AddConstant(MakeJitConstant("WI_SLICE_SIZE_STATIC", GetSliceSize(params)));
jit.AddConstant(MakeJitConstant("INDICES_LAST_DIM", GetIndicesLastDim(params)));
if (!params.fused_ops.empty()) {
@@ -155,9 +156,11 @@ bool GatherNDKernelRef::Validate(const Params& p, const optional_params& o) cons
return false;
}
for (uint8_t i = 0; i < batch_dims; i++) {
if (input_dims[i] != indices_dims[i]) {
return false;
if (!params.inputs[0].is_dynamic()) {
for (uint8_t i = 0; i < batch_dims; i++) {
if (input_dims[i] != indices_dims[i]) {
return false;
}
}
}
@@ -177,13 +180,35 @@ KernelsData GatherNDKernelRef::GetKernelsData(const Params& params, const option
KernelData kd = KernelData::Default<gather_nd_params>(params);
gather_nd_params& newParams = *static_cast<gather_nd_params*>(kd.params.get());
auto dispatchData = SetDefault(newParams, options);
auto dispatchData = SetDefault(newParams);
auto cldnn_jit = GetJitConstants(newParams);
auto entry_point = GetEntryPoint(kernelName, newParams.layerID, params, options);
auto jit = CreateJit(kernelName, cldnn_jit, entry_point);
auto& kernel = kd.kernels[0];
FillCLKernelData(kernel, dispatchData, params.engineInfo, kernelName, jit, entry_point, "", false, false, 2, GetFusedPrimitiveInputsCount(params));
kd.update_dispatch_data_func = [this](const Params& params, KernelData& kd) {
const auto& prim_params = static_cast<const gather_nd_params&>(params);
auto dispatchData = SetDefault(prim_params);
OPENVINO_ASSERT(kd.kernels.size() == 1, "[GPU] Invalid kernels size for update dispatch data func");
kd.kernels[0].params.workGroups.global = dispatchData.gws;
kd.kernels[0].params.workGroups.local = dispatchData.lws;
kd.kernels[0].skip_execution = KernelData::SkipKernelExecution(prim_params);
};
FillCLKernelData(kernel,
dispatchData,
params.engineInfo,
kernelName,
jit,
entry_point,
"",
false,
false,
2,
GetFusedPrimitiveInputsCount(params),
1,
newParams.has_dynamic_tensors());
return { kd };
}
src/plugins/intel_gpu/src/kernel_selector/kernels/gather/gather_nd_kernel_ref.h
+1 -1
View File
@@ -32,7 +32,7 @@ public:
GatherNDKernelRef() : KernelBaseOpenCL("gather_nd_ref") {}
virtual ~GatherNDKernelRef() {}
virtual JitConstants GetJitConstants(const gather_nd_params& params) const;
virtual CommonDispatchData SetDefault(const gather_nd_params& params, const optional_params&) const;
virtual CommonDispatchData SetDefault(const gather_nd_params& params) const;
KernelsData GetKernelsData(const Params& params, const optional_params& options) const override;
ParamsKey GetSupportedKey() const override;
std::vector<FusedOpType> GetSupportedFusedOps() const override {
src/plugins/intel_gpu/tests/functional/single_layer_tests/dynamic/gather_nd.cpp
+163
View File
@@ -0,0 +1,163 @@
// Copyright (C) 2023 Intel Corporation
// SPDX-License-Identifier: Apache-2.0
//
#include "shared_test_classes/single_layer/gather_nd.hpp"
#include "shared_test_classes/base/ov_subgraph.hpp"
#include "ie_precision.hpp"
#include "ngraph_functions/builders.hpp"
#include "common_test_utils/ov_tensor_utils.hpp"
#include <string>
using namespace ngraph;
using namespace InferenceEngine;
using namespace ov::test;
namespace GPULayerTestsDefinitions {
struct GatherNDShapeParams {
InputShape inputShapes;
InputShape targetShapes;
int batch_dims;
};
typedef std::tuple<
GatherNDShapeParams,
ElementType, // Network precision
bool // Is const Indices
> GatherNDGPUTestParams;
class GatherNDGPUTest : public testing::WithParamInterface<GatherNDGPUTestParams>,
virtual public ov::test::SubgraphBaseTest {
public:
static std::string getTestCaseName(testing::TestParamInfo<GatherNDGPUTestParams> obj) {
GatherNDShapeParams Shapes;
ElementType netPrecision;
bool isIndicesConstant;
std::tie(Shapes, netPrecision, isIndicesConstant) = obj.param;
std::ostringstream result;
result << "IS=(";
result << CommonTestUtils::partialShape2str({Shapes.inputShapes.first}) << "_";
for (size_t i = 0lu; i < Shapes.inputShapes.second.size(); i++) {
result << "{";
result << CommonTestUtils::vec2str(Shapes.inputShapes.second[i]) << "_";
result << "}_";
}
result << "TS=(";
result << CommonTestUtils::partialShape2str({Shapes.targetShapes.first}) << "_";
for (size_t i = 0lu; i < Shapes.targetShapes.second.size(); i++) {
result << "{";
result << CommonTestUtils::vec2str(Shapes.targetShapes.second[i]) << "_";
result << "}_";
}
result << "batchDims=" << Shapes.batch_dims << "_";
result << "netPrc=" << netPrecision << "_";
result << "constIdx=" << (isIndicesConstant ? "True" : "False") << "_";
return result.str();
}
protected:
void SetUp() override {
GatherNDShapeParams Shapes;
ElementType netPrecision;
bool isIndicesConstant;
const ElementType intInputsPrecision = ElementType::i32;
std::tie(Shapes, netPrecision, isIndicesConstant) = this->GetParam();
const int batchDims = Shapes.batch_dims;
targetDevice = CommonTestUtils::DEVICE_GPU;
std::shared_ptr<ov::Node> indicesNode;
std::shared_ptr<ov::Node> gather_ndNode;
if (isIndicesConstant) {
init_input_shapes({Shapes.inputShapes});
} else { // Not being tested because currently parameter targetshape is not supported
init_input_shapes({Shapes.inputShapes, Shapes.targetShapes});
}
ngraph::ParameterVector params{std::make_shared<ov::op::v0::Parameter>(netPrecision, inputDynamicShapes[0])};
params.back()->set_friendly_name("data");
if (isIndicesConstant) {
int64_t idx_range = INT64_MAX;
for (size_t i = 0; i < Shapes.inputShapes.second.size(); ++i) {
for (size_t j = 0; j < Shapes.inputShapes.second[i].size(); ++j) {
idx_range = std::min(static_cast<int64_t>(Shapes.inputShapes.second[i][j]), idx_range);
}
}
indicesNode = ngraph::builder::makeConstant<int64_t>(
ngraph::element::i64,
Shapes.targetShapes.second[0],
{},
true,
idx_range - 1,
0);
} else {
params.push_back(std::make_shared<ov::op::v0::Parameter>(intInputsPrecision, inputDynamicShapes[1]));
params.back()->set_friendly_name("indices");
}
auto paramOuts =
ngraph::helpers::convert2OutputVector(ngraph::helpers::castOps2Nodes<ov::op::v0::Parameter>(params));
gather_ndNode = std::make_shared<ov::op::v8::GatherND>(paramOuts[0],
isIndicesConstant ? indicesNode : paramOuts[1],
batchDims);
ngraph::ResultVector results{std::make_shared<ngraph::opset4::Result>(gather_ndNode)};
function = std::make_shared<ngraph::Function>(results, params, "GatherND");
}
};
TEST_P(GatherNDGPUTest, CompareWithRefs) {
SKIP_IF_CURRENT_TEST_IS_DISABLED()
run();
}
namespace {
const std::vector<ov::element::Type_t> netPrecisions = {
ov::element::f32,
ov::element::f16,
ov::element::i32
};
const std::vector<GatherNDShapeParams> dynamicInputShapeConstTargetShape = {
{
ov::test::InputShape(ov::PartialShape({-1, -1}), {{2, 2}}),
ov::test::InputShape(ov::PartialShape({}), {{2, 1}}),
1
},
{
ov::test::InputShape(ov::PartialShape({-1, -1, -1}), {{2, 3, 4}, {3, 4, 5}}),
ov::test::InputShape(ov::PartialShape({}), {{2, 1}}),
0
},
{
ov::test::InputShape(ov::PartialShape({-1, -1, -1, -1}), {{2, 7, 8, 9}, {2, 7, 4, 8}}),
ov::test::InputShape(ov::PartialShape({}), {{2, 1}}),
1
},
{
ov::test::InputShape(ov::PartialShape({-1, -1, -1, -1, -1}), {{15, 12, 20, 15, 2}, {15, 12, 18, 7, 17}}),
ov::test::InputShape(ov::PartialShape({}), {{15, 12, 5, 9, 1, 3}}),
2
},
{
ov::test::InputShape(ov::PartialShape({-1, -1, -1, -1, -1, -1}), {{4, 3, 2, 5, 5, 2}, {4, 3, 2, 5, 7, 2}}),
ov::test::InputShape(ov::PartialShape({}), {{4, 3, 2, 5, 6, 2}}),
4
}
};
INSTANTIATE_TEST_SUITE_P(smoke_dynamic_input_shapes_const_target_shapes, GatherNDGPUTest,
::testing::Combine(
::testing::ValuesIn(dynamicInputShapeConstTargetShape), // input shapes
::testing::ValuesIn(netPrecisions), // network precision
::testing::Values(true)), // is const indices
GatherNDGPUTest::getTestCaseName);
} // namespace
} // namespace GPULayerTestsDefinitions
src/plugins/intel_gpu/tests/unit/test_cases/gather_nd_gpu_test.cpp
+173
View File
@@ -7,6 +7,7 @@
#include <intel_gpu/primitives/input_layout.hpp>
#include <intel_gpu/primitives/gather_nd.hpp>
#include "gather_nd_inst.h"
using namespace cldnn;
using namespace ::tests;
@@ -780,3 +781,175 @@ TEST(gather_nd_gpu_fp16, export_import) {
DoTestV5(engine,input0, input1, expected_results, indices_rank, batch_dims, format::bfyx, { 3, 1, 1, 1 }, true);
DoTestV8(engine, input0, input1, expected_results, indices_rank, batch_dims, format::bfyx, { 3, 1, 1, 1 }, true);
}
TEST(gather_nd_gpu_fp16, dynamic_r4) {
auto& engine = get_test_engine();
ov::Shape in1_shape = {3, 3, 2, 3};
ov::Shape in2_shape = {3, 2, 2, 3};
const int batch_dims = 0;
auto in1_layout = layout{ov::PartialShape::dynamic(in1_shape.size()), data_types::f16, format::bfyx};
auto in2_layout = layout{ov::PartialShape(in2_shape), data_types::f16, format::bfyx};
auto input1 = engine.allocate_memory(layout{ov::PartialShape(in1_shape), data_types::f16, format::bfyx}); // data
auto input2 = engine.allocate_memory(layout{ov::PartialShape(in2_shape), data_types::f16, format::bfyx}); // Indexes
set_values(input1, {
FLOAT16(1), FLOAT16(2), FLOAT16(3), FLOAT16(4), FLOAT16(5), FLOAT16(6),
FLOAT16(7), FLOAT16(8), FLOAT16(9), FLOAT16(10), FLOAT16(11), FLOAT16(12),
FLOAT16(13), FLOAT16(14), FLOAT16(15), FLOAT16(16), FLOAT16(17), FLOAT16(18),
FLOAT16(19), FLOAT16(20), FLOAT16(21), FLOAT16(22), FLOAT16(23), FLOAT16(24),
FLOAT16(25), FLOAT16(26), FLOAT16(27), FLOAT16(28), FLOAT16(29), FLOAT16(30),
FLOAT16(31), FLOAT16(32), FLOAT16(33), FLOAT16(34), FLOAT16(35), FLOAT16(36),
FLOAT16(41), FLOAT16(42), FLOAT16(43), FLOAT16(44), FLOAT16(45), FLOAT16(46),
FLOAT16(51), FLOAT16(52), FLOAT16(53), FLOAT16(54), FLOAT16(55), FLOAT16(56),
FLOAT16(61), FLOAT16(62), FLOAT16(63), FLOAT16(64), FLOAT16(65), FLOAT16(66),
});
set_values(input2, {
FLOAT16(2), FLOAT16(0), FLOAT16(0), FLOAT16(2), FLOAT16(2), FLOAT16(0),
FLOAT16(1), FLOAT16(0), FLOAT16(0), FLOAT16(1), FLOAT16(1), FLOAT16(0),
FLOAT16(1), FLOAT16(0), FLOAT16(1), FLOAT16(1), FLOAT16(1), FLOAT16(1),
FLOAT16(2), FLOAT16(0), FLOAT16(0), FLOAT16(2), FLOAT16(1), FLOAT16(0),
FLOAT16(1), FLOAT16(1), FLOAT16(1), FLOAT16(0), FLOAT16(1), FLOAT16(1),
FLOAT16(1), FLOAT16(2), FLOAT16(1), FLOAT16(0), FLOAT16(2), FLOAT16(1),
});
std::vector<float> expected_results = {
FLOAT16(41), FLOAT16(42), FLOAT16(43), FLOAT16(61), FLOAT16(62), FLOAT16(63),
FLOAT16(19), FLOAT16(20), FLOAT16(21), FLOAT16(25), FLOAT16(26), FLOAT16(27),
FLOAT16(22), FLOAT16(23), FLOAT16(24), FLOAT16(28), FLOAT16(29), FLOAT16(30),
FLOAT16(41), FLOAT16(42), FLOAT16(43), FLOAT16(51), FLOAT16(52), FLOAT16(53),
FLOAT16(28), FLOAT16(29), FLOAT16(30), FLOAT16(10), FLOAT16(11), FLOAT16(12),
FLOAT16(34), FLOAT16(35), FLOAT16(36), FLOAT16(16), FLOAT16(17), FLOAT16(18),
};
auto expected_fmt = format::bfyx;
const tensor expected_ts = { 3, 2, 3, 2 };
uint8_t dim_size = 4;
topology topology;
topology.add(input_layout("input1", in1_layout));
topology.add(input_layout("input2", in2_layout));
topology.add(gather_nd("gather_nd", input_info("input1"), input_info("input2"), static_cast<uint8_t>(in1_shape.size()), static_cast<uint8_t>(in2_shape.size()), batch_dims, true));
ExecutionConfig config = get_test_default_config(engine);
config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
network network(engine, topology, config);
network.set_input_data("input1", input1);
network.set_input_data("input2", input2);
auto inst = network.get_primitive("gather_nd");
auto impl = inst->get_impl();
ASSERT_TRUE(impl != nullptr);
ASSERT_TRUE(impl->is_dynamic());
auto outputs = network.execute();
auto output = outputs.at("gather_nd").get_memory();
// Compare output shape
auto output_format = output->get_layout().format;
auto output_shape = output->get_layout().get_tensor();
ASSERT_EQ(expected_fmt, output_format);
for (int32_t i = 0; i < dim_size; i++)
{
ASSERT_EQ(expected_ts.sizes()[i], output_shape.sizes()[i]);
}
cldnn::mem_lock<uint16_t> output_ptr(output, get_test_stream());
ASSERT_EQ(expected_results.size(), output_ptr.size());
for (size_t i = 0; i < expected_results.size(); ++i) {
EXPECT_EQ(expected_results[i], half_to_float(output_ptr[i])) << i;
}
}
TEST(gather_nd_gpu_fp16, dynamic_r5) {
auto& engine = get_test_engine();
ov::Shape in1_shape = {2, 3, 3, 2, 2};
ov::Shape in2_shape = {2, 3, 3, 2, 1};
const int batch_dims = 4;
auto in1_layout = layout{ov::PartialShape::dynamic(in1_shape.size()), data_types::f16, format::bfzyx};
auto in2_layout = layout{ov::PartialShape(in2_shape), data_types::f16, format::bfzyx};
auto input1 = engine.allocate_memory(layout{ov::PartialShape(in1_shape), data_types::f16, format::bfzyx}); // data
auto input2 = engine.allocate_memory(layout{ov::PartialShape(in2_shape), data_types::f16, format::bfzyx}); // Indexes
set_values(input1, {
FLOAT16(11), FLOAT16(12), FLOAT16(13), FLOAT16(14), FLOAT16(15), FLOAT16(16), FLOAT16(17), FLOAT16(18), FLOAT16(19), FLOAT16(10), FLOAT16(21), FLOAT16(18),
FLOAT16(21), FLOAT16(22), FLOAT16(23), FLOAT16(24), FLOAT16(25), FLOAT16(26), FLOAT16(27), FLOAT16(28), FLOAT16(29), FLOAT16(20), FLOAT16(27), FLOAT16(28),
FLOAT16(31), FLOAT16(32), FLOAT16(33), FLOAT16(34), FLOAT16(35), FLOAT16(36), FLOAT16(37), FLOAT16(38), FLOAT16(39), FLOAT16(30), FLOAT16(31), FLOAT16(30),
FLOAT16(11), FLOAT16(12), FLOAT16(13), FLOAT16(14), FLOAT16(15), FLOAT16(16), FLOAT16(17), FLOAT16(18), FLOAT16(19), FLOAT16(10), FLOAT16(17), FLOAT16(18),
FLOAT16(21), FLOAT16(22), FLOAT16(23), FLOAT16(24), FLOAT16(25), FLOAT16(26), FLOAT16(27), FLOAT16(28), FLOAT16(29), FLOAT16(20), FLOAT16(27), FLOAT16(28),
FLOAT16(31), FLOAT16(32), FLOAT16(33), FLOAT16(34), FLOAT16(35), FLOAT16(36), FLOAT16(37), FLOAT16(38), FLOAT16(39), FLOAT16(30), FLOAT16(29), FLOAT16(30),
});
set_values(input2, {
FLOAT16(1), FLOAT16(1), FLOAT16(1), FLOAT16(1), FLOAT16(1), FLOAT16(1),
FLOAT16(0), FLOAT16(0), FLOAT16(0), FLOAT16(0), FLOAT16(0), FLOAT16(0),
FLOAT16(1), FLOAT16(0), FLOAT16(0), FLOAT16(1), FLOAT16(1), FLOAT16(0),
FLOAT16(1), FLOAT16(1), FLOAT16(1), FLOAT16(1), FLOAT16(1), FLOAT16(1),
FLOAT16(0), FLOAT16(0), FLOAT16(0), FLOAT16(0), FLOAT16(0), FLOAT16(0),
FLOAT16(1), FLOAT16(0), FLOAT16(0), FLOAT16(1), FLOAT16(1), FLOAT16(0),
});
std::vector<float> expected_results = {
FLOAT16(12), FLOAT16(14), FLOAT16(16), FLOAT16(18), FLOAT16(10), FLOAT16(18),
FLOAT16(21), FLOAT16(23), FLOAT16(25), FLOAT16(27), FLOAT16(29), FLOAT16(27),
FLOAT16(32), FLOAT16(33), FLOAT16(35), FLOAT16(38), FLOAT16(30), FLOAT16(31),
FLOAT16(12), FLOAT16(14), FLOAT16(16), FLOAT16(18), FLOAT16(10), FLOAT16(18),
FLOAT16(21), FLOAT16(23), FLOAT16(25), FLOAT16(27), FLOAT16(29), FLOAT16(27),
FLOAT16(32), FLOAT16(33), FLOAT16(35), FLOAT16(38), FLOAT16(30), FLOAT16(29),
};
auto expected_fmt = format::bfyx;
const tensor expected_ts = { 2, 3, 2, 3 };
uint8_t dim_size = 4;
topology topology;
topology.add(input_layout("input1", in1_layout));
topology.add(input_layout("input2", in2_layout));
topology.add(gather_nd("gather_nd", input_info("input1"), input_info("input2"), static_cast<uint8_t>(in1_shape.size()), static_cast<uint8_t>(in2_shape.size()), batch_dims, false));
ExecutionConfig config = get_test_default_config(engine);
config.set_property(ov::intel_gpu::allow_new_shape_infer(true));
network network(engine, topology, config);
network.set_input_data("input1", input1);
network.set_input_data("input2", input2);
auto inst = network.get_primitive("gather_nd");
auto impl = inst->get_impl();
ASSERT_TRUE(impl != nullptr);
ASSERT_TRUE(impl->is_dynamic());
auto outputs = network.execute();
auto output = outputs.at("gather_nd").get_memory();
// Compare output shape
auto output_format = output->get_layout().format;
auto output_shape = output->get_layout().get_tensor();
ASSERT_EQ(expected_fmt, output_format);
for (int32_t i = 0; i < dim_size; i++)
{
ASSERT_EQ(expected_ts.sizes()[i], output_shape.sizes()[i]);
}
cldnn::mem_lock<uint16_t> output_ptr(output, get_test_stream());
ASSERT_EQ(expected_results.size(), output_ptr.size());
for (size_t i = 0; i < expected_results.size(); ++i) {
EXPECT_EQ(expected_results[i], half_to_float(output_ptr[i])) << i;
}
}