[AIMIGRAPHX-1002] Match nop reshapes with broadcasted inputs

src/include/migraphx/matcher.hpp

@@ -1129,6 +1129,26 @@ auto same_shape(Ms... ms)
     return all_of(same_shape(ms)...);
 }
 
+template <class M>
+auto same_lens(M m)
+{
+    return make_basic_fun_matcher(
+        [=](matcher_context& ctx, instruction_ref ins) -> optional<instruction_ref> {
+            auto i = m.match(ctx, ins);
+            if(not i)
+                return nullopt;
+            if(shape::same_lens((*i)->get_shape(), ins->get_shape()))
+                return ins;
+            return nullopt;
+        });
+}
+
+template <class... Ms>
+auto same_lens(Ms... ms)
+{
+    return all_of(same_lens(ms)...);
+}
+
 template <class... Ms>
 auto skip_broadcasts(Ms... ms)
 {

src/simplify_reshapes.cpp

@@ -593,7 +593,7 @@ struct find_nop_reshapes
     auto matcher() const
     {
         // clang-format off
-        static const std::unordered_set<std::string> names = {
+        static const std::unordered_set<std::string> shape_names = {
             "flatten",
             "reshape",
             "contiguous",
@@ -608,14 +608,18 @@ struct find_nop_reshapes
             "slice",
             "step",
             "transpose",
+        };
+        static const std::unordered_set<std::string> lens_names = {
             "reduce_mean",
             "reduce_max",
             "reduce_min",
             "reduce_sum",
             "reduce_prod",
         };
-
-       return match::name(names)(match::same_shape(match::arg(0)));
+        // clang-format on
+        auto shape_match = match::name(shape_names)(match::same_shape(match::arg(0)));
+        auto lens_match  = match::name(lens_names)(match::same_lens(match::arg(0)));
+        return match::any_of(shape_match, lens_match);
     }
 
     void apply(module& m, const match::matcher_result& mr) const

test/simplify_reshapes_test.cpp

@@ -5059,4 +5059,30 @@ TEST_CASE(slice_reshape_multibroadcast_rebase_axis)
     EXPECT(m1.get_output_shapes() == m2.get_output_shapes());
 }
 
+TEST_CASE(broadcast_nop_reduce_mean)
+{
+    migraphx::module m1;
+    {
+        auto lit = m1.add_literal(
+            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {1}}, {42}});
+        auto bcast = m1.add_instruction(
+            migraphx::make_op("broadcast", {{"axis", 0}, {"out_lens", {1, 1, 8}}}), lit);
+        auto reduce_mean =
+            m1.add_instruction(migraphx::make_op("reduce_mean", {{"axes", {0}}}), bcast);
+        m1.add_return({reduce_mean});
+    }
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto lit = m2.add_literal(
+            migraphx::literal{migraphx::shape{migraphx::shape::float_type, {1}}, {42}});
+        auto bcast = m2.add_instruction(
+            migraphx::make_op("broadcast", {{"axis", 0}, {"out_lens", {1, 1, 8}}}), lit);
+        m2.add_return({bcast});
+    }
+
+    EXPECT(m1.sort() == m2.sort());
+}
+
 int main(int argc, const char* argv[]) { test::run(argc, argv); }