reorder_slice_add_mul matcher

src/simplify_algebra.cpp

@@ -412,24 +412,83 @@
     }
 };
 
+
+
+auto fusable_split()
+{
+    return match::make_basic_pred_matcher([](instruction_ref ins) {
+
+        std::vector<instruction_ref> slices;
+
+        for (const auto& output : ins->outputs())
+        {
+            if (output->name() == "slice")
+            {
+                slices.push_back(output);
+            }
+        }
+
+        if (slices.empty())
+            return false;
+
+        std::vector<instruction_ref> add_instructions;
+        for (const auto& slice : slices)
+        {
+            bool used_by_add = false;
+            for (auto& user : slice->outputs())
+            {
+                if (user->name() == "add")
+                {
+                    used_by_add = true;
+                    add_instructions.push_back(user);
+                }
+            }
+            if (!used_by_add)
+                return false;
+        }        
+
+        bool any_add_followed_by_mul = false;
+        for (const auto& add_ins : add_instructions)
+        {
+            for (auto& user : add_ins->outputs())
+            {
+                if (user->name() == "mul")
+                {
+                    any_add_followed_by_mul = true;
+                    ins->outputs();
+                }
+            }
+        }
+
+        if (!any_add_followed_by_mul)
+            return false;
+
+        return true;
+    });
+}
+
+
 // ******************************
 //  a * (x + b) => a * x + a * b
 // ******************************
 // When a * (x + b) is followed by another add of constant, then the
 // additional add can be const folded. Also, better fusions can be applied
 // when the add comes after.
+
 struct find_mul_add
 {
     auto matcher() const
     {
         return match::name("mul")(match::either_arg(0, 1)(
-            match::name("add")(
-                match::either_arg(0, 1)(
-                    match::any().bind("x"),
-                    match::any_of(conv_const_weights(), match::is_constant()).bind("b")),
-                match::none_of(match::args(match::is_constant(), match::is_constant())),
-                match::used_once()),
-            match::is_constant().bind("a")));
+                match::name("add")(
+                    match::either_arg(0, 1)(
+                        match::none_of(match::name("slice")(match::arg(0)(
+                            fusable_split().bind("slc")
+                        ))).bind("x"),
+                        match::any_of(conv_const_weights(), match::is_constant()).bind("b")),
+                    match::none_of(match::args(match::is_constant(), match::is_constant())),
+                    match::used_once()),
+                match::is_constant().bind("a")));
     }
 
     void apply(module& m, const match::matcher_result& r) const
@@ -439,13 +498,48 @@
         auto b_ins = r.instructions["b"];
         auto x_ins = r.instructions["x"];
         assert(x_ins != b_ins);
-
+    
         auto ax_ins = m.insert_instruction(ins, make_op("mul"), a_ins, x_ins);
         auto ab_ins = m.insert_instruction(ins, make_op("mul"), a_ins, b_ins);
         m.replace_instruction(ins, make_op("add"), ax_ins, ab_ins);
+
     }
 };
 
+
+struct find_slice_add_mul
+{
+
+        auto matcher() const
+        {        
+            return match::name("mul")(match::either_arg(0, 1)(
+                match::name("add")(
+                    match::either_arg(0, 1)(
+                        match::name("slice")(match::arg(0)(
+                            fusable_split().bind("slc")
+                        )).bind("x"),
+                        match::any_of(conv_const_weights(), match::is_constant()).bind("b")),
+                    match::none_of(match::args(match::is_constant(), match::is_constant())),
+                    match::used_once()),
+                match::is_constant().bind("a")));
+        }
+    void apply(module& m, const match::matcher_result& r) const
+    {
+
+        auto ins    = r.result;
+        auto a_ins  = r.instructions["a"];
+        auto b_ins  = r.instructions["b"];
+        auto x_ins  = r.instructions["x"];
+
+        assert(x_ins != b_ins);
+
+        auto ax_ins = m.insert_instruction(ins, make_op("add"), x_ins, b_ins);
+        m.replace_instruction(ins, make_op("mul"), ax_ins, a_ins);
+
+    }
+};
+
+
 struct find_dot_add
 {
     auto matcher() const
@@ -1569,6 +1663,7 @@
 
         auto outputs = ins->outputs();
         group_by(outputs.begin(), outputs.end(), each, pred);
+
     }
 };
 
@@ -1981,6 +2076,7 @@
                             find_dot_slice{},
                             find_dot_mul{},
                             find_mul_add{},
+                            find_slice_add_mul{},
                             find_unit_ops{},
                             find_neg_unit_ops{},
                             eliminate_zero_point{},

test/simplify_algebra_test.cpp

@@ -509,6 +509,7 @@
     EXPECT(m1 == m2);
 }
 
