Merge branch 'broadcast-concat-dot' into sd-opt

src/normalize_attributes.cpp

@@ -40,10 +40,12 @@ inline namespace MIGRAPHX_INLINE_NS {
  *
  * See normalize_attribute.hpp for explaining the options.
  */
+template <class Message>
 auto tune_attribute(const std::vector<int64_t>& vec,
                     const std::vector<int64_t>& axes,
                     const value& val,
-                    const std::vector<std::size_t>& lens)
+                    const std::vector<std::size_t>& lens,
+                    Message m)
 {
     std::vector<int64_t> result(vec);
     int64_t n_rank                                 = lens.size();
@@ -84,14 +86,14 @@ auto tune_attribute(const std::vector<int64_t>& vec,
         {
             if(not std::equal(result.begin(), result.end(), max_vals.begin(), std::less_equal<>{}))
             {
-                MIGRAPHX_THROW("TUNE_VECTOR: value out of range!");
+                MIGRAPHX_THROW(m() + "value out of range!");
             }
         }
         else
         {
             if(not std::equal(result.begin(), result.end(), max_vals.begin(), std::less<>{}))
             {
-                MIGRAPHX_THROW("TUNE_VECTOR: value out of range!");
+                MIGRAPHX_THROW(m() + "value out of range!");
             }
         }
     }
@@ -124,14 +126,14 @@ auto tune_attribute(const std::vector<int64_t>& vec,
             if(not std::equal(
                    min_vals.begin(), min_vals.end(), result.begin(), std::less_equal<>{}))
             {
-                MIGRAPHX_THROW("TUNE_VECTOR: attribute out of range!");
+                MIGRAPHX_THROW(m() + "attribute out of range!");
             }
         }
         else
         {
             if(not std::equal(result.begin(), result.end(), min_vals.begin(), std::less<>{}))
             {
-                MIGRAPHX_THROW("TUNE_VECTOR: attribute out of range!");
+                MIGRAPHX_THROW(m() + "attribute out of range!");
             }
         }
     }
@@ -193,7 +195,8 @@ bool normalize_attributes(operation& op, const std::vector<std::size_t>& lens)
         const auto& key = rv.get_key();
         if(val.contains(key))
         {
-            auto vv = val.at(key).without_key();
+            auto message = [&] { return op.name() + ": " + key + ": "; };
+            auto vv      = val.at(key).without_key();
             if(vv.is_array())
             {
                 std::vector<int64_t> axes;
@@ -202,7 +205,7 @@ bool normalize_attributes(operation& op, const std::vector<std::size_t>& lens)
                     axes = val.at("axes").without_key().to_vector<int64_t>();
                 }
                 auto vec    = vv.to_vector<int64_t>();
-                auto result = tune_attribute(vec, axes, rv.without_key(), lens);
+                auto result = tune_attribute(vec, axes, rv.without_key(), lens, message);
                 val[key]    = result;
                 op.from_value(val);
                 val   = op.to_value();
@@ -211,7 +214,7 @@ bool normalize_attributes(operation& op, const std::vector<std::size_t>& lens)
             else
             {
                 auto num    = vv.to<int64_t>();
-                auto result = tune_attribute({num}, {num}, rv.without_key(), lens);
+                auto result = tune_attribute({num}, {num}, rv.without_key(), lens, message);
                 val[key]    = result.front();
                 op.from_value(val);
                 val   = op.to_value();

src/simplify_algebra.cpp

@@ -516,12 +516,72 @@ struct find_inner_broadcast
     }
 };
 
+struct find_dot_broadcast
+{
+    auto matcher() const
+    {
+        return match::name("dot")(match::all_of[match::inputs()](match::broadcast()));
+    }
+
+    void apply(module& m, const match::matcher_result& r) const
+    {
+        auto ins = r.result;
+        auto a   = ins->inputs()[0];
+        auto b   = ins->inputs()[1];
+        if(a->get_operator().name() != b->get_operator().name())
+            return;
+        if(ins->get_shape().lens().size() < 3)
+            return;
+        auto nbatch_axes      = ins->get_shape().lens().size() - 2;
+        const auto& a_strides = a->get_shape().strides();
+        const auto& b_strides = b->get_shape().strides();
+        // Find leading batch axes that are broadcasted
+        auto p =
+            std::mismatch(a_strides.begin(),
+                          a_strides.begin() + nbatch_axes,
+                          b_strides.begin(),
+                          b_strides.begin() + nbatch_axes,
+                          [](auto astride, auto bstride) { return astride == 0 and bstride == 0; });
+        auto naxes = p.first - a_strides.begin();
+        assert(naxes <= nbatch_axes);
+        std::vector<std::size_t> axes(naxes);
+        std::iota(axes.begin(), axes.end(), 0);
+
+        auto insert_squeeze = [&](instruction_ref b_ins) -> instruction_ref {
+            auto input = b_ins->inputs()[0];
+            std::vector<std::size_t> lens(b_ins->get_shape().lens().begin() + naxes,
+                                          b_ins->get_shape().lens().end());
+            if(b_ins->name() == "multibroadcast")
+            {
+                return m.insert_instruction(
+                    ins, make_op("multibroadcast", {{"out_lens", lens}}), input);
+            }
+            else if(b_ins->name() == "broadcast")
+            {
+                auto v    = b_ins->get_operator().to_value();
+                auto axis = v.at("axis").to<std::size_t>() - naxes;
+                return m.insert_instruction(
+                    ins, make_op("broadcast", {{"axis", axis}, {"out_lens", lens}}), input);
+            }
+            assert(false);
+            return m.end();
+        };
+        auto a1        = insert_squeeze(a);
+        auto b1        = insert_squeeze(b);
+        auto dot       = m.insert_instruction(ins, make_op("dot"), a1, b1);
+        auto broadcast = m.insert_instruction(
+            ins, make_op("multibroadcast", {{"out_lens", ins->get_shape().lens()}}), dot);
+        m.replace_instruction(ins, broadcast);
+    }
+};
+
 struct find_concat_op
 {
     auto matcher() const
     {
         return match::name("concat")(match::any_of[match::inputs()](
-            match::any_of(match::pointwise(), match::name("broadcast")), match::used_once()));
+            match::any_of(match::pointwise(), match::name("broadcast", "multibroadcast")),
+            match::used_once()));
     }
 
     template <class Iterator>
@@ -540,7 +600,8 @@ struct find_concat_op
 
     static bool is_valid_op(const operation& op)
     {
-        return op.name() == "broadcast" or op.attributes().contains("pointwise");
+        return contains({"broadcast", "multibroadcast"}, op.name()) or
+               op.attributes().contains("pointwise");
     }
 
     void apply(module& m, const match::matcher_result& r) const
@@ -568,6 +629,16 @@ struct find_concat_op
                 op               = b;
                 iaxis            = 0;
             }
+            else if(op.name() == "multibroadcast")
+            {
+                shape bshape = (*start)->get_shape();
+                auto input   = (*start)->inputs()[0];
+                if(iaxis >= bshape.strides().size() or bshape.strides()[iaxis] == 0)
+                    return {start, last};
+                op.from_value({{"out_lens", get_output_lens(start, last, iaxis)}});
+                auto delta = bshape.lens().size() - input->get_shape().lens().size();
+                iaxis -= delta;
+            }
 
             std::vector<instruction_ref> concats;
             for(std::size_t i = 0; i < x->inputs().size(); i++)
@@ -1388,6 +1459,7 @@ void simplify_algebra::apply(module& m) const
     {
         match::find_matches(m,
                             find_inner_broadcast{},
+                            find_dot_broadcast{},
                             find_double_add_lit_broadcast{},
                             find_add_lit_broadcast{},
                             find_add_convs{},

test/simplify_algebra_test.cpp

@@ -3003,6 +3003,38 @@ TEST_CASE(reorder_slice_ins_deps)
     EXPECT(m == create_module());
 }
 
+TEST_CASE(dot_broadcast_different_rank)
+{
+    migraphx::module m1;
+    {
+        auto x  = m1.add_parameter("x", {migraphx::shape::float_type, {768}});
+        auto y  = m1.add_parameter("y", {migraphx::shape::float_type, {768, 3072}});
+        auto xb = m1.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {2, 384, 768}}}), x);
+        auto yb = m1.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {2, 768, 3072}}}), y);
+        auto dot = m1.add_instruction(migraphx::make_op("dot"), xb, yb);
+        m1.add_return({dot});
+    };
+
+    migraphx::module m2;
+    {
+        auto x = m2.add_parameter("x", {migraphx::shape::float_type, {768}});
+        auto y = m2.add_parameter("y", {migraphx::shape::float_type, {768, 3072}});
+        auto xb =
+            m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {384, 768}}}), x);
+        auto yb =
+            m2.add_instruction(migraphx::make_op("multibroadcast", {{"out_lens", {768, 3072}}}), y);
+        auto dot       = m2.add_instruction(migraphx::make_op("dot"), xb, yb);
+        auto broadcast = m2.add_instruction(
+            migraphx::make_op("multibroadcast", {{"out_lens", {2, 384, 3072}}}), dot);
+        m2.add_return({broadcast});
+    };
+
+    run_pass(m1);
+    EXPECT(m1.sort() == m2.sort());
+}
+
 TEST_CASE(dot_fusion_reshape)
 {
     migraphx::module m1;