Add gemm_softmax_gemm

requirements.txt

@@ -28,4 +28,4 @@ ROCmSoftwarePlatform/half@rocm-5.6.0
 pybind/pybind11@d159a563383d10c821ba7b2a71905d1207db6de4 --build
 msgpack/msgpack-c@cpp-3.3.0 -DMSGPACK_BUILD_TESTS=Off
 sqlite3@3.17 -DCMAKE_POSITION_INDEPENDENT_CODE=On
-ROCmSoftwarePlatform/composable_kernel@5172ec5280f14974beee2acf1af1db3b2670244c -DCK_BUILD_JIT_LIB=On -DCMAKE_POSITION_INDEPENDENT_CODE=On
+# ROCmSoftwarePlatform/composable_kernel@5172ec5280f14974beee2acf1af1db3b2670244c -DCK_BUILD_JIT_LIB=On -DCMAKE_POSITION_INDEPENDENT_CODE=On

src/targets/gpu/fuse_ck.cpp

@@ -65,21 +65,62 @@ struct ck_gemm
             return r;
         return r.with_type(mods.front()->get_output_shapes().front().type());
     }
+
+    static bool is_ck_supported_type(shape::type_t t)
+    {
+        return contains({shape::half_type, shape::int8_type, shape::int32_type}, t);
+    }
 };
 MIGRAPHX_REGISTER_OP(ck_gemm);
 