+
 TEST_CASE(simplify_conv_add)
 {
     migraphx::shape s{migraphx::shape::float_type, {1, 3, 32, 32}};
@@ -4133,6 +4134,46 @@
     EXPECT(m1.sort() == m2.sort());
 }
 
+TEST_CASE(my_optim)
+{
+    migraphx::shape as{migraphx::shape::float_type, {1, 77, 2304}};
+
+    migraphx::module m2;
+    {
+        auto a       = m2.add_parameter("b", as); 
+
+        auto slice_a = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {768}}}), a);
+        auto slice_b = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {768}}, {"ends", {1536}}}), a);
+        auto slice_c = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {1}}, {"starts", {1536}}, {"ends", {2304}}}), a);    
+
+        auto one = m2.add_literal(
+            migraphx::generate_literal({migraphx::shape::float_type, {1, 77, 768}}, 3.0f));
+        auto two = m2.add_literal(
+            migraphx::generate_literal({migraphx::shape::float_type, {1, 77, 768}}, 2.0f));
+        auto three = m2.add_literal(
+            migraphx::generate_literal({migraphx::shape::float_type, {1, 77, 768}}, 4.0f));
+        auto four = m2.add_literal(
+            migraphx::generate_literal({migraphx::shape::float_type, {1, 77, 768}}, 2.0f));
+
+
+        auto add1 = m2.add_instruction(migraphx::make_op("add"), one, slice_a);
+        auto mul1 = m2.add_instruction(migraphx::make_op("mul"), add1, four);
+
+        auto add2 = m2.add_instruction(migraphx::make_op("add"), two, slice_b);
+        auto add3 = m2.add_instruction(migraphx::make_op("add"), three, slice_c);
+
+        m2.add_return({mul1});
+    };
+    migraphx::module m1 = m2;
+    run_pass(m2);
+
+    EXPECT(m1.sort() == m2.sort());
+
+}
+
 TEST_CASE(dot_slice_ab)
 {
     migraphx::shape as{migraphx::shape::float_type, {2, 256, 32}};
@@ -4195,6 +4236,61 @@
     EXPECT(m1.sort() == m2.sort());
 }
 
+
+TEST_CASE(complex_graph_operations)
+{
+    migraphx::module m;
+
+    auto x_0 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 0));
+    auto x_1 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 1));
+    auto x_2 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 2));
+    auto x_3 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 3));
+    auto x_4 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 4));
+    auto x_5 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 5));
+    auto x_6 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 6));
+    auto x_7 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 7));
+    auto x_8 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 8));
+    auto x_9 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 9));
+    auto x_10 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 10));
+    auto x_11 = m.add_literal(migraphx::generate_literal({migraphx::shape::float_type, {64, 512}}, 11));
+    auto x_12 = m.add_literal(migraphx::literal({migraphx::shape::float_type, {1}}, {0.125}));
+
+
+    auto p_x = m.add_parameter("x", {migraphx::shape::float_type, {64, 64}});
+
+
+    auto x_14 = m.add_instruction(
+        migraphx::make_op("multibroadcast", {{"out_lens", {64, 512}}}), x_12);
+
+
+    auto x_15 = m.add_instruction(migraphx::make_op("dot"), p_x, x_11);
+    auto x_16 = m.add_instruction(migraphx::make_op("add"), x_15, x_10);
+    auto x_17 = m.add_instruction(migraphx::make_op("add"), x_16, x_9);
+    auto x_18 = m.add_instruction(migraphx::make_op("add"), x_17, x_8);
+
+    auto x_19 = m.add_instruction(migraphx::make_op("dot"), p_x, x_7);
+    auto x_20 = m.add_instruction(migraphx::make_op("add"), x_19, x_6);
+    auto x_21 = m.add_instruction(migraphx::make_op("add"), x_20, x_5);
+    auto x_22 = m.add_instruction(migraphx::make_op("add"), x_21, x_4);
+
+    auto x_23 = m.add_instruction(migraphx::make_op("dot"), p_x, x_3);
+    auto x_24 = m.add_instruction(migraphx::make_op("add"), x_23, x_2);
+    auto x_25 = m.add_instruction(migraphx::make_op("add"), x_24, x_1);
+    auto x_26 = m.add_instruction(migraphx::make_op("add"), x_25, x_0);
+
+    auto x_27 = m.add_instruction(migraphx::make_op("mul"), x_26, x_14);
+
+    m.add_return({x_18, x_22, x_27});
+
+    run_pass(m);
+
+    EXPECT(m.get_output_shapes().size() == 3);
+
+
+}
+
+
+
 TEST_CASE(dot_slice_not_applicable_1)
 {
     migraphx::shape as{migraphx::shape::float_type, {2, 256, 32}};