Remove Standard Shape requirement for ArgOps (#1042)

allow nonstd shape for the arg ops, non-standard shapes include broadcast, slice and transpose

src/include/migraphx/op/argmax.hpp

@@ -35,7 +35,7 @@ struct argmax
 
     shape normalize_compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1).standard();
+        check_shapes{inputs, *this}.has(1);
         auto lens = inputs[0].lens();
 
         lens[axis] = 1;

src/include/migraphx/op/argmin.hpp

@@ -35,7 +35,7 @@ struct argmin
 
     shape normalize_compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1).standard();
+        check_shapes{inputs, *this}.has(1);
         auto lens = inputs[0].lens();
 
         lens[axis] = 1;

src/targets/gpu/argmax.cpp

@@ -9,7 +9,7 @@ namespace gpu {
 
 shape hip_argmax::compute_shape(const std::vector<shape>& inputs) const
 {
-    check_shapes{inputs, *this}.has(2).standard();
+    check_shapes{inputs, *this}.has(2);
     return op.normalize_compute_shape({inputs.at(0)});
 }
 

src/targets/gpu/argmin.cpp

@@ -9,7 +9,7 @@ namespace gpu {
 
 shape hip_argmin::compute_shape(const std::vector<shape>& inputs) const
 {
-    check_shapes{inputs, *this}.has(2).standard();
+    check_shapes{inputs, *this}.has(2);
     return op.normalize_compute_shape({inputs.at(0)});
 }
 

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

@@ -76,8 +76,9 @@ void arg_op(Op op, hipStream_t stream, const argument& result, const argument& a
     size_t batch_item_num = batch_lens[axis];
     batch_lens[axis]      = 1;
     migraphx::shape batch_shape{arg_shape.type(), batch_lens};
+    migraphx::shape std_arg_shape{arg_shape.type(), arg_shape.lens()};
 
-    hip_visit_all(arg, arg_shape, batch_shape)([&](auto input, auto arg_s, auto batch_s) {
+    hip_visit_all(arg, std_arg_shape, batch_shape)([&](auto input, auto arg_s, auto batch_s) {
         auto* output = device_cast(result.get<int64_t>().data());
         using type   = device_type<std::remove_cv_t<typename decltype(input)::value_type>>;
         // use one block for items in one batch.

test/ref_ops_nonstd_shape_test.cpp

+#include <iostream>
+#include <vector>
+#include <migraphx/literal.hpp>
+#include <migraphx/instruction.hpp>
+#include <migraphx/ref/target.hpp>
+#include <migraphx/verify.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/auto_contiguous.hpp>
+#include <migraphx/pass_manager.hpp>
+#include "test.hpp"
+
+TEST_CASE(argmax_test_nonstd_shape)
+{
+    migraphx::program p;
+    auto* mm                = p.get_main_module();
+    std::vector<float> data = {1.2255,  1.6834,  -2.0305, -0.3221, 0.4701,  0.2583, 0.7545, 2.5758,
+                               -1.6849, 0.0928,  0.9022,  -0.8765, -0.4090, 0.9301, 2.0724, -1.5706,
+                               0.4867,  -0.1493, 0.6957,  -0.2179, 0.7142,  0.7177, 0.0183, 1.3497};
+    migraphx::shape data_shape{migraphx::shape::float_type, {2, 3, 4}};
+    auto dl = mm->add_literal(migraphx::literal{data_shape, data});
+    auto dl_trans =
+        mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 2, 0}}}), dl);
+    mm->add_instruction(migraphx::make_op("argmax", {{"axis", -3}}), dl_trans);
+    auto p_uncompiled = p;
+    p.compile(migraphx::ref::target{});
+    auto result   = p.eval({}).back();
+    auto res_gold = p_uncompiled.eval({}).back();
+    std::vector<int64_t> result_vec;
+    result.visit([&](auto output) { result_vec.assign(output.begin(), output.end()); });
+    std::vector<int64_t> res_gold_vec;
+    res_gold.visit([&](auto output) { res_gold_vec.assign(output.begin(), output.end()); });
+    EXPECT(migraphx::verify_range(result_vec, res_gold_vec));
+}
+
+TEST_CASE(argmin_test_nonstd_shape)
+{
+    migraphx::program p;
+    auto* mm                = p.get_main_module();
+    std::vector<float> data = {1.2255,  1.6834,  -2.0305, -0.3221, 0.4701,  0.2583, 0.7545, 2.5758,
+                               -1.6849, 0.0928,  0.9022,  -0.8765, -0.4090, 0.9301, 2.0724, -1.5706,
+                               0.4867,  -0.1493, 0.6957,  -0.2179, 0.7142,  0.7177, 0.0183, 1.3497};
+    migraphx::shape data_shape{migraphx::shape::float_type, {2, 3, 4}};
+    auto dl = mm->add_literal(migraphx::literal{data_shape, data});
+    auto dl_trans =
+        mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {1, 2, 0}}}), dl);
+    mm->add_instruction(migraphx::make_op("argmin", {{"axis", -1}}), dl_trans);
+    auto p_uncompiled = p;
+    p.compile(migraphx::ref::target{});
+    auto result   = p.eval({}).back();
+    auto res_gold = p_uncompiled.eval({}).back();
+    std::vector<int64_t> result_vec;
+    result.visit([&](auto output) { result_vec.assign(output.begin(), output.end()); });
+    std::vector<int64_t> res_gold_vec;
+    res_gold.visit([&](auto output) { res_gold_vec.assign(output.begin(), output.end()); });
+    EXPECT(migraphx::verify_range(result_vec, res_gold_vec));
+}
+
+int main(int argc, const char* argv[]) { test::run(argc, argv); }

