Add a step to unsqeeze axis (#1242)

Instead of just unsqueezing to an axis of 1 a step can be set to use instead. So instead of unsqueezing {3, 12} to {3, 1, 12} a step of 2 will unsqeeze to {3, 2, 6} instead

src/include/migraphx/op/unsqueeze.hpp

@@ -42,11 +42,12 @@ namespace op {
 struct unsqueeze
 {
     std::vector<int64_t> axes;
+    std::vector<int64_t> steps;
 
     template <class Self, class F>
     static auto reflect(Self& self, F f)
     {
-        return pack(f(self.axes, "axes"));
+        return pack(f(self.axes, "axes"), f(self.steps, "steps"));
     }
 
     value attributes() const
@@ -73,6 +74,9 @@ struct unsqueeze
                 MIGRAPHX_THROW("UNSQUEEZE: Input must be a scalar");
         }
 
+        if(steps.size() > axes.size())
+            MIGRAPHX_THROW("UNSQUEEZE: Steps provided with no axis");
+
         std::size_t new_size = old_lens.size() + axes.size();
 
         std::vector<std::size_t> new_lens(new_size);
@@ -80,16 +84,27 @@ struct unsqueeze
         std::size_t p = 0;
         for(auto i : range(new_size))
         {
-            if(std::find(axes.begin(), axes.end(), i) != axes.end())
+            auto axis_idx = std::find(axes.begin(), axes.end(), i) - axes.begin();
+            if(axis_idx < axes.size())
             {
-                new_lens[i] = 1;
-                if(p == 0) // unsqueeze on the first axes
+                std::int64_t step = 1;
+                if(axis_idx < steps.size())
+                    step = steps[axis_idx];
+                if(step == 0)
+                    MIGRAPHX_THROW("UNSQUEEZE: step must be non-zero");
+                new_lens[i] = step;
+                if(p < old_strides.size())
                 {
-                    new_strides[i] = old_lens[0] * old_strides[0];
+                    if((old_lens[p] % step) != 0)
+                        MIGRAPHX_THROW("UNSQUEEZE: Axis dimenstion is not divisible by step");
+                    old_lens[p] /= step;
+                    new_strides[i] = old_strides[p] * old_lens[p];
                 }
-                else // unsqueeze on middle or last axes
+                else
                 {
-                    new_strides[i] = (p < old_strides.size()) ? old_strides[p - 1] : 1;
+                    if(step != 1)
+                        MIGRAPHX_THROW("UNSQUEEZE: Step must be 1 for extra axes");
+                    new_strides[i] = 1;
                 }
             }
             else

src/simplify_reshapes.cpp

@@ -272,7 +272,7 @@ struct find_concat_transpose
 {
     auto matcher() const
     {
-        return match::name("concat")(match::all_of[match::inputs()](match::transpose_shape()));
+        return match::name("concat")(match::all_of[match::inputs()](match::name("transpose")));
     }
 
     void apply(module& m, const match::matcher_result& mr) const

test/op_shape_test.cpp

@@ -1533,15 +1533,46 @@ TEST_CASE(test_squeeze_wrong_axis)
 
 TEST_CASE(test_unsqueeze)
 {
-    migraphx::shape s1{migraphx::shape::float_type, {4, 3, 3}};
-    migraphx::shape s2{migraphx::shape::float_type, {4, 3, 1, 3}};
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {4, 5, 1, 3}};
     expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2}}}), s1);
 }
 
+TEST_CASE(test_unsqueeze_step)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 12}};
+    migraphx::shape s2{migraphx::shape::float_type, {4, 5, 2, 6}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2}}, {"steps", {2}}}), s1);
+}
+
+TEST_CASE(test_unsqueeze_step_non_divisable)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 3}};
+    throws_shape(migraphx::make_op("unsqueeze", {{"axes", {2}}, {"steps", {2}}}), s1);
+}
+
+TEST_CASE(test_unsqueeze_step_zero)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 12}};
+    throws_shape(migraphx::make_op("unsqueeze", {{"axes", {2}}, {"steps", {0}}}), s1);
+}
+
+TEST_CASE(test_unsqueeze_step_at_end)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 12}};
+    throws_shape(migraphx::make_op("unsqueeze", {{"axes", {3}}, {"steps", {2}}}), s1);
+}
+
+TEST_CASE(test_unsqueeze_mismatch_step_axis)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 12}};
+    throws_shape(migraphx::make_op("unsqueeze", {{"axes", {2}}, {"steps", {2, 3}}}), s1);
+}
+
 TEST_CASE(test_unsqueeze_negative_axis)
 {
-    migraphx::shape s1{migraphx::shape::float_type, {4, 3, 3}};
-    migraphx::shape s2{migraphx::shape::float_type, {4, 3, 1, 3}};
+    migraphx::shape s1{migraphx::shape::float_type, {4, 5, 3}};
+    migraphx::shape s2{migraphx::shape::float_type, {4, 5, 1, 3}};
     expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {-2}}}), s1);
 }
 
@@ -1567,21 +1598,28 @@ TEST_CASE(test_unsqueeze_scalar_tensor2)
 TEST_CASE(test_unsqueeze_transpose)
 {
     migraphx::shape s1{migraphx::shape::float_type, {4, 4, 3}, {12, 1, 4}};
-    migraphx::shape s2{migraphx::shape::float_type, {4, 4, 1, 3}, {12, 1, 1, 4}};
+    migraphx::shape s2{migraphx::shape::float_type, {4, 4, 1, 3}, {12, 1, 12, 4}};
     expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2}}}), s1);
 }
 
+TEST_CASE(test_unsqueeze_transpose_step)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {4, 4, 6}, {24, 1, 4}};
+    migraphx::shape s2{migraphx::shape::float_type, {4, 4, 2, 3}, {24, 1, 12, 4}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2}}, {"steps", {2}}}), s1);
+}
+
 TEST_CASE(test_unsqueeze_multibroadcast)
 {
     migraphx::shape s1{migraphx::shape::float_type, {2, 3, 4}, {0, 1, 0}};
-    migraphx::shape s2{migraphx::shape::float_type, {2, 3, 1, 4}, {0, 1, 1, 0}};
+    migraphx::shape s2{migraphx::shape::float_type, {2, 3, 1, 4}, {0, 1, 0, 0}};
     expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2}}}), s1);
 }
 
 TEST_CASE(test_unsqueeze_slice)
 {
     migraphx::shape s1{migraphx::shape::float_type, {2, 3, 4}, {108, 36, 1}};
-    migraphx::shape s2{migraphx::shape::float_type, {2, 3, 1, 4}, {108, 36, 36, 1}};
+    migraphx::shape s2{migraphx::shape::float_type, {2, 3, 1, 4}, {108, 36, 4, 1}};
     expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2}}}), s1);
 }
 
@@ -1613,6 +1651,27 @@ TEST_CASE(test_unsqueeze_multiple_axes_2)
     expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {0, 1}}}), s1);
 }
 
+TEST_CASE(test_unsqueeze_multiple_axes_3)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {3, 4, 5}};
+    migraphx::shape s2{migraphx::shape::float_type, {3, 4, 1, 5, 1, 1}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2, 4, 5}}}), s1);
+}
+
+TEST_CASE(test_unsqueeze_multiple_axes_4)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {3, 4, 5}};
+    migraphx::shape s2{migraphx::shape::float_type, {3, 4, 1, 5, 1, 1}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {5, 4, 2}}}), s1);
+}
+
+TEST_CASE(test_unsqueeze_multiple_axes_step)
+{
+    migraphx::shape s1{migraphx::shape::float_type, {3, 4, 10}};
+    migraphx::shape s2{migraphx::shape::float_type, {3, 4, 2, 5, 1, 1}};
+    expect_shape(s2, migraphx::make_op("unsqueeze", {{"axes", {2, 4, 5}}, {"steps", {2}}}), s1);
+}
+
 TEST_CASE(transpose_shape)
 {
     migraphx::shape input{migraphx::shape::float_type, {2, 2}};

test/simplify_reshapes_test.cpp

@@ -1141,6 +1141,38 @@ TEST_CASE(transpose_contiguous_reshape_binary_broadcast)
     EXPECT(m1 == m2);
 }
 
+TEST_CASE(transpose_unsqueeze_concat)
+{
+    migraphx::module m1;
+    {
+        auto l0 = m1.add_parameter("0", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
+        auto lt0 =
+            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), l0);
+        auto l1 = m1.add_parameter("1", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
+        auto lt1 =
+            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), l1);
+        auto l2 = m1.add_parameter("2", migraphx::shape{migraphx::shape::float_type, {1, 2, 1, 1}});
+        auto lt2 =
+            m1.add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 3, 1}}}), l2);
+        std::vector<migraphx::instruction_ref> args{lt0, lt1, lt2};
+        std::vector<migraphx::instruction_ref> unsqueezed_args;
+        int64_t axis = 3;
+
+        std::transform(
+            args.begin(),
+            args.end(),
+            std::back_inserter(unsqueezed_args),
+            [&](migraphx::instruction_ref arg) {
+                return m1.add_instruction(migraphx::make_op("unsqueeze", {{"axes", {axis}}}), arg);
+            });
+        m1.add_instruction(migraphx::make_op("concat", {{"axis", axis}}), unsqueezed_args);
+    }
+    // TODO: This could be simplified to a single transpose after concat
+    migraphx::module m2 = m1;
+    run_pass(m1);
+    EXPECT(m1 == m2);
+}
+
 TEST_CASE(transpose_slice)
 {
     migraphx::module m1;