Cyphers/26mig [MLIR] (#3716)

* Hotfix for negative axes in unsqueeze op (#3705)

* Hotfix for negative axes in unsqueeze op

* Review fix I

* Enable Gather with negative indices. (#3701)

* Enable Gather with negative indices.

* Address PR feedback.

* Remove GOE from Gather unit tests.

* merge error

* Disable test for MLIR

src/ngraph/frontend/onnx_import/op/unsqueeze.cpp

@@ -17,6 +17,7 @@
 #include "ngraph/op/fused/unsqueeze.hpp"
 #include "ngraph/op/constant.hpp"
 #include "squeeze.hpp"
+#include "utils/common.hpp"
 
 namespace ngraph
 {
@@ -30,8 +31,11 @@ namespace ngraph
                 {
                     auto data = node.get_ng_inputs().at(0);
                     auto axes = node.get_attribute_value<std::vector<std::int64_t>>("axes", {});
+                    const auto expanded_rank = data->get_shape().size() + axes.size();
+                    std::vector<std::size_t> valid_axes =
+                        common::validate_axes(node, axes, expanded_rank);
                     auto axes_node = std::make_shared<ngraph::op::Constant>(
-                        element::i64, Shape{axes.size()}, axes);
+                        element::i64, Shape{valid_axes.size()}, valid_axes);
                     return {std::make_shared<ngraph::op::Unsqueeze>(data, axes_node)};
                 }
 

src/ngraph/runtime/cpu/kernel/gather.hpp

@@ -68,6 +68,7 @@ namespace ngraph
                 {
                     Eigen::array<Eigen::Index, Rank1> in_dims;
                     Eigen::array<Eigen::Index, Rank2> out_dims;
+                    auto axis_length = inputs_shape[axis];
 
                     for (size_t i = 0; i < Rank1; i++)
                     {
@@ -84,6 +85,7 @@ namespace ngraph
                         static_cast<ElementType*>(inputs), in_dims);
 
                     auto indices_ptr = static_cast<IndicesType*>(indices);
+                    IndicesType index_value;
                     auto indices_rank = indices_shape.size();
                     auto outer_loop_num = 1;
                     for (size_t i = 0; i < axis; i++)
@@ -123,7 +125,10 @@ namespace ngraph
                             // at axis
                             in_extents[axis] = 1;
                             // at axis, get the value from indices arg
-                            in_offsets[axis] = indices_ptr[0];
+                            index_value = indices_ptr[0];
+                            // take care of negative indices
+                            in_offsets[axis] =
+                                index_value >= 0 ? index_value : index_value + axis_length;
 
                             // before axis
                             for (size_t r = 0; r < axis; r++)
@@ -195,7 +200,10 @@ namespace ngraph
                             }
                             // at axis, get the value from indices arg
                             int k = i % num_indices;
-                            in_offsets[axis] = indices_ptr[k];
+                            index_value = indices_ptr[k];
+                            // take care of negative indices
+                            in_offsets[axis] =
+                                index_value >= 0 ? index_value : index_value + axis_length;
 
                             // indices_from_indices_arg depends on indices_shape and k.
                             // suppose the inputs has shape {3, 3, 3}, indices has shape {2, 2}, and

src/ngraph/runtime/intelgpu/unit_test.manifest

@@ -64,6 +64,7 @@ gather_2d_indices_axis_1_2d_input
 gather_scalar_indices_no_axis_2d_input
 gather_1d_indices_no_axis_1d_input
 gather_2d_indices_no_axis_2d_input
+gather_2d_negative_and_positive_indices_no_axis_2d_input
 gather_3d_indices_no_axis_2d_input
 gather_4d_indices_no_axis_2d_input
 gemm

src/ngraph/runtime/plaidml/unit_test.manifest

@@ -199,6 +199,7 @@ gather_4d_indices_no_axis_uint8
 gather_4d_indices_no_axis_2d_input
 gather_3d_indices_no_axis_2d_input
 gather_2d_indices_no_axis_2d_input
+gather_2d_negative_and_positive_indices_no_axis_2d_input
 gather_1d_indices_no_axis_1d_input
 gather_scalar_indices_no_axis_2d_input
 gather_2d_indices_axis_1_2d_input

src/ngraph/runtime/reference/gather_nd.hpp

@@ -83,6 +83,8 @@ namespace ngraph
                     for (size_t i = 0; i < slice_rank; i++)
                     {
                         U index = indices[indices_index];
+                        // take care of negative indices
+                        index = index >= 0 ? index : index + params_shape[i];
                         params_start_corner[i] = index;
                         params_end_corner[i] = index + 1;
                         indices_index++;

test/backend/gather.in.cpp

@@ -44,7 +44,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_4d_indices_no_axis_uint8)
     auto P = make_shared<op::Parameter>(element::u8, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -77,7 +77,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_4d_indices_no_axis_2d_input)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -113,7 +113,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_3d_indices_no_axis_2d_input)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -144,7 +144,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_2d_indices_no_axis_2d_input)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -162,6 +162,33 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_2d_indices_no_axis_2d_input)
                                   MIN_FLOAT_TOLERANCE_BITS));
 }
 
+NGRAPH_TEST(${BACKEND_NAME},
+            MLIR_DISABLE_TEST(gather_2d_negative_and_positive_indices_no_axis_2d_input))
+{
+    Shape params_shape{3, 2};
+    Shape indices_shape{2, 2};
+    Shape out_shape{2, 2, 2};
+    auto P = make_shared<op::Parameter>(element::f32, params_shape);
+    auto I = make_shared<op::Parameter>(element::i32, indices_shape);
+    auto G = make_shared<op::Gather>(P, I);
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
+
+    auto backend = runtime::Backend::create("${BACKEND_NAME}");
+
+    // Create some tensors for input/output
+    auto p = backend->create_tensor(element::f32, params_shape);
+    copy_data(p, vector<float>{1.0f, 1.1f, 2.0f, 2.1f, 3.0f, 3.1f});
+    auto i = backend->create_tensor(element::i32, indices_shape);
+    copy_data(i, vector<int32_t>{0, -2, 1, 2});
+    auto result = backend->create_tensor(element::f32, out_shape);
+
+    auto c = backend->compile(f);
+    c->call_with_validate({result}, {p, i});
+    EXPECT_TRUE(test::all_close_f((vector<float>{1.0f, 1.1f, 2.0f, 2.1f, 2.0f, 2.1f, 3.0f, 3.1f}),
+                                  read_vector<float>(result),
+                                  MIN_FLOAT_TOLERANCE_BITS));
+}
+
 NGRAPH_TEST(${BACKEND_NAME}, gather_1d_indices_no_axis_1d_input)
 {
     Shape params_shape{3};
@@ -170,7 +197,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_1d_indices_no_axis_1d_input)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -195,7 +222,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_scalar_indices_no_axis_2d_input)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -220,7 +247,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_2d_indices_axis_1_2d_input)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I, 1);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -246,7 +273,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_1d_indices_axis_2_4d_input)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I, 2);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -277,7 +304,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_scalar_indices_axis_1_2d_input)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I, 1);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -302,7 +329,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_single_indices)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -327,7 +354,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_scalar_from_2d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -352,7 +379,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_1d_from_2d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -378,7 +405,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_scalar_from_3d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -403,7 +430,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_1d_from_3d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -429,7 +456,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_2d_from_3d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -455,7 +482,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_scalar_from_2d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -480,7 +507,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_1d_from_2d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -506,7 +533,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_scalar_from_3d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -532,7 +559,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_1d_from_3d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -558,7 +585,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_nd_batch_2d_from_3d)
     auto P = make_shared<op::Parameter>(element::f32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::GatherND>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -584,7 +611,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_int8)
     auto P = make_shared<op::Parameter>(element::i8, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -610,7 +637,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_int16)
     auto P = make_shared<op::Parameter>(element::i16, params_shape);
     auto I = make_shared<op::Parameter>(element::i64, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -636,7 +663,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_int32)
     auto P = make_shared<op::Parameter>(element::i32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -662,7 +689,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_int64)
     auto P = make_shared<op::Parameter>(element::i64, params_shape);
     auto I = make_shared<op::Parameter>(element::i64, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -688,7 +715,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_uint8)
     auto P = make_shared<op::Parameter>(element::u8, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -714,7 +741,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_uint16)
     auto P = make_shared<op::Parameter>(element::u16, params_shape);
     auto I = make_shared<op::Parameter>(element::i64, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -740,7 +767,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_uint32)
     auto P = make_shared<op::Parameter>(element::u32, params_shape);
     auto I = make_shared<op::Parameter>(element::i32, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -766,7 +793,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_uint64)
     auto P = make_shared<op::Parameter>(element::u64, params_shape);
     auto I = make_shared<op::Parameter>(element::i64, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");
 
@@ -792,7 +819,7 @@ NGRAPH_TEST(${BACKEND_NAME}, gather_no_axis_bool)
     auto P = make_shared<op::Parameter>(element::boolean, params_shape);
     auto I = make_shared<op::Parameter>(element::i64, indices_shape);
     auto G = make_shared<op::Gather>(P, I);
-    auto f = make_shared<Function>(make_shared<op::GetOutputElement>(G, 0), ParameterVector{P, I});
+    auto f = make_shared<Function>(G, ParameterVector{P, I});
 
     auto backend = runtime::Backend::create("${BACKEND_NAME}");