#ifndef GPU_INTEL_CONV_JIT_ZERO_OUT_HPP
#define GPU_INTEL_CONV_JIT_ZERO_OUT_HPP
#include "gpu/intel/jit/codegen/kernel_ext.hpp"
#include "gpu/intel/jit/ir/kernel_desc.hpp"
#include "gpu/intel/jit/ir/kernel_info.hpp"
#include "ngen.hpp"
namespace dnnl {
namespace impl {
namespace gpu {
namespace intel {
namespace conv {
namespace jit {
using namespace intel::jit;
class zero_out_kernel_desc_t : public kernel_desc_base_t {
public:
static const size_t bytes_per_thr;
zero_out_kernel_desc_t() = default;
zero_out_kernel_desc_t(int regs, int simd, bool dpas)
: regs_(regs), simd_(simd), dpas_(dpas) {}
std::string kernel_name() const override;
dsl::kernel::options_t options(const impl::engine_t *engine) const override;
compute::range_t local_range() const override;
void init_kernel_iface(dsl::kernel::iface_t &kernel_iface) const override;
void init_kernel_info(kernel_info_t &kernel_info,
const kernel_params_base_t ¶ms,
const impl::engine_t *engine) const override;
status_t create_kernel(compute::kernel_t &kernel, primitive_t *primitive,
impl::engine_t *engine) const override;
status_t create_generator(
const intel::engine_t &engine, compute::kernel_t &kernel) const;
serialization_stream_t serialize() const override;
static zero_out_kernel_desc_t deserialize(const serialization_stream_t &s);
#if __cplusplus >= 202002L
bool operator==(const zero_out_kernel_desc_t &) const = default;
#endif
static compute::nd_range_t nd_range(int simd, size_t size);
private:
int regs_ = 0;
int simd_ = 0;
bool dpas_ = false;
};
class zero_out_kernel_params_t : public kernel_params_base_t {
public:
zero_out_kernel_params_t() = default;
zero_out_kernel_params_t(size_t size) : size(size) {}
size_t size = 0;
};
template <ngen::HW hw>
class zero_out_kernel_t : public ir::ir_to_ngen_generator_t<generator_t<hw>> {
public:
IR_TO_NGEN_GENERATOR_FORWARD(generator_t<hw>)
using base_type = ir::ir_to_ngen_generator_t<generator_t<hw>>;
zero_out_kernel_t(const dsl::kernel::options_t &options,
const kernel_info_t &kernel_info, const impl::engine_t *engine)
: zero_out_kernel_t(zero_out_kernel_desc_t(options.regs(),
options.simd(), options.require_dpas()),
engine) {}
zero_out_kernel_t(
const kernel_desc_base_t &_desc, const impl::engine_t *engine)
: base_type(get_kernel_iface(_desc), _desc.options(engine),
debug_config_t {GENERATOR_NAME, GENERATOR_LINE}) {
requireLocalID(3);
requireLocalSize();
requireGRF(options().regs());
requireSIMD(options().simd());
requireBarrier();
externalName(_desc.kernel_name());
newArgument(kernel_iface()[0].template as<var_t>().name,
ir::to_ngen(kernel_iface()[0].type()));
newArgument(kernel_iface()[1].template as<var_t>().name,
ngen::ExternalArgumentType::GlobalPtr,
ngen::GlobalAccessType::Stateless);
finalizeInterface();
generate_prologue();
ra().claim(getLocalSize(0));
ra().claim(getLocalID(0));
int simd_size = getSIMD();
bool use_lsc = (hw >= ngen::HW::XeHPG);
auto size = getArgument(kernel_iface()[0].template as<var_t>().name);
ra().claim(size);
auto ptr = getArgument(kernel_iface()[1].template as<var_t>().name);
ra().claim(ptr);
auto global_id = ra().template alloc_sub<uint32_t>();
auto off0 = ra().template alloc_sub<uint32_t>();
const int bytes_per_thr
= into<int>(zero_out_kernel_desc_t::bytes_per_thr);
if (base_type::emu_strategy_.emulate64) {
base_type::emu_state_.temp[0] = ra().alloc();
base_type::emu_state_.temp[1] = ra().alloc();
}
mul(1, global_id, r0.ud(1), getLocalSize(0).uw());
add(1, global_id, global_id, getLocalID(0));
shl(1, off0, global_id, math::ilog2q(bytes_per_thr / simd_size));
int grf_size = ngen::GRF::bytes(hw);
int bytes_per_store = 16;
int uw_size = sizeof(uint16_t);
int ud_size = sizeof(uint32_t);
int uq_size = sizeof(uint64_t);
auto zero = ra().alloc_range(bytes_per_store * ud_size / grf_size);
auto off_vec = ra().alloc_range(bytes_per_thr * ud_size / grf_size);
auto off_vec_q_strided
= ra().alloc_range(bytes_per_thr * uq_size / grf_size);
auto ptr_vec = ra().alloc_range(bytes_per_thr * uq_size / grf_size);
for (int i = 0; i < bytes_per_store * ud_size; i += 64) {
auto z = ir::get_subregister(hw, ngen::DataType::ud, zero, i);
mov(16, z, 0);
}
auto idx_vec = ra().alloc().uw(0);
mov(8, idx_vec(1), ngen::Immediate::uv(0, 1, 2, 3, 4, 5, 6, 7));
mov(8, idx_vec(2), idx_vec(1));
for (int i = 16; i < grf_size / uw_size; i += 16) {
mov(8, idx_vec.uw(i)(2), idx_vec(2));
}
ir::reg_buf_t dst, src0, src1;
for (int i = 0; i < bytes_per_thr; i += 8) {
auto off_sub_vec = ir::get_subregister(
hw, ngen::DataType::ud, off_vec, i)(1);
this->eadd3(8,
ir::ngen_operand_t(ir::reg_buf_data_t(dst, off_sub_vec)),
ir::ngen_operand_t(ir::reg_buf_data_t(
src1, idx_vec.uw((i % grf_size) * 2)(2))),
ir::ngen_operand_t(ir::reg_buf_data_t(src0, off0)),
ir::ngen_operand_t(i));
auto ptr_sub_vec = ir::get_subregister(
hw, ngen::DataType::uq, ptr_vec, i)(1);
auto off_sub_vec_q_strided = ir::get_subregister(
hw, ngen::DataType::ud, off_vec_q_strided, i * 2)(2);
emov(8, off_sub_vec_q_strided, off_sub_vec);
eadd(8, ptr_sub_vec, ptr, off_sub_vec_q_strided);
}
for (int i = 0; i < bytes_per_thr; i += bytes_per_store) {
auto off_sub_vec = ir::get_subregister(
hw, ngen::DataType::ud, off_vec, i)(1);
cmp(16 | lt | f0[0], off_sub_vec, size);
auto h = ir::get_subregister(hw, ngen::DataType::uq, ptr_vec, i);
if (use_lsc) {
std::unique_ptr<ngen::DataSpecLSC> lsc_spec;
lsc_spec = utils::make_unique<ngen::DataSpecLSC>(
ngen::scattered(ngen::DataSizeLSC::D8U32, 1));
store.ugm(16 | f0[0], *lsc_spec, A64, h, zero[0]);
} else {
store(16 | f0[0], ngen::scattered_byte(), A64, h, zero[0]);
}
}
generate_epilogue();
}
private:
static dsl::kernel::iface_t get_kernel_iface(
const kernel_desc_base_t &desc) {
dsl::kernel::iface_t iface(desc.kernel_name());
desc.init_kernel_iface(iface);
return iface;
}
};
} } } } } }
#endif