Merge branch 'gather_operator' into extend_gemm_op

src/include/migraphx/operators.hpp

@@ -774,35 +774,12 @@ struct gather
         return {type, lens};
     }
 
-    template <typename V>
-    std::vector<std::size_t> compute_data_index(const V& indices,
-                                                const int axis_index,
-                                                const std::vector<std::size_t>& out_idx) const
-    {
-        auto data_idx = out_idx;
-        std::size_t index{};
-        if(!indices.get_shape().scalar())
-        {
-            auto start_it = data_idx.begin() + axis_index;
-            auto end_it   = data_idx.begin() + axis_index + indices.get_shape().lens().size();
-            std::vector<std::size_t> ind_idx(start_it, end_it);
-            data_idx.erase(start_it, end_it);
-            index = indices(ind_idx.begin(), ind_idx.end());
-        }
-        else
-        {
-            index = indices.front();
-        }
-        data_idx.insert(data_idx.begin() + axis_index, index);
-
-        return data_idx;
-    }
-
     argument compute(const shape& output_shape, std::vector<argument> args) const
     {
         argument result{output_shape};
         // negative axis means counting dimensions from back
-        int axis_index = (axis < 0) ? (args[0].get_shape().lens().size() + axis) : axis;
+        int axis_index =
+            (axis < 0) ? static_cast<int>(args[0].get_shape().lens().size() + axis) : axis;
 
         // max dimension in axis
         visit_all(result, args[0])([&](auto output, auto data) {
@@ -813,9 +790,13 @@ struct gather
                 }
                 else
                 {
-                    shape_for_each(output.get_shape(), [&](const auto& out_idx) {
-                        auto data_idx = compute_data_index(indices, axis_index, out_idx);
-                        output(out_idx.begin(), out_idx.end()) =
+                    auto out_lens        = data.get_shape().lens();
+                    out_lens[axis_index] = indices.get_shape().elements();
+                    migraphx::shape out_comp_shape{data.get_shape().type(), out_lens};
+                    shape_for_each(out_comp_shape, [&](const auto& out_idx) {
+                        auto data_idx        = out_idx;
+                        data_idx[axis_index] = indices[data_idx[axis_index]];
+                        output[out_comp_shape.index(out_idx.begin(), out_idx.end())] =
                             data(data_idx.begin(), data_idx.end());
                     });
                 }

src/targets/gpu/device/gather.cpp

@@ -23,28 +23,19 @@ argument gather(hipStream_t stream,
             const auto* indices_ptr = device_cast(indices.data());
             auto* out_ptr           = device_cast(output.data());
             const auto* in_ptr      = device_cast(input.data());
-            if(output_shape.scalar())
-            {
-                gs_launch(stream, 1)(
-                    [=](auto i) { out_ptr[i] = in_ptr[static_cast<int>(indices_ptr[0])]; });
-            }
-            else
-            {
-                // if indices are a scalar, output has one dim smaller than input
-                auto& input_shape = args[0].get_shape();
-                auto lens         = input_shape.lens();
-                lens[axis_index]  = args[1].get_shape().elements();
-                migraphx::shape out_comp_shape{output_shape.type(), lens};
-                visit_tensor_size(out_comp_shape.lens().size(), [&](auto n_out_dim) {
-                    hip_tensor_descriptor<n_out_dim> desc_input(input_shape);
-                    hip_tensor_descriptor<n_out_dim> desc_output(out_comp_shape);
-                    gs_launch(stream, nelements)([=](auto ii) {
-                        auto in_idx        = desc_output.multi(ii);
-                        in_idx[axis_index] = indices_ptr[in_idx[axis_index]];
-                        out_ptr[ii]        = in_ptr[desc_input.linear(in_idx)];
-                    });
+            auto& input_shape       = args[0].get_shape();
+            auto lens               = input_shape.lens();
+            lens[axis_index]        = args[1].get_shape().elements();
+            migraphx::shape out_comp_shape{output_shape.type(), lens};
+            visit_tensor_size(out_comp_shape.lens().size(), [&](auto n_out_dim) {
+                hip_tensor_descriptor<n_out_dim> desc_input(input_shape);
+                hip_tensor_descriptor<n_out_dim> desc_output(out_comp_shape);
+                gs_launch(stream, nelements)([=](auto ii) {
+                    auto in_idx        = desc_output.multi(ii);
+                    in_idx[axis_index] = indices_ptr[in_idx[axis_index]];
+                    out_ptr[ii]        = in_ptr[desc_input.linear(in_idx)];
                 });
-            }
+            });
         });
     });
 

test/onnx/constant_scalar.onnx

+shape-gather-example:O
+2value"Constant*
+value**

Bconst_tensor 
constantb
+
z
+
+

+
B
\ No newline at end of file

test/onnx/onnx_test.cpp

@@ -521,6 +521,15 @@ TEST_CASE(constant_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(constant_test_scalar)
+{
+    migraphx::program p;
+    p.add_literal(migraphx::literal{migraphx::shape{migraphx::shape::int32_type, {1}, {0}}, {1}});
+    auto prog = migraphx::parse_onnx("constant_scalar.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(constant_fill_test)
 {
     {