Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "170ae09602b724db69fbbc06f12abf7ce2a23165"
Commit f3fcfcc7 authored by Alan Turner's avatar Alan Turner
Browse files

Fix fusion pass and add tuning

parent 3133fd79
Expand all Show whitespace changes
Inline Side-by-side
Showing with 1067 additions and 2651 deletions
src/targets/gpu/CMakeLists.txt
View file @ f3fcfcc7
......@@ -90,6 +90,7 @@ add_library(migraphx_gpu
device_name.cpp
elu.cpp
fuse_ck.cpp
fuse_ck_gemm_softmax_gemm.cpp
fuse_mlir.cpp
fuse_ops.cpp
gather.cpp
......
src/targets/gpu/fuse_ck.cpp
View file @ f3fcfcc7
......@@ -49,43 +49,43 @@ struct ck_gemm
};
MIGRAPHX_REGISTER_OP(ck_gemm);
struct ck_gemm_scale_bias_softmax_gemm
{
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_scale_bias_softmax_gemm");
}
shape compute_shape(std::vector<shape> inputs, const std::vector<module_ref>& mods) const
{
check_shapes{inputs, *this}.same_ndims();
// if(mods.size() != 1)
// MIGRAPHX_THROW("should have one submodule.");
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)
{
// std::cout << input << std::endl;
check_gemm_shape(input);
}
return op.compute_shape({op.compute_shape({a, b}), b1});
}
};
MIGRAPHX_REGISTER_OP(ck_gemm_scale_bias_softmax_gemm);
// struct ck_gemm_scale_bias_softmax_gemm
// {
// 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_scale_bias_softmax_gemm");
// }
// shape compute_shape(std::vector<shape> inputs, const std::vector<module_ref>& mods) const
// {
// check_shapes{inputs, *this}.same_ndims();
// // if(mods.size() != 1)
// // MIGRAPHX_THROW("should have one submodule.");
// 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)
// {
// // std::cout << input << std::endl;
// check_gemm_shape(input);
// }
// return op.compute_shape({op.compute_shape({a, b}), b1});
// }
// };
// MIGRAPHX_REGISTER_OP(ck_gemm_scale_bias_softmax_gemm);
namespace {
......@@ -156,38 +156,38 @@ struct find_ck_gemm
struct find_ck_gemm_scale_bias_softmax_gemm
{
auto matcher() const
{
auto gemm1 =
match::skip(match::name("contiguous"))(match::name("dot")(is_ck_gemm().bind("gemm1")));
auto pw =
match::name("pointwise")(match::any_of[match::inputs()](gemm1)).bind("scale_bias");
auto softmax = match::name("softmax")(match::any_of[match::inputs()](pw)).bind("softmax");
return match::name("dot")(is_ck_gemm().bind("gemm2"))(
match::any_of[match::inputs()](softmax));
}
// auto matcher() const
// {
// auto gemm1 =
// match::skip(match::name("contiguous"))(match::name("dot")(is_ck_gemm().bind("gemm1")));
// auto pw =
// match::name("pointwise")(match::any_of[match::inputs()](gemm1)).bind("scale_bias");
// auto softmax = match::name("softmax")(match::any_of[match::inputs()](pw)).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
{
std::cout << "Matched" << std::endl;
auto ins = r.result;
auto gemm2_ins = r.instructions["gemm2"];
auto sm_ins = r.instructions["softmax"];
auto pw_ins = r.instructions["scale_bias"];
auto gemm1_ins = r.instructions["gemm1"];
gemm2_ins->debug_print();
sm_ins->debug_print();
pw_ins->debug_print();
gemm1_ins->debug_print();
auto inputs = gemm1_ins->inputs(); // A, B
inputs.push_back(gemm2_ins->inputs().back()); // B1
// inputs.push_back(pw_ins->inputs().back()); // C
mpm.get_module().replace_instruction(
ins, ck_gemm_scale_bias_softmax_gemm{gemm2_ins->get_operator()}, inputs);
}
// void apply(module_pass_manager& mpm, const match::matcher_result& r) const
// {
// std::cout << "Matched" << std::endl;
// auto ins = r.result;
// auto gemm2_ins = r.instructions["gemm2"];
// auto sm_ins = r.instructions["softmax"];
// auto pw_ins = r.instructions["scale_bias"];
// auto gemm1_ins = r.instructions["gemm1"];
// gemm2_ins->debug_print();
// sm_ins->debug_print();
// pw_ins->debug_print();
// gemm1_ins->debug_print();
// auto inputs = gemm1_ins->inputs(); // A, B
// inputs.push_back(gemm2_ins->inputs().back()); // B1
// // inputs.push_back(pw_ins->inputs().back()); // C
// mpm.get_module().replace_instruction(
// ins, ck_gemm_scale_bias_softmax_gemm{gemm2_ins->get_operator()}, inputs);
// }
// auto matcher() const
// {
......@@ -223,11 +223,11 @@ struct find_ck_gemm_scale_bias_softmax_gemm
void fuse_ck::apply(module_pass_manager& mpm) const
{
// mpm.get_module().debug_print();
match::find_matches(mpm, find_ck_gemm_scale_bias_softmax_gemm{});
// if(not enabled(MIGRAPHX_DISABLE_CK_GEMM_FUSION{}))
// match::find_matches(mpm, find_ck_gemm_pointwise{});
// if(not enabled(MIGRAPHX_DISABLE_CK_GEMM{}))
// match::find_matches(mpm, find_ck_gemm{});
// match::find_matches(mpm, find_ck_gemm_scale_bias_softmax_gemm{});
if(not enabled(MIGRAPHX_DISABLE_CK_GEMM_FUSION{}))
match::find_matches(mpm, find_ck_gemm_pointwise{});
if(not enabled(MIGRAPHX_DISABLE_CK_GEMM{}))
match::find_matches(mpm, find_ck_gemm{});
}
} // namespace gpu
......
src/targets/gpu/fuse_ck_gemm_softmax_gemm.cpp 0 → 100644
View file @ f3fcfcc7
#include <migraphx/gpu/fuse_ck_gemm_softmax_gemm.hpp>
#include <migraphx/matcher.hpp>
#include <migraphx/pass_manager.hpp>
#include <migraphx/make_op.hpp>
#include <migraphx/register_op.hpp>
#include <migraphx/env.hpp>
namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
struct module;
namespace gpu {
struct gemm_softmax_gemm_gemm
{
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 gemm_softmax_gemm_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});
}
};
MIGRAPHX_REGISTER_OP(gemm_softmax_gemm_gemm);
namespace {
MIGRAPHX_PRED_MATCHER(is_ck_gemm, instruction_ref ins)
{
if(ins->name() != "dot")
return false;
auto a = ins->inputs().front()->get_shape();
auto b = ins->inputs().back()->get_shape();
if(a.lens().back() > 2048)
return false;
return true;
}
struct find_gemm_softmax_gemm_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 add =
match::name("add")(match::any_of[match::inputs()](mul));
auto softmax = match::name("softmax")(match::any_of[match::inputs()](add)).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"];
auto inputs = gemm1_ins->inputs(); // A, B
inputs.push_back(gemm2_ins->inputs().back()); // B1
mpm.get_module().replace_instruction(
ins, gemm_softmax_gemm_gemm{gemm2_ins->get_operator()}, inputs);
}
};
} // namespace
void fuse_ck_gemm_softmax_gemm::apply(module_pass_manager& mpm) const
{
match::find_matches(mpm, find_gemm_softmax_gemm_gemm{});
}
} // namespace gpu
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx
src/targets/gpu/include/migraphx/gpu/fuse_ck_gemm_softmax_gemm.hpp 0 → 100644
View file @ f3fcfcc7
#ifndef MIGRAPHX_GUARD_GPU_FUSE_CK_GEMM_SOFTMAX_GEMM_HPP
#define MIGRAPHX_GUARD_GPU_FUSE_CK_GEMM_SOFTMAX_GEMM_HPP
#include <migraphx/config.hpp>
#include <migraphx/gpu/context.hpp>
namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
struct module_pass_manager;
namespace gpu {
struct fuse_ck_gemm_softmax_gemm
{
context* ctx = nullptr;
std::string name() const { return "gpu::fuse_ck_gemm_softmax_gemm"; }
void apply(module_pass_manager& mpm) const;
};
} // namespace gpu
} // namespace MIGRAPHX_INLINE_NS
} // namespace migraphx
#endif // MIGRAPHX_GUARD_GPU_FUSE_CK_HPP
src/targets/gpu/jit/ck_gemm_softmax_gemm.cpp
View file @ f3fcfcc7
......@@ -39,7 +39,7 @@
#include <migraphx/file_buffer.hpp>
const std::vector<std::string>&
get_instance(std::size_t i, const std::function<bool(const std::vector<std::string>&)>& pred);
get_gsg_instance(std::size_t i, const std::function<bool(const std::vector<std::string>&)>& pred);
namespace migraphx {
inline namespace MIGRAPHX_INLINE_NS {
......@@ -62,34 +62,12 @@ namespace migraphx {
${preamble}
template <ck::index_t... Is>
using S = ck::Sequence<Is...>;
using Row = ck::tensor_layout::gemm::RowMajor;
using Col = ck::tensor_layout::gemm::ColumnMajor;
using F16 = ck::half_t;
using F32 = float;
using PassThrough = ck_passthrough;
static constexpr auto GemmDefault = ck::tensor_operation::device::GemmSpecialization::Default;
using AElementOp = PassThrough;
using B0ElementOp = PassThrough;
using Acc0ElementOp = ck_scale;//ck::tensor_operation::element_wise::Scale;
using B1ElementOp = PassThrough;
using CElementOp = PassThrough;//ck_add;//ck::tensor_operation::element_wise::Add;
using gemm = CK_DeviceBatchedGemmSoftmaxGemm_Xdl_CShuffle< Row, Col, Row, Row, F16, F16, F16, F16, F32, F16, AElementOp, B0ElementOp, Acc0ElementOp, B1ElementOp, CElementOp, GemmDefault, 1, 256, 256, 128, 32, 64, 32, 8, 8, 2, 32, 32, 2, 4, 2, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<4, 64, 1>, S<1, 0, 2>, S<1, 0, 2>, 2, 8, 8, true, S<16, 16, 1>, S<0, 2, 1>, S<0, 2, 1>, 1, 4, 2, false, 1, 2, S<1, 32, 1, 8>, 8, false, std::ratio<1, 8>>;
extern "C" {
__global__ void ${kernel}(${params})
{
// transform_args(make_tensors(), rotate_last())(${args})([](auto... xs) {
// ck_gemm_softmax_gemm<CK_DeviceGemmMultipleD<${instance}>, ${blocks_per_batch}>(xs...);
// });
transform_args(make_tensors(), rotate_last())(${args})([](auto... xs) {
ck_gemm_softmax_gemm<gemm, ${blocks_per_batch}>(xs...);
ck_gemm_softmax_gemm<CK_DeviceBatchedGemmSoftmaxGemm_Xdl_CShuffle<${instance}>, ${blocks_per_batch}>(xs...);
});
}
......@@ -112,13 +90,13 @@ struct instance
std::size_t get_pb(std::size_t i) const
{
assert(i < 4);
assert(i <= 4);
return int_at(block_size_index + 1 + i);
}
std::array<std::size_t, 3> get_pad(const std::array<std::size_t, 3>& config) const
std::array<std::size_t, 4> get_pad(const std::array<std::size_t, 4>& config) const
{
std::array<std::size_t, 3> result{};
std::array<std::size_t, 4> result{};
for(auto i : range(config.size()))
{
result[i] = int_div_ceil(config[i], get_pb(i)) * get_pb(i) - config[i];
......@@ -126,33 +104,16 @@ struct instance
return result;
}
std::size_t get_grid_size(const std::array<std::size_t, 3>& config) const
{
return int_div_ceil(config[0], get_pb(0)) * int_div_ceil(config[1], get_pb(1));
}
void set_ds_layout(const std::string& s)
{
assert(params[2] == "ck::Tuple<>");
params[2] = s;
}
void set_ds_type(const std::string& s)
{
assert(params[8] == "ck::Tuple<>");
params[8] = s;
}
void set_ds_op(const std::string& s)
std::size_t get_grid_size(const std::array<std::size_t, 4>& config) const
{
assert(params[12] == "ck_passthrough");
params[12] = s;
return int_div_ceil(config[0], get_pb(0)) * int_div_ceil(config[3], get_pb(3));
}
void set_gemm(const std::string& s)
{
assert(params[13] == "ck::tensor_operation::device::GemmSpecialization::Default");
params[13] = s;
assert(params[15] == "ck::tensor_operation::device::GemmSpecialization::Default" or
params[15] == "ck::tensor_operation::device::GemmSpecialization::MNKOPadding");
params[15] = s;
}
std::string str() const { return join_strings(params, ","); }
......@@ -179,12 +140,12 @@ 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!" << std::endl;
std::cout << "*********** Warning: No CK GSG tuning!" << 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 << "*********** Warning: CK GSG tuning missing for config!" << std::endl;
return 4;
}
return it->second;
......@@ -194,8 +155,12 @@ struct ck_gemm_softmax_gemm_compiler : compiler<ck_gemm_softmax_gemm_compiler>
{
static std::string get_layout(const shape& s)
{
return s.transposed() ? "ck::tensor_layout::gemm::ColumnMajor"
: "ck::tensor_layout::gemm::RowMajor";
if (not s.transposed())
return "ck::tensor_layout::gemm::RowMajor";
auto lens = s.lens();
return lens[lens.size() - 1] > lens[lens.size() - 2] ?
"ck::tensor_layout::gemm::ColumnMajor" : "ck::tensor_layout::gemm::RowMajor";
}
static std::string get_type(const shape& s)
......@@ -222,6 +187,7 @@ struct ck_gemm_softmax_gemm_compiler : compiler<ck_gemm_softmax_gemm_compiler>
{
auto a_shape = inputs[0];
auto b_shape = inputs[1];
auto b1_shape = inputs[2];
auto c_shape = inputs.back();
auto m = a_shape.lens()[0];
auto k = a_shape.lens()[1];
......@@ -229,48 +195,39 @@ struct ck_gemm_softmax_gemm_compiler : compiler<ck_gemm_softmax_gemm_compiler>
auto rank = a_shape.lens().size();
// std::array<char, 3> keys{'M', 'N', 'K'};
// std::array<std::size_t, 3> config{
// c_shape.lens()[rank - 2], c_shape.lens().back(), a_shape.lens().back()};
// auto tuning_val = v.get("tuning_val", get_tuning_for({a_shape, b_shape, c_shape}));
// auto ip = instance{get_instance(tuning_val, [&](const auto& x) -> bool {
// return get_layout(a_shape) == x[0] and get_layout(b_shape) == x[1] and
// get_layout(c_shape) == x[3] and get_type(a_shape) == x[4] and
// get_type(b_shape) == x[5] and get_type(c_shape) == x[9];
// })};
// assert(inputs.size() < 4 or v.contains("post"));
// if(v.contains("post"))
// {
// ip.set_ds_layout(ck_tuple(inputs.begin() + 2, inputs.end() - 1, &get_layout));
// ip.set_ds_type(ck_tuple(inputs.begin() + 2, inputs.end() - 1, &get_type));
// ip.set_ds_op(v.at("post").to<std::string>());
// }
// auto padding = ip.get_pad(config);
// std::string gemm_type;
// for(auto i : range(padding.size()))
// {
// if(padding[i] != 0)
// gemm_type += keys[i];
// }
// if(gemm_type.empty())
// gemm_type = "Default";
// else
// gemm_type += "Padding";
// ip.set_gemm("ck::tensor_operation::device::GemmSpecialization::" + gemm_type);
auto gemm1_nperblock = 64;
auto gemm01_mperblock = 256;
auto blocks_per_batch = int_div_ceil(m, gemm01_mperblock) *
int_div_ceil(n, gemm1_nperblock); // ip.get_grid_size(config);
std::array<char, 4> keys{'M', 'N', 'K', 'O'};
// config (m0, n0, k0, n1)
std::array<std::size_t, 4> config{
c_shape.lens()[rank - 2], b_shape.lens()[rank - 2], a_shape.lens().back(), c_shape.lens().back()};
auto tuning_val = v.get("tuning_val", get_tuning_for({a_shape, b_shape, b1_shape, c_shape}));
auto ip = instance{get_gsg_instance(tuning_val, [&](const auto& x) -> bool {
return get_layout(a_shape) == x[0] and get_layout(b_shape) == x[1] and
get_layout(c_shape) == x[3] and get_type(a_shape) == x[4] and
get_type(b_shape) == x[5] and get_type(c_shape) == x[9];
})};
auto padding = ip.get_pad(config);
std::string gemm_type;
for(auto i : range(padding.size()))
{
if(padding[i] != 0)
gemm_type += keys[i];
}
if(gemm_type.empty())
gemm_type = "Default";
else
gemm_type += "Padding";
ip.set_gemm("ck::tensor_operation::device::GemmSpecialization::" + gemm_type);
auto blocks_per_batch = ip.get_grid_size(config);
auto batch_count = std::accumulate(c_shape.lens().rbegin() + 2,
c_shape.lens().rend(),
std::size_t{1},
std::multiplies<std::size_t>());
hip_compile_options options;
auto block_size = 256; // ip.get_block_size();
auto block_size = ip.get_block_size();
auto grid_size = batch_count * blocks_per_batch;
options.set_launch_params(v, grid_size * block_size, block_size);
options.inputs = inputs;
......@@ -282,7 +239,7 @@ struct ck_gemm_softmax_gemm_compiler : compiler<ck_gemm_softmax_gemm_compiler>
options.params += " -DMIGRAPHX_CK_CHECK=1";
auto src = interpolate_string(ck_gemm_softmax_gemm_kernel,
{{"instance", "" /* ip.str() */},
{{"instance", ip.str()},
{"params", enum_params(inputs.size(), "void * private_p")},
{"args", enum_params(inputs.size(), "private_p")},
{"blocks_per_batch", to_string(blocks_per_batch)},
......@@ -302,7 +259,6 @@ struct ck_gemm_softmax_gemm_compiler : compiler<ck_gemm_softmax_gemm_compiler>
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";
}
......@@ -310,7 +266,7 @@ struct ck_gemm_softmax_gemm_compiler : compiler<ck_gemm_softmax_gemm_compiler>
return action_decorate(replace(compile_op(ctx, shapes, v)), [=] {
if(enabled(MIGRAPHX_LOG_CK_GEMM{}))
{
std::vector<shape> gemm_shapes{shapes[0], shapes[1], shapes.back()};
std::vector<shape> gemm_shapes{shapes[0], shapes[1], shapes[2], shapes.back()};
std::cout << "ck_gemm_softmax_gemm: " << to_json_string(to_value(gemm_shapes))
<< std::endl;
}
......
src/targets/gpu/jit/ck_gsg_instances.cpp
View file @ f3fcfcc7
This diff is collapsed.
src/targets/gpu/kernels/include/migraphx/kernels/ck_gemm_softmax_gemm.hpp
View file @ f3fcfcc7
......@@ -69,7 +69,7 @@ __device__ void ck_gemm_softmax_gemm_matrix(C c, A a, B b, B1 b1)
ck::make_tuple(ck::Sequence<0, 2>{}, ck::Sequence<1>{}));
constexpr const auto b_shape = get_shape_c<B>{};
constexpr const auto n = b_shape.lens[0]; // col-major
constexpr const auto n = b_shape.lens[1];
constexpr const auto sb = b_shape.strides[1]; // col-major
constexpr const auto BK1 = gemm.get_BK1();
constexpr const auto BK0 = k / BK1;
......@@ -85,9 +85,9 @@ __device__ void ck_gemm_softmax_gemm_matrix(C c, A a, B b, B1 b1)
ck::make_tuple(ck::Sequence<0, 2>{}, ck::Sequence<1>{}));
constexpr const auto b1_shape = get_shape_c<B1>{};
constexpr const auto k1 = b1_shape.lens[0]; // row-major
constexpr const auto n1 = b1_shape.lens[1]; // row-major
constexpr const auto sb1 = b1_shape.strides[0]; // rowl-major
constexpr const auto k1 = b1_shape.lens[0];
constexpr const auto n1 = b1_shape.lens[1];
constexpr const auto sb1 = b1_shape.strides[0]; // row-major
constexpr const auto B1K1 = gemm.get_B1K1();
constexpr const auto B1K0 = k1 / B1K1;
......@@ -139,11 +139,11 @@ __device__ void ck_gemm_softmax_gemm_matrix(C c, A a, B b, B1 b1)
__shared__ char p_shared[GridwiseGemm::GetSharedMemoryNumberOfByte()];
static_assert(GridwiseGemm::CheckValidity(a_grid_desc_ak0_m_ak1,
b_grid_desc_bk0_n_bk1,
b1_grid_desc_bk0_n_bk1,
c_grid_desc_m_n,
block_2_ctile_map));
// static_assert(GridwiseGemm::CheckValidity(a_grid_desc_ak0_m_ak1,
// b_grid_desc_bk0_n_bk1,
// b1_grid_desc_bk0_n_bk1,
// c_grid_desc_m_n,
// block_2_ctile_map));
GridwiseGemm::template Run<HasMainKBlockLoop>(to_ck_const_pointer(a.data()),
to_ck_const_pointer(b.data()),
to_ck_const_pointer(b1.data()),
......
src/targets/gpu/target.cpp
View file @ f3fcfcc7
......@@ -56,6 +56,7 @@
#include <migraphx/gpu/context.hpp>
#include <migraphx/gpu/device_name.hpp>
#include <migraphx/gpu/fuse_ck.hpp>
#include <migraphx/gpu/fuse_ck_gemm_softmax_gemm.hpp>
#include <migraphx/gpu/fuse_mlir.hpp>
#include <migraphx/gpu/fuse_ops.hpp>
#include <migraphx/gpu/prefuse_ops.hpp>
......@@ -131,6 +132,8 @@ std::vector<pass> target::get_passes(migraphx::context& gctx, const compile_opti
simplify_reshapes{},
propagate_constant{},
dead_code_elimination{},
fuse_ck_gemm_softmax_gemm{&ctx},
dead_code_elimination{},
enable_pass(not enabled(MIGRAPHX_DISABLE_POINTWISE_FUSION{}), fuse_pointwise{}),
dead_code_elimination{},
fuse_mlir{&ctx},
......
test/verify/0ck_gemm_softmax_gemm.cpp
View file @ f3fcfcc7
......@@ -31,15 +31,39 @@ 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, {16, 12, 384, 64}};
// migraphx::shape m2_shape{migraphx::shape::half_type, {16, 12, 384, 384}};
// auto m2_elements = 16 * 12 * 384 * 384;
// 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});
// // a = one;
// // b = one;
// // b1 = one;
// 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);
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);
size_t batch = 2;
migraphx::shape m1_shape{migraphx::shape::half_type, {batch, 384, 2304}};
migraphx::shape m2_shape{migraphx::shape::half_type, {batch, 12, 384, 384}};
auto m2_elements = batch * 12 * 384 * 384;
auto g = mm->add_parameter("1", 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);
......@@ -47,7 +71,13 @@ struct ck_gemm_softmax_gemm : verify_program<ck_gemm_softmax_gemm>
std::vector<float> ones(m2_elements, 1);
auto one = mm->add_literal(migraphx::literal{m2_shape, ones});
g = mm->add_instruction(migraphx::make_op("reshape", {{"dims", {batch, 384, 36, 64}}}), g);
g = mm->add_instruction(migraphx::make_op("transpose", {{"permutation", {0, 2, 1, 3}}}), g);
auto a = mm->add_instruction(migraphx::make_op("slice", {{"axes", {1}}, {"starts", {0}}, {"ends", {12}}}), g);
auto b = mm->add_instruction(migraphx::make_op("slice", {{"axes", {1}}, {"starts", {12}}, {"ends", {24}}}), g);
auto b1 = mm->add_instruction(migraphx::make_op("slice", {{"axes", {1}}, {"starts", {24}}, {"ends", {36}}}), g);
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);
......
tools/tune_ck.py
View file @ f3fcfcc7
import os, json, subprocess, tempfile, sys, argparse, contextlib
ck_function = -1
@contextlib.contextmanager
def tmp_file(dump=None):
......@@ -20,6 +21,9 @@ def pretty_print(obj):
def run_driver(b):
print(b)
outfile = open("temp2.json", "w")
json.dump(b, outfile)
outfile.close()
with tmp_file(lambda tf: json.dump(b, tf)) as tf:
cp = subprocess.run('./bin/gpu-driver {}'.format(tf),
capture_output=True,
......@@ -45,7 +49,7 @@ def get_device_time(s):
def benchmark_ck(config, tuning):
try:
b = {
b0 = {
'settings': {
'iterations': 100
},
......@@ -56,6 +60,18 @@ def benchmark_ck(config, tuning):
'inputs': config
}
}
b1 = {
'settings': {
'iterations': 100
},
'compile_op': {
'name': 'ck_gemm_softmax_gemm',
'check': True,
'tuning_val': tuning,
'inputs': config
}
}
b = b0 if (ck_function == 0) else b1
for line in run_driver(b):
dtime = get_device_time(line)
print(dtime)
......@@ -72,17 +88,26 @@ def benchmark(config, size):
def parse_log(f):
for line in open(f).readlines():
line = line.strip()
if not line.startswith('ck_gemm:'):
continue
global ck_function
if line.startswith('ck_gemm:'):
line = line[len('ck_gemm:'):].strip()
config = json.loads(line)
ck_function = 0
yield config
if line.startswith('ck_gemm_softmax_gemm:'):
line = line[len('ck_gemm_softmax_gemm:'):].strip()
config = json.loads(line)
ck_function = 1
yield config
def benchmark_log(f, n):
result = []
for config in parse_log(f):
tuned = benchmark(config, n)
logs = parse_log(f)
for config in logs:
additional_tv = ck_function * 2
tuned = benchmark(config, n + additional_tv)
print("Tuned:", tuned)
result.append([config, tuned])
return result
......
tools/tune_ck.sh 0 → 100755
View file @ f3fcfcc7
#!/bin/bash
MODEL=$1
LOG="ck_bbc.log"
TUNING_DB="ck_bbc.json"
rm $LOG
touch $LOG
for N in 1 16 32 64
do
MIGRAPHX_LOG_CK_GEMM=1 ./bin/driver run $MODEL -g --fill1 input_ids --input-dim @input_ids $N 384 | grep 'ck_gemm.*: \[{' | sort -u >> $LOG
done
python3 ../tools/tune_ck.py -n 16 -l $LOG -o $TUNING_DB
\ No newline at end of file
tools/tune_ck_gsg.py 0 → 100644
View file @ f3fcfcc7
import os, json, subprocess, tempfile, sys, argparse, contextlib
@contextlib.contextmanager
def tmp_file(dump=None):
tmp_name = None
try:
with tempfile.NamedTemporaryFile(mode='w+', delete=False) as f:
tmp_name = f.name
if dump:
dump(f)
yield tmp_name
finally:
os.unlink(tmp_name)
def pretty_print(obj):
print(json.dumps(obj, indent=2))
def run_driver(b):
print(b)
with tmp_file(lambda tf: json.dump(b, tf)) as tf:
cp = subprocess.run('./bin/gpu-driver {}'.format(tf),
capture_output=True,
check=True,
shell=True)
for line in cp.stdout.decode().split("\n"):
s = line.strip()
if not s:
continue
if not ']: ' in s:
continue
yield s.split(']: ')[1].strip()
def convert_to_float(s):
return s[:-2]
def get_device_time(s):
fields = s.split(',')
return convert_to_float(fields[-1].strip())
def benchmark_ck(config, tuning):
try:
b = {
'settings': {
'iterations': 100
},
'compile_op': {
'name': 'ck_gemm_softmax_gemm',
'check': True,
'tuning_val': tuning,
'inputs': config
}
}
for line in run_driver(b):
dtime = get_device_time(line)
print(dtime)
return float(dtime)
except:
return sys.float_info.max
def benchmark(config, size):
times = [benchmark_ck(config, i) for i in range(size)]
return times.index(min(times))
def parse_log(f):
for line in open(f).readlines():
line = line.strip()
if not line.startswith('ck_gemm_softmax_gemm:'):
continue
line = line[len('ck_gemm_softmax_gemm:'):].strip()
config = json.loads(line)
yield config
def benchmark_log(f, n):
result = []
for config in parse_log(f):
tuned = benchmark(config, n)
print("Tuned:", tuned)
result.append([config, tuned])
return result
def parse_args():
parser = argparse.ArgumentParser(description="Simple tuner for CK gemms")
parser.add_argument('--log',
'-l',
type=str,
metavar='file',
help='Path to logfile')
parser.add_argument('--out',
'-o',
type=str,
metavar='file',
help='Output json file to save tunings')
parser.add_argument('-n', type=int, help='Number of instances to tune')
args = parser.parse_args()
return args
def run(args):
tuned = benchmark_log(args.log, args.n)
json.dump(tuned, open(args.out, 'w+'))
run(parse_args())
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment