#include "gpu/intel/jit/pass/pass.hpp"
#include "gemmstone/../../dsl/ir/pass/trace.hpp"
namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace jit {
class external_var_visitor_t : public scope_visitor_t {
public:
void _visit(const var_t &obj) override {
if (!is_expr_defined(obj)) external_vars.insert(obj);
}
object_eq_set_t<expr_t> external_vars;
};
stmt_t inject_external_var_let(const stmt_t &_stmt, ir_context_t &ir_ctx) {
ir::trace_start();
auto stmt = _stmt;
external_var_visitor_t v;
v.visit(stmt);
std::vector<expr_t> external_vars(
v.external_vars.begin(), v.external_vars.end());
std::sort(external_vars.begin(), external_vars.end(),
[&](const expr_t &a, const expr_t &b) {
return a.as<var_t>().name < b.as<var_t>().name;
});
for (auto &var : external_vars)
stmt = let_t::make(var, {}, stmt);
ir::trace_pass("inject_external_var_let", stmt, ir_ctx);
return stmt;
}
class store_splitter_t : public ir_mutator_t {
public:
store_splitter_t(const dsl::hw_t &hw) : hw_(hw) {}
object_t _mutate(const store_t &obj) override {
int elems = obj.value.type().elems();
int elem_size = obj.value.type().base().size();
int stride = (obj.has_default_stride() ? 1 : obj.stride / elem_size);
int store_size = elem_size * stride * elems;
const auto grf_size = hw_.grf_size();
if (store_size <= 2 * grf_size) return ir_mutator_t::_mutate(obj);
int step = 2 * grf_size / (stride * elem_size);
stmt_t new_stmt;
for (int i = 0; i < elems; i += step) {
int cur_elems = std::min(step, elems - i);
gpu_assert(math::is_pow2(cur_elems));
int off = i * stride * elem_size;
auto store = store_t::make(obj.buf, obj.off + off,
split_expr(obj.value, i, i + cur_elems), obj.stride);
new_stmt = new_stmt.append(store);
}
return std::move(new_stmt);
}
private:
static expr_t split_expr(const expr_t &e, int beg, int end) {
auto *shuffle = e.as_ptr<shuffle_t>();
if (shuffle) return shuffle_t::make(shuffle, beg, end);
auto *binary = e.as_ptr<binary_op_t>();
if (binary) {
auto a = split_expr(binary->a, beg, end);
auto b = split_expr(binary->b, beg, end);
return binary_op_t::make(binary->op_kind, a, b);
}
auto *load = e.as_ptr<load_t>();
if (load) {
int stride = load->stride;
if (load->has_default_stride()) stride = load->type.base().size();
return load_t::make(load->type.with_elems(end - beg), load->buf,
load->off + beg * stride, load->stride);
}
gpu_error_not_expected();
return expr_t();
}
dsl::hw_t hw_;
};
stmt_t split_wide_stores(const stmt_t &s, ir_context_t &ir_ctx) {
ir::trace_start();
auto ret = store_splitter_t(ir_ctx.hw()).mutate(s);
ir::trace_pass("split_wide_stores", ret, ir_ctx);
return ret;
}
class if_condition_fixer_t : public ir_mutator_t {
public:
if_condition_fixer_t(int simd_size)
: simd_size_(simd_size), in_cond_(false) {}
object_t _mutate(const if_t &obj) override {
auto _new_obj = ir_mutator_t::_mutate(obj);
auto &new_obj = _new_obj.as<if_t>();
flag_setter_t in_cond(&in_cond_, true);
auto cond = mutate(new_obj.cond);
return if_t::make(cond, new_obj.body, new_obj.else_body);
}
object_t _mutate(const binary_op_t &obj) override {
if (!in_cond_) return obj;
auto broadcast = [&](const expr_t &operand) {
object_t ret;
if (is_cmp_op(obj.op_kind) && obj.type.elems() == 1)
ret = shuffle_t::make_broadcast(operand, simd_size_);
else
ret = mutate(operand);
return ret;
};
auto a = broadcast(obj.a);
auto b = broadcast(obj.b);
return binary_op_t::make(obj.op_kind, a, b);
}
private:
struct flag_setter_t {
flag_setter_t(bool *flag, bool value) : flag(flag), old(*flag) {
*flag = value;
}
~flag_setter_t() { *flag = old; }
bool *flag;
bool old;
};
int simd_size_;
bool in_cond_;
};
stmt_t fixup_if_conditions(const stmt_t &s, ir_context_t &ir_ctx) {
ir::trace_start();
auto ret = if_condition_fixer_t(ir_ctx.options().simd()).mutate(s);
ir::trace_pass("fixup_if_conditions", ret, ir_ctx);
return ret;
}
class int64_expr_optimizer_t : public ir_mutator_t {
public:
#define HANDLE_IR_OBJECT(type) \
object_t _mutate(const type &obj) override { return mutate_expr(obj); }
HANDLE_EXPR_IR_OBJECTS()
#undef HANDLE_IR_OBJECT
private:
template <typename T>
object_t mutate_expr(const T &obj) {
auto new_obj = ir_mutator_t::_mutate(obj);
if (auto *binary = new_obj.template as_ptr<binary_op_t>()) {
if (binary->op_kind == op_kind_t::_add) {
new_obj = simplify_64_bit_add(new_obj);
}
}
return new_obj;
}
};
stmt_t optimize_int64_exprs(const stmt_t &s, ir_context_t &ir_ctx) {
ir::trace_start();
auto ret = int64_expr_optimizer_t().mutate(s);
ir::trace_pass("optimize_int64_exprs", ret, ir_ctx);
return ret;
}
} } } } }