-namespace {
 
-bool is_ck_supported_type(shape::type_t t)
+struct ck_gemm_softmax_gemm
 {
-    return contains({shape::half_type, shape::int8_type, shape::int32_type}, t);
-}
+    operation op = make_op("dot");
+
+    template <class Self, class F>
+    static auto reflect(Self& self, F f)
+    {
+        return pack(f(self.op, "op"));
+    }
+
+    std::string name() const { return "gpu::ck_gemm_softmax_gemm"; }
+
+    void check_gemm_shape(const shape& s) const
+    {
+        if(not contains(range(s.strides().rbegin(), s.strides().rbegin() + 3), 1))
+            MIGRAPHX_THROW("Invalid shape for ck_gemm_softmax_gemm");
+    }
+
+    shape compute_shape(std::vector<shape> inputs, const std::vector<module_ref>& mods) const
+    {
+        check_shapes{inputs, *this}.same_ndims();
+        if(inputs.size() < 2)
+            MIGRAPHX_THROW("should have at least two inputs.");
+        auto a  = inputs[0];
+        auto b  = inputs[1];
+        auto b1 = inputs[2];
+        for(const auto& input : inputs)
+        {
+            check_gemm_shape(input);
+        }
+        return op.compute_shape({op.compute_shape({a, b}), b1});
+    }
+
+    static bool is_ck_supported_type(shape::type_t t)
+    {
+        return contains({shape::half_type}, t);
+    }
+};
+MIGRAPHX_REGISTER_OP(ck_gemm_softmax_gemm);
+
+namespace {
 
 MIGRAPHX_PRED_MATCHER(is_ck_gemm, instruction_ref ins)
 {
     if(ins->name() != "dot" and ins->name() != "quant_dot")
         return false;
-    if(not is_ck_supported_type(ins->get_shape().type()))
+    if(not ck_gemm::is_ck_supported_type(ins->get_shape().type()))
         return false;
     auto a = ins->inputs().front()->get_shape();
     auto b = ins->inputs().back()->get_shape();
@@ -99,9 +140,38 @@ MIGRAPHX_PRED_MATCHER(is_ck_gemm, instruction_ref ins)
     // Skipping GEMMs with a K dimension greater than 2048 is a course-grained strategy
     // to avoid poor-performing GEMM kernels from CK
     // To-do: Investigate a more precise strategy
-    return k <= 2048;
+    return true;//k <= 2048;
 }
 
+struct find_ck_gemm_softmax_gemm
+{
+    auto matcher() const
+    {
+        auto gemm1 =
+            match::skip(match::name("contiguous"))(match::name("dot")(is_ck_gemm().bind("gemm1")));
+        auto mul     = match::name("mul")(match::any_of[match::inputs()](gemm1)).bind("scale");
+        auto softmax = match::name("softmax")(match::any_of[match::inputs()](mul)).bind("softmax");
+        return match::name("dot")(is_ck_gemm().bind("gemm2"))(
+            match::any_of[match::inputs()](softmax));
+    }
+
+    void apply(module_pass_manager& mpm, const match::matcher_result& r) const
+    {
+        auto ins       = r.result;
+        auto gemm2_ins = r.instructions["gemm2"];
+        auto gemm1_ins = r.instructions["gemm1"];
+
+        // if (not ck_gemm_softmax_gemm::is_ck_supported_type(gemm1_ins->get_shape().type()))
+        //     return;
+        
+        auto inputs = gemm1_ins->inputs();            // A, B
+        inputs.push_back(gemm2_ins->inputs().back()); // B1
+
+        mpm.get_module().replace_instruction(
+            ins, ck_gemm_softmax_gemm{gemm2_ins->get_operator()}, inputs);
+    }
+};
+
 struct find_ck_gemm_pointwise
 {
     // Find a gemm followed by a pointwise operation.
@@ -127,7 +197,11 @@ struct find_ck_gemm_pointwise
            ins->get_shape().type() != gemm_ins->get_shape().type())
             return;
         if(std::any_of(ins->inputs().begin(), ins->inputs().end(), [](auto input) {
-               return not is_ck_supported_type(input->get_shape().type());
+               return not ck_gemm::is_ck_supported_type(input->get_shape().type());
+           }))
+            return;
+        if(std::any_of(ins->inputs().begin(), ins->inputs().end(), [](auto input) {
+               return not input->inputs().empty() and input->inputs().front()->name() == "capture";
            }))
             return;
         assert(gemm_it != inputs.end());
@@ -152,7 +226,7 @@ struct find_ck_gemm_pointwise
 
 struct find_ck_gemm
 {
-    auto matcher() const { return match::name("dot")(is_ck_gemm().bind("gemm")); }
+    auto matcher() const { return match::name("dot", "quant_dot")(is_ck_gemm().bind("gemm")); }
 
     void apply(module_pass_manager& mpm, const match::matcher_result& r) const
     {
@@ -165,6 +239,7 @@ struct find_ck_gemm
 
 void fuse_ck::apply(module_pass_manager& mpm) const
 {
+    match::find_matches(mpm, find_ck_gemm_softmax_gemm{});
     match::find_matches(mpm, find_ck_gemm_pointwise{});
     match::find_matches(mpm, find_ck_gemm{});
 }

src/targets/gpu/jit/ck_gemm_softmax_gemm.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#include <fstream>
+#include <migraphx/filesystem.hpp>
+#include <migraphx/gpu/compiler.hpp>
+#include <migraphx/make_op.hpp>
+#include <migraphx/gpu/context.hpp>
+
+#include <migraphx/env.hpp>
+#include <migraphx/file_buffer.hpp>
+#include <migraphx/gpu/compile_gen.hpp>
+#include <migraphx/gpu/compile_hip.hpp>
+#include <migraphx/gpu/compile_hip_code_object.hpp>
+#include <migraphx/module.hpp>
+#include <migraphx/ranges.hpp>
+#include <migraphx/reduce_dims.hpp>
+#include <migraphx/stringutils.hpp>
+
+#include "ck/host/device_batched_gemm_softmax_gemm.hpp"
+
+namespace migraphx {
+inline namespace MIGRAPHX_INLINE_NS {
+
+namespace gpu {
+
+using namespace migraphx::gpu::gen; // NOLINT
+
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_LOG_CK_GEMM);
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_CK_TUNING);
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_CK_TUNING_VALUE);
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_CK_DEBUG);
+MIGRAPHX_DECLARE_ENV_VAR(MIGRAPHX_TUNE_CK);
+
+// NOLINTNEXTLINE
+static const char* const ck_gemm_softmax_gemm_kernel = R"__migraphx__(
+#include <args.hpp>
+#include <migraphx/kernels/ck_gemm_softmax_gemm.hpp>
+#include <migraphx/kernels/pointwise.hpp>
+#include <migraphx/kernels/ops.hpp>
+#include <${include}>
+
+namespace migraphx {
+
+${preamble}
+
+extern "C" {
+
+MIGRAPHX_GLOBAL void ${kernel}(${params})
+{
+    transform_args(make_tensors(), rotate_last())(${args})([](auto... xs) {
+        ck_gemm_softmax_gemm<${solution}, ${blocks_per_batch}>(xs...);
+    });
+}
+
+}
+
+} // namespace migraphx
+
+)__migraphx__";
+
+// NOLINTNEXTLINE
+static const char* const disable_warning_pragma = R"__migraphx__(
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Weverything"
+${content}
+#pragma clang diagnostic pop
+)__migraphx__";
+
+template <class P>
+static std::string ck_disable_warnings(P p)
+{
+    return interpolate_string(disable_warning_pragma,
+                              {{"content", std::string{p.first, p.second}}});
+}
+
+static std::unordered_map<std::string, std::string> create_ck_header_strings()
+{
+    std::unordered_map<std::string, std::string> result;
+    auto ck_headers = ck::host::GetHeaders();
+
+    std::transform(
+        ck_headers.begin(), ck_headers.end(), std::inserter(result, result.begin()), [&](auto&& p) {
+            return std::make_pair(p.first, ck_disable_warnings(p.second));
+        });
+    return result;
+}
+
+static std::vector<src_file> create_ck_headers()
+{
+    static const auto& header_strings = create_ck_header_strings();
+    std::vector<src_file> srcs;
+    std::transform(
+        header_strings.begin(), header_strings.end(), std::back_inserter(srcs), [&](auto&& p) {
+            return src_file{fs::path{p.first},
+                            {p.second.data(), p.second.data() + p.second.size()}};
+        });
+    return srcs;
+}
+
+static const std::vector<src_file>& ck_headers()
+{
+    static const auto& headers = create_ck_headers();
+    return headers;
+}
+
+static bool transposed_matrix(const shape& s) { return s.strides().back() != 1; }
+
+using tuning_entry = std::pair<std::vector<shape>, size_t>;
+static std::vector<tuning_entry> read_tuning(const std::string& s)
+{
+    if(not fs::exists(s))
+        return {};
+    return from_value<std::vector<tuning_entry>>(from_json_string(read_string(s)));
+}
+
+static float matrix_distance(const shape& x, const shape& y)
+{
+    if(x.type() != y.type())
+        return std::numeric_limits<float>::max();
+    if(transposed_matrix(x) != transposed_matrix(y))
+        return std::numeric_limits<float>::max();
+    auto sum_squared = std::inner_product(x.lens().rbegin(),
+                                          x.lens().rbegin() + 2,
+                                          y.lens().rbegin(),
+                                          0,
+                                          std::plus<>{},
+                                          [](auto a, auto b) { return (a - b) * (a - b); });
+    return std::sqrt(sum_squared);
+}
+
+static std::size_t get_tuning_for(const std::vector<shape>& inputs)
+{
+    static auto tuning = read_tuning(string_value_of(MIGRAPHX_CK_TUNING{}, ""));
+    if(tuning.empty())
+    {
+        std::cout << "*********** Warning: No CK tuning! for config:" << std::endl;
+        std::cout << "  " << inputs[0] << std::endl;
+        std::cout << "  " << inputs[1] << std::endl;
+        std::cout << "  " << inputs[2] << std::endl;
+    }
+    auto it = std::find_if(
+        tuning.begin(), tuning.end(), [&](const auto& p) { return p.first == inputs; });
+    if(it == tuning.end())
+    {
+        std::cout << "*********** Warning: CK tuning missing for config!" << std::endl;
+        std::cout << "  " << inputs[0] << std::endl;
+        std::cout << "  " << inputs[1] << std::endl;
+        std::cout << "  " << inputs[2] << std::endl;
+        std::vector<std::pair<float, std::size_t>> w;
+        std::transform(tuning.begin(), tuning.end(), std::back_inserter(w), [&](const auto& p) {
+            if(inputs.size() < 3 or p.first.size() < 3)
+                MIGRAPHX_THROW("Invalid CK config");
+            auto avg_distance = std::inner_product(
+                p.first.begin(),
+                p.first.begin() + 3,
+                inputs.begin(),
+                0.0f,
+                std::plus<>{},
+                [](const auto& x, const auto& y) { return matrix_distance(x, y) / 3.0f; });
+            return std::make_pair(avg_distance, p.second);
+        });
+        std::sort(w.begin(), w.end());
+        std::size_t default_value = 4;
+        if(not w.empty())
+            default_value = w.front().second;
+        auto tuning_val = value_of(MIGRAPHX_CK_TUNING_VALUE{}, default_value);
+        std::cout << "*********** Warning: CK try tuning: " << tuning_val << std::endl;
+        return tuning_val;
+    }
+    return it->second;
+}
+
+struct ck_gemm_softmax_gemm_compiler : compiler<ck_gemm_softmax_gemm_compiler>
+{
+    static std::string get_layout(const shape& s)
+    {
+        return transposed_matrix(s) ? "ck::tensor_layout::gemm::ColumnMajor"
+                                    : "ck::tensor_layout::gemm::RowMajor";
+    }
+
+    static ck::host::DataType get_type(const shape& s)
+    {
+        if(s.type() == shape::half_type)
+            return ck::host::DataType::Half;
+        else if(s.type() == shape::float_type)
+            return ck::host::DataType::Float;
+        else if(s.type() == shape::int8_type)
+            return ck::host::DataType::Int8;
+        else if(s.type() == shape::int32_type)
+            return ck::host::DataType::Int32;
+        MIGRAPHX_THROW("Unsupported ck type");
+    }
+
+    template <class Iterator, class F>
+    static std::string ck_tuple(Iterator start, Iterator last, F f)
+    {
+        std::vector<std::string> s;
+        std::transform(start, last, std::back_inserter(s), f);
+        return "ck::Tuple<" + join_strings(s, ",") + ">";
+    }
+
+    static std::vector<shape> adjust_inputs(std::vector<shape> inputs, bool& swap_inputs)
+    {
+        swap_inputs  = false;
+        auto c_shape = inputs.back();
+        if(not transposed_matrix(c_shape))
+            return inputs;
+        std::vector<int64_t> perm(c_shape.lens().size());
+        std::iota(perm.begin(), perm.end(), 0);
+        std::swap(perm[perm.size() - 1], perm[perm.size() - 2]);
+        std::transform(inputs.begin(), inputs.end(), inputs.begin(), [&](shape s) {
+            return reorder_shape(s, perm);
+        });
+        swap_inputs = true;
+        return inputs;
+    }
+
+    static std::size_t get_batch_count(const shape& s)
+    {
+        return std::accumulate(
+            s.lens().rbegin() + 2, s.lens().rend(), std::size_t{1}, std::multiplies<std::size_t>());
+    }
+
+    static void fold_batch_dims(shape& s)
+    {
+        auto lens = s.lens();
+        if(lens.size() <= 2)
+            return;
+        auto batch_count = get_batch_count(s);
+        auto m1          = lens.at(lens.size() - 2);
+        auto m2          = lens.at(lens.size() - 1);
+        if(transposed_matrix(s))
+            s = shape{s.type(), {m1, m2 * batch_count}};
+        else
+            s = shape{s.type(), {m1 * batch_count, m2}};
+    }
+
+    static void remove_batch_dims(shape& s)
+    {
+        auto lens = s.lens();
+        if(lens.size() <= 2)
+            return;
+        auto m1 = lens.at(lens.size() - 2);
+        auto m2 = lens.at(lens.size() - 1);
+        s       = shape{s.type(), {m1, m2}};
+    }
+
+    std::vector<std::string> names() const { return {"ck_gemm_softmax_gemm", "gpu::ck_gemm_softmax_gemm"}; }
+
+    static bool standard_batch(const shape& s)
+    {
+        if(s.lens().size() < 3)
+            return true;
+        std::vector<std::size_t> lens(s.lens().begin(), s.lens().end() - 2);
+        std::vector<std::size_t> strides(s.strides().begin(), s.strides().end() - 2);
+        auto base = *(s.lens().end() - 2) * *(s.lens().end() - 1);
+        std::transform(strides.begin(), strides.end(), strides.begin(), [&](auto stride) {
+            return stride / base;
+        });
+        return shape{s.type(), lens, strides}.standard();
+    }
+
+    bool can_fold_batch(const std::vector<shape>& inputs) const
+    {
+        const auto& b_shape = inputs[1];
+        if(std::any_of(inputs.begin() + 2, inputs.end() - 1, [](auto input) {
+               return not standard_batch(input);
+           }))
+            return false;
+        const auto& b_strides = b_shape.strides();
+        return std::all_of(
+            b_strides.begin(), b_strides.end() - 2, [](auto stride) { return stride == 0; });
+    }
+
+    ck::host::device_batched_gemm_softmax_gemm::Problem create_problem(const std::vector<shape>& inputs,
+                                                             const value& v) const
+    {
+        const auto& a_shape = inputs[0];
+        const auto& b_shape = inputs[1];
+        const auto& b1_shape = inputs[2];
+        const auto& c_shape = inputs.back();
+
+        // cppcheck-suppress unreadVariable
+        auto rank = a_shape.ndim();
+
+        auto batch_count = get_batch_count(c_shape);
+        auto m           = c_shape.lens()[rank - 2];
+        m                = can_fold_batch(inputs) ? m * batch_count : m;
+        auto n           = c_shape.lens().back();
+        auto k           = a_shape.lens().back();
+        auto o           = c_shape.lens().back();
+
+        const bool trans_a = transposed_matrix(a_shape);
+        const bool trans_b = transposed_matrix(b_shape);
+        const bool trans_b1 = transposed_matrix(b1_shape);
+        const bool trans_c = transposed_matrix(c_shape);
+        const auto a_type  = get_type(a_shape);
+        const auto b_type  = get_type(b_shape);
+        const auto b1_type  = get_type(b1_shape);
+        const auto c_type  = get_type(c_shape);
+        const auto scale = 1.0f;
+
+        std::string ck_passthrough = "ck_passthrough";
+        std::string cde_op         = ck_passthrough;
+        
+        /// update params after adding to jitlib
+        return ck::host::device_batched_gemm_softmax_gemm::Problem{m,
+                                                         n,
+                                                         k,
+                                                         o,
+                                                         trans_a,
+                                                         trans_b,
+                                                         trans_b1,
+                                                         trans_c,
+                                                         a_type,
+                                                         b_type,
+                                                         b1_type,
+                                                         c_type,
+                                                         ck_passthrough,
+                                                         ck_passthrough,
+                                                         ck_passthrough,
+                                                         ck_passthrough,
+                                                         scale};
+    }
+
+    operation compile_op(context& ctx, const std::vector<shape>& inputs, const value& v) const
+    {
+        const auto& a_shape = inputs[0];
+        const auto& b_shape = inputs[1];
+        const auto& c_shape = inputs.back();
+        /// update for 4-arg lookup?
+        auto tuning_value   = v.get("tuning_value", 4);
+        if(not v.contains("tuning_value"))
+            tuning_value = get_tuning_for({a_shape, b_shape, c_shape});
+        auto batch_count = get_batch_count(c_shape);
+        auto problem     = create_problem(inputs, v);
+
+        const auto include_header   = problem.GetIncludeHeader();
+        const auto solutions        = problem.GetSolutions(ctx.get_current_device().get_gfx_name());
+        const auto& solution        = solutions.at(tuning_value);
+        const auto template_str     = solution.template_str;
+        const auto blocks_per_batch = solution.grid_size;
+        const auto block_size       = solution.block_size;
+
+        hip_compile_options options;
+        options.additional_src_files = ck_headers();
+        auto grid_size = can_fold_batch(inputs) ? blocks_per_batch : batch_count * blocks_per_batch;
+        options.set_launch_params(v, grid_size * block_size, block_size);
+        options.inputs         = inputs;
+        options.output         = c_shape;
+        options.kernel_name    = v.get("kernel", "ck_gemm_softmax_gemm_kernel");
+        options.virtual_inputs = inputs;
+        if(can_fold_batch(inputs))
+        {
+            auto vinputs = inputs;
+            fold_batch_dims(vinputs[0]);
+            remove_batch_dims(vinputs[1]);
+            std::for_each(vinputs.begin() + 2, vinputs.end(), fold_batch_dims);
+            options.virtual_inputs = vinputs;
+        }
+
+        if(v.get("check", false) or enabled(MIGRAPHX_CK_DEBUG{}))
+            options.params += " -DMIGRAPHX_CK_CHECK=1";
+
+        auto src = interpolate_string(ck_gemm_softmax_gemm_kernel,
+                                      {{"solution", template_str},
+                                       {"include", include_header},
+                                       {"params", enum_params(inputs.size(), "void * private_p")},
+                                       {"args", enum_params(inputs.size(), "private_p")},
+                                       {"blocks_per_batch", to_string(blocks_per_batch)},
+                                       {"preamble", v.get("preamble", std::string{})},
+                                       {"kernel", options.kernel_name}});
+
+        return compile_hip_code_object(src, options);
+    }
+
+    value create_settings(instruction_ref ins, const operation& op) const
+    {
+        auto v      = op.to_value();
+        v["kernel"] = "ck_gemm_softmax_gemm_kernel";
+        if(not ins->module_inputs().empty())
+        {
+            auto* pm      = ins->module_inputs().front();
+            v["preamble"] = generate_pointwise(*pm, "post_ck_gemm_softmax_gemm_function") +
+                            "\nMIGRAPHX_LIFT_CLASS(post_ck_gemm_softmax_gemm, post_ck_gemm_softmax_gemm_function);";
+            v["post"]   = "ck_function_adaptor<post_ck_gemm_softmax_gemm>";
+            v["kernel"] = "ck_gemm_softmax_gemm_" + generate_name_from_ops(*pm) + "_kernel";
+        }
+        return v;
+    }
+
+    compiler_replace
+    compile(context& ctx, instruction_ref ins, const operation& op, const value& solution) const
+    {
+        auto shapes = to_shapes(ins->inputs());
+        auto v      = create_settings(ins, op);
+        if(not solution.is_null())
+            v["tuning_value"] = solution;
+        return {compile_op(ctx, shapes, v),
+                [=](module& m, instruction_ref ins2, const operation& code_object) {
+                    if(enabled(MIGRAPHX_LOG_CK_GEMM{}))
+                    {
+                        std::vector<shape> gemm_shapes{
+                            shapes[0], shapes[1], shapes.back().with_type(shapes[0].type())};
+                        std::cout << "gpu::ck_gemm_softmax_gemm: " << to_json_string(to_value(gemm_shapes))
+                                  << std::endl;
+                    }
+                    m.replace_instruction(ins2, code_object, ins2->inputs());
+                }};
+    }
+
+    optional<tuning_config>
+    get_tuning_config(context& ctx, instruction_ref ins, const operation& op, bool exhaustive) const
+    {
+        if(not exhaustive and not enabled(MIGRAPHX_TUNE_CK{}))
+            return nullopt;
+        tuning_config tc;
+        auto shapes    = to_shapes(ins->inputs());
+        auto problem   = create_problem(shapes, create_settings(ins, op));
+        auto solutions = problem.GetSolutions(ctx.get_current_device().get_gfx_name());
+        tc.solutions.resize(solutions.size());
+        std::iota(tc.solutions.begin(), tc.solutions.end(), 0);
+        std::vector<shape> gemm_shapes{shapes[0], shapes[1], shapes.back()};
+        tc.problem = to_value(gemm_shapes);
+        return tc;
+    }
+};
+
+} // namespace gpu
+} // namespace MIGRAPHX_INLINE_NS
+} // namespace migraphx

