#ifndef GPU_INTEL_JIT_IR_BLOCKING_HPP
#define GPU_INTEL_JIT_IR_BLOCKING_HPP
#include "gpu/intel/jit/ir/problem.hpp"
#include "gpu/intel/logging.hpp"
#include <set>
namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace jit {
class blocking_t {
public:
int simd() const { return simd_; }
const tile_t &loop() const { return loop_; }
const tile_t &thread_group() const { return thread_group_; }
const tile_t &iter() const { return iter_; }
dim_t loop_dim(const pvar_t &d) const { return loop_[d]; }
dim_t thread_group_dim(const pvar_t &d) const { return thread_group_[d]; }
dim_t iter_dim(const pvar_t &d) const { return iter_[d]; }
void set_simd(int simd) { simd_ = simd; }
void set_loop(const pvar_t &d, dim_t value) { loop_[d] = value; }
void set_thread_group(const pvar_t &d, dim_t value) {
thread_group_[d] = value;
}
void set_iter(const pvar_t &d, dim_t value) { iter_[d] = value; }
bool is_empty() const {
return loop_.is_empty() && thread_group_.is_empty() && iter_.is_empty();
}
bool is_spatial() const {
for (const auto &d : {pvars::iw, pvars::ow}) {
if (iter_.has(d) && iter_[d] != 1) return true;
}
return false;
}
void unset(const pvar_t &d) {
if (loop_.has(d)) loop_[d] = 1;
if (thread_group_.has(d)) thread_group_[d] = 1;
if (iter_.has(d)) iter_[d] = 1;
}
bool operator==(const blocking_t &other) const {
return (loop_ == other.loop_) && (thread_group_ == other.thread_group_)
&& (iter_ == other.iter_);
}
void stringify(std::ostream &out) const {
out << "simd=" << simd_;
out << " l=";
out << loop_.str();
out << " T=";
out << thread_group_.str();
out << " i=";
out << iter_.str();
}
void parse(std::istream &in) {
stream_match(in, "simd=");
simd_ = stream_parse<int>(in);
stream_match(in, "l=");
loop_.parse(in);
stream_match(in, "T=");
thread_group_.parse(in);
stream_match(in, "i=");
iter_.parse(in);
}
size_t get_hash() const {
return ir_utils::get_hash(loop_, thread_group_, iter_);
}
std::string str(bool csv = false) const {
ostringstream_t oss;
if (csv) {
oss << simd_;
oss << "," << loop_;
oss << "," << thread_group_;
oss << "," << iter_;
} else {
oss << "simd=" << simd_;
oss << " l=" << loop_;
oss << " T=" << thread_group_;
oss << " i=" << iter_;
}
return oss.str();
}
float get_efficiency(const tile_t &shape) const {
double ret = 1;
for (auto &d : shape) {
dim_t loop = loop_.get(d, 1);
dim_t tg = thread_group_.get(d, 1);
dim_t iter = iter_.get(d, 1);
dim_t size = shape[d];
dim_t size_padded = utils::rnd_up(size, loop * tg * iter);
if (size_padded != size) ret *= double(size) / size_padded;
}
return (float)ret;
}
bool operator<(const blocking_t &other) const {
if (simd_ != other.simd_) return simd_ < other.simd_;
auto tile_compare = [](const tile_t &a, const tile_t &b) {
auto a_it = a.begin(), b_it = b.begin();
auto a_end = a.end(), b_end = b.end();
while (a_it != a_end && b_it != b_end) {
if (a_it.idx() != b_it.idx()) return a_it.idx() < b_it.idx();
if (a_it.value() != b_it.value())
return a_it.value() < b_it.value();
a_it++;
b_it++;
}
return a_it == a_end && b_it != b_end;
};
if (loop_ != other.loop_) return tile_compare(loop_, other.loop_);
if (thread_group_ != other.thread_group_)
return tile_compare(thread_group_, other.thread_group_);
return tile_compare(iter_, other.iter_);
}
XE_DEFINE_DUMP()
private:
int simd_ = 0;
tile_t loop_;
tile_t thread_group_;
tile_t iter_;
};
enum class tile_flags_t : uint32_t {
undef = 0,
loop = (1 << 0),
thread_group = (1 << 1),
iter = (1 << 2),
loop_span = (1 << 3),
loop_iter_unroll = (1 << 4),
};
GPU_DEFINE_BIT_MASK_ENUM_OPS(tile_flags_t)
struct div_info_t {
int iter_unit = 1;
int unroll_unit = 1;
void set_iter_unit(int new_unit) {
iter_unit = math::lcm(iter_unit, new_unit);
}
void set_unroll_unit(int new_unit) {
unroll_unit = math::lcm(unroll_unit, new_unit);
}
bool is_iter_ok(dim_t blk) const {
if (iter_unit != 1 && blk % iter_unit != 0) return false;
if (iter_unit != 1 && !math::is_pow2(blk)) return false;
return true;
}
};
struct tile_info_t {
tile_info_t() = default;
tile_info_t(const pvar_t &dim) : dim(dim) {}
void add(tile_flags_t f) { flags = flags | f; }
void remove(tile_flags_t f) { flags = flags & ~f; }
void set_iter_unit(int unit) { div_info.set_iter_unit(unit); }
void set_unroll_unit(int unit) { div_info.set_unroll_unit(unit); }
void set_min_iter_block(int block, int pow2_block = 0) {
min_iter_blk = block;
if (pow2_block != 0) min_iter_pow2_blk = pow2_block;
}
std::vector<int> iter_blocks(dim_t size) const;
std::vector<int> thread_group_blocks(dim_t size) const;
std::vector<dim_t> loop_blocks(dim_t size, int iter_blk) const;
static bool block_ok(dim_t size, int blk, int target_eff) {
dim_t size_padded = utils::rnd_up(size, blk);
double eff = size / (double)size_padded;
return eff * 100 >= target_eff;
}
static std::vector<dim_t> get_factors(dim_t n);
static std::vector<dim_t> get_loop_blocks(dim_t n);
pvar_t dim;
tile_flags_t flags = tile_flags_t::undef;
div_info_t div_info;
int min_iter_blk = default_min_iter_blk;
int min_iter_pow2_blk = default_min_iter_pow2_blk;
int max_iter_blk = default_max_iter_blk;
int max_thread_group_blk = default_max_thread_group_blk;
static const int default_min_iter_blk = 6;
static const int default_min_iter_pow2_blk = 8;
static const int default_max_iter_blk = 64;
static const int default_max_thread_group_blk = 16;
};
enum class level_t {
undef = 0,
loop,
thread_group,
iter,
};
class level_tile_t {
public:
bool has(level_t level) const {
switch (level) {
case level_t::loop: return loop != 0;
case level_t::thread_group: return thread_group != 0;
case level_t::iter: return iter != 0;
default: gpu_error_not_expected();
}
return false;
}
std::string str() const {
if (utils::everyone_is(0, loop, thread_group, iter)) return "x";
ostringstream_t oss;
if (loop != 0) oss << "l" << loop;
if (thread_group != 0) oss << "T" << thread_group;
if (iter != 0) oss << "i" << iter;
return oss.str();
}
XE_DEFINE_DUMP()
dim_t loop = 0;
int thread_group = 0;
int iter = 0;
};
void get_level_tiles(
dim_t size, const tile_info_t &info, std::vector<level_tile_t> &ret);
class level_tile_set_t {
public:
level_tile_set_t(const std::vector<std::vector<level_tile_t>> &tiles,
const std::vector<int> &deps, const std::vector<pvar_t> &dims)
: tiles_(tiles), deps_(deps), dims_(dims) {}
int count() const;
std::vector<blocking_t> product(int simd) const;
std::vector<blocking_t> sample(int target,
const std::function<bool(const blocking_t &)> &is_ok, int simd,
int tries_mult_bound = 5) const;
private:
static void set(
blocking_t &blk, const pvar_t &dim, const level_tile_t &tile);
void product_impl(int idx, std::vector<int> &cur_idxs, blocking_t &blk,
std::vector<blocking_t> &ret) const;
std::vector<level_tile_t> sample(ir_utils::fast_random_t &r) const;
std::vector<std::vector<level_tile_t>> tiles_;
std::vector<int> deps_;
std::vector<pvar_t> dims_;
};
class blocking_scheme_t {
public:
virtual ~blocking_scheme_t() = default;
blocking_scheme_t() = default;
blocking_scheme_t(const std::string &s) {
gpu_assert(s[s.length() - 1] == ']');
auto parts = gpu_utils::split(s.substr(0, s.length() - 1), "],");
const std::string colon = ":";
const std::string comma = ",";
for (auto &p : parts) {
auto p_parts = gpu_utils::split(p, colon);
auto &key = p_parts[0];
auto &vec = p_parts[1];
gpu_assert(vec[0] == '[');
auto s_dims
= gpu_utils::split(vec.substr(1, vec.length() - 1), comma);
for (auto &s : s_dims)
set(key, s);
}
}
virtual level_tile_set_t make_level_tile_set(
const tile_t &padded_shape) const {
const auto all_dims = dims();
const int ndims = int(all_dims.size());
const std::vector<int> deps(ndims, -1);
std::vector<std::vector<level_tile_t>> tiles(ndims);
for (int i = 0; i < ndims; i++) {
auto &d = all_dims[i];
get_level_tiles(padded_shape[d], tile_info(d), tiles[i]);
}
return level_tile_set_t(tiles, deps, all_dims);
}
tile_info_t &tile_info(const pvar_t &d) {
auto it = tile_infos_.find(d);
if (it != tile_infos_.end()) return it->second;
auto &info = tile_infos_[d];
info = tile_info_t(d);
return info;
}
const tile_info_t &tile_info(const pvar_t &d) const {
return tile_infos_.at(d);
}
std::vector<pvar_t> dims() const {
std::set<pvar_t> dims;
for (auto *t : {&loop_, &thread_group_, &iter_}) {
for (auto &d : t->keys()) {
dims.insert(d);
}
}
return std::vector<pvar_t>(dims.begin(), dims.end());
}
std::string str() const {
ostringstream_t oss;
oss << "l:" << loop_;
oss << " T:" << thread_group_;
oss << " i:" << iter_;
return oss.str();
}
XE_DEFINE_DUMP()
private:
void set(const std::string &s_tile, const std::string &_s_dim) {
gpu_assert(!_s_dim.empty());
bool no_min_check = (_s_dim[0] == '#');
const auto &s_dim = no_min_check ? _s_dim.substr(1) : _s_dim;
auto d = pvar_t(s_dim);
if (no_min_check) gpu_assert(s_tile == "i");
if (s_tile == "i") {
add_iter_dim(d);
if (no_min_check) tile_info(d).set_min_iter_block(1);
} else if (s_tile == "T") {
add_thread_group_dim(d);
} else if (s_tile == "l") {
add_loop_dim(d);
} else if (s_tile == "ls") {
add_loop_dim_with_span(d);
} else if (s_tile == "li") {
add_loop_dim_with_iter_unroll(d);
} else {
gpu_error_not_expected() << s_tile;
}
}
void add_loop_dim(const pvar_t &d) {
loop_[d] = 1;
auto &info = tile_info(d);
info.add(tile_flags_t::loop);
}
void add_loop_dim_with_span(const pvar_t &d) {
add_loop_dim(d);
tile_info(d).add(tile_flags_t::loop_span);
}
void add_loop_dim_with_iter_unroll(const pvar_t &d) {
add_loop_dim(d);
tile_info(d).add(tile_flags_t::loop_iter_unroll);
}
void add_thread_group_dim(const pvar_t &d) {
thread_group_[d] = 1;
auto &info = tile_info(d);
info.add(tile_flags_t::thread_group);
}
void add_iter_dim(const pvar_t &d) {
iter_[d] = 1;
auto &info = tile_info(d);
info.add(tile_flags_t::iter);
}
protected:
tile_t loop_;
tile_t thread_group_;
tile_t iter_;
std::map<pvar_t, tile_info_t> tile_infos_;
};
template <class blocking_scheme_kind>
class blocking_scheme_list_impl_t {
public:
blocking_scheme_list_impl_t() : blocking_scheme_list_impl_t(0) {}
blocking_scheme_list_impl_t(int tune_level) : tune_level_(tune_level) {}
void add(bool filter, const blocking_scheme_kind &scheme) {
if ((tune_level_ == 0) && !filter) return;
schemes_.push_back(scheme);
}
const std::vector<blocking_scheme_kind> &get() const { return schemes_; }
private:
int tune_level_;
std::vector<blocking_scheme_kind> schemes_;
};
using blocking_scheme_list_t = blocking_scheme_list_impl_t<blocking_scheme_t>;
class blocking_checker_t {
public:
virtual ~blocking_checker_t() = default;
virtual void reset_checks() = 0;
virtual bool relax_checks() = 0;
virtual bool is_ok(const blocking_t &blk) const = 0;
};
class blocking_generator_t {
public:
blocking_generator_t(int vec_size, blocking_checker_t &chk,
const std::vector<level_tile_set_t> &level_tile_sets) {
for (auto &ts : level_tile_sets)
generate_all(vec_size, chk, ts);
}
std::vector<blocking_t> blockings() const {
return std::vector<blocking_t>(blockings_.begin(), blockings_.end());
}
private:
void generate_all(int vec_size, blocking_checker_t &chk,
const level_tile_set_t &level_tile_set);
void generate_sample(int vec_size, const blocking_checker_t &chk,
const level_tile_set_t &level_tile_set);
std::set<blocking_t> blockings_;
};
class blocking_params_t {
public:
static const int bufs_hint_undef = -1;
blocking_params_t() = default;
blocking_params_t(
const blocking_t &blocking, int bufs_hint = bufs_hint_undef)
: blocking_(blocking), bufs_hint_(bufs_hint) {}
int id() const { return id_; }
int bufs_hint() const { return bufs_hint_; }
bool is_empty() const { return blocking_.is_empty(); }
const blocking_t &blocking() const { return blocking_; }
void set_id(int id) { id_ = id; }
void stringify(std::ostream &out) const {
out << blocking_.str();
if (bufs_hint_ != -1) out << " bufs=" << bufs_hint_;
}
void parse(std::istream &in) {
id_ = -1;
blocking_.parse(in);
if (stream_try_match(in, "bufs=")) {
bufs_hint_ = stream_parse<int>(in);
}
}
std::string str(bool csv = false) const {
ostringstream_t oss;
if (csv) {
oss << blocking_.str(csv);
oss << "," << bufs_hint_;
} else {
oss << "cfg=\"";
oss << blocking_.str(csv);
if (bufs_hint_ == 0) oss << " s=x0 p=x0";
oss << "\"";
}
return oss.str();
}
XE_DEFINE_DUMP()
static std::vector<std::string> csv_keys() {
return {"simd", "loop", "tg", "iter", "bufs_hint"};
}
private:
int id_ = -1;
blocking_t blocking_;
int bufs_hint_ = bufs_hint_undef;
};
class prim_config_t;
class params_generator_t {
public:
params_generator_t() = default;
params_generator_t(const blocking_params_t ¶ms);
params_generator_t(int tune_level, int simd_size, blocking_checker_t &chk,
const std::vector<level_tile_set_t> &level_tile_sets,
const blocking_params_t ¶ms);
params_generator_t(int tune_level, int simd_size, blocking_checker_t &chk,
const std::vector<level_tile_set_t> &level_tile_sets, int idx = -1);
const std::vector<blocking_params_t> ¶ms_vec() const {
return params_vec_;
}
std::vector<blocking_params_t> ¶ms_vec() { return params_vec_; }
bool is_empty() const { return params_vec_.empty(); }
bool is_valid() const { return cur_idx_ < configs(); }
void move_next() { cur_idx_++; }
int cur_index() const { return cur_idx_; }
void set_cur_index(int idx) {
gpu_assert(idx < configs());
cur_idx_ = idx;
}
const blocking_params_t &cur_params() const { return at(cur_idx_); }
const blocking_params_t &at(int idx) const {
gpu_assert(idx >= 0 && idx < configs());
return params_vec_[idx];
}
void set_params(prim_config_t &cfg);
int configs() const { return (int)params_vec_.size(); }
template <typename KeyFuncT>
void sort(int beg, int end, const KeyFuncT &key_func) {
gpu_assert(beg >= 0 && beg < configs());
gpu_assert(end >= beg && end <= configs());
std::sort(params_vec_.begin() + beg, params_vec_.begin() + end,
[&](const blocking_params_t &a, const blocking_params_t &b) {
auto a_eval = key_func(a);
auto b_eval = key_func(b);
if (a_eval == b_eval) return a.id() < b.id();
return a_eval < b_eval;
});
}
template <typename PredicateFuncT>
void remove_if(const PredicateFuncT &func) {
gpu_assert(cur_idx_ == -1);
params_vec_.erase(
std::remove_if(params_vec_.begin(), params_vec_.end(), func),
params_vec_.end());
}
void shuffle(size_t seed);
void print_all() const {
using namespace ir_utils;
std::vector<std::string> headers = {};
table_t table("List of configs", headers);
for (int i = 0; i < configs(); i++) {
auto ¶ms = params_vec_[i];
gpu_trace() << "params #" << i << ": " << params;
}
}
private:
static void assign_ids(std::vector<blocking_params_t> &vec);
static void append_params(std::vector<blocking_params_t> &vec,
const blocking_params_t ¶ms);
static void append_params(std::vector<blocking_params_t> &vec,
const std::vector<level_tile_set_t> &level_tile_sets,
blocking_checker_t &chk, int tune_level, int simd_size);
std::vector<blocking_params_t> params_vec_;
int cur_idx_ = 0;
};
enum class tiler_mode_t {
undef,
env_config,
env_tiler,
lookup,
model,
tune,
default_mode = lookup
};
std::string to_string(tiler_mode_t mode);
struct tiler_params_t {
tiler_mode_t mode = tiler_mode_t::default_mode;
bool do_list = false;
int tune_iters = 0;
int env_params_idx = -1;
};
const tiler_params_t &tiler_params();
class tile_to_vec_t {
public:
tile_to_vec_t() = default;
tile_to_vec_t(const std::vector<std::vector<tile_t>> &tiles,
const std::vector<int> &ids = {});
float dist(int id0, int id1) const {
auto &v0 = vecs_[id0];
auto &v1 = vecs_[id1];
float ret = 0;
for (int i = 0; i < (int)v0.size(); i++) {
ret += float(std::abs(v0[i] - v1[i]));
}
return ret;
}
private:
struct indexed_tile_t {
struct indexed_dim_t {
indexed_dim_t() = default;
indexed_dim_t(const pvar_t &dim) : dim_(dim) {}
bool is_empty() const { return values_.empty(); }
const pvar_t &dim() const { return dim_; }
void add(dim_t value) { values_.emplace(value, dim_idx::invalid); }
void finalize() {
dim_idx_t idx = 0;
add(1);
for (auto &kv : values_) {
kv.second = idx++;
}
}
dim_idx_t to_index(dim_t value) const {
auto it = values_.find(value);
gpu_assert(it != values_.end());
return it->second;
}
pvar_t dim_;
std::map<dim_t, dim_idx_t> values_;
};
void add(const pvar_t &d, dim_t value) {
if (dim_mappers_.count(d) == 0) {
dim_mappers_[d] = indexed_dim_t(d);
}
dim_mappers_[d].add(value);
}
void add(const tile_t &t) {
for (auto &d : t) {
add(d, t[d]);
}
}
void finalize() {
for (auto &kv : dim_mappers_)
if (!kv.second.is_empty()) kv.second.finalize();
}
dim_idx_t to_index(const pvar_t &d, dim_t value) const {
return dim_mappers_.at(d).to_index(value);
}
std::vector<dim_idx_t> to_index(const tile_t &t) const {
std::vector<dim_idx_t> ret;
for (auto &kv : dim_mappers_) {
auto &m = kv.second;
if (m.is_empty()) continue;
ret.push_back(to_index(m.dim(), t.get(m.dim(), 1)));
}
return ret;
}
std::unordered_map<pvar_t, indexed_dim_t> dim_mappers_;
};
std::vector<std::vector<int>> vecs_;
};
class tune_data_t {
public:
void add_time(int id, uint64_t nsec) {
resize(id + 1);
auto &p = points_[id];
p.id = id;
p.nsec = std::min(p.nsec, nsec);
if (p.repeats == 0) reported_points_++;
p.repeats++;
if (nsec < best_point_.nsec) best_point_ = p;
}
int best_id() const { return best_point_.id; }
uint64_t nsec(int id) const { return points_[id].nsec; }
std::vector<int> best_ids(int n) const {
auto sorted_points = points_;
std::sort(sorted_points.begin(), sorted_points.end(),
[&](const bench_point_t &a, const bench_point_t &b) {
return a.nsec < b.nsec;
});
std::vector<int> ret;
for (int i = 0; i < std::min((int)sorted_points.size(), n); i++) {
auto &p = sorted_points[i];
if (p.id == -1) break;
ret.push_back(p.id);
}
return ret;
}
int reported_points() const { return reported_points_; }
void resize(int new_size) {
int size = (int)points_.size();
if (new_size <= size) return;
points_.resize(new_size);
for (int i = size; i < new_size; i++) {
points_[i].id = i;
}
}
private:
static const uint64_t max_nsec_ = std::numeric_limits<uint64_t>::max();
struct bench_point_t {
int id = -1;
int repeats = 0;
uint64_t nsec = max_nsec_;
bool is_ok() const { return nsec != max_nsec_; }
};
std::vector<bench_point_t> points_;
int reported_points_ = 0;
bench_point_t best_point_;
};
} } } } }
#endif