Merge branch 'develop' of https://github.com/ROCmSoftwarePlatform/AMDMIGraphX into pad_op

src/include/migraphx/operators.hpp

@@ -671,49 +671,58 @@ struct as_shape
 
 struct gather
 {
-    std::size_t axis = 0;
+    int axis = 0;
     std::string name() const { return "gather"; }
 
     shape compute_shape(std::vector<shape> inputs) const
     {
         check_shapes{inputs, *this}.has(2);
         auto lens = inputs[0].lens();
-        if(axis >= lens.size())
+        int n_dim = static_cast<int>(lens.size());
+        if(axis >= n_dim || axis < -n_dim)
         {
-            MIGRAPHX_THROW("Gather, axis is out of range.");
+            MIGRAPHX_THROW("Gather: axis is out of range.");
         }
+
+        // negative axis means counting dimensions from back
+        int axis_index = (axis < 0) ? (n_dim + axis) : axis;
+
         auto type        = inputs[0].type();
-        lens[axis] = inputs[1].elements();
+        lens[axis_index] = inputs[1].elements();
 
         return {type, lens};
     }
 
     template <class T>
     void compute_index(const T& out_idx,
+                       const int axis_index,
                        const std::vector<std::size_t>& vec_indices,
                        const std::size_t max_dim,
                        T& in_idx) const
     {
         in_idx          = out_idx;
-        std::size_t idx = vec_indices.at(out_idx[axis]);
+        std::size_t idx = vec_indices.at(out_idx[axis_index]);
         if(idx >= max_dim)
         {
             MIGRAPHX_THROW("Gather: indices are out of range in input tensor");
         }
-        in_idx[axis] = idx;
+        in_idx[axis_index] = 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) ? (output_shape.lens().size() + axis) : axis;
+
         // max dimension in axis
-        std::size_t max_dim = args[0].get_shape().lens()[axis];
+        std::size_t max_dim = args[0].get_shape().lens()[axis_index];
         std::vector<std::size_t> vec_indices;
         args[1].visit([&](auto indices) { vec_indices.assign(indices.begin(), indices.end()); });
         visit_all(result, args[0])([&](auto output, auto input) {
             std::vector<std::size_t> in_idx;
             shape_for_each(output.get_shape(), [&](const auto& idx) {
-                this->compute_index(idx, vec_indices, max_dim, in_idx);
+                this->compute_index(idx, axis_index, vec_indices, max_dim, in_idx);
                 output(idx.begin(), idx.end()) = input(in_idx.begin(), in_idx.end());
             });
         });

src/onnx/onnx.cpp

@@ -24,7 +24,8 @@ struct onnx_parser
 {
     using attribute_map = std::unordered_map<std::string, onnx::AttributeProto>;
     using node_map      = std::unordered_map<std::string, onnx::NodeProto>;
-    using op_func = std::function<instruction_ref(attribute_map, std::vector<instruction_ref>)>;
+    using op_func =
+        std::function<std::vector<instruction_ref>(attribute_map, std::vector<instruction_ref>)>;
     node_map nodes;
     std::unordered_map<std::string, instruction_ref> instructions;
     program prog    = program();
@@ -89,6 +90,15 @@ struct onnx_parser
 
     template <class F>
     void add_op(std::string name, F f)
+    {
+        ops.emplace(name, [=](auto&&... xs) {
+            return std::vector<instruction_ref>{f(std::forward<decltype(xs)>(xs)...)};
+        });
+    }
+
+    // Multi output op
+    template <class F>
+    void add_multi_op(std::string name, F f)
     {
         ops.emplace(name, f);
     }
@@ -96,7 +106,7 @@ struct onnx_parser
     template <class F>
     void add_mem_op(std::string name, F f)
     {
-        ops.emplace(name, [=](auto&&... xs) {
+        add_op(name, [=](auto&&... xs) {
             return std::mem_fn(f)(*this, name, std::forward<decltype(xs)>(xs)...);
         });
     }
@@ -104,7 +114,7 @@ struct onnx_parser
     template <class T>
     void add_binary_op(std::string name, T x)
     {
-        ops.emplace(name, [this, x](attribute_map attributes, std::vector<instruction_ref> args) {
+        add_op(name, [this, x](attribute_map attributes, std::vector<instruction_ref> args) {
             if(args.size() != 2)
                 MIGRAPHX_THROW("binary operators should have 2 operands");
             if(contains(attributes, "broadcast") and contains(attributes, "axis"))
@@ -173,7 +183,7 @@ struct onnx_parser
     template <class T>
     void add_generic_op(std::string name, T x)
     {
-        ops.emplace(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
+        add_op(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
             return prog.add_instruction(x, args);
         });
     }
@@ -181,7 +191,7 @@ struct onnx_parser
     template <class T>
     void add_variadic_op(std::string name, T x)
     {
-        ops.emplace(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
+        add_op(name, [this, x](attribute_map, std::vector<instruction_ref> args) {
             return std::accumulate(std::next(args.begin()),
                                    args.end(),
                                    args.front(),
@@ -361,7 +371,7 @@ struct onnx_parser
     instruction_ref
     parse_gather(const std::string&, attribute_map attributes, std::vector<instruction_ref> args)
     {
-        std::size_t axis = 0;
+        int axis = 0;
         if(contains(attributes, "axis"))
         {
             axis = parse_value(attributes.at("axis")).at<int>();
@@ -666,7 +676,7 @@ struct onnx_parser
         }
         else
         {
-            throw std::runtime_error("Failed reading");
+            MIGRAPHX_THROW("Failed reading onnx file.");
         }
     }
 
@@ -696,7 +706,7 @@ struct onnx_parser
         }
         for(auto&& p : nodes)
         {
-            this->parse_node(get_name(p.second));
+            this->parse_node(p.first);
         }
     }
 
@@ -712,23 +722,37 @@ struct onnx_parser
             {
                 if(nodes.count(input) > 0)
                 {
-                    auto&& iname = get_name(nodes.at(input));
-                    assert(name != iname);
-                    this->parse_node(iname);
-                    args.push_back(instructions.at(iname));
+                    assert(name != input);
+                    this->parse_node(input);
+                    args.push_back(instructions.at(input));
                 }
                 else
                 {
                     args.push_back(instructions.at(input));
                 }
             }
+            std::vector<instruction_ref> result;
             if(ops.count(node.op_type()) == 0)
             {
-                instructions[name] = prog.add_instruction(unknown{node.op_type()}, args);
+                result.push_back(prog.add_instruction(unknown{node.op_type()}, args));
             }
             else
             {
-                instructions[name] = ops[node.op_type()](get_attributes(node), args);
+                result = ops[node.op_type()](get_attributes(node), args);
+            }
+            // Even no output nodes produce output in migraphx
+            if(node.output().empty() and result.size() == 1)
+            {
+                instructions[name] = result.front();
+            }
+            else
+            {
+                assert(node.output().size() >= result.size());
+                std::transform(result.begin(),
+                               result.end(),
+                               node.output().begin(),
+                               std::inserter(instructions, instructions.end()),
+                               [](auto&& x, auto&& y) { return std::make_pair(y, x); });
             }
         }
     }
@@ -743,25 +767,24 @@ struct onnx_parser
         return result;
     }
 
-    static std::string get_name(const onnx::NodeProto& node)
-    {
-        if(node.name().empty())
-        {
-            std::string generated = "migraphx_unnamed_node";
-            return std::accumulate(node.output().begin(),
-                                   node.output().end(),
-                                   generated,
-                                   [](auto x, auto y) { return x + "_" + y; });
-        }
-        return node.name();
-    }
-
     static node_map get_nodes(const onnx::GraphProto& graph)
     {
         std::unordered_map<std::string, onnx::NodeProto> result;
+        std::size_t n = 0;
         for(auto&& node : graph.node())
         {
-            result[get_name(node)] = node;
+            if(node.output().empty())
+            {
+                if(node.name().empty())
+                {
+                    result["migraphx_unamed_node_" + std::to_string(n)] = node;
+                    n++;
+                }
+                else
+                {
+                    result[node.name()] = node;
+                }
+            }
             for(auto&& output : node.output())
             {
                 result[output] = node;

src/targets/gpu/device/gather.cpp

@@ -14,8 +14,9 @@ namespace device {
 argument gather(hipStream_t stream,
                 const migraphx::shape& output_shape,
                 std::vector<migraphx::argument> args,
-                std::size_t axis)
+                int axis)
 {
+    int axis_index = (axis < 0) ? (axis + output_shape.lens().size()) : axis;
     visit_all(args.back(), args[0])([&](auto output, auto input) {
         std::size_t nelements = output_shape.elements();
         args[1].visit([&](auto indices) {
@@ -27,7 +28,7 @@ argument gather(hipStream_t stream,
                 hip_tensor_descriptor<ndim> desc_output(output.get_shape());
                 gs_launch(stream, nelements)([=](auto i) {
                     auto lens        = desc_output.multi(i);
-                    lens[axis] = indices_ptr[lens[axis]];
+                    lens[axis_index] = indices_ptr[lens[axis_index]];
                     outptr[i]        = inptr[desc_input.linear(lens)];
                 });
             });

src/targets/gpu/include/migraphx/gpu/device/gather.hpp

@@ -13,7 +13,7 @@ namespace device {
 argument gather(hipStream_t stream,
                 const migraphx::shape& output_shape,
                 std::vector<migraphx::argument> args,
-                std::size_t axis);
+                int axis);
 
 } // namespace device
 } // namespace gpu

test/cpu_ops_test.cpp

@@ -113,7 +113,7 @@ TEST_CASE(gather_test)
         migraphx::shape s_indices{migraphx::shape::int32_type, {1, 2}};
         std::vector<int> indices{0, 2};
         auto a1  = p.add_literal(migraphx::literal{s_indices, indices});
-        std::size_t axis = 0;
+        int axis = 0;
         p.add_instruction(migraphx::op::gather{axis}, a0, a1);
         p.compile(migraphx::cpu::target{});
         auto result = p.eval({});
@@ -133,7 +133,27 @@ TEST_CASE(gather_test)
         migraphx::shape s_indices{migraphx::shape::int32_type, {1, 2}};
         std::vector<int> indices{0, 2};
         auto a1  = p.add_literal(migraphx::literal{s_indices, indices});
-        std::size_t axis = 1;
+        int axis = 1;
+        p.add_instruction(migraphx::op::gather{axis}, a0, a1);
+        p.compile(migraphx::cpu::target{});
+        auto result = p.eval({});
+        std::vector<float> res_data(4 * 5);
+        std::vector<float> golden = {0.5f, 2.5f, 3.5f, 5.5f, 6.5f, 8.5f};
+        result.visit([&](auto output) { res_data.assign(output.begin(), output.end()); });
+        EXPECT(migraphx::verify_range(res_data, golden));
+    }
+
+    {
+        migraphx::program p;
+
+        std::vector<float> data(3 * 3);
+        std::iota(data.begin(), data.end(), 0.5);
+        migraphx::shape s{migraphx::shape::float_type, {3, 3}};
+        auto a0 = p.add_literal(migraphx::literal{s, data});
+        migraphx::shape s_indices{migraphx::shape::int32_type, {1, 2}};
+        std::vector<int> indices{0, 2};
+        auto a1  = p.add_literal(migraphx::literal{s_indices, indices});
+        int axis = -1;
         p.add_instruction(migraphx::op::gather{axis}, a0, a1);
         p.compile(migraphx::cpu::target{});
         auto result = p.eval({});

test/gpu/miopen.cpp

@@ -963,7 +963,23 @@ struct test_gather
         std::vector<int> indices{1, 2, 2, 1};
         auto a0  = p.add_parameter("data", s);
         auto a1  = p.add_literal(migraphx::literal{s_indices, indices});
-        std::size_t axis = 0;
+        int axis = 0;
+        p.add_instruction(migraphx::op::gather{axis}, a0, a1);
+        return p;
+    }
+};
+
+struct test_gather_neg_axis
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        migraphx::shape s{migraphx::shape::float_type, {3, 3}};
+        migraphx::shape s_indices{migraphx::shape::int32_type, {2, 2}};
+        std::vector<int> indices{1, 2, 2, 1};
+        auto a0  = p.add_parameter("data", s);
+        auto a1  = p.add_literal(migraphx::literal{s_indices, indices});
+        int axis = -1;
         p.add_instruction(migraphx::op::gather{axis}, a0, a1);
         return p;
     }
@@ -1111,4 +1127,6 @@ int main()
     verify_program<test_conv_bn_relu_pooling>();
     verify_program<test_conv_bn_relu_pooling2>();
     verify_program<test_slice>();
+    verify_program<test_gather>();
+    verify_program<test_gather_neg_axis>();
 }

test/onnx/onnx_test.cpp

@@ -417,7 +417,7 @@ TEST_CASE(gather_test)
     migraphx::program p;
     auto l0  = p.add_parameter("data", migraphx::shape{migraphx::shape::float_type, {3, 4, 5, 6}});
     auto l1  = p.add_parameter("indices", migraphx::shape{migraphx::shape::int32_type, {2, 3}});
-    std::size_t axis = 1;
+    int axis = 1;
     p.add_instruction(migraphx::op::gather{axis}, l0, l1);
     auto prog = migraphx::parse_onnx("gather_test.onnx");
 
@@ -432,7 +432,7 @@ TEST_CASE(shape_gather_test)
         p.add_literal(migraphx::shape{migraphx::shape::int64_type, {3}}, l0->get_shape().lens());
     migraphx::shape const_shape{migraphx::shape::int32_type, {1}};
     auto l2  = p.add_literal(migraphx::literal{const_shape, {1}});
-    std::size_t axis = 0;
+    int axis = 0;
     p.add_instruction(migraphx::op::gather{axis}, l1, l2);
     auto prog = migraphx::parse_onnx("shape_gather.onnx");
 

test/op_shape_test.cpp

@@ -217,7 +217,7 @@ TEST_CASE(gather)
     {
         migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
         migraphx::shape indices{migraphx::shape::int32_type, {2, 3}};
-        std::size_t axis = 1;
+        int axis = 1;
         expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 6, 4, 5}},
                      migraphx::op::gather{axis},
                      input,
@@ -227,7 +227,24 @@ TEST_CASE(gather)
     {
         migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
         migraphx::shape indices{migraphx::shape::int32_type, {2, 3}};
-        std::size_t axis = 4;
+        int axis = -4;
+        expect_shape(migraphx::shape{migraphx::shape::float_type, {6, 3, 4, 5}},
+                     migraphx::op::gather{axis},
+                     input,
+                     indices);
+    }
+
+    {
+        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
+        migraphx::shape indices{migraphx::shape::int32_type, {2, 3}};
+        int axis = 4;
+        throws_shape(migraphx::op::gather{axis}, input, indices);
+    }
+
+    {
+        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
+        migraphx::shape indices{migraphx::shape::int32_type, {2, 3}};
+        int axis = -5;
         throws_shape(migraphx::op::gather{axis}, input, indices);
     }
 }