test/verify/test_arg_ops.cpp

@@ -2,34 +2,92 @@
 #include "verify_program.hpp"
 #include <migraphx/program.hpp>
 #include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
 #include <migraphx/op/argmax.hpp>
 #include <migraphx/op/argmin.hpp>
 
-template <class T, int Axis>
-struct test_arg_ops : verify_program<test_arg_ops<T, Axis>>
+template <class T, int Axis, int NonStdShape>
+struct test_arg_ops : verify_program<test_arg_ops<T, Axis, NonStdShape>>
 {
     migraphx::program create_program() const
     {
         migraphx::program p;
         auto* mm = p.get_main_module();
-        migraphx::shape s{migraphx::shape::float_type, {2, 3, 4, 1025}};
+        migraphx::shape s{migraphx::shape::float_type, {2, 1, 4, 1025}};
         auto param = mm->add_parameter("data", s);
+        switch(NonStdShape)
+        {
+        case 0:
+            param = mm->add_instruction(
+                migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), param);
+            break;
+        case 1:
+            param = mm->add_instruction(
+                migraphx::make_op("multibroadcast", {{"out_lens", {2, 3, 4, 1025}}}), param);
+            break;
+        case 2:
+            param = mm->add_instruction(
+                migraphx::make_op("slice", {{"axes", {2}}, {"starts", {1}}, {"ends", {3}}}), param);
+            break;
+        default: break;
+        }
         mm->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::argmax, -1>;
-template struct test_arg_ops<migraphx::op::argmax, -2>;
-
-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>;
-template struct test_arg_ops<migraphx::op::argmin, -3>;
-template struct test_arg_ops<migraphx::op::argmin, -4>;
+// transpose argmax tests
+template struct test_arg_ops<migraphx::op::argmax, 0, 0>;
+template struct test_arg_ops<migraphx::op::argmax, 1, 0>;
+template struct test_arg_ops<migraphx::op::argmax, 2, 0>;
+template struct test_arg_ops<migraphx::op::argmax, 3, 0>;
+template struct test_arg_ops<migraphx::op::argmax, -1, 0>;
+template struct test_arg_ops<migraphx::op::argmax, -2, 0>;
+// transpose argmin tests
+template struct test_arg_ops<migraphx::op::argmin, 0, 0>;
+template struct test_arg_ops<migraphx::op::argmin, 1, 0>;
+template struct test_arg_ops<migraphx::op::argmin, 2, 0>;
+template struct test_arg_ops<migraphx::op::argmin, 3, 0>;
+template struct test_arg_ops<migraphx::op::argmin, -3, 0>;
+template struct test_arg_ops<migraphx::op::argmin, -4, 0>;
+// broadcast argmax tests
+template struct test_arg_ops<migraphx::op::argmax, 0, 1>;
+template struct test_arg_ops<migraphx::op::argmax, 1, 1>;
+template struct test_arg_ops<migraphx::op::argmax, 2, 1>;
+template struct test_arg_ops<migraphx::op::argmax, 3, 1>;
+template struct test_arg_ops<migraphx::op::argmax, -1, 1>;
+template struct test_arg_ops<migraphx::op::argmax, -2, 1>;
+// broadcast argmin tests
+template struct test_arg_ops<migraphx::op::argmin, 0, 1>;
+template struct test_arg_ops<migraphx::op::argmin, 1, 1>;
+template struct test_arg_ops<migraphx::op::argmin, 2, 1>;
+template struct test_arg_ops<migraphx::op::argmin, 3, 1>;
+template struct test_arg_ops<migraphx::op::argmin, -3, 1>;
+template struct test_arg_ops<migraphx::op::argmin, -4, 1>;
+// slice argmax tests
+template struct test_arg_ops<migraphx::op::argmax, 0, 2>;
+template struct test_arg_ops<migraphx::op::argmax, 1, 2>;
+template struct test_arg_ops<migraphx::op::argmax, 2, 2>;
+template struct test_arg_ops<migraphx::op::argmax, 3, 2>;
+template struct test_arg_ops<migraphx::op::argmax, -1, 2>;
+template struct test_arg_ops<migraphx::op::argmax, -2, 2>;
+// slice argmin tests
+template struct test_arg_ops<migraphx::op::argmin, 0, 2>;
+template struct test_arg_ops<migraphx::op::argmin, 1, 2>;
+template struct test_arg_ops<migraphx::op::argmin, 2, 2>;
+template struct test_arg_ops<migraphx::op::argmin, 3, 2>;
+template struct test_arg_ops<migraphx::op::argmin, -3, 2>;
+template struct test_arg_ops<migraphx::op::argmin, -4, 2>;
+// default case, standard shape argmax tests
+template struct test_arg_ops<migraphx::op::argmax, 0, 3>;
+template struct test_arg_ops<migraphx::op::argmax, 1, 3>;
+template struct test_arg_ops<migraphx::op::argmax, 2, 3>;
+template struct test_arg_ops<migraphx::op::argmax, 3, 3>;
+template struct test_arg_ops<migraphx::op::argmax, -1, 3>;
+template struct test_arg_ops<migraphx::op::argmax, -2, 3>;
+// default case, standard shape argmin tests
+template struct test_arg_ops<migraphx::op::argmin, 0, 3>;
+template struct test_arg_ops<migraphx::op::argmin, 1, 3>;
+template struct test_arg_ops<migraphx::op::argmin, 2, 3>;
+template struct test_arg_ops<migraphx::op::argmin, 3, 3>;
+template struct test_arg_ops<migraphx::op::argmin, -3, 3>;
+template struct test_arg_ops<migraphx::op::argmin, -4, 3>;