[GPU] Use int64_t type for axis in GatherElements (#12317)

This commit is contained in:
Roman Lyamin authored and GitHub committed 2022-08-11 14:07:08 +04:00
1 parent fe4e875586
commit e98cdcf496
5 files changed
+71 -102

No files matched your search

@@ -19,15 +19,6 @@ namespace cldnn {
struct gather_elements : public primitive_base<gather_elements> {
CLDNN_DECLARE_PRIMITIVE(gather_elements)
enum gather_elements_axis {
along_b,
along_f,
along_x,
along_y,
along_z,
along_w
};
/// @brief Constructs gather_elements primitive.
/// @param id This primitive id.
/// @param data Input data primitive id.
@@ -40,7 +31,7 @@ struct gather_elements : public primitive_base<gather_elements> {
const primitive_id& indices,
const format& output_format,
const tensor& output_shape,
const gather_elements_axis axis,
const int64_t axis,
const primitive_id& ext_prim_id = "",
const padding& output_padding = padding())
: primitive_base(id, {data, indices}, ext_prim_id, output_padding), output_format(output_format), output_shape(output_shape), axis(axis) {}
@@ -51,7 +42,7 @@ struct gather_elements : public primitive_base<gather_elements> {
tensor output_shape;
/// @brief Which axis to gather on.
gather_elements_axis axis;
int64_t axis;
};
/// @}
/// @}
@@ -14,22 +14,35 @@ using namespace cldnn;
namespace cldnn {
namespace ocl {
kernel_selector::gather_elements_axis convert_axis(gather_elements::gather_elements_axis axis) {
static inline kernel_selector::gather_elements_axis convert_axis(int64_t axis, size_t rank) {
if (axis < 0) {
axis += rank;
}
switch (axis) {
case gather_elements::along_x:
return kernel_selector::gather_elements_axis::X;
case gather_elements::along_y:
return kernel_selector::gather_elements_axis::Y;
case gather_elements::along_z:
return kernel_selector::gather_elements_axis::Z;
case gather_elements::along_w:
return kernel_selector::gather_elements_axis::W;
case gather_elements::along_f:
return kernel_selector::gather_elements_axis::FEATURE;
case gather_elements::along_b:
return kernel_selector::gather_elements_axis::BATCH;
default:
return kernel_selector::gather_elements_axis::BATCH;
case 0: return kernel_selector::gather_elements_axis::BATCH;
case 1: return kernel_selector::gather_elements_axis::FEATURE;
case 2:
if (rank == 6)
return kernel_selector::gather_elements_axis::W;
else if (rank == 5)
return kernel_selector::gather_elements_axis::Z;
else
return kernel_selector::gather_elements_axis::Y;
case 3:
if (rank == 6)
return kernel_selector::gather_elements_axis::Z;
else if (rank == 5)
return kernel_selector::gather_elements_axis::Y;
else
return kernel_selector::gather_elements_axis::X;
case 4:
if (rank == 6)
return kernel_selector::gather_elements_axis::Y;
else
return kernel_selector::gather_elements_axis::X;
case 5: return kernel_selector::gather_elements_axis::X;
default: IE_THROW() << "Incorrect gather_elements axis.";
}
}
@@ -48,7 +61,8 @@ public:
auto gather_elements_optional_params =
get_default_optional_params<kernel_selector::gather_elements_optional_params>(arg.get_program());
gather_elements_params.axis = convert_axis(prim->axis);
size_t rank = arg.get_output_layout().get_rank();
gather_elements_params.axis = convert_axis(prim->axis, rank);
gather_elements_params.inputs.push_back(convert_data_tensor(impl_param.input_layouts[1]));
@@ -13,50 +13,26 @@
namespace ov {
namespace intel_gpu {
static cldnn::gather_elements::gather_elements_axis GetGatherAxis(int axis, unsigned rank) {
if (axis < 0)
axis += rank;
if (axis < 0 || axis >= static_cast<int32_t>(rank))
IE_THROW() << "GatherElements axis is not correspond to number of dimensions";
// Difference in dimension ordering between IE and GPU plugin,
// reverse spatial dimensions after batch and feature.
unsigned cldnn_axis = axis;
if (axis >= 2) {
auto spatial_axis = axis - 2;
// Default and minimum number of dimensions is 4
auto spatial_size = std::max(rank, 4u) - 2;
cldnn_axis = spatial_size - spatial_axis - 1 + 2;
}
switch (cldnn_axis) {
case 0: return cldnn::gather_elements::gather_elements_axis::along_b;
case 1: return cldnn::gather_elements::gather_elements_axis::along_f;
case 2: return cldnn::gather_elements::gather_elements_axis::along_x;
case 3: return cldnn::gather_elements::gather_elements_axis::along_y;
case 4: return cldnn::gather_elements::gather_elements_axis::along_z;
case 5: return cldnn::gather_elements::gather_elements_axis::along_w;
default: IE_THROW() << "Unsupported GatherElements axis: " << axis;
}
return cldnn::gather_elements::gather_elements_axis::along_f; // shouldn't get here
}
static void CreateGatherElementsOp(Program& p, const std::shared_ptr<ngraph::op::v6::GatherElements>& op) {
p.ValidateInputs(op, {2});
auto inputPrimitives = p.GetInputPrimitiveIDs(op);
std::string layerName = layer_type_name_ID(op);
size_t rank = op->get_input_shape(0).size();
int32_t axis = static_cast<int32_t>(op->get_axis());
auto outLayout = cldnn::format::get_default_format(op->get_output_shape(0).size());
size_t rank = op->get_input_shape(0).size();
int64_t axis = op->get_axis();
if (axis < 0)
axis += rank;
OPENVINO_ASSERT(axis >= 0 && axis < static_cast<int64_t>(rank),
"GatherElements axis is not correspond to number of dimensions");
auto primitive = cldnn::gather_elements(layerName,
inputPrimitives[0],
inputPrimitives[1],
outLayout,
tensor_from_dims(op->get_output_shape(0)),
GetGatherAxis(axis, rank),
axis,
op->get_friendly_name());
p.AddPrimitive(primitive);
@@ -29,7 +29,7 @@ struct gather_elements_test_params {
format output_format;
tensor output_shape;
cldnn::gather_elements::gather_elements_axis axis;
int64_t axis;
data_types default_type;
format default_format;
@@ -50,22 +50,10 @@ public:
}
size_t get_axis_dim(gather_elements_test_params& p) {
switch (p.axis) {
case cldnn::gather_elements::gather_elements_axis::along_x:
return p.input_shape.spatial[0];
case cldnn::gather_elements::gather_elements_axis::along_y:
return p.input_shape.spatial[1];
case cldnn::gather_elements::gather_elements_axis::along_z:
return p.input_shape.spatial[2];
case cldnn::gather_elements::gather_elements_axis::along_w:
return p.input_shape.spatial[3];
case cldnn::gather_elements::gather_elements_axis::along_f:
return p.input_shape.feature[0];
case cldnn::gather_elements::gather_elements_axis::along_b:
return p.input_shape.batch[0];
default:
return 1;
}
size_t rank = p.input_format.dimension();
auto default_fmt = format::get_default_format(rank);
auto dims = p.input_shape.sizes(default_fmt);
return dims[p.axis];
}
layout get_input_layout(gather_elements_test_params& p) {
@@ -89,28 +77,28 @@ public:
/* ----------------------------------------------------------------------------------------------------- */
/* ------------------------------------------ GatherElements cases ------------------------------------- */
/* ----------------------------------------------------------------------------------------------------- */
#define CASE_GATHER_ELEMENTS_FP16_4D_1 data_types::f16, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, cldnn::gather_elements::gather_elements_axis::along_y, data_types::f16, format::bfyx
#define CASE_GATHER_ELEMENTS_FP16_4D_2 data_types::f16, format::bfyx, { 3, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, cldnn::gather_elements::gather_elements_axis::along_b, data_types::f16, format::bfyx
#define CASE_GATHER_ELEMENTS_FP16_4D_3 data_types::f16, format::bfyx, { 1, 3, 2, 9 }, format::bfyx, { 1, 3, 5, 9 }, format::bfyx, { 1, 3, 5, 9 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f16, format::bfyx
#define CASE_GATHER_ELEMENTS_FP16_4D_1 data_types::f16, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, 2, data_types::f16, format::bfyx
#define CASE_GATHER_ELEMENTS_FP16_4D_2 data_types::f16, format::bfyx, { 3, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, 0, data_types::f16, format::bfyx
#define CASE_GATHER_ELEMENTS_FP16_4D_3 data_types::f16, format::bfyx, { 1, 3, 2, 9 }, format::bfyx, { 1, 3, 5, 9 }, format::bfyx, { 1, 3, 5, 9 }, 3, data_types::f16, format::bfyx
#define CASE_GATHER_ELEMENTS_FP16_5D_1 data_types::f16, format::bfzyx, { 3, 2, 5, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f16, format::bfzyx
#define CASE_GATHER_ELEMENTS_FP16_5D_2 data_types::f16, format::bfzyx, { 5, 4, 7, 4, 4 }, format::bfzyx, { 5, 4, 7, 4, 3 }, format::bfzyx, { 5, 4, 7, 4, 3 }, cldnn::gather_elements::gather_elements_axis::along_z, data_types::f16, format::bfzyx
#define CASE_GATHER_ELEMENTS_FP16_5D_1 data_types::f16, format::bfzyx, { 3, 2, 5, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, 4, data_types::f16, format::bfzyx
#define CASE_GATHER_ELEMENTS_FP16_5D_2 data_types::f16, format::bfzyx, { 5, 4, 7, 4, 4 }, format::bfzyx, { 5, 4, 7, 4, 3 }, format::bfzyx, { 5, 4, 7, 4, 3 }, 2, data_types::f16, format::bfzyx
#define CASE_GATHER_ELEMENTS_FP16_6D_1 data_types::f16, format::bfwzyx, { 5, 4, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, cldnn::gather_elements::gather_elements_axis::along_f, data_types::f16, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP16_6D_2 data_types::f16, format::bfwzyx, { 2, 1, 2, 3, 2, 1 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, cldnn::gather_elements::gather_elements_axis::along_w, data_types::f16, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP16_6D_3 data_types::f16, format::bfwzyx, { 2, 2, 3, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f16, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP16_6D_1 data_types::f16, format::bfwzyx, { 5, 4, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, 1, data_types::f16, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP16_6D_2 data_types::f16, format::bfwzyx, { 2, 1, 2, 3, 2, 1 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, 2, data_types::f16, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP16_6D_3 data_types::f16, format::bfwzyx, { 2, 2, 3, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, 5, data_types::f16, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP32_4D_1 data_types::f32, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, cldnn::gather_elements::gather_elements_axis::along_y, data_types::f32, format::bfyx
#define CASE_GATHER_ELEMENTS_FP32_4D_2 data_types::f32, format::bfyx, { 3, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, cldnn::gather_elements::gather_elements_axis::along_b, data_types::f32, format::bfyx
#define CASE_GATHER_ELEMENTS_FP32_4D_3 data_types::f32, format::bfyx, { 1, 3, 2, 9 }, format::bfyx, { 1, 3, 5, 9 }, format::bfyx, { 1, 3, 5, 9 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f32, format::bfyx
#define CASE_GATHER_ELEMENTS_FP32_4D_1 data_types::f32, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, format::bfyx, { 3, 7, 9, 8 }, 2, data_types::f32, format::bfyx
#define CASE_GATHER_ELEMENTS_FP32_4D_2 data_types::f32, format::bfyx, { 3, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, format::bfyx, { 2, 2, 8, 3 }, 0, data_types::f32, format::bfyx
#define CASE_GATHER_ELEMENTS_FP32_4D_3 data_types::f32, format::bfyx, { 1, 3, 2, 9 }, format::bfyx, { 1, 3, 5, 9 }, format::bfyx, { 1, 3, 5, 9 }, 3, data_types::f32, format::bfyx
#define CASE_GATHER_ELEMENTS_FP32_5D_1 data_types::f32, format::bfzyx, { 3, 2, 5, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f32, format::bfzyx
#define CASE_GATHER_ELEMENTS_FP32_5D_2 data_types::f32, format::bfzyx, { 5, 4, 7, 4, 4 }, format::bfzyx, { 5, 4, 7, 4, 3 }, format::bfzyx, { 5, 4, 7, 4, 3 }, cldnn::gather_elements::gather_elements_axis::along_z, data_types::f32, format::bfzyx
#define CASE_GATHER_ELEMENTS_FP32_5D_1 data_types::f32, format::bfzyx, { 3, 2, 5, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, format::bfzyx, { 3, 2, 2, 2, 3 }, 4, data_types::f32, format::bfzyx
#define CASE_GATHER_ELEMENTS_FP32_5D_2 data_types::f32, format::bfzyx, { 5, 4, 7, 4, 4 }, format::bfzyx, { 5, 4, 7, 4, 3 }, format::bfzyx, { 5, 4, 7, 4, 3 }, 2, data_types::f32, format::bfzyx
#define CASE_GATHER_ELEMENTS_FP32_6D_1 data_types::f32, format::bfwzyx, { 5, 4, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, cldnn::gather_elements::gather_elements_axis::along_f, data_types::f32, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP32_6D_2 data_types::f32, format::bfwzyx, { 2, 1, 2, 3, 2, 1 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, cldnn::gather_elements::gather_elements_axis::along_w, data_types::f32, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP32_6D_3 data_types::f32, format::bfwzyx, { 2, 2, 3, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, cldnn::gather_elements::gather_elements_axis::along_x, data_types::f32, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP32_6D_1 data_types::f32, format::bfwzyx, { 5, 4, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, format::bfwzyx, { 5, 2, 6, 7, 8, 2 }, 1, data_types::f32, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP32_6D_2 data_types::f32, format::bfwzyx, { 2, 1, 2, 3, 2, 1 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, format::bfwzyx, { 2, 1, 2, 3, 2, 3 }, 2, data_types::f32, format::bfwzyx
#define CASE_GATHER_ELEMENTS_FP32_6D_3 data_types::f32, format::bfwzyx, { 2, 2, 3, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, format::bfwzyx, { 2, 2, 6, 4, 4, 2 }, 5, data_types::f32, format::bfwzyx
class gather_elements_quantize : public GatherElementsPrimitiveFusingTest {};
TEST_P(gather_elements_quantize, basic) {
@@ -21,7 +21,7 @@ inline void DoTest(engine& engine,
const cldnn::memory::ptr& input1, // indices
const std::vector<float>& expected_results,
const tensor& output_tensor,
const cldnn::gather_elements::gather_elements_axis axis) {
const int64_t axis) {
topology topology;
topology.add(input_layout("InputData", input0->get_layout()));
topology.add(input_layout("InputIndices", input1->get_layout()));
@@ -45,7 +45,7 @@ inline void DoTest(engine& engine,
TEST(gather_elements_gpu_fp16, d3283_i2283_a0) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_b;
auto axis = 0;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfyx, { 3, 2, 8, 3 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfyx, { 2, 2, 8, 3 } }); // indices
@@ -106,7 +106,7 @@ TEST(gather_elements_gpu_fp16, d3283_i2283_a0) {
TEST(gather_elements_gpu_fp16, d2235_i2235_a3) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_x;
auto axis = 3;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfyx, { 2, 2, 3, 5 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfyx, { 2, 2, 3, 5 } }); // indices
set_values(input0, {
@@ -184,7 +184,7 @@ TEST(gather_elements_gpu_fp16, d2235_i2235_a3) {
TEST(gather_elements_gpu_fp16, d1329_i1359_an1) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_x;
auto axis = 3;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfyx, { 1, 3, 2, 9 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfyx, { 1, 3, 5, 9 } }); // indices
set_values(input0, {
@@ -283,7 +283,7 @@ TEST(gather_elements_gpu_fp16, d1329_i1359_an1) {
TEST(gather_elements_gpu_fp16, d12853_i12923_a3) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_y;
auto axis = 3;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 1, 2, 8, 5, 3 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 1, 2, 8, 2, 3 } }); // indices
@@ -356,7 +356,7 @@ TEST(gather_elements_gpu_fp16, d12853_i12923_a3) {
TEST(gather_elements_gpu_fp16, d25441_i22441_an4) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_f;
auto axis = 1;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 2, 5, 4, 4, 1 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 2, 2, 4, 4, 1 } }); // indices
@@ -448,7 +448,7 @@ TEST(gather_elements_gpu_fp16, d25441_i22441_an4) {
TEST(gather_elements_gpu_fp16, d32843_i12843_a0) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_b;
auto axis = 0;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 3, 2, 8, 4, 3 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfzyx, { 1, 2, 8, 4, 3 } }); // indices
@@ -588,7 +588,7 @@ TEST(gather_elements_gpu_fp16, d32843_i12843_a0) {
TEST(gather_elements_gpu_fp16, d223442_i226442_a5) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_x;
auto axis = 5;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 2, 2, 3, 4, 4, 2 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 2, 2, 6, 4, 4, 2 } }); // indices
@@ -991,7 +991,7 @@ TEST(gather_elements_gpu_fp16, d223442_i226442_a5) {
TEST(gather_elements_gpu_fp16, d124251_i124221_an3) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_z;
auto axis = 3;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 1, 2, 4, 2, 5, 1 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 1, 2, 4, 2, 2, 1 } }); // indices
@@ -1046,7 +1046,7 @@ TEST(gather_elements_gpu_fp16, d124251_i124221_an3) {
TEST(gather_elements_gpu_fp16, d233113_i233115_a2) {
auto& engine = get_test_engine();
auto axis = cldnn::gather_elements::gather_elements_axis::along_w;
auto axis = 2;
auto input0 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 2, 3, 3, 1, 1, 3 } }); // data
auto input1 = engine.allocate_memory({ data_types::f16, format::bfwzyx, { 2, 3, 3, 1, 1, 5 } }); // indices