#include "gpu/intel/jit/ir/v2/send.hpp"
namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace jit {
namespace v2 {
bool process_coef_y_stride(
plane_t &plane, expr_t &u, const expr_t &v, const prover_t &prover) {
auto reduce_if_divisible = [](expr_t &a, const expr_t &b) {
auto args = op_split(op_kind_t::_mul, a);
for (auto &arg : args) {
if (arg.is_equal(b)) {
arg = expr_t();
a = op_combine(op_kind_t::_mul, args);
return true;
}
}
return false;
};
if (reduce_if_divisible(u, plane.y_stride)) return true;
if (auto *op = u.as_ptr<unary_op_t>()) {
auto a = op->a;
return process_coef_y_stride(plane, a, v, prover);
}
for (auto &size_val : {std::make_pair(u, 0), std::make_pair(v, 1)}) {
auto dim = to_pvar(size_val.first);
if (!dim.is_undef()) {
if (!prover.require((plane.y_stride == 1)
| (size_val.first == size_val.second)))
return false;
return true;
}
}
gpu_error_not_expected() << "Cannot make " << u << " x " << v
<< " divisible by " << plane.y_stride;
return false;
}
bool adjust_for_non_unit_y_stride(
plane_t &plane, const coord_t &coord, const prover_t &prover) {
if (plane.y_stride.is(1)) return true;
auto y_stride_dim = to_pvar(plane.y_stride);
if (y_stride_dim.is_undef()) return false;
auto _y = to_linear(plane.y);
auto &y = _y.as<linear_t>();
auto to_size = [&](const expr_t &idx) {
for (auto &d : coord) {
if (coord[d].is_equal(idx)) return var(d);
}
return expr_t();
};
auto c = y.c;
auto u_vec = y.u_vec;
if (!process_coef_y_stride(plane, c, expr_t(), prover)) return false;
for (int i = 0; i < y.nargs(); i++) {
if (!process_coef_y_stride(
plane, u_vec[i], to_size(y.v_vec[i]), prover))
return false;
}
if (!prover.require(plane.H % plane.y_stride == 0)) return false;
plane.y = linear_t::to_expr(c, u_vec, y.v_vec);
plane.H /= plane.y_stride;
return true;
}
send_2d_desc_t::send_2d_desc_t(const view_t &view, const send_params_t ¶ms,
const prover_t &prover) {
auto plane = view.plane();
if (!params.hint_2d) return;
if (!plane) return;
if (!adjust_for_non_unit_y_stride(plane, view.coord(), prover)) return;
auto &hint = params.hint_2d;
hw = params.hw;
address = params.address;
op = params.op;
type = view.type();
transpose = hint.transpose;
vnni = hint.vnni;
W = plane.W;
H = plane.H;
P = plane.P;
w = hint.width;
h = hint.height;
c = 1;
w_rcount = ir_utils::safe_div(plane.w, w);
h_rcount = ir_utils::safe_div(plane.h, h);
w_dim = plane.w_dim;
h_dim = plane.h_dim;
base = get_2d_base(view);
x_base = plane.x;
y_base = plane.y;
try_promote_count();
is_valid = is_supported(view, prover);
}
} } } } } }