reorder_slice_add_mul matcher

src/simplify_algebra.cpp

@@ -422,14 +422,46 @@
 {
     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")));
+        return match::name("mul")(
+            match::none_of[match::outputs()](match::name("convolution")),
+            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")));
+    }
+
+    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("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);
+    }
+};
+/* Delete this later
+struct find_mul_add
+{
+    auto matcher() const
+    {
+        return match::name("mul")(
+            match::none_of[match::outputs()](match::name("convolution")),
+            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")));
     }
 
     void apply(module& m, const match::matcher_result& r) const
@@ -445,6 +477,38 @@
         m.replace_instruction(ins, make_op("add"), ax_ins, ab_ins);
     }
 };
+*/
+
+struct find_scalar_mul_conv
+{
+    auto matcher() const
+    {
+        return match::name("mul")(
+            match::either_arg(0, 1)(
+                match::is_constant().bind("scalar"),
+                match::name("convolution").bind("conv")
+            ));
+    }
+
+    void apply(module& m, const match::matcher_result& r) const
+    {
+        auto ins      = r.result;
+        auto scalar   = r.instructions["scalar"];
+        auto conv_ins = r.instructions["conv"];
+
+        // Get the convol's input and weights
+        auto conv_input   = conv_ins->inputs().front();  // input to conv
+        auto conv_weights = conv_ins->inputs().back();   // weights of the conv
+
+        auto scaled_weights = m.insert_instruction(ins, make_op("mul"), scalar, conv_weights);
+
+        // new conv with modified weights
+        auto new_conv = m.insert_instruction(ins, conv_ins->get_operator(), conv_input, scaled_weights);
+
+        m.replace_instruction(ins, new_conv);
+    }
+};
+
 
 struct find_dot_add
 {
@@ -1981,6 +2045,7 @@
                             find_dot_slice{},
                             find_dot_mul{},
                             find_mul_add{},
+                            find_scalar_mul_conv{},
                             find_unit_ops{},
                             find_neg_unit_ops{},
                             eliminate_zero_point{},

src/simplify_reshapes.cpp

@@ -744,7 +744,8 @@ struct find_unary_shape_transforms
     {
         return ins->name() == "@literal" or
                ins->get_operator().attributes().contains("pointwise") or
-               contains(ins->name(), "reduce");
+               contains(ins->name(), "reduce") or
+               ins->get_operator().attributes().contains("scalar"); //For scalar operation
     }
 
     void apply(module& m, const match::matcher_result& mr) const