src/targets/gpu/kernels/include/migraphx/kernels/ck_gemm_softmax_gemm.hpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+#ifndef MIGRAPHX_GUARD_KERNELS_CK_GEMM_HPP
+#define MIGRAPHX_GUARD_KERNELS_CK_GEMM_HPP
+
+#include <migraphx/kernels/index.hpp>
+#include <migraphx/kernels/algorithm.hpp>
+#include <migraphx/kernels/integral_constant.hpp>
+#include <migraphx/kernels/tensor_view.hpp>
+#include <migraphx/kernels/ck.hpp>
+#include <migraphx/kernels/gemm_batcher.hpp>
+
+namespace migraphx {
+
+// In CK, the B matrix is ordered as N,K instead of K,N
+template <class Dims>
+constexpr auto ck_transposeb_dims(Dims dims)
+{
+    return unpack(dims, [](auto k, auto n) { return make_const_array(n, k); });
+}
+
+template <class Tensor>
+using ck_transposeb = decltype(make_shape(ck_transposeb_dims(get_shape_c<Tensor>{}.lens),
+                                          ck_transposeb_dims(get_shape_c<Tensor>{}.strides)));
+
+template <class G, class C, class A, class B, class B1>
+__device__ void ck_gemm_softmax_gemm_matrix(C c, A a, B b, B1 b1)
+{
+    constexpr auto desc = G::make_descriptor(to_ck_tensor<A>(),
+                                             to_ck_tensor<ck_transposeb<B>>(),
+                                             to_ck_tensor<ck_transposeb<B1>>(),
+                                             to_ck_tensor<C>());
+
+    static_assert(desc.IsValid(), "Invalid ck gemm.");
+
+    G::Run(desc,
+           to_ck_const_pointer(a.data()),
+           to_ck_const_pointer(b.data()),
+           to_ck_const_pointer(b1.data()),
+           to_ck_pointer(c.data()));
+}
+
+template <class G, index_int BlocksPerBatch, class... Ts>
+__device__ void ck_gemm_softmax_gemm(Ts... xs)
+{
+    gemm_batch_args(make_index(), _c<BlocksPerBatch>, xs...)(
+        [](auto... ys) { ck_gemm_softmax_gemm_matrix<G>(ys...); });
+}
+
+} // namespace migraphx
+#endif

