Merge pull request #292 from ROCmSoftwarePlatform/argmax_min

onnx operators argmax/argmin

test/gpu/miopen.cpp

@@ -615,6 +615,30 @@ template struct test_softmax<1, migraphx::shape::half_type>;
 template struct test_softmax<2, migraphx::shape::half_type>;
 template struct test_softmax<3, migraphx::shape::half_type>;
 
+template <class T, int Axis>
+struct test_arg_ops : verify_program<test_arg_ops<T, Axis>>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 1025}};
+        auto param = p.add_parameter("data", s);
+        p.add_instruction(T{Axis}, param);
+
+        return p;
+    }
+};
+
+template struct test_arg_ops<migraphx::op::argmax, 0>;
+template struct test_arg_ops<migraphx::op::argmax, 1>;
+template struct test_arg_ops<migraphx::op::argmax, 2>;
+template struct test_arg_ops<migraphx::op::argmax, 3>;
+
+template struct test_arg_ops<migraphx::op::argmin, 0>;
+template struct test_arg_ops<migraphx::op::argmin, 1>;
+template struct test_arg_ops<migraphx::op::argmin, 2>;
+template struct test_arg_ops<migraphx::op::argmin, 3>;
+
 struct test_conv : verify_program<test_conv>
 {
     migraphx::program create_program() const
@@ -3348,22 +3372,6 @@ struct test_lstm_bidirct_default_actv2 : verify_program<test_lstm_bidirct_defaul
     }
 };
 
-template <int Axis>
-struct test_logsoftmax_1 : verify_program<test_logsoftmax_1<Axis>>
-{
-    migraphx::program create_program() const
-    {
-        migraphx::program p;
-        migraphx::shape s{migraphx::shape::float_type, {3}};
-        auto param = p.add_parameter("0", s);
-        p.add_instruction(migraphx::op::logsoftmax{Axis}, param);
-
-        return p;
-    }
-};
-
-template struct test_logsoftmax_1<0>;
-
 template <int Axis, migraphx::shape::type_t T>
 struct test_logsoftmax : verify_program<test_logsoftmax<Axis, T>>
 {

test/onnx/onnx_test.cpp

@@ -784,6 +784,28 @@ TEST_CASE(logsoftmax)
     EXPECT(p == prog);
 }
 
+TEST_CASE(argmax)
+{
+    migraphx::program p;
+    auto l0  = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {3, 4, 5, 6}});
+    auto ins = p.add_instruction(migraphx::op::argmax{2}, l0);
+    p.add_instruction(migraphx::op::squeeze{{2}}, ins);
+    auto prog = migraphx::parse_onnx("argmax_test.onnx");
+
+    EXPECT(p == prog);
+}
+
+TEST_CASE(argmin)
+{
+    migraphx::program p;
+    auto l0  = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {3, 4, 5, 6}});
+    auto ins = p.add_instruction(migraphx::op::argmin{3}, l0);
+    p.add_instruction(migraphx::op::squeeze{{3}}, ins);
+    auto prog = migraphx::parse_onnx("argmin_test.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(no_pad_test)
 {
     migraphx::program p;
@@ -794,6 +816,38 @@ TEST_CASE(no_pad_test)
     EXPECT(p == prog);
 }
 
+TEST_CASE(reducesum_test1)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {3, 4, 5, 6}});
+    auto l1 = p.add_instruction(migraphx::op::reduce_sum{{2}}, l0);
+    p.add_instruction(migraphx::op::squeeze{{2}}, l1);
+    auto prog = migraphx::parse_onnx("reducesum_test1.onnx");
+
+    EXPECT(p == prog);
+}
+
+TEST_CASE(reducesum_test2)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {3, 4, 5, 6}});
+    auto l1 = p.add_instruction(migraphx::op::reduce_sum{{2, 3}}, l0);
+    p.add_instruction(migraphx::op::squeeze{{2, 3}}, l1);
+    auto prog = migraphx::parse_onnx("reducesum_test2.onnx");
+
+    EXPECT(p == prog);
+}
+
+TEST_CASE(reducesum_test3)
+{
+    migraphx::program p;
+    auto l0 = p.add_parameter("x", migraphx::shape{migraphx::shape::float_type, {3, 4, 5, 6}});
+    p.add_instruction(migraphx::op::reduce_sum{{2, 3}}, l0);
+    auto prog = migraphx::parse_onnx("reducesum_test3.onnx");
+
+    EXPECT(p == prog);
+}
+
 TEST_CASE(clip_test)
 {
     migraphx::program p;

test/onnx/reducesum_test3.onnx

+reducesum-example:}
+1
+
x
y"	ReduceSum*
+axes@@
*
+keepdims

test_reducesumZ
+
x
+

+
+
+
+b
+
y
+

+
+
+

+
B

test/op_shape_test.cpp

@@ -390,6 +390,38 @@ TEST_CASE(softmax) { test_softmax_variations<migraphx::op::softmax>(); }
 
 TEST_CASE(logsoftmax) { test_softmax_variations<migraphx::op::logsoftmax>(); }
 
+template <class T>
+void test_argop_var()
+{
+    {
+        migraphx::shape input{migraphx::shape::half_type, {2, 3, 4, 5}};
+        expect_shape(migraphx::shape{migraphx::shape::int64_type, {1, 3, 4, 5}}, T{0}, input);
+    }
+
+    {
+        migraphx::shape input{migraphx::shape::half_type, {2, 3, 4, 5}};
+        expect_shape(migraphx::shape{migraphx::shape::int64_type, {2, 1, 4, 5}}, T{1}, input);
+    }
+
+    {
+        migraphx::shape input{migraphx::shape::half_type, {2, 3, 4, 5}};
+        expect_shape(migraphx::shape{migraphx::shape::int64_type, {2, 3, 1, 5}}, T{2}, input);
+    }
+
+    {
+        migraphx::shape input{migraphx::shape::half_type, {2, 3, 4, 5}};
+        expect_shape(migraphx::shape{migraphx::shape::int64_type, {2, 3, 4, 1}}, T{3}, input);
+    }
+
+    {
+        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
+        throws_shape(T{4}, input);
+    }
+}
+
+TEST_CASE(argmax) { test_argop_var<migraphx::op::argmax>(); }
+TEST_CASE(argmin) { test_argop_var<migraphx::op::argmin>(); }
+
 // 2 inputs arguments
 TEST_CASE(matmul)
 {