#include "gpu/intel/conv/jit/v2/planner/search.hpp"
#include "common/profiler.hpp"
#include "gpu/intel/conv/jit/v2/model.hpp"
#include "gpu/intel/conv/jit/v2/plan.hpp"
#include "gpu/intel/conv/jit/v2/plan_registry.hpp"
#include "gpu/intel/conv/jit/v2/planner/bench.hpp"
#include "gpu/intel/conv/jit/v2/planner/model_fit.hpp"
#include "gpu/intel/jit/ir/blocking.hpp"
#include "gpu/intel/jit/utils/utils.hpp"
#include "oneapi/dnnl/dnnl.hpp"
#include <random>
#include <initializer_list>
#include <unordered_map>
#include <unordered_set>
namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace conv {
namespace jit {
namespace v2 {
namespace planner {
enum class tile_flags_t : uint32_t {
undef = 0,
loop = (1 << 0),
thread_group = (1 << 1),
iter = (1 << 2),
};
tile_flags_t operator&(tile_flags_t a, tile_flags_t b) {
auto _a = static_cast<uint32_t>(a);
auto _b = static_cast<uint32_t>(b);
return static_cast<tile_flags_t>(_a & _b);
}
tile_flags_t operator|(tile_flags_t a, tile_flags_t b) {
auto _a = static_cast<uint32_t>(a);
auto _b = static_cast<uint32_t>(b);
return static_cast<tile_flags_t>(_a | _b);
}
tile_flags_t operator~(tile_flags_t a) {
auto _a = static_cast<uint32_t>(a);
return static_cast<tile_flags_t>(~_a);
}
bool any(tile_flags_t a) {
return a != tile_flags_t::undef;
}
inline std::string to_string(tile_flags_t flags) {
ostringstream_t oss;
if (any(flags & tile_flags_t::loop)) oss << "l";
if (any(flags & tile_flags_t::thread_group)) oss << "t";
if (any(flags & tile_flags_t::iter)) oss << "i";
return oss.str();
}
inline std::vector<int> pow_range(int a, int b, int step) {
std::vector<int> ret;
for (int i = a; i <= b; i *= step)
ret.push_back(i);
return ret;
}
struct tile_info_t {
pvar_t dim;
tile_flags_t flags = tile_flags_t::undef;
tile_info_t() = default;
tile_info_t(const pvar_t &dim) : dim(dim) {}
void add(tile_flags_t f) { flags = flags | f; }
std::vector<int> iter_tiles() const {
if (!any(flags & tile_flags_t::iter)) return {1};
return pow_range(8, 64, 2);
}
std::vector<int> thread_group_tiles() const {
if (!any(flags & tile_flags_t::thread_group)) return {1};
return pow_range(1, 16, 2);
}
std::string str() const {
ostringstream_t oss;
oss << dim << ": " << to_string(flags);
return oss.str();
}
XE_DEFINE_DUMP()
};
class tile_scheme_t {
public:
tile_scheme_t() = default;
tile_scheme_t(const std::string &s) {
std::vector<std::string> parts;
std::vector<size_t> key_idxs;
std::string cur;
int beg = -1;
for (int i = 0; i <= (int)s.length(); i++) {
char c = (i <= (int)s.length() ? s[i] : ' ');
bool is_id = (std::isdigit(c) || std::isalpha(c) || c == '_');
if (beg != -1 && is_id) continue;
if (beg == -1 && is_id) {
beg = i;
} else if (beg != -1 && !is_id) {
parts.push_back(s.substr(beg, i - beg));
beg = -1;
}
if (c == '=') key_idxs.push_back(parts.size() - 1);
}
key_idxs.push_back(parts.size());
std::unordered_map<std::string, std::vector<std::string>> k2v;
for (size_t i = 0; i < key_idxs.size() - 1; i++) {
size_t cur = key_idxs[i];
size_t next = key_idxs[i + 1];
for (size_t j = cur + 1; j < next; j++) {
set(parts[cur], parts[j]);
}
}
}
void unset(const pvar_t &dim) { tile_infos_.unset(dim); }
std::vector<pvar_t> dims() const { return tile_infos_.keys(); }
const tile_info_t &tile_info(const pvar_t &dim) const {
return tile_infos_.at(dim);
}
private:
void set(const std::string &key, const std::string &value) {
if (key == "iter") {
auto dim = pvar_t(value);
tile_infos_[dim].add(tile_flags_t::iter);
} else if (key == "tg") {
auto dim = pvar_t(value);
tile_infos_[dim].add(tile_flags_t::thread_group);
} else {
gpu_error_not_expected();
}
}
pvar_map_t<tile_info_t> tile_infos_;
};
struct dim_tile_t {
int loop = 0;
int tg = 0;
int iter = 0;
std::string str() const {
ostringstream_t oss;
if (loop != 0) oss << "l" << loop;
if (tg != 0) oss << "t" << tg;
if (iter != 0) oss << "i" << iter;
return oss.str();
}
XE_DEFINE_DUMP()
};
std::ostream &operator<<(std::ostream &out, const dim_tile_t &tile) {
out << tile.str();
return out;
}
struct tiling_desc_t {
tile_t iter;
tile_t thread_group;
void set(const pvar_t &dim, const dim_tile_t &tile) {
if (tile.iter != 1) iter[dim] = tile.iter;
if (tile.tg != 1) thread_group[dim] = tile.tg;
}
void unset(const pvar_t &dim) {
iter.unset(dim);
thread_group.unset(dim);
}
std::string str() const {
ostringstream_t oss;
oss << "iter: " << iter.str();
oss << " thread_group: " << thread_group.str();
return oss.str();
}
XE_DEFINE_DUMP()
};
class dim_tile_set_t {
public:
dim_tile_set_t(const tile_scheme_t &scheme) : dims_(scheme.dims()) {
for (auto &d : scheme.dims()) {
auto &d_tiles = tiles_[d];
d_tiles = get_dim_tiles(scheme, d);
}
}
std::vector<tiling_desc_t> create_tiling_descs() const {
std::vector<tiling_desc_t> ret;
tiling_desc_t tiling_desc;
std::vector<int> cur_idxs(dims_.size());
product_impl(0, cur_idxs, tiling_desc, ret);
return ret;
}
private:
void product_impl(int idx, std::vector<int> &cur_idxs,
tiling_desc_t &tiling_desc, std::vector<tiling_desc_t> &ret) const {
if (idx == (int)dims_.size()) {
ret.push_back(tiling_desc);
return;
}
auto &v = tiles_.at(dims_[idx]);
for (int i = 0; i < (int)v.size(); i++) {
cur_idxs[idx] = i;
tiling_desc.set(dims_[idx], v[i]);
product_impl(idx + 1, cur_idxs, tiling_desc, ret);
tiling_desc.unset(dims_[idx]);
}
}
static std::vector<dim_tile_t> get_dim_tiles(
const tile_scheme_t &scheme, const pvar_t &dim) {
std::vector<dim_tile_t> ret;
auto &info = scheme.tile_info(dim);
auto iter_tiles = info.iter_tiles();
auto tg_tiles = info.thread_group_tiles();
for (int iter : iter_tiles) {
for (int tg : tg_tiles) {
dim_tile_t tile;
tile.iter = iter;
tile.tg = tg;
ret.push_back(tile);
}
}
return ret;
}
std::vector<pvar_t> dims_;
pvar_map_t<std::vector<dim_tile_t>> tiles_;
};
struct search_params_t {
kernel_desc_t base_desc;
bool is_iter_set = false;
bool is_tg_set = false;
bool is_prefetch_set = false;
search_params_t(
const kernel_desc_t &_base_desc, const parse_result_t &parse_result)
: base_desc(_base_desc) {
is_iter_set = parse_result.is_set("--iter");
is_tg_set = parse_result.is_set("--tg");
is_prefetch_set = parse_result.is_set("--prefetch");
}
search_params_t(const planner_params_t ¶ms)
: search_params_t(params.desc, params.parse_result) {}
};
std::vector<tile_scheme_t> get_tile_schemes(const search_params_t ¶ms) {
std::vector<tile_scheme_t> schemes;
if (params.base_desc.prop == prop_kind::forward) {
schemes.emplace_back("tg=[ic], iter=[mb,g,oc,ic]");
schemes.emplace_back("tg=[ic], iter=[ow,g,oc,ic]");
schemes.emplace_back("tg=[oc,mb], iter=[mb,g,oc,ic]");
schemes.emplace_back("tg=[oc,mb], iter=[ow,g,oc,ic]");
schemes.emplace_back("tg=[oc,ow], iter=[mb,g,oc,ic]");
schemes.emplace_back("tg=[oc,ow], iter=[ow,g,oc,ic]");
} else if (params.base_desc.prop == prop_kind::backward_data) {
schemes.emplace_back("tg=[ic,iw], iter=[mb,g,oc,ic]");
schemes.emplace_back("tg=[ic,mb], iter=[mb,g,oc,ic]");
schemes.emplace_back("tg=[ic,iw], iter=[iw,g,oc,ic]");
} else if (params.base_desc.prop == prop_kind::backward_weights) {
schemes.emplace_back("tg=[oc,ic], iter=[mb,g,oc,ic]");
schemes.emplace_back("tg=[oc,ic], iter=[ow,g,oc,ic]");
} else {
gpu_error_not_expected();
}
for (auto &s : schemes) {
if (params.base_desc.is_dw) {
s.unset(pvars::ic);
s.unset(pvars::oc);
} else {
s.unset(pvars::g);
}
}
return schemes;
}
class search_kernel_desc_group_t {
public:
search_kernel_desc_group_t() = default;
search_kernel_desc_group_t(const prb_reqs_t &reqs) : reqs_(reqs) {}
const prb_reqs_t &reqs() const { return reqs_; }
const std::vector<kernel_desc_t> &descs() const { return descs_; }
void add_desc(const kernel_desc_t &desc) {
gpu_assert(desc.reqs().str() == reqs_.str())
<< "Reqs mismatch:\n"
<< desc.cmd_str() << "\ndesc.reqs:" << desc.reqs().str()
<< "\nreqs:\n"
<< reqs_.str();
if (descs_.empty()) {
is_dw_ = desc.is_dw;
} else {
gpu_assert(desc.is_dw == is_dw_);
}
descs_.push_back(desc);
}
bench_input_params_t bench_input_params(
int nprbs, const dsl::hw_t &hw) const {
if (descs_.empty()) return bench_input_params_t();
auto &kd = descs_.front();
bench_input_params_t params;
params.hw = hw;
params.prop = kd.prop;
params.src_tag = kd.src_tag;
params.wei_tag = kd.wei_tag;
params.dst_tag = kd.dst_tag;
params.reqs = reqs_;
params.is_dw = is_dw_;
params.nprbs = nprbs;
return params;
}
private:
prb_reqs_t reqs_;
std::vector<kernel_desc_t> descs_;
bool is_dw_ = false;
};
bench_data_set_t bench_kernel_desc_group(const bench_manager_t &bench_mger,
const search_kernel_desc_group_t &desc_group, int nprbs, int max_descs);
class kernel_search_manager_t {
public:
static const int bench_nprbs = 50;
static const int registry_top_k = 8;
static const int max_descs = 256;
kernel_search_manager_t(
const bench_manager_t &bench_mger, const search_params_t ¶ms)
: bench_mger_(bench_mger), params_(params) {}
void search() {
std::cout << "Starting kernel search" << std::endl;
auto ®istry = plan_registry();
auto desc_groups = gen_desc_groups();
for (auto &dg : desc_groups) {
auto bench_data_set = bench_kernel_desc_group(
bench_mger_, dg, bench_nprbs, max_descs);
auto best = bench_data_set.find_best(registry_top_k);
for (auto &bd : best) {
auto entry = prepare_plan_registry_entry(
bench_mger_, bd.kernel_desc);
registry.set(entry);
}
}
std::cout << "Kernel search completed" << std::endl;
}
private:
std::vector<search_kernel_desc_group_t> gen_desc_groups() const {
std::unordered_set<std::string> seen;
std::vector<kernel_desc_t> descs;
for (auto &s : get_tile_schemes(params_)) {
dim_tile_set_t tile_set(s);
auto tiling_descs = tile_set.create_tiling_descs();
for (auto &td : tiling_descs) {
auto d = params_.base_desc;
if (!params_.is_tg_set) d.thread_group_tile = td.thread_group;
if (!params_.is_iter_set) d.iter_tile = td.iter;
auto d_key = jit::stringify(d);
if (seen.count(d_key) > 0) continue;
seen.insert(std::move(d_key));
if (!create_plan(d, bench_mger_.hw())) {
std::cout << d.brief_str() << ": \033[1;31mFAIL\033[0m"
<< std::endl;
continue;
}
std::cout << d.brief_str() << ": \033[1;32mOK\033[0m"
<< std::endl;
descs.push_back(std::move(d));
}
}
gpu_info() << "gen_desc_groups(): descs.size() = " << descs.size();
std::unordered_map<std::string, search_kernel_desc_group_t> desc_groups;
std::vector<int> prefetch_dists;
if (params_.is_prefetch_set) {
prefetch_dists.push_back(params_.base_desc.prefetch.dist);
} else {
prefetch_dists.push_back(1);
prefetch_dists.push_back(3);
}
for (auto &d : descs) {
auto ret = desc_groups.emplace(
d.reqs().str(), search_kernel_desc_group_t(d.reqs()));
ret.first->second.add_desc(d);
for (int dist : prefetch_dists) {
auto _d = d;
_d.prefetch = prefetch_desc_t(dist, true, true);
if (!create_plan(_d, bench_mger_.hw())) {
std::cout << d.brief_str() << ": \033[1;31mFAIL\033[0m"
<< std::endl;
continue;
}
std::cout << _d.brief_str() << ": \033[1;32mOK\033[0m"
<< std::endl;
ret.first->second.add_desc(_d);
}
}
std::vector<search_kernel_desc_group_t> ret;
ret.reserve(desc_groups.size());
for (auto &kv : desc_groups) {
ret.push_back(kv.second);
}
std::cout << "Generated " << ret.size()
<< " kernel descriptor groups\n";
return ret;
}
static std::vector<tile_t> generate_iter_outer_tiles(
const kernel_desc_t &desc) {
std::vector<tile_t> tiles = {tile_t()};
for (auto &d : desc.iter_tile) {
auto bmnk = to_gemm(d, desc.prop);
if (!utils::one_of(std::move(bmnk), pvars::m, pvars::n)) continue;
for (int outer : {2, 4}) {
if (desc.iter_tile.at(d) % outer != 0) continue;
tile_t tile_outer;
tile_outer[d] = outer;
tiles.push_back(std::move(tile_outer));
}
}
return tiles;
}
const bench_manager_t &bench_mger_;
search_params_t params_;
};
class search_sequence_t {
public:
search_sequence_t(const std::vector<kernel_desc_t> &descs, int max_entries)
: max_entries_(max_entries) {
std::vector<std::vector<tile_t>> tiles;
pvar_t prefetch_dim("p");
for (int i = 0; i < (int)descs.size(); i++) {
auto &d = descs[i];
entries_.emplace_back(i, d);
std::vector<tile_t> d_tiles;
auto iter = to_gemm(d.iter_tile, d.prop);
auto tg = to_gemm(d.thread_group_tile, d.prop);
d_tiles.push_back(std::move(iter));
d_tiles.push_back(std::move(tg));
tile_t prefetch_tile;
prefetch_tile[prefetch_dim] = d.prefetch.dist;
d_tiles.push_back(std::move(prefetch_tile));
tiles.push_back(std::move(d_tiles));
}
tile_to_vec_ = tile_to_vec_t(tiles);
entry_it_ = entries_.begin();
std::default_random_engine rng(0);
std::shuffle(entries_.begin(), entries_.end(), rng);
}
explicit operator bool() const {
return entry_idx_ < max_entries_ && entry_it_ != entries_.end();
}
std::pair<int, kernel_desc_t> next() {
gpu_assert((bool)*this);
auto &e = *entry_it_;
++entry_it_;
return std::make_pair(e.id, e.desc);
}
void update(const bench_data_set_t &data_set) {
entry_idx_++;
if (batch_entry_idx_++ < rescore_period_) return;
batch_entry_idx_ = 0;
const int nbest = 5;
auto best_ids = data_set.find_best_ids(nbest);
std::unordered_map<int, float> min_dists;
for (auto it = entry_it_; it != entries_.end(); ++it) {
min_dists[it->id] = std::numeric_limits<float>::max();
for (auto &id : best_ids) {
min_dists[it->id] = std::min(
min_dists[it->id], tile_to_vec_.dist(it->id, id));
}
}
std::sort(entry_it_, entries_.end(),
[&](const entry_t &a, const entry_t &b) {
return min_dists[a.id] < min_dists[b.id];
});
}
private:
struct entry_t {
int id = -1;
kernel_desc_t desc;
entry_t(int id, const kernel_desc_t &desc) : id(id), desc(desc) {}
};
static const int rescore_period_ = 16;
std::vector<entry_t> entries_;
std::vector<entry_t>::iterator entry_it_;
tile_to_vec_t tile_to_vec_;
int batch_entry_idx_ = 0;
int entry_idx_ = 0;
int max_entries_ = 0;
};
bench_data_set_t bench_kernel_desc_group(const bench_manager_t &bench_mger,
const search_kernel_desc_group_t &desc_group, int nprbs,
int max_descs) {
bench_runner_t runner(
bench_mger, desc_group.bench_input_params(nprbs, bench_mger.hw()));
bench_data_set_t bd_set;
search_sequence_t seq(desc_group.descs(), max_descs);
while (seq) {
auto seq_next = seq.next();
int kernel_desc_id = seq_next.first;
auto &kernel_desc = seq_next.second;
auto bd = runner.bench(kernel_desc);
if (!bd) continue;
bd.id = kernel_desc_id;
bd_set.add(bd);
seq.update(bd_set);
}
return bd_set;
}
std::string merge_cmd_lines(const std::string &recipe_line,
const parse_result_t &cmd_parse_result) {
auto &iface = kernel_desc_t::parse_iface();
kernel_desc_t recipe_desc;
parse_result_t recipe_parse_result;
iface.parse(recipe_line, recipe_desc, &recipe_parse_result);
bool is_first = true;
ostringstream_t oss;
for (auto &kv : cmd_parse_result.args()) {
auto &name = kv.first;
;
auto &value = kv.second;
if (!is_first) oss << " ";
oss << name << "=" << value;
is_first = false;
}
for (auto &kv : recipe_parse_result.args()) {
auto &name = kv.first;
auto &value = kv.second;
if (cmd_parse_result.args().count(name) > 0) continue;
if (!is_first) oss << " ";
oss << name << "=" << value;
is_first = false;
}
return oss.str();
}
void auto_search(
const bench_manager_t &bench_mger, const planner_params_t ¶ms) {
std::vector<const char *> recipes = {
"--hw xehpc --prop fwd --src axb:s8 --wei axcb:s8 --dst axb:s8 --fma dpas --simd 16 --regs 256 --2d 1",
"--hw xehpc --prop fwd --src axb:s8 --wei axcb:s8 --dst axb:s8 --fma dpas --simd 16 --regs 256",
"--hw xehpc --prop fwd --src axb:bf16 --wei axcb:bf16 --dst axb:bf16 --fma dpas --simd 16 --regs 256 --2d 1",
"--hw xehpc --prop fwd --src axb:bf16 --wei axcb:bf16 --dst axb:bf16 --fma dpas --simd 16 --regs 256",
"--hw xehpc --prop fwd --src axb:f32 --wei axcb:f32 --dst axb:f32 --fma mad --simd 32 --regs 128 --2d 1",
"--hw xehpc --prop fwd --src axb:f32 --wei axcb:f32 --dst axb:f32 --fma mad --simd 32 --regs 128",
"--hw xehpc --prop bwd_d --src axb:bf16 --wei axbc:bf16 --dst axb:bf16 --fma dpas --simd 16 --regs 256 --2d 1",
"--hw xehpc --prop bwd_d --src axb:bf16 --wei axbc:bf16 --dst axb:bf16 --fma dpas --simd 16 --regs 256",
"--hw xehpc --prop bwd_d --src axb:f32 --wei axbc:f32 --dst axb:f32 --fma mad --simd 32 --regs 128 --2d 1",
"--hw xehpc --prop bwd_d --src axb:f32 --wei axbc:f32 --dst axb:f32 --fma mad --simd 32 --regs 128",
"--hw xehpc --prop bwd_w --src axb:bf16 --wei axcb:bf16 --dst axb:bf16 --fma dpas --simd 16 --regs 256 --2d 1",
"--hw xehpc --prop bwd_w --src axb:bf16 --wei axcb:bf16 --dst axb:bf16 --fma dpas --simd 16 --regs 256",
"--hw xehpc --prop bwd_w --src axb:f32 --wei axcb:f32 --dst axb:f32 --fma mad --simd 16 --regs 128 --2d 1",
"--hw xehpc --prop bwd_w --src axb:f32 --wei axcb:f32 --dst axb:f32 --fma mad --simd 16 --regs 128",
"--hw xehpc --dw 1 --prop fwd --src axb:s8 --wei axcb:s8 --dst axb:s8 --fma mad --simd 32 --regs 128",
"--hw xehpc --dw 1 --prop fwd --src axb:bf16 --wei axcb:bf16 --dst axb:bf16 --fma mad --simd 32 --regs 128",
"--hw xehpc --dw 1 --prop fwd --src axb:f32 --wei axcb:f32 --dst axb:f32 --fma mad --simd 32 --regs 128",
"--hw xehpc --dw 1 --prop bwd_d --src axb:bf16 --wei axbc:bf16 --dst axb:bf16 --fma mad --simd 32 --regs 128",
"--hw xehpc --dw 1 --prop bwd_d --src axb:f32 --wei axbc:f32 --dst axb:f32 --fma mad --simd 32 --regs 128",
"--hw xehpc --dw 1 --prop bwd_w --src axb:bf16 --wei axcb:bf16 --dst axb:bf16 --fma mad --simd 16 --regs 128",
"--hw xehpc --dw 1 --prop bwd_w --src axb:f32 --wei axcb:f32 --dst axb:f32 --fma mad --simd 32 --regs 128",
};
auto &iface = kernel_desc_t::parse_iface();
double t = get_msec();
std::unordered_set<std::string> seen;
for (const char *_r : recipes) {
std::string line = merge_cmd_lines(_r, params.parse_result);
if (seen.count(line) > 0) continue;
seen.insert(line);
kernel_desc_t desc;
parse_result_t parse_result;
iface.parse(line, desc, &parse_result);
kernel_search_manager_t mger(
bench_mger, search_params_t(desc, parse_result));
mger.search();
}
t = get_msec() - t;
std::cout << "Kernel search done, took: " << t / 1e3 << " sec" << std::endl;
}
void search(const bench_manager_t &bench_mger, const planner_params_t ¶ms) {
switch (params.mode) {
case planner_mode_t::search: {
kernel_search_manager_t mger(bench_mger, search_params_t(params));
mger.search();
break;
}
case planner_mode_t::auto_search:
auto_search(bench_mger, params);
break;
default: gpu_error_not_expected();
}
}
} } } } } } } }