test/onnx/gen_onnx.py

@@ -7799,3 +7799,82 @@ def where_mixed_test():
                                  outputs=['z'])
 
     return ([node], [c, x, y], [z])
+
+
+@onnx_test()
+def gemm_softmax_gemm_test():
+    a = helper.make_tensor_value_info('a', TensorProto.FLOAT16, [1, 1])
+    b = helper.make_tensor_value_info('b', TensorProto.FLOAT16, [1, 1])
+    # c = helper.make_tensor_value_info('c', TensorProto.FLOAT16, [1, 1])
+    b1 = helper.make_tensor_value_info('b1', TensorProto.FLOAT16, [1, 1])
+    # bias = helper.make_tensor_value_info('bias', TensorProto.FLOAT16, [1, 1])
+    out = helper.make_tensor_value_info('out', TensorProto.FLOAT16, [1, 1])
+
+    scale_array = np.array([1])
+    bias_array = np.array([0])
+
+    scale_tensor = helper.make_tensor(name='scale',
+                                    data_type=TensorProto.FLOAT16,
+                                    dims=[1, 1],
+                                    vals=[1])
+    bias_tensor = helper.make_tensor(name='bias',
+                                    data_type=TensorProto.FLOAT16,
+                                    dims=[1, 1],
+                                    vals=[0])
+
+    gemm1 = onnx.helper.make_node('MatMul',
+                                 inputs=['a', 'b'],
+                                 outputs=['gemm1_out'])
+    mul1 = onnx.helper.make_node('Mul',
+                                 inputs=['gemm1_out', 'scale'],
+                                 outputs=['mul1_out'])
+    add1 = onnx.helper.make_node('Add',
+                                 inputs=['mul1_out', 'bias'],
+                                 outputs=['add1_out'])
+    softmax = onnx.helper.make_node('Softmax',
+                                 inputs=['add1_out'],
+                                 outputs=['softmax_out'])
+    gemm2 = onnx.helper.make_node('MatMul',
+                                 inputs=['softmax_out', 'b1'],
+                                 outputs=['out'])
+
+
+    return ([gemm1, mul1, add1, softmax, gemm2], [a, b, b1], [out], [scale_tensor, bias_tensor])
+
+
+
+@onnx_test()
+def old_gemm_softmax_gemm_test():
+    a = helper.make_tensor_value_info('a', TensorProto.FLOAT16, [1, 1])
+    b = helper.make_tensor_value_info('b', TensorProto.FLOAT16, [1, 1])
+    c = helper.make_tensor_value_info('c', TensorProto.FLOAT16, [1, 1])
+    b1 = helper.make_tensor_value_info('b1', TensorProto.FLOAT16, [1, 1])
+    bias = helper.make_tensor_value_info('bias', TensorProto.FLOAT16, [1, 1])
+    out = helper.make_tensor_value_info('out', TensorProto.FLOAT16, [1, 1])
+
+    scale_array = np.array([(1/8)])
+
+    scale_tensor = helper.make_tensor('scale',
+                                    TensorProto.FLOAT16,
+                                    [1, 1],
+                                    [1])
+
+    gemm1 = onnx.helper.make_node('MatMul',
+                                 inputs=['a', 'b'],
+                                 outputs=['gemm1_out'])
+    mul1 = onnx.helper.make_node('Mul',
+                                 inputs=['gemm1_out', 'scale'],
+                                 outputs=['mul1_out'])
+    add1 = onnx.helper.make_node('Add',
+                                 inputs=['mul1_out', 'c'],
+                                 outputs=['add1_out'])
+    softmax = onnx.helper.make_node('Softmax',
+                                 inputs=['add1_out'],
+                                 outputs=['softmax_out'])
+    gemm2 = onnx.helper.make_node('MatMul',
+                                 inputs=['softmax_out', 'b1'],
+                                 outputs=['out'])
+
+
+    return ([gemm1, mul1, add1, softmax, gemm2], [a, b, c, b1, bias], [out], [scale_tensor])
+

