Fuse transposes in pointwise and reduce fusions

src/include/migraphx/rewrite_reshapes.hpp

@@ -33,6 +33,7 @@
 #include <migraphx/simplify_reshapes.hpp>
 #include <migraphx/eliminate_common_subexpression.hpp>
 #include <migraphx/dead_code_elimination.hpp>
+#include <migraphx/shape_transform_descriptor.hpp>
 
 namespace migraphx {
 inline namespace MIGRAPHX_INLINE_NS {
@@ -72,18 +73,19 @@ struct rewrite_reshapes
 
         auto matcher() const
         {
-            auto reshape =
-                match::name("reshape", "squeeze", "unsqueeze", "flatten")(match::used_once());
-            auto skip_contiguous_broadcast =
-                match::skip(match::name("contiguous", "multibroadcast")(match::used_once()));
-            auto skip_contiguous_broadcast_arg = [&](auto... ms) {
-                return match::arg(0)(skip_contiguous_broadcast(ms...));
-            };
+            auto reshapes          = match::name("reshape",
+                                        "squeeze",
+                                        "unsqueeze",
+                                        "flatten",
+                                        "transpose",
+                                        "contiguous",
+                                        "multibroadcast",
+                                        "broadcast")(match::used_once());
             auto pointwise         = match::name(op1)(match::used_once());
-            auto reshape_pointwise =
-                reshape(skip_contiguous_broadcast_arg(pointwise.bind("x"))).bind("reshape");
-            return match::name(op2)(match::any_of[match::inputs()](
-                skip_contiguous_broadcast(reshape_pointwise).bind("input")));
+            auto reshapes_pointwise =
+                reshapes(match::arg(0)(match::skip(reshapes())(pointwise.bind("x"))));
+            return match::name(op2)(
+                match::any_of[match::inputs()](reshapes_pointwise.bind("input")));
         }
 
         template <class F>
@@ -100,6 +102,12 @@ struct rewrite_reshapes
             return last;
         }
 
+        template <class F>
+        static bool any_input_of(instruction_ref start, instruction_ref last, F f)
+        {
+            return find_input_if(start, last, f) != last;
+        }
+
         static bool match_input(instruction_ref ins, instruction_ref x_ins)
         {
             if(ins->inputs().empty())
@@ -120,65 +128,76 @@ struct rewrite_reshapes
             return result;
         }
 
+        static bool is_broadcast(instruction_ref ins) { return ins->name() == "multibroadcast"; }
+
         void apply(module_pass_manager& mpm, const match::matcher_result& r) const
         {
             auto ins         = r.result;
             auto x_ins       = r.instructions["x"];
-            auto reshape_ins = r.instructions["reshape"];
             auto input_ins   = r.instructions["input"];
 
-            const auto has_broadcast_before_reshape = is_broadcasted(reshape_ins, x_ins);
-            const auto has_broadcast_after_reshape  = is_broadcasted(input_ins, reshape_ins);
-            if(not has_broadcast_before_reshape.has_value())
-                return;
-            if(not has_broadcast_after_reshape.has_value())
+            // If its just a broadcast then skip
+            if(not any_input_of(input_ins, x_ins, [](instruction_ref x) {
+                   return not contains({"multibroadcast", "broadcast", "contiguous"}, x->name());
+               }))
                 return;
-            if(*has_broadcast_after_reshape and *has_broadcast_before_reshape)
-                return;
-            const bool has_broadcast =
-                *has_broadcast_after_reshape or *has_broadcast_before_reshape;
 
             auto dims1 = T::base_dims(ins);
             auto dims2 = T::base_dims(x_ins);
 
             if(elements(dims1) != elements(dims2))
                 return;
 
-            auto cd = common_dims::compute(T::base_dims(ins), T::base_dims(x_ins));
-            if(cd.dims.empty())
-                return;
+            std::vector<operation> ops;
+            auto next_ins = input_ins;
+            while(next_ins != x_ins)
+            {
+                ops.push_back(next_ins->get_operator());
+                next_ins = next_ins->inputs().front();
+            }
+            assert(next_ins == x_ins);
+            std::reverse(ops.begin(), ops.end());
 
-            if(ins->name() != "pointwise" and not T::supports(ins, cd.dims, cd.axes_map1))
+            auto desc =
+                shape_transform_descriptor::create(x_ins->get_shape().lens(), ops).rebase(dims2);
+            if(desc.empty())
                 return;
-            if(x_ins->name() != "pointwise" and not T::supports(x_ins, cd.dims, cd.axes_map2))
-                return;
-
-            auto reshape_input = [&](const auto& ins_to_insert) {
-                return [&](auto input) {
-                    auto dims = cd.get_dimensions_for(input->get_shape().lens());
-                    return mpm.get_module().insert_instruction(
-                        ins_to_insert, make_op("reshape", {{"dims", dims}}), input);
+            auto cdims         = desc.common_dims();
+            auto reshape_input = [&](const auto& ins_to_insert, auto generate) {
+                return [&, generate](auto input) {
+                    auto gops  = std::invoke(generate, desc, input->get_shape().lens());
+                    auto start = input;
+                    for(const auto& op : gops)
+                    {
+                        start = mpm.get_module().insert_instruction(ins_to_insert, op, start);
+                    }
+                    return start;
                 };
             };
             auto x_inputs = x_ins->inputs();
             std::transform(
-                x_inputs.begin(), x_inputs.end(), x_inputs.begin(), reshape_input(x_ins));
-            auto new_x_ins = insert(mpm, x_ins, x_inputs, cd.axes_map2);
-            if(has_broadcast)
+                x_inputs.begin(),
+                x_inputs.end(),
+                x_inputs.begin(),
+                reshape_input(x_ins, &shape_transform_descriptor::generate_common_from_src));
+            auto new_x_ins = insert(mpm, x_ins, x_inputs, desc.common_axes_map_from_src());
+            if(new_x_ins->get_shape().lens() != cdims)
             {
                 new_x_ins = mpm.get_module().insert_instruction(
-                    x_ins, make_op("multibroadcast", {{"out_lens", cd.dims}}), new_x_ins);
+                    x_ins, make_op("multibroadcast", {{"out_lens", cdims}}), new_x_ins);
             }
 
             auto inputs = ins->inputs();
             std::transform(inputs.begin(), inputs.end(), inputs.begin(), [&](auto input) {
                 if(input == input_ins)
                     return new_x_ins;
-                return reshape_input(ins)(input);
+                return reshape_input(ins,
+                                     &shape_transform_descriptor::generate_common_from_dst)(input);
             });
-            auto pw = insert(mpm, ins, inputs, cd.axes_map1);
-            mpm.get_module().replace_instruction(
-                ins, make_op("reshape", {{"dims", ins->get_shape().lens()}}), pw);
+            auto pw = insert(mpm, ins, inputs, desc.common_axes_map_from_dst());
+            auto rins =
+                reshape_input(ins, &shape_transform_descriptor::generate_dst_from_common)(pw);
+            mpm.get_module().replace_instruction(ins, rins);
         }
 
         static bool same_dims(instruction_ref ins)

src/include/migraphx/shape_transform_descriptor.hpp

@@ -74,6 +74,11 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
     shape_transform_descriptor() = default;
     explicit shape_transform_descriptor(const std::vector<std::size_t>& dims);
 
+    static shape_transform_descriptor create(const std::vector<std::size_t>& dims,
+                                             const std::vector<operation>& ops);
+
+    shape_transform_descriptor rebase(const std::vector<std::size_t>& dims) const;
+
     bool apply(const std::vector<operation>& ops);
     bool apply_reshape(const std::vector<std::size_t>& rdims);
     bool apply_reshape_impl(const std::vector<std::size_t>& rdims);
@@ -84,6 +89,22 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
     std::size_t elements() const;
     std::vector<operation> generate() const;
 
+    bool has_broadcast() const;
+    void flatten_broadcast();
+
+    std::vector<std::size_t> common_dims(const std::vector<std::size_t>& input_dims = {}) const;
+    std::vector<operation>
+    generate_common_from_src(const std::vector<std::size_t>& input_dims = {}) const;
+    std::vector<operation>
+    generate_common_from_dst(const std::vector<std::size_t>& input_dims = {}) const;
+    std::vector<operation>
+    generate_dst_from_common(const std::vector<std::size_t>& input_dims = {}) const;
+    std::vector<std::vector<std::size_t>> common_axes_map_from_src() const;
+    std::vector<std::vector<std::size_t>> common_axes_map_from_dst() const;
+
+    bool empty() const;
+    std::vector<std::size_t> lens() const;
+
     struct MIGRAPHX_EXPORT dimension
     {
         void simplify();
@@ -105,6 +126,9 @@ struct MIGRAPHX_EXPORT shape_transform_descriptor
 
             void add_split_axis(std::size_t i);
 
+            void expose();
+            void hide();
+
             MIGRAPHX_EXPORT friend bool operator==(const sub& x, const sub& y);
             MIGRAPHX_EXPORT friend bool operator!=(const sub& x, const sub& y);
             MIGRAPHX_EXPORT friend std::ostream& operator<<(std::ostream& os, const sub& x);