Merge branch 'dyn_squeeze' of github.com:ROCmSoftwarePlatform/AMDMIGraphX into dyn_unsqueeze

src/include/migraphx/op/squeeze.hpp

@@ -29,6 +29,7 @@
 #include <migraphx/config.hpp>
 #include <migraphx/value.hpp>
 #include <migraphx/op/normalize_attribute.hpp>
+#include <migraphx/dyn_output.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -54,14 +55,51 @@ struct squeeze
     std::string name() const { return "squeeze"; }
     shape normalize_compute_shape(std::vector<shape> inputs) const
     {
-        check_shapes{inputs, *this}.has(1);
+        check_shapes{inputs, *this, true}.has(1);
         auto input_shape = inputs[0];
+        if(input_shape.dynamic())
+        {
+            std::vector<shape::dynamic_dimension> one_dyn_dims{{1, 1, 0}, {1, 1, 1}};
+            if(std::any_of(axes.begin(), axes.end(), [&](auto axis) {
+                   return not contains(one_dyn_dims, input_shape.dyn_dims()[axis]);
+               }))
+            {
+                MIGRAPHX_THROW(
+                    "SQUEEZE: dynamic axis dimension should be equal to {1, 1, 0} or {1, 1, 1}");
+            }
+            std::vector<shape::dynamic_dimension> dyn_dims = {};
+            if(axes.empty())
+            {
+                for(auto i : range(input_shape.ndim()))
+                {
+                    auto dd = input_shape.dyn_dims()[i];
+                    if(not contains(one_dyn_dims, dd))
+                    {
+                        dyn_dims.push_back(dd);
+                    }
+                }
+            }
+            else
+            {
+                for(auto i : range(input_shape.ndim()))
+                {
+                    if(std::find(axes.begin(), axes.end(), i) == axes.end())
+                    {
+                        dyn_dims.push_back(input_shape.dyn_dims()[i]);
+                    }
+                }
+            }
+            return {input_shape.type(), dyn_dims};
+        }
+        else
+        {
             auto type        = input_shape.type();
             auto old_lens    = input_shape.lens();
             auto old_strides = input_shape.strides();
-        if(std::any_of(axes.begin(), axes.end(), [&](auto axis) { return old_lens[axis] != 1; }))
+            if(std::any_of(
+                   axes.begin(), axes.end(), [&](auto axis) { return old_lens[axis] != 1; }))
             {
-            MIGRAPHX_THROW("squeeze axis dimension should be equal to 1");
+                MIGRAPHX_THROW("SQUEEZE: static axis dimension should be equal to 1");
             }
             std::vector<std::size_t> new_lens;
             std::vector<std::size_t> new_strides;
@@ -96,10 +134,11 @@ struct squeeze
                 return shape{type, new_lens, new_strides};
             }
         }
+    }
 
-    argument compute(shape output_shape, std::vector<argument> args) const
+    argument compute(const dyn_output& dyn_out, std::vector<argument> args) const
     {
-        return args[0].reshape(output_shape);
+        return args[0].reshape(dyn_out.computed_shape);
     }
     std::ptrdiff_t output_alias(const std::vector<shape>&) const { return 0; }
 };

test/op_shape_test.cpp

@@ -1997,6 +1997,30 @@ TEST_CASE(test_squeeze_all)
     expect_shape(s2, migraphx::make_op("squeeze", {{"axes", {0}}}), s1);
 }
 
+TEST_CASE(test_squeeze_dyn)
+{
+    migraphx::shape s1{migraphx::shape::float_type,
+                       {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {1, 1, 0}, {3, 3, 0}}};
+    migraphx::shape s2{migraphx::shape::float_type, {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {3, 3, 0}}};
+    expect_shape(s2, migraphx::make_op("squeeze", {{"axes", {3}}}), s1);
+
+    migraphx::shape s3{migraphx::shape::float_type, {{1, 4, 0}, {3, 3, 0}, {3, 3, 0}}};
+    expect_shape(s3, migraphx::make_op("squeeze"), s1);
+
+    throws_shape(migraphx::make_op("squeeze", {{"axes", {0}}}), s1);
+}
+
+TEST_CASE(test_squeeze_dyn_neg_axes)
+{
+    migraphx::shape s1{migraphx::shape::float_type,
+                       {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {1, 1, 0}, {3, 3, 0}}};
+    migraphx::shape s2{migraphx::shape::float_type, {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {3, 3, 0}}};
+    expect_shape(s2, migraphx::make_op("squeeze", {{"axes", {-2}}}), s1);
+
+    migraphx::shape s3{migraphx::shape::float_type, {{1, 4, 0}, {3, 3, 0}, {3, 3, 0}}};
+    expect_shape(s3, migraphx::make_op("squeeze", {{"axes", {-2, -4}}}), s1);
+}
+
 TEST_CASE(test_squeeze_transpose)
 {
     migraphx::shape s1{migraphx::shape::float_type, {4, 4, 1}, {4, 1, 4}};

test/ref_ops_test.cpp

@@ -7084,6 +7084,25 @@ TEST_CASE(squeeze_test)
     }
 }
 
+TEST_CASE(squeeze_dyn_test)
+{
+    migraphx::program p;
+    auto* mm = p.get_main_module();
+    migraphx::shape s1{migraphx::shape::float_type,
+                       {{1, 4, 0}, {1, 1, 0}, {3, 3, 0}, {1, 1, 0}, {3, 3, 0}}};
+    auto p0 = mm->add_parameter("x", s1);
+    mm->add_instruction(migraphx::make_op("squeeze", {{"axes", {1}}}), p0);
+    p.compile(migraphx::ref::target{});
+
+    std::vector<float> input_data(4 * 3 * 3);
+    migraphx::parameter_map params0;
+    migraphx::shape input_fixed_shape0{migraphx::shape::float_type, {4, 1, 3, 1, 3}};
+    params0["x"] = migraphx::argument(input_fixed_shape0, input_data.data());
+    auto result  = p.eval(params0).back();
+    migraphx::shape s2{migraphx::shape::float_type, {4, 3, 1, 3}};
+    EXPECT(result.get_shape() == s2);
+}
+
 TEST_CASE(step_test)
 {
     {