test/onnx/old_gemm_softmax_gemm_test.onnx

+old_gemm_softmax_gemm_test:
+
+
a
+
b	gemm1_out"MatMul
+!
+	gemm1_out
+scalemul1_out"Mul
+
+mul1_out
+
cadd1_out"Add
+ 
+add1_outsoftmax_out"Softmax
+
+softmax_out
+b1out"MatMulold_gemm_softmax_gemm_test*


+*

BscaleZ
+
a
+
+
+

+
Z
+
b
+
+
+

+
Z
+
c
+
+
+

+
Z
+b1
+
+
+

+
Z
+bias
+
+
+

+
b
+out
+
+
+

+
B
\ No newline at end of file

test/verify/ck_gemm_softmax_gemm.cpp

+/*
+ * The MIT License (MIT)
+ *
+ * Copyright (c) 2015-2022 Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * Permission is hereby granted, free of charge, to any person obtaining a copy
+ * of this software and associated documentation files (the "Software"), to deal
+ * in the Software without restriction, including without limitation the rights
+ * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ * copies of the Software, and to permit persons to whom the Software is
+ * furnished to do so, subject to the following conditions:
+ *
+ * The above copyright notice and this permission notice shall be included in
+ * all copies or substantial portions of the Software.
+ *
+ * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
+ * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ * THE SOFTWARE.
+ */
+
+#include "verify_program.hpp"
+#include <migraphx/program.hpp>
+#include <migraphx/generate.hpp>
+#include <migraphx/make_op.hpp>
+
+struct ck_gemm_softmax_gemm : verify_program<ck_gemm_softmax_gemm>
+{
+    migraphx::program create_program() const
+    {
+        migraphx::program p;
+        auto* mm = p.get_main_module();
+        migraphx::shape m1_shape{migraphx::shape::half_type, {1, 12, 256, 256}};
+        migraphx::shape m2_shape{migraphx::shape::half_type, {1, 12, 256, 256}};
+        auto m2_elements = 1 * 12 * 256 * 256;
+        auto a   = mm->add_parameter("1", m1_shape);
+        auto b   = mm->add_parameter("2", m1_shape);
+        auto b1   = mm->add_parameter("3", m1_shape);
+        auto c   = mm->add_parameter("4", m1_shape);
+        std::vector<float> eights(m2_elements, 0.125);
+        auto eight = mm->add_literal(migraphx::literal{m2_shape, eights});
+        std::vector<float> zeros(m2_elements, 0);
+        auto zero = mm->add_literal(migraphx::literal{m2_shape, zeros});
+        std::vector<float> ones(m2_elements, 1);
+        auto one = mm->add_literal(migraphx::literal{m2_shape, ones});
+
+        b = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 1, 3, 2}}}), b);
+        auto gemm1 = mm->add_instruction(migraphx::make_op("dot"), a, b);
+        auto scale = mm->add_instruction(migraphx::make_op("mul"), gemm1, eight);
+        auto bias = mm->add_instruction(migraphx::make_op("add"), scale, zero);
+        auto softmax = mm->add_instruction(migraphx::make_op("softmax", {{"axis", -1}}), bias);
+        mm->add_instruction(migraphx::make_op("dot"), softmax, b1);
+
+        return p;
+    }
+};