#include "gpu/intel/jit/eltwise_injector.hpp"
#include "common/impl_registration.hpp"
#include "gpu/intel/jit/utils/type_bridge.hpp"
#include <limits>
namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace jit {
using namespace ngen;
template <typename ngen_generator_t>
int eltwise_injector_f32_t<ngen_generator_t>::min_scratch_regs() {
using namespace alg_kind;
if (is_fwd_) {
switch ((int)alg_) {
case eltwise_elu:
case eltwise_elu_use_dst_for_bwd: return 1;
case eltwise_exp:
case eltwise_exp_use_dst_for_bwd: return 0;
case eltwise_gelu_erf: return 4;
case eltwise_hardsigmoid: return 0;
case eltwise_hardswish: return 1;
case eltwise_log: return 0;
case eltwise_mish: return 4;
case eltwise_pow: return 1;
case eltwise_relu:
case eltwise_relu_use_dst_for_bwd: return 1;
case eltwise_abs: return 0;
case eltwise_soft_relu: return 1;
case eltwise_sqrt:
case eltwise_sqrt_use_dst_for_bwd: return 0;
case eltwise_square: return 0;
case eltwise_swish: return 1;
case eltwise_tanh:
case eltwise_tanh_use_dst_for_bwd:
return (hw() < gpu_xe3p_35_10) ? 2 : 1;
case eltwise_round: return 0;
case eltwise_linear: return 0;
case eltwise_clip:
case eltwise_clip_v2:
case eltwise_clip_v2_use_dst_for_bwd: return 0;
case eltwise_gelu_tanh: return 2;
case eltwise_logistic:
case eltwise_logistic_use_dst_for_bwd: return 0;
case eltwise_stochastic_round: return 6;
case eltwise_mx_scale: return hw() >= gpu_xe_hpc ? 1 : 2;
default: assert(!"unsupported eltwise algorithm");
}
} else {
switch (alg_) {
case eltwise_relu: return 1;
case eltwise_abs: return 1;
case eltwise_square: return 0;
case eltwise_linear: return 0;
case eltwise_clip: return 1;
case eltwise_gelu_tanh: return (hw() < gpu_xe3p_35_10) ? 2 : 1;
default: assert(!"unsupported eltwise algorithm");
}
}
return 0;
}
template <typename ngen_generator_t>
int eltwise_injector_f32_t<ngen_generator_t>::preferred_scratch_regs() {
using namespace alg_kind;
if (is_fwd_) {
switch (alg_) {
case eltwise_elu:
case eltwise_elu_use_dst_for_bwd: return 8;
case eltwise_gelu_erf: return 8;
case eltwise_hardswish: return 8;
case eltwise_mish: return 8;
case eltwise_relu:
case eltwise_relu_use_dst_for_bwd: return (alpha_ == 0.f) ? 1 : 8;
case eltwise_tanh: return 8;
case eltwise_gelu_tanh: return 8;
case eltwise_soft_relu: return 8;
case eltwise_swish: return 8;
default: break;
}
} else {
switch (alg_) {
case eltwise_gelu_tanh: return 8;
default: break;
}
}
return min_scratch_regs();
}
template <typename ngen_generator_t>
int eltwise_injector_f32_t<ngen_generator_t>::max_batch_size() {
using namespace alg_kind;
auto ss = scratch_.getLen();
if (is_fwd_) {
switch (alg_) {
case eltwise_relu:
case eltwise_relu_use_dst_for_bwd:
if (alpha_ == 0.)
break;
else
return ss;
case eltwise_elu:
case eltwise_elu_use_dst_for_bwd:
case eltwise_hardswish:
case eltwise_pow:
case eltwise_soft_relu:
case eltwise_swish: return ss;
case eltwise_tanh:
case eltwise_mish:
case eltwise_gelu_erf: return ss / min_scratch_regs();
case eltwise_gelu_tanh: return ss & ~1;
default: break;
}
} else {
switch (alg_) {
case eltwise_gelu_tanh: return ss / 2;
default: break;
}
}
return 128;
}
template <typename ngen_generator_t>
int eltwise_injector_f32_t<ngen_generator_t>::phase_count(alg_kind_t alg) {
using namespace alg_kind;
if (is_fwd_) {
switch (alg) {
case eltwise_elu:
case eltwise_elu_use_dst_for_bwd: return 5;
case eltwise_exp:
case eltwise_exp_use_dst_for_bwd: return 2;
case eltwise_gelu_erf: return 25;
case eltwise_hardsigmoid: return 4;
case eltwise_hardswish: return 5;
case eltwise_log: return 2;
case eltwise_mish:
return phase_count(alg_kind::eltwise_soft_relu)
+ phase_count(alg_kind::eltwise_tanh) + 1;
case eltwise_pow: return 6;
case eltwise_relu:
case eltwise_relu_use_dst_for_bwd: return (alpha_ == 0) ? 1 : 2;
case eltwise_soft_relu: return 10;
case eltwise_swish: return (hw() < gpu_xe3p_35_10) ? 5 : 3;
case eltwise_tanh:
case eltwise_tanh_use_dst_for_bwd:
return (hw() < gpu_xe3p_35_10) ? 6 : 1;
case eltwise_linear: return (beta_ == 0) ? 1 : 2;
case eltwise_clip:
case eltwise_clip_v2:
case eltwise_clip_v2_use_dst_for_bwd: return 2;
case eltwise_gelu_tanh: return 8;
case eltwise_logistic:
case eltwise_logistic_use_dst_for_bwd:
return (hw() < gpu_xe3p_35_10) ? 4 : 1;
default: break;
}
} else {
switch (alg) {
case eltwise_abs: return 2;
case eltwise_clip: return 4;
case eltwise_gelu_tanh: return (hw() < gpu_xe3p_35_10) ? 14 : 8;
default: break;
}
}
return 1;
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::relu_zero_ns_prepare_fwd() {
h->mov(1, scratch_[0].f(0), 0.f);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::relu_zero_ns_compute_fwd(
int simd, const ngen::GRF &r) {
h->csel(simd | le | f0[0], r, scratch_[0].f(0), r, r);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::relu_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off) {
auto temp = scratch_[off].f();
switch (phase) {
case 0: h->mul(simd, temp, r, alpha_); break;
case 1: h->csel(simd | le | f0[0], r, temp, r, r); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::abs_compute_fwd(
int simd, const ngen::GRF &r) {
h->mov(simd, r, abs(r));
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::soft_relu_compute_fwd_inner(
int simd, const ngen::GRF &input, const ngen::GRF &temp,
const ngen::GRF &dest, int phase, int off, float alpha) {
const float exp_overflow_bound = 88.72283172607421875f;
const float log2e = 1.44269502162933349609375f;
const float reciproc_log2e = 1.f / log2e; switch (phase) {
case 0: h->mul(simd, temp, input, alpha); break;
case 1: h->add(simd, dest, input, -exp_overflow_bound); break;
case 2: h->csel(simd | le | f0[0], dest, dest, temp, dest); break;
case 3: h->mul(simd, temp, temp, log2e); break;
case 4: h->exp(simd, temp, temp); break;
case 5: h->add(simd, temp, temp, 1.f); break;
case 6: h->log(simd, temp, temp); break;
case 7: h->mul(simd, temp, temp, reciproc_log2e); break;
case 8: h->csel(simd | le | f0[0], temp, temp, dest, dest); break;
case 9: h->mul(simd, dest, temp, 1.f / alpha); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::soft_relu_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off) {
auto temp = scratch_[off].f();
soft_relu_compute_fwd_inner(simd, r, temp, r, phase, off, alpha_);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::sqrt_compute_fwd(
int simd, const ngen::GRF &r) {
h->sqt(simd, r, r);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::square_compute_fwd(
int simd, const ngen::GRF &r) {
h->mul(simd, r, r, r);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::tanh_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off, int batch) {
if (hw() < gpu_xe3p_35_10) {
const float log2e = 1.44269502162933349609375f; auto one_half = scratch_[0].f(7);
auto a = scratch_[off + batch].f();
switch (phase) {
case 0: h->mul(simd, a, abs(r), 2.f * log2e); break;
case 1: h->exp(simd, a, a); break;
case 2: h->mad(simd, a, one_half, a, one_half); break;
case 3: h->inv(simd, a, a); break;
case 4: h->add(simd, a, -a, 1.f); break;
case 5: h->csel(simd | ge | f0[0], r, a, -a, r); break;
default: assert(!"invalid phase");
}
} else {
switch (phase) {
case 0: h->tanh(simd, r, r); break;
default: assert(!"invalid phase");
}
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::tanh_compute_fwd_compat(
int simd, const ngen::GRF &r, int phase, int off, int batch) {
auto k = scratch_[0].f(4);
auto l = scratch_[0].f(5);
auto m = scratch_[0].f(6);
auto n = scratch_[0].f(7);
auto a = scratch_[off + batch].f();
switch (phase) {
case 0: h->mad(simd, a, m, r, r); break;
case 1: h->mad(simd, a, n, a, a); break;
case 2: h->inv(simd, a, a); break;
case 3: h->mul(simd, a, a, r); break;
case 4: h->mad(simd, r, k, r, r); break;
case 5: h->mad(simd, r, l, r, r); break;
case 6: h->mul(simd, r, r, 0.0519867f); break; case 7: h->mul(simd | sat, r, r, abs(a)); break;
case 8: h->csel(simd | ge | f0[0], r, r, -r, a); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::round_compute_fwd(
int simd, const ngen::GRF &r) {
h->rnde(simd, r, r);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::mx_scale_compute_fwd(int simd,
const ngen::GRF &r, int off, const ngen::Subregister &seed, int nreg,
const ngen::DataType dst_dt, int phase) {
assert(simd == 32);
assert(utils::one_of(dst_dt, ngen::DataType::bf8, ngen::DataType::hf8,
static_cast<ngen::DataType>(0x5A)));
int scale_off = (phase / nreg * 4);
int grf_bytes = GRF::bytes(hw());
int scale_reg_incr = scale_off / grf_bytes;
auto scale_dst = scale_reg_incr ? ngen::GRF(seed.getBase() + scale_reg_incr)
.ub(scale_off % grf_bytes)
: seed.ub(scale_off);
ngen::GRF r_alt(off);
auto max = scratch_[0].f();
float zero_e8m0 = 5.877472e-39f, inv_fmax_bf8 = 0.000030517578125f,
inv_fmax_hf8 = 0.00390625f, inv_fmax_f4_e2m1 = 0.25f;
auto fmax = (dst_dt == ngen::DataType::bf8) ? Immediate::f(inv_fmax_bf8)
: (dst_dt == ngen::DataType::hf8) ? Immediate::f(inv_fmax_hf8)
: Immediate::f(inv_fmax_f4_e2m1);
h->add(16, r.ud(0)(1), r.ud(0)(1), Immediate::ud(0x80800000));
h->add(16, r_alt.ud(0)(1), r_alt.ud(0)(1), Immediate::ud(0x80800000));
h->sel(16 | ge, max.f(0)(1), abs(r.f(0)(1)), abs(r_alt.f(0)(1)));
h->add(16, r.d(0)(1), r.d(0)(1), Immediate::d(0x7F800000));
h->add(16, r_alt.d(0)(1), r_alt.d(0)(1), Immediate::d(0x7F800000));
h->sel(8 | ge, max.ud(0)(1), max.ud(8)(1), max.ud(0)(1));
h->sel(4 | ge, max.ud(0)(1), max.ud(4)(1), max.ud(0)(1));
h->sel(2 | ge, max.ud(0)(1), max.ud(2)(1), max.ud(0)(1));
h->sel(1 | ge, max.ud(0)(1), max.ud(1)(1), max.ud(0)(1));
h->add(1, max.d(0)(1), max.d(0)(1), Immediate::d(0x7F800000));
h->and_(1, max.ud(0), max.ud(0), Immediate::ud(0x7F800000));
h->mul(1, max.f(0), max.f(0), fmax);
h->and_(1, max.ud(0), max.ud(0), Immediate::ud(0x7F800000));
h->sel(1 | ge, max.f(0)(1), max.f(0), Immediate::f(zero_e8m0));
h->inv(1, max.f(1), max.f(0));
h->mul(16, r, r, max.f(1)(0));
h->mul(16, r_alt.f(0)(1), r_alt.f(0)(1), max.f(1)(0));
h->shr(1, max.ud(0), max.ud(0), 23);
h->mov(1, scale_dst.ub(0), max.ub(0));
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::sround_compute_fwd(int simd,
const ngen::GRF &r, int phase, const ngen::Subregister &seed,
const ngen::DataType dst_dt, int off) {
auto bias = scratch_[0].ud();
auto u_r = r[0].ud()(8, 8, 1);
auto u_f = r[0].f()(8, 8, 1);
int base_idx = off * simd;
h->template mov<uint16_t>(8, bias.uw(0)(1), Immediate::uv(0x76543210));
h->template mov<uint16_t>(8, bias.uw(8)(1), Immediate::uv(0xfedcba98));
auto imm = Immediate::ud(base_idx);
if (hw() >= gpu_xe_hpc)
h->add(16, bias.ud(), bias.uw(), imm);
else {
auto extra = scratch_[1].ud();
h->add(8, extra.ud(0)(1), bias.uw(8)(8, 8, 1), imm);
h->add(8, bias.ud(0)(1), bias.uw(0)(8, 8, 1), imm);
}
const uint32_t dst_dt_digits = dnnl::impl::types::digits<uint32_t>(
convert_ngen_type_to_dnnl(dst_dt));
assert(dst_dt_digits <= 24);
data_type_t dnnl_t = convert_ngen_type_to_dnnl(dst_dt);
const float f_min = types::min_value<float>(dnnl_t);
const float max = types::max_value<float>(dnnl_t);
const float lowest = types::lowest_value<float>(dnnl_t);
auto bia_scratch = scratch_[4].ud();
h->and_(simd | h->nz | f0[0], h->null.ud(), ~u_r, 0x7F800000);
const int truncation_mask = (0xffffffff << (24 - dst_dt_digits));
philox_4x32(simd, seed, bias);
if (getBytes(dst_dt) == 2) {
h->mov(simd, bia_scratch.ud(0), bias.ub(0)(simd, simd, 1));
} else {
h->mov(simd, bia_scratch.ud(0), bias.uw(0)(simd, simd, 1));
}
h->and_(simd, bia_scratch, bia_scratch, ~truncation_mask);
h->add(simd | f0[0], u_r, u_r, bia_scratch);
h->and_(simd | f0[0], u_r, u_r, truncation_mask);
h->max_(simd, u_f, u_f, lowest);
h->min_(simd, u_f, u_f, max);
h->cmp(simd | lt | f0[0], abs(u_f), f_min);
h->mov(simd | f0[0], u_r, 0);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::philox_4x32(
int simd, const ngen::Subregister &seed, const ngen::GRF &bias) {
auto sround_seed = seed;
auto ctr = bias.ud(0);
auto key = scratch_[2].ud();
auto ctr_mul = scratch_[3].ud();
auto offs = scratch_[4].ud();
auto off_inc = scratch_[5].uw();
auto addr = h->indirect[h->a0].ud(0)(0, 1, 0);
if (hw() >= gpu_xe_hpc)
h->mov(4, key.uq(0)(4, 4, 1), uint64_t(0xBB67AE859E3779B9uLL));
else {
h->mov(4, key.ud(0)(4, 4, 2), uint32_t(0x9E3779B9u));
h->mov(4, key.ud(1)(4, 4, 2), uint32_t(0xBB67AE85u));
}
if (hw() < gpu_xe3p_35_10) {
h->template mov<uint16_t>(4, offs.uw(16)(1), Immediate::uv(0x00009988));
h->template mul<uint32_t>(
4, offs.ud(8)(1), key.ud(0)(4, 4, 1), offs.uw(16)(4, 4, 1));
h->add(4, offs.ud(8)(1), offs.ud(8)(4, 4, 1), sround_seed);
h->template mov<uint16_t>(8, offs.uw(8)(1), Immediate::uv(0x77665544));
h->template mov<uint16_t>(8, offs.uw(0)(1), Immediate::uv(0x33221100));
h->template mul<uint32_t>(
8, key.ud(8)(1), key.ud(0)(8, 8, 1), offs.uw(8)(8, 8, 1));
h->template mul<uint32_t>(
8, key.ud(0)(1), key.ud(0)(8, 8, 1), offs.uw(0)(8, 8, 1));
h->add(16, key, key, sround_seed);
} else {
h->template mov<uint16_t>(4, offs.ud(0)(1), Immediate::uv(0x00009988));
h->template mul<uint32_t>(
4, offs.ud(0)(1), key.ud(0)(4, 4, 1), offs.ud(0)(4, 4, 1));
h->template mov<uint16_t>(8, offs.ud(8)(1), Immediate::uv(0x77665544));
h->template mov<uint16_t>(8, key.ud(8)(8, 8, 1), key.ud(0)(8, 8, 1));
h->template mul<uint32_t>(
8, key.ud(8)(1), key.ud(8)(8, 8, 1), offs.ud(8)(8, 8, 1));
h->template mov<uint16_t>(8, offs.ud(0)(1), Immediate::uv(0x33221100));
h->template mul<uint32_t>(
8, key.ud(0)(1), key.ud(0)(8, 8, 1), offs.ud(0)(8, 8, 1));
h->add(16, key, key, sround_seed);
}
h->mov(4, ctr_mul.ud(2)(4), 0xCD9E8D57);
h->mov(4, ctr_mul.ud(0)(4), 0xD2511F53);
auto ctr_base_sub = offs.uw(8)(8, 8, 1);
h->mov(8, ctr_base_sub,
(ctr.getBase() * GRF::bytes(hw())) + ctr.getOffset());
h->template mov<uint16_t>(8, off_inc.uw(0)(1), Immediate::uv(0x56741230));
h->template mul<uint16_t>(8, off_inc.uw(0)(1), off_inc.uw(0)(8, 8, 1), 4);
auto philox_round = [&](ngen::Subregister &ctr, ngen::GRF &ctr_mul,
ngen::GRF &key, int idx) {
h->add(8, h->a0, ctr_base_sub, off_inc);
h->template movi<uint32_t>(8, ctr.ud(0)(1), addr);
h->add(8, h->a0, h->a0, 32);
h->template movi<uint32_t>(8, ctr.ud(8)(1), addr);
h->mov(4, ctr_mul.ud(1)(4), key.ud(idx + 1)(0, 1, 0));
h->mov(4, ctr_mul.ud(3)(4), key.ud(idx)(0, 1, 0));
h->xor_(8, ctr_mul.ud(1)(2), ctr_mul.ud(1)(8, 4, 2),
ctr.ud(1)(8, 4, 2));
auto ctrLo = ctr.ud(0)(8, 4, 2);
auto ctrHi = ctr.ud(1)(8, 4, 2);
const auto grf_size = ngen::GRF::bytes(hw());
const int esize = grf_size / 8; const int steps = utils::div_up(8, esize);
if (hw() < gpu_xe3p_35_10) {
auto acc = h->acc0.retype(DataType::ud);
for (int i = 0, off = 0; i < steps; ++i, off += 2 * esize)
h->mul(esize, acc[i](2), ctr.ud(off)(8, 4, 2),
ctr_mul.uw(off)(8, 2, 4));
h->mach(8, ctrLo, ctr.ud(0)(8, 4, 2), ctr_mul.ud(0)(8, 4, 2));
h->mov(8, ctrHi, ctrLo);
for (int i = 0, off = 0; i < steps; ++i, off += 2 * esize)
h->mov(esize, ctr.ud(off)(2), acc[i](2));
} else {
for (int i = 0, off = 0; i < steps; ++i, off += 2 * esize)
h->mul(esize, ctr.uq(off)(1), ctr.ud(off)(8, 4, 2),
ctr_mul.ud(off)(8, 4, 2));
}
h->xor_(8, ctr.ud(1)(2), ctr.ud(1)(8, 4, 2), ctr_mul.ud(1)(8, 4, 2));
if (idx == 0) {
h->template mov<uint16_t>(
8, off_inc.uw(0)(1), Immediate::uv(0x65472103));
h->template mul<uint16_t>(
8, off_inc.uw(0)(1), off_inc.uw(0)(8, 8, 1), 4);
}
};
philox_round(ctr, ctr_mul, key, 0);
philox_round(ctr, ctr_mul, key, 2);
philox_round(ctr, ctr_mul, key, 4);
philox_round(ctr, ctr_mul, key, 6);
philox_round(ctr, ctr_mul, key, 8);
philox_round(ctr, ctr_mul, key, 10);
philox_round(ctr, ctr_mul, key, 12);
philox_round(ctr, ctr_mul, key, 14);
philox_round(ctr, ctr_mul, offs, 8);
philox_round(ctr, ctr_mul, offs, 10);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::swish_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off) {
auto temp = scratch_[off].f();
if (hw() < gpu_xe3p_35_10) {
const float log2e = 1.442695f; switch (phase) {
case 0: h->mul(simd, temp, r, -1.f * log2e * alpha_); break;
case 1: h->exp(simd, temp, temp); break;
case 2: h->add(simd, temp, temp, 1.f); break;
case 3: h->inv(simd, temp, temp); break;
case 4: h->mul(simd, r, r, temp); break;
default: assert(!"invalid phase");
}
} else {
switch (phase) {
case 0: h->mul(simd, temp, r, alpha_); break;
case 1: h->sigm(simd, temp, temp); break;
case 2: h->mul(simd, r, r, temp); break;
default: assert(!"invalid phase");
}
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::linear_compute_fwd(
int simd, const ngen::GRF &r, int phase) {
switch (phase) {
case 0: h->mul(simd, r, r, alpha_); break;
case 1: h->add(simd, r, r, beta_); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::clip_compute_fwd(
int simd, const ngen::GRF &r, int phase, float alpha, float beta) {
switch (phase) {
case 0: h->max_(simd, r, r, alpha); break;
case 1: h->min_(simd, r, r, beta); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::gelu_tanh_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off) {
const float k = 0.044715f;
const float sqrt_2_over_pi = 0.7978845f; const float log2e = 1.442695f;
int msimd = simd;
if (hw() == gpu_xe_hp)
msimd = 16;
auto a = scratch_[off].f();
switch (phase) {
case 0: h->mul(simd, a, r, r); break;
case 1: h->mul(simd, a, a, k); break;
case 2: h->mad(simd, a, r, a, r); break;
case 3: h->mul(simd, a, a, -2 * sqrt_2_over_pi * log2e); break;
case 4: h->exp(msimd, a, a); break;
case 5: h->add(simd, a, a, 1.0f); break;
case 6: h->inv(msimd, a, a); break;
case 7: h->mul(simd, r, a, r); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::logistic_compute_fwd(
int simd, const ngen::GRF &r, int phase) {
if (hw() < gpu_xe3p_35_10) {
const float log2e = 1.442695f; switch (phase) {
case 0: h->mul(simd, r, r, -1.f * log2e); break;
case 1: h->exp(simd, r, r); break;
case 2: h->add(simd, r, r, 1.f); break;
case 3: h->inv(simd, r, r); break;
default: assert(!"invalid phase");
}
} else {
switch (phase) {
case 0: h->sigm(simd, r, r); break;
default: assert(!"invalid phase");
}
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::relu_prepare_bwd() {
auto neg_slope = scratch_[0].f(0);
auto pos_slope = scratch_[0].f(4);
h->mov(1, neg_slope, alpha_);
h->mov(1, pos_slope, 1.f);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::relu_compute_bwd(
int simd, const ngen::GRF &r) {
auto neg_slope = scratch_[0].f(0);
auto pos_slope = scratch_[0].f(4);
h->csel(simd | le | f0[0], r, neg_slope, pos_slope, r);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::abs_prepare_bwd() {
auto neg_one = scratch_[0].f(0);
auto pos_one = scratch_[0].f(4);
h->mov(1, neg_one, -1.f);
h->mov(1, pos_one, 1.f);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::clip_prepare_bwd() {
auto pos_inf_imm = Immediate(std::numeric_limits<float>::infinity());
auto zero = scratch_[0].f(0);
auto one = scratch_[0].f(1);
auto pos_inf = scratch_[0].f(2);
h->mov(1, zero, 0.f);
h->mov(1, one, 1.f);
h->mov(1, pos_inf, pos_inf_imm);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::tanh_prepare_fwd() {
if (hw() < gpu_xe3p_35_10) {
auto one_half = scratch_[0].f(7);
h->mov(1, one_half, 0.5f);
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::tanh_prepare_fwd_compat() {
auto k = scratch_[0].f(4);
auto l = scratch_[0].f(5);
auto m = scratch_[0].f(6);
auto n = scratch_[0].f(7);
h->mov(1, k, 77.0954f); h->mov(1, l, -4435.55f); h->mov(1, m, 17.06396f); h->mov(1, n, -212.7724f); }
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::abs_compute_bwd(
int simd, const ngen::GRF &r, int phase) {
auto neg_one = scratch_[0].f(0);
auto pos_one = scratch_[0].f(4);
switch (phase) {
case 0: h->csel(simd | lt | f0[0], r, neg_one, r, r); break;
case 1: h->csel(simd | gt | f0[0], r, pos_one, r, r); break;
default: break;
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::square_compute_bwd(
int simd, const ngen::GRF &r) {
h->add(simd, r, r, r);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::linear_compute_bwd(
int simd, const ngen::GRF &r) {
h->mov(simd, r, alpha_);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::clip_compute_bwd(
int simd, const ngen::GRF &r, int phase, float alpha, float beta) {
auto zero = scratch_[0].f(0);
auto one = scratch_[0].f(1);
auto pos_inf = scratch_[0].f(2);
switch (phase) {
case 0: h->add(simd, r, r, -alpha); break;
case 1: h->csel(simd | le | f0[0], r, pos_inf, r, r); break;
case 2: h->add(simd, r, r, alpha - beta); break;
case 3: h->csel(simd | le | f0[0], r, one, zero, r); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::gelu_tanh_compute_bwd(
int simd, const ngen::GRF &r, int phase, int off, int batch) {
const float k = 0.044715f;
const float sqrt_2_over_pi = 0.7978845f; const float log2e = 1.442695f;
int msimd = simd;
if (hw() == gpu_xe_hp) msimd = 16;
auto a = scratch_[off].f();
if (hw() < gpu_xe3p_35_10) {
auto b = scratch_[off + batch].f();
switch (phase) {
case 0: h->mul(simd, a, r, r); break;
case 1: h->mul(simd, b, a, 3.0f * k); break;
case 2: h->mul(simd, a, a, k); break;
case 3: h->mad(simd, a, r, a, r); break;
case 4: h->mad(simd, b, r, b, r); break;
case 5: h->mul(simd, a, a, -2 * sqrt_2_over_pi * log2e); break;
case 6: h->mul(simd, b, b, 2 * sqrt_2_over_pi); break;
case 7: h->exp(msimd, a, a); break;
case 8: h->add(simd, r, a, 1.0f); break;
case 9: h->inv(msimd, r, r); break;
case 10: h->mul(simd, a, a, r); break;
case 11: h->mul(simd, a, a, b); break;
case 12: h->add(simd, a, a, 1.0f); break;
case 13: h->mul(simd, r, r, a); break;
default: assert(!"invalid phase");
}
} else {
switch (phase) {
case 0: h->mul(simd, a, r, k); break;
case 1: h->mad(simd, a, half(1.f), a, r); break;
case 2: h->mul(simd, a, a, r); break;
case 3: h->mul(simd, a, a, sqrt_2_over_pi); break;
case 4: h->tanh(simd, r, a); break;
case 5: h->mad(simd, a, a, a, -r); break;
case 6: h->mad(simd, r, half(0.5f), r, half(0.5f)); break;
case 7: h->mad(simd, r, r, r, a); break;
default: assert(!"invalid phase");
}
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::elu_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off) {
auto temp = scratch_[off].f();
const float log2e = 1.442695f; switch (phase) {
case 0: h->mul(simd, temp, r, log2e); break;
case 1: h->exp(simd, temp, temp); break;
case 2: h->add(simd, temp, temp, -1.f); break;
case 3: h->mul(simd, temp, temp, alpha_); break;
case 4: h->csel(simd | le | f0[0], r, temp, r, r); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::exp_compute_fwd(
int simd, const ngen::GRF &r, int phase) {
const float log2e = 1.442695f; switch (phase) {
case 0: h->mul(simd, r, r, log2e); break;
case 1: h->exp(simd, r, r); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::gelu_erf_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off, int batch) {
auto temp = scratch_[off].f();
auto at_accum = scratch_[off + batch].f();
auto tpow = scratch_[off + 2 * batch].f();
auto temp2 = scratch_[off + 3 * batch].f();
const float log2e = 1.442695f; const float reciproc_sqrt_2 = 0.707106769084930419921875f; const float p = 0.3275911f;
const float a1 = 0.254829592f;
const float a2 = -0.284496736f;
const float a3 = 1.421413741f;
const float a4 = -1.453152027f;
const float a5 = 1.061405429f;
switch (phase) {
case 0: h->mul(simd, temp, abs(r), reciproc_sqrt_2); break;
case 1: h->mul(simd, temp, temp, p); break;
case 2: h->add(simd, temp, temp, 1.f); break;
case 3: h->inv(simd, temp, temp); break;
case 4: h->mul(simd, at_accum, temp, a1); break;
case 5: h->mul(simd, tpow, temp, temp); break;
case 6: h->mul(simd, temp2, tpow, a2); break;
case 7: h->add(simd, at_accum, temp2, at_accum); break;
case 8: h->mul(simd, tpow, tpow, temp); break;
case 9: h->mul(simd, temp2, tpow, a3); break;
case 10: h->add(simd, at_accum, temp2, at_accum); break;
case 11: h->mul(simd, tpow, tpow, temp); break;
case 12: h->mul(simd, temp2, tpow, a4); break;
case 13: h->add(simd, at_accum, temp2, at_accum); break;
case 14: h->mul(simd, tpow, tpow, temp); break;
case 15: h->mul(simd, temp2, tpow, a5); break;
case 16: h->add(simd, at_accum, temp2, at_accum); break;
case 17: h->mul(simd, temp, r, r); break;
case 18: h->mul(simd, temp, temp, -log2e * 0.5f); break;
case 19: h->exp(simd, temp, temp); break;
case 20: h->mul(simd, temp, temp, at_accum); break;
case 21: h->mul(simd, temp, temp, r); break;
case 22: h->mul(simd, temp, temp, 0.5f); break;
case 23: h->add(simd, temp2, r, -temp); break;
case 24: h->csel(simd | le | f0[0], r, temp, temp2, r); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::hardsigmoid_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off) {
switch (phase) {
case 0: h->mul(simd, r, r, alpha_); break;
case 1: h->add(simd, r, r, beta_); break;
case 2: h->min_(simd, r, r, 1.f); break;
case 3: h->max_(simd, r, r, 0.f); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::hardswish_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off) {
auto temp = scratch_[off].f();
switch (phase) {
case 0: h->mul(simd, temp, r, alpha_); break;
case 1: h->add(simd, temp, temp, beta_); break;
case 2: h->min_(simd, temp, temp, 1.f); break;
case 3: h->max_(simd, temp, temp, 0.f); break;
case 4: h->mul(simd, r, r, temp); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::log_compute_fwd(
int simd, const ngen::GRF &r, int phase) {
const float reciproc_log2e = 1.f / 1.442695f; switch (phase) {
case 0: h->log(simd, r, r); break;
case 1: h->mul(simd, r, r, reciproc_log2e); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::mish_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off, int batch) {
auto temp = scratch_[off + batch].f();
auto temp2 = scratch_[off + 2 * batch].f();
const int srelu_phases = phase_count(alg_kind::eltwise_soft_relu);
const int tanh_phases = phase_count(alg_kind::eltwise_tanh);
if (phase < srelu_phases)
soft_relu_compute_fwd_inner(simd, r, temp, temp2, phase, off, 1.f);
if (phase >= srelu_phases && phase < srelu_phases + tanh_phases) {
if (use_tanh_compat())
tanh_compute_fwd_compat(
simd, temp2, phase - srelu_phases, off, batch);
else
tanh_compute_fwd(simd, temp2, phase - srelu_phases, off, batch);
}
if (phase == srelu_phases + tanh_phases) h->mul(simd, r, r, temp2);
if (phase > srelu_phases + tanh_phases) assert(!"invalid phase");
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::pow_compute_fwd(
int simd, const ngen::GRF &r, int phase, int off) {
auto temp = scratch_[off].f();
switch (phase) {
case 0:
if (float((long long int)beta_) == beta_) {
h->mov(simd, temp, abs(r));
} else {
h->mov(simd, temp, r);
}
break;
case 1: h->log(simd, temp, temp); break;
case 2: h->mul(simd, temp, temp, beta_); break;
case 3: h->exp(simd, temp, temp); break;
case 4:
if (((long long int)beta_) & 0x1)
h->csel(simd | lt | f0[0], temp, -temp, temp, r);
break;
case 5: h->mul(simd, r, temp, alpha_); break;
default: assert(!"invalid phase");
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::compute(const int *grfs,
int ngrf, const int seed, const int off, const ngen::DataType dt) {
using namespace alg_kind;
auto bmax = max_batch_size();
auto phases = phase_count(alg_);
for (int idx0 = 0; idx0 < ngrf; idx0 += bmax) {
auto batch = nstl::min(ngrf - idx0, bmax);
for (int phase = 0; phase < phases; phase++) {
for (int ii = 0, nreg = 0; ii < batch; ii += nreg) {
auto grf0 = grfs[idx0 + ii];
auto base = GRF(grf0).f();
nreg = 1;
if (ii + 1 < batch)
if (grf0 + 1 == grfs[idx0 + ii + 1]) nreg = 2;
int simd = nreg * GRF::bytes(hw()) / sizeof(float);
auto grf0_t = grfs[idx0 + (ii / 2)];
auto base_t = GRF(grf0_t).f();
auto grf1 = grfs[idx0 + off + (ii / 2)];
if (is_fwd_) {
switch ((int)alg_) {
case eltwise_elu:
case eltwise_elu_use_dst_for_bwd:
elu_compute_fwd(simd, base, phase, ii);
break;
case eltwise_exp:
case eltwise_exp_use_dst_for_bwd:
exp_compute_fwd(simd, base, phase);
break;
case eltwise_gelu_erf:
gelu_erf_compute_fwd(simd, base, phase, ii, batch);
break;
case eltwise_hardsigmoid:
hardsigmoid_compute_fwd(simd, base, phase, ii);
break;
case eltwise_hardswish:
hardswish_compute_fwd(simd, base, phase, ii);
break;
case eltwise_log:
log_compute_fwd(simd, base, phase);
break;
case eltwise_mish:
mish_compute_fwd(simd, base, phase, ii, batch);
break;
case eltwise_pow:
pow_compute_fwd(simd, base, phase, ii);
break;
case eltwise_relu:
case eltwise_relu_use_dst_for_bwd:
if (alpha_ == 0.f)
relu_zero_ns_compute_fwd(simd, base);
else
relu_compute_fwd(simd, base, phase, ii);
break;
case eltwise_abs: abs_compute_fwd(simd, base); break;
case eltwise_soft_relu:
soft_relu_compute_fwd(simd, base, phase, ii);
break;
case eltwise_sqrt:
case eltwise_sqrt_use_dst_for_bwd:
sqrt_compute_fwd(simd, base);
break;
case eltwise_square:
square_compute_fwd(simd, base);
break;
case eltwise_tanh:
case eltwise_tanh_use_dst_for_bwd:
if (use_tanh_compat())
tanh_compute_fwd_compat(
simd, base, phase, ii, batch);
else
tanh_compute_fwd(simd, base, phase, ii, batch);
break;
case eltwise_round:
round_compute_fwd(simd, base);
break;
case eltwise_swish:
swish_compute_fwd(simd, base, phase, ii);
break;
case eltwise_linear:
linear_compute_fwd(simd, base, phase);
break;
case eltwise_clip:
case eltwise_clip_v2:
case eltwise_clip_v2_use_dst_for_bwd:
clip_compute_fwd(simd, base, phase, alpha_, beta_);
break;
case eltwise_gelu_tanh:
gelu_tanh_compute_fwd(simd, base, phase, ii);
break;
case eltwise_logistic:
case eltwise_logistic_use_dst_for_bwd:
logistic_compute_fwd(simd, base, phase);
break;
case eltwise_stochastic_round:
sround_compute_fwd(simd, base, phase,
GRF(seed).ud(off), dt, ii);
break;
case eltwise_mx_scale:
mx_scale_compute_fwd(simd, base_t, grf1,
GRF(seed).ud(0), nreg, dt, ii);
break;
default: assert(!"unsupported eltwise algorithm");
}
} else {
switch (alg_) {
case eltwise_relu: relu_compute_bwd(simd, base); break;
case eltwise_abs:
abs_compute_bwd(simd, base, phase);
break;
case eltwise_square:
square_compute_bwd(simd, base);
break;
case eltwise_linear:
linear_compute_bwd(simd, base);
break;
case eltwise_clip:
clip_compute_bwd(simd, base, phase, alpha_, beta_);
break;
case eltwise_gelu_tanh:
gelu_tanh_compute_bwd(simd, base, phase, ii, batch);
break;
default: assert(!"unsupported eltwise algorithm");
}
}
if (phase == phases - 1 && scale_ != 1.f) {
h->mul(simd, base, base, scale_);
}
}
}
}
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::compute(
const ngen::GRFRange ®s, const int seed, const int off,
const ngen::DataType dt) {
int grfs[ngen::GRF::maxRegs()];
for (int i = 0; i < regs.getLen(); i++)
grfs[i] = regs.getBase() + i;
compute(grfs, regs.getLen(), seed, off, dt);
}
template <typename ngen_generator_t>
void eltwise_injector_f32_t<ngen_generator_t>::prepare() {
using namespace alg_kind;
assert(scratch_.getLen() >= min_scratch_regs());
if (is_fwd_) {
switch (alg_) {
case eltwise_relu:
case eltwise_relu_use_dst_for_bwd:
if (alpha_ == 0.f) relu_zero_ns_prepare_fwd();
break;
case eltwise_mish:
case eltwise_tanh:
if (use_tanh_compat())
tanh_prepare_fwd_compat();
else
tanh_prepare_fwd();
break;
default: break;
}
} else {
switch (alg_) {
case eltwise_relu: relu_prepare_bwd(); break;
case eltwise_abs: abs_prepare_bwd(); break;
case eltwise_clip: clip_prepare_bwd(); break;
default: break;
}
}
}
template <ngen::HW hw>
using code_gen = ngen::BinaryCodeGenerator<hw>;
REG_XELP_ISA(template struct eltwise_injector_f32_t<code_gen<gpu_xe_lp>>);
REG_XEHP_ISA(template struct eltwise_injector_f32_t<code_gen<gpu_xe_hp>>);
REG_XEHPG_ISA(template struct eltwise_injector_f32_t<code_gen<gpu_xe_hpg>>);
REG_XEHPC_ISA(template struct eltwise_injector_f32_t<code_gen<gpu_xe_hpc>>);
REG_XE2_ISA(template struct eltwise_injector_f32_t<code_gen<gpu_xe2>>);
REG_XE3_ISA(template struct eltwise_injector_f32_t<code_gen<gpu_xe3>>);
REG_XE3P_ISA(template struct eltwise_injector_f32_t<code_gen<gpu_xe3p_35_10>>);
REG_XE3P_ISA(template struct eltwise_injector_f32_t<code_gen<gpu_xe3p_35_11>>);
REG_XE3P_ISA(
template struct eltwise_injector_f32_t<code_gen<gpu_xe3p_35_unknown>>);
#ifdef NGEN_ASM
template struct eltwise_injector_f32_t<ngen::AsmCodeGenerator>;
#endif
} } } } }