Fix concat transpose

src/simplify_reshapes.cpp

@@ -249,7 +249,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/simplify_reshapes_test.cpp

@@ -1118,4 +1118,37 @@ 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);
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }