Fuse gemm add with pointwise fusions

src/targets/gpu/fuse_ops.cpp

@@ -931,6 +931,67 @@ struct find_gemm_add
     }
 };
 
+auto pointwise_name(const std::string& s)
+{
+    return precompile_name("pointwise")(match::make_basic_pred_matcher([=](auto ins) {
+        module_ref pm = ins->module_inputs().front();
+        auto n = std::count_if(pm->begin(), pm->end(), [&](auto& i) {
+            return i.name() == s;
+        });
+        if (n != 1)
+            return false;
+        return std::all_of(pm->begin(), pm->end(), [&](auto& i) {
+            return starts_with(i.name(), "@") or i.name() == s;
+        });
+    }));
+}
+
+struct find_gemm_pointwise
+{
+    auto matcher() const
+    {
+        return pointwise_name("add")(
+            match::nargs(3),
+            match::all_of[match::inputs()](match::standard_shape()),
+            match::either_arg(0, 1)(match::used_once().bind("c"),
+                                    match::name("gpu::gemm")(match::nargs(3)).bind("gemm")));
+    }
+
+    void apply(module& m, const match::matcher_result& r) const
+    {
+        auto ins      = r.result;
+        auto gemm_ins = r.instructions["gemm"];
+        auto c_ins    = r.instructions["c"];
+
+        auto gemm = any_cast<rocblas_gemm<op::dot>>(gemm_ins->get_operator());
+
+        // Already fused gemm
+        if(not float_equal(gemm.beta, 0))
+            return;
+
+        if(std::any_of(ins->inputs().begin(), ins->inputs().end(), [](auto i) {
+               return not i->get_shape().standard();
+           }))
+            return;
+
+        auto inputs = gemm_ins->inputs();
+        inputs.pop_back();
+
+        auto copy_ins = c_ins;
+
+        // Insert copy
+        if(ins == m.end() or c_ins->outputs().size() > 1 or c_ins->inputs().empty())
+        {
+            copy_ins = m.insert_instruction(ins, hip_copy{}, c_ins, ins->inputs().back());
+        }
+        inputs.push_back(copy_ins);
+        inputs.push_back(copy_ins);
+
+        gemm.beta = 1;
+        m.replace_instruction(ins, gemm, inputs);
+    }
+};
+
 struct find_commutative_broadcast
 {
     auto matcher() const
@@ -967,7 +1028,7 @@ void fuse_ops::apply(module& m) const
                         find_add_unary{"gpu::tanh", hip_add_tanh{}, hip_triadd_tanh{}},
                         find_add_clip{});
     run_passes(m, {dead_code_elimination{}});
-    match::find_matches(m, find_triadd_layernorm{}, find_gemm_add{}, find_commutative_broadcast{});
+    match::find_matches(m, find_triadd_layernorm{}, find_gemm_add{}, find_gemm_pointwise{}, find_commutative_broadcast{});
 }
 
 } // namespace gpu