reorder_slice_add_mul matcher

src/simplify_algebra.cpp

@@ -412,6 +412,21 @@
     }
 };
 
+auto fusable_split(const std::string& name)
+{    
+    return match::make_basic_pred_matcher([&](instruction_ref ins) {
+        return all_of(ins->outputs(), [&](instruction_ref slice) {
+            if(slice->name() != "slice")
+                return false;
+            return any_of(slice->outputs(), [&](instruction_ref x) {
+                if (x->name() == name)
+                return true;
+            });
+        });
+    });
+}
+
+
 // ******************************
 //  a * (x + b) => a * x + a * b
 // ******************************
@@ -421,15 +436,22 @@
 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")));
+    {   
+        auto slice_1 = match::skip(match::name("slice"))(match::none_of(fusable_split("add")));
+
+        return match::name("mul")(
+            match::either_arg(0, 1)(
+                match::name("add")(
+                    match::either_arg(0, 1)(
+                        slice_1.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
@@ -443,6 +465,46 @@
         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);
+
+    }
+};
+
+// When a slice is followed by a*(x+b), a⋅(x+b), this matcher performs 
+// the add first, followed by the mul. 
+// If multiple slices originate from the same instruction and are followed by add,
+// all the adds can be const folded and performed before the slicing.
+struct find_slice_add_mul
+{
+        auto matcher() const
+        {        
+            auto slice_1 = match::name("slice")(match::arg(0)(fusable_split("add")));
+
+            return match::name("mul")(
+                match::either_arg(0, 1)(
+                    match::name("add")(
+                        match::either_arg(0, 1)(
+                            slice_1.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);
+
     }
 };
 
@@ -1569,6 +1631,7 @@
 
         auto outputs = ins->outputs();
         group_by(outputs.begin(), outputs.end(), each, pred);
+
     }
 };
 
@@ -1981,6 +2044,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,63 @@
     EXPECT(m1.sort() == m2.sort());
 }
 
+TEST_CASE(skip_mul_add_for_horizontal_add_fusion)
+{
+    migraphx::shape as{migraphx::shape::float_type, {1, 77, 2304}};
+
+    migraphx::module m1;
+    {
+        auto a = m1.add_parameter("a", as); 
+
+        auto slice_a = m1.add_instruction(
+            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {768}}}), a);
+        auto slice_b = m1.add_instruction(
+            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {768}}, {"ends", {1536}}}), a);
+        auto slice_c = m1.add_instruction(
+            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {1536}}, {"ends", {2304}}}), a);    
+
+        auto one = m1.add_literal(migraphx::generate_literal(
+            {migraphx::shape::float_type, {1, 77, 768}}, 3.0f));
+        auto two = m1.add_literal(migraphx::generate_literal(
+            {migraphx::shape::float_type, {1, 77, 768}}, 2.0f));
+        auto three = m1.add_literal(migraphx::generate_literal(
+            {migraphx::shape::float_type, {1, 77, 768}}, 4.0f));
+
+        auto add1 = m1.add_instruction(migraphx::make_op("add"), slice_a, one);
+
+        auto add2 = m1.add_instruction(migraphx::make_op("add"), slice_b, two);
+        auto add3 = m1.add_instruction(migraphx::make_op("add"), slice_c, three);
+
+        m1.add_return({add1});
+
+    };
+
+    run_pass(m1);
+
+    migraphx::module m2;
+    {
+        auto a = m2.add_parameter("a", as); 
+
+        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 concat = m2.add_instruction(migraphx::make_op("concat", {{"axis", 2}}), one, two, three);
+
+        auto add = m2.add_instruction(migraphx::make_op("add"), a, concat);
+
+        auto slice_a = m2.add_instruction(
+            migraphx::make_op("slice", {{"axes", {2}}, {"starts", {0}}, {"ends", {768}}}), add);
+
+        m2.add_return({slice_a});
+    };
+
+    EXPECT(m1.sort() == m2.sort());
+}
+
 TEST_CASE(dot_slice_ab)
 {
     migraphx::shape as{migraphx::shape::float_type, {2, 256, 32}};
@@ -4195,6 +4253,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}};
@@ -4209,7 +4322,8 @@
                               {{"axes", {1, 2}}, {"starts", {64, 64}}, {"ends", {128, 128}}}),
             dot);
         auto slice2 = m1.add_instruction(
-            migraphx::make_op("slice", {{"axes", {0}}, {"starts", {0}}, {"ends", {1}}}), dot);
+            migraphx::make_op("slice", {{"axes", {0}}, {"starts"
+            , {0}}, {"ends", {1}}}), dot);
 
         m1.add_return({slice1, slice2});
     };