diff --git a/src/plugins/intel_cpu/src/nodes/mkldnn_topk_node.cpp b/src/plugins/intel_cpu/src/nodes/mkldnn_topk_node.cpp index d368a5bc3f1..09ee24531a3 100644 --- a/src/plugins/intel_cpu/src/nodes/mkldnn_topk_node.cpp +++ b/src/plugins/intel_cpu/src/nodes/mkldnn_topk_node.cpp @@ -2,21 +2,1787 @@ // SPDX-License-Identifier: Apache-2.0 // -#include +#include "mkldnn_topk_node.h" + +#include + +#include +#include +#include +#include +#include "emitters/jit_load_store_emitters.hpp" +#include "ie_parallel.hpp" +#include +#include +#include + +#include +#include +#include "common/cpu_memcpy.h" #include -#include -#include -#include "ie_parallel.hpp" -#include "mkldnn_topk_node.h" -#include "utils/general_utils.h" - -#if defined(HAVE_SSE) || defined(HAVE_AVX2) || defined(HAVE_AVX512F) -#include -#endif +using namespace mkldnn; using namespace MKLDNNPlugin; using namespace InferenceEngine; +using namespace mkldnn::impl; +using namespace mkldnn::impl::cpu::x64; +using namespace mkldnn::impl::utils; +using namespace Xbyak; + +#define GET_OFF(field) offsetof(jit_topk_call_args, field) + +#define vmm_mask Vmm(0) +#define vmm_tmp Vmm(1) +#define vmm_val(i) Vmm(2 * (i) + 2) +#define vmm_idx(i) Vmm(2 * (i) + 3) +#define vmm_val_l Vmm(2) +#define vmm_idx_l Vmm(3) +#define vmm_val_r Vmm(4) +#define vmm_idx_r Vmm(5) + +#define xmm_mask Xmm(0) +#define xmm_tmp Xmm(1) +#define xmm_val(i) Xmm(2 * (i) + 2) +#define xmm_idx(i) Xmm(2 * (i) + 3) +#define xmm_val_l Xmm(2) +#define xmm_idx_l Xmm(3) +#define xmm_val_r Xmm(4) +#define xmm_idx_r Xmm(5) + +#define xmm_val_p Xmm(6) +#define xmm_idx_p Xmm(7) + +#define JMP_TO_LABEL(label) \ + if (isa == cpu::x64::avx512_common) { \ + kmovw(reg_tmp_32, k_mask); \ + } else { \ + uni_vmovmskps(reg_tmp_32, xmm_mask); \ + } \ + and_(reg_tmp_32, 0x1); \ + cmp(reg_tmp_32, 0x0); \ + je(label, T_NEAR); + +static inline bool isFloatCompatible(memory::data_type type) { + return memory::data_type::f32 == type || memory::data_type::bf16 == type; +} + +template +struct jit_uni_topk_kernel_f32 : public jit_uni_topk_kernel, public jit_generator { + DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_uni_topk_kernel_f32) + + explicit jit_uni_topk_kernel_f32(jit_topk_config_params jcp) + : jit_uni_topk_kernel(jcp), jit_generator() {} + + void create_ker() override { + jit_generator::create_kernel(); + ker_ = (decltype(ker_))jit_ker(); + } + + void generate() override { + load_emitter.reset(new jit_load_emitter(this, isa)); + store_emitter.reset(new jit_store_emitter(this, isa)); + + this->preamble(); + + mov(reg_src, ptr[reg_params + GET_OFF(src)]); + mov(reg_dst, ptr[reg_params + GET_OFF(dst)]); + mov(reg_dst_idx, ptr[reg_params + GET_OFF(index)]); + mov(reg_prc, ptr[reg_params + GET_OFF(process)]); + mov(reg_prc_idx, ptr[reg_params + GET_OFF(process_index)]); + mov(reg_work_amount, ptr[reg_params + GET_OFF(work_amount)]); + + bool shape_agnostic_alg = jcp_.algorithm == TopKAlgorithm::topk_heap_sort || + (jcp_.algorithm == TopKAlgorithm::topk_bubble_sort && !jcp_.bubble_inplace); + + if (!shape_agnostic_alg) + mov(reg_table, l_table); + + data_type = MKLDNNExtensionUtils::IEPrecisionToDataType(jcp_.precision); + if (jcp_.layout == TopKLayoutType::topk_blocked && jcp_.topk_innermost) + blk_stride = jcp_.sort_stride * jcp_.blk_size; + + if (jcp_.mode_max) { + cmp_flg = _cmp_lt_os; // if val[left] < val[right], set mask 1, swap + heap_cmp_flg = _cmp_nle_us; // min heap is used for max topk, if a > b, set mask 1, swap + } else { + cmp_flg = _cmp_nle_us; // if val[left] > val[right], set mask 1, swap + heap_cmp_flg = _cmp_lt_os; // max heap is used for min topk, if a < b, set mask 1, swap + } + + if (isa == cpu::x64::avx512_common) + uni_vpxor(vmm_zero, vmm_zero, vmm_zero); + + load_pool_gpr_idxs = {static_cast(reg_load_store_mask.getIdx()), static_cast(reg_load_table.getIdx())}; + store_pool_gpr_idxs = {static_cast(reg_load_store_mask.getIdx())}; + store_pool_vec_idxs = {static_cast(vmm_zero.getIdx())}; + + topk_loop(); + + this->postamble(); + + load_emitter->emit_data(); + store_emitter->emit_data(); + + if (!shape_agnostic_alg) + prepare_idx_table(); + } + +private: + using Vmm = typename conditional3::type; + size_t vlen = cpu_isa_traits::vlen; + mkldnn::memory::data_type data_type; + + Xbyak::Address table_val(int index) { return ptr[reg_table + index * vlen]; } + Xbyak::Address table_bubble_block_idx(int index) { return ptr[reg_bubble_block_idx + index * vlen]; } + Xbyak::Address table_bubble_seq_idx(int index) { return ptr[reg_bubble_seq_idx + index * sizeof(int)]; } + Xbyak::Address table_heap_seq_idx(int index) { return ptr[reg_heap_seq_idx + index * sizeof(int)]; } + + Xbyak::Reg64 reg_src = r8; + Xbyak::Reg64 reg_dst = r9; + Xbyak::Reg64 reg_dst_idx = r10; + Xbyak::Reg64 reg_prc = r11; + Xbyak::Reg64 reg_prc_idx = r12; + Xbyak::Reg64 reg_work_amount = r13; + Xbyak::Reg64 reg_table = r14; + Xbyak::Reg64 reg_params = abi_param1; + Xbyak::Reg64 reg_i = rax; + Xbyak::Reg64 reg_aux = rdx; + Xbyak::Reg64 reg_aux_idx = rbx; + + Xbyak::Reg8 reg_tmp_8 = r15b; + Xbyak::Reg32 reg_tmp_32 = r15d; + Xbyak::Reg64 reg_tmp_64 = r15; + + Xbyak::Reg64 reg_load_table = rbp; + Xbyak::Reg64 reg_load_store_mask = rsi; + + // ================================================ for shape_agnostic_alg ================================================ + // *** for both heap sort and bubble sort *** + Xbyak::Reg64 reg_tmp = reg_aux_idx; + + // *** for heap sort only *** + Xbyak::Reg64 reg_j = reg_i; // save reg_i by rsp before using reg_j + Xbyak::Reg64 reg_offset = reg_load_table; // reuse reg_load_table after finish using load/store_emiter + Xbyak::Reg64 reg_offset_idx = reg_load_store_mask; // reuse reg_load_store_mask after finish using load/store_emiter + Xbyak::Reg64 reg_heap_seq_idx = reg_table; + Xbyak::Reg64 reg_heap_axis_dim = reg_work_amount; + Xbyak::Reg64 reg_heap_top_k = reg_prc; // save reg_top_k by rsp before using reg_prc + Xbyak::Reg64 reg_heap_k_sub_step = reg_heap_top_k; + Xbyak::Reg64 reg_zero = reg_offset; // save reg_zero by rsp before using reg_offset, also refer to reg_offset + Xbyak::Reg64 reg_end = reg_prc_idx; // save reg_heap_outer_aux by rsp before using reg_prc_idx + Xbyak::Reg64 reg_heap_outer_aux = reg_prc_idx; + Xbyak::Reg64 reg_i_sub_1 = reg_i; // denotes i-1 + Xbyak::Reg64 reg_heap_k_sub_1 = reg_heap_top_k; // denotes k-1 + Xbyak::Reg64 reg_heapify_end = reg_heap_axis_dim; // save reg_heap_axis_dim by rsp before using reg_inner_end + Xbyak::Reg64 reg_heapify_i = reg_src; // save reg_src by rsp before using reg_heapify_i + Xbyak::Reg64 reg_heapify_valid = reg_heap_seq_idx; // save reg_heap_seq_idx by rsp before using reg_heapify_valid + Xbyak::Reg64 reg_heapify_tmp = reg_params; // save reg_params by rsp before using reg_heapify_tmp + + // *** for bubble sort only *** + Xbyak::Reg64 reg_bubble_seq_idx = reg_table; + Xbyak::Reg64 reg_bubble_block_idx = reg_prc; + Xbyak::Reg64 reg_bubble_axis_dim = reg_prc_idx; + Xbyak::Reg64 reg_block_l = reg_bubble_block_idx; // save reg_bubble_block_idx by rsp before using reg_l + Xbyak::Reg64 reg_block_r = reg_bubble_axis_dim; // save reg_bubble_axis_dim by rsp before using reg_r + Xbyak::Reg64 reg_seq_l = reg_load_table; // blocked layout on channel + Xbyak::Reg64 reg_seq_r = reg_prc; // blocked layout on channel + Xbyak::Reg64 reg_offset_l = reg_i; // save reg_i by rsp before using reg_offset_l + Xbyak::Reg64 reg_offset_r = reg_prc_idx; // save reg_prc_idx by rsp before using reg_offset_r + Xbyak::Reg64 reg_bubble_block_top_k = reg_bubble_seq_idx; + Xbyak::Reg64 reg_bubble_block_k_sub_1 = reg_bubble_block_top_k; + Xbyak::Reg64 reg_bubble_seq_top_k = reg_load_store_mask; + Xbyak::Reg64 reg_bubble_seq_k_sub_1 = reg_bubble_seq_top_k; + Xbyak::Reg64 reg_block_sort_stride = reg_aux; // by vector + Xbyak::Reg64 reg_block_sort_stride_byte = reg_block_sort_stride; // by vector + Xbyak::Reg64 reg_seq_tmp = reg_seq_l; // blocked layout on channel + Xbyak::Reg64 reg_seq_sort_stride = reg_work_amount;// blocked layout on channel + Xbyak::Reg64 reg_blk_stride = reg_seq_sort_stride; // blocked layout on channel, denotes reg_seq_sort_stride * jcp_.blk_size + Xbyak::Reg64 reg_sub_idx = reg_bubble_block_idx; // blocked layout on channel + // ======================================================================================================================== + + Vmm vmm_zero = Vmm(0); // vmm_zero represents Vmm(0) when isa is avx512_common, otherwise vmm_mask represents Vmm(0) + + const Xbyak::Opmask k_mask = Xbyak::Opmask(1); + const int step = vlen / sizeof(float); + const int tail = jcp_.work_amount % step; + const int topk_tail = jcp_.top_k % step; + + int blk_stride; // stride of channel blocks at the same space coordinate, only used in blocked layout with topk on channel + unsigned char cmp_flg; + unsigned char heap_cmp_flg; + + Xbyak::Label l_table; + + std::unique_ptr load_emitter = nullptr; + std::unique_ptr store_emitter = nullptr; + + std::vector store_pool_gpr_idxs; + std::vector load_pool_gpr_idxs; + std::vector store_pool_vec_idxs; + + inline void topk_loop() { + if (jcp_.algorithm == TopKAlgorithm::topk_bubble_sort) { + if (jcp_.layout == TopKLayoutType::topk_blocked && jcp_.topk_innermost) { + if (jcp_.top_k == 1) { + topk_bubble_BLK_on_channel_horiz(); + } else { + topk_bubble_BLK_on_channel_verti(); + } + } else { + topk_bubble_vector(); + } + } else if (jcp_.algorithm == TopKAlgorithm::topk_bitonic_sort) { + if (jcp_.layout == TopKLayoutType::topk_blocked && jcp_.topk_innermost) { + topk_bitonic_BLK_on_channel(); + } else { + topk_bitonic_vector(); + } + } else if (jcp_.algorithm == TopKAlgorithm::topk_heap_sort) { + topk_heap_sorting(); + } + } + + inline void topk_bitonic_vector() { + Xbyak::Label topk_main_loop_label; + Xbyak::Label topk_main_loop_end_label; + L(topk_main_loop_label); + { + cmp(reg_work_amount, step); + jl(topk_main_loop_end_label, T_NEAR); + + topk_bitonic(step); + + add(reg_src, step * jcp_.data_size); + add(reg_dst, step * jcp_.data_size); + add(reg_dst_idx, step * sizeof(int)); + sub(reg_work_amount, step); + + jmp(topk_main_loop_label, T_NEAR); + } + L(topk_main_loop_end_label); + + // tail + if (tail) { + Xbyak::Label topk_tail_loop_end_label; + cmp(reg_work_amount, tail); + jl(topk_tail_loop_end_label, T_NEAR); + + topk_bitonic(tail); + + L(topk_tail_loop_end_label); + } + } + + inline void topk_bitonic(int elt_num) { + // src => prc + for (int i = 0; i < jcp_.axis_dim; i++) { + load_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_tmp.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num, i * jcp_.sort_stride * jcp_.data_size), + {}, {load_pool_gpr_idxs}); + store_emitter->emit_code({static_cast(vmm_tmp.getIdx())}, {static_cast(reg_prc.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, elt_num, i * jcp_.sort_stride * jcp_.data_size), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + + load_emitter->emit_code({static_cast(reg_table.getIdx())}, {static_cast(vmm_tmp.getIdx())}, + std::make_shared(Precision::I32, Precision::I32, elt_num, i * vlen), + {}, {load_pool_gpr_idxs}); + store_emitter->emit_code({static_cast(vmm_tmp.getIdx())}, {static_cast(reg_prc_idx.getIdx())}, + std::make_shared(Precision::I32, Precision::I32, elt_num, i * jcp_.sort_stride * sizeof(int)), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + } + + // sort + bitonic_sort_vector(elt_num); + if (jcp_.sort_index) { + bitonic_sort_vector(elt_num, false); + } + + // prc => dst + for (int i = 0; i < jcp_.top_k; i++) { + load_emitter->emit_code({static_cast(reg_prc.getIdx())}, {static_cast(vmm_tmp.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num, i * jcp_.sort_stride * jcp_.data_size), + {}, {load_pool_gpr_idxs}); + store_emitter->emit_code({static_cast(vmm_tmp.getIdx())}, {static_cast(reg_dst.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, elt_num, i * jcp_.sort_stride * jcp_.data_size), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + + load_emitter->emit_code({static_cast(reg_prc_idx.getIdx())}, {static_cast(vmm_tmp.getIdx())}, + std::make_shared(Precision::I32, Precision::I32, elt_num, i * jcp_.sort_stride * sizeof(int)), + {}, {load_pool_gpr_idxs}); + store_emitter->emit_code({static_cast(vmm_tmp.getIdx())}, {static_cast(reg_dst_idx.getIdx())}, + std::make_shared(Precision::I32, Precision::I32, elt_num, i * jcp_.sort_stride * sizeof(int)), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + } + } + + // src memory layout: (N) * (CB * H * W * blk_size) + // prc memory layout: (C) * (N * H * W) + // topk_bitonic_BLK_on_channel: sort (C) * (N * H * W / blk_size * blk_size) elements + // sort (C) * (N * H * W % blk_size) elements in the rear + inline void topk_bitonic_BLK_on_channel() { + Xbyak::Label topk_main_loop_label; + Xbyak::Label topk_main_loop_end_label; + L(topk_main_loop_label); + { + cmp(reg_work_amount, step); + jl(topk_main_loop_end_label, T_NEAR); + + // src => prc + bitonic_BLK_on_channel_load(step); + + // sort + bitonic_sort_vector(step); + if (jcp_.sort_index) { + bitonic_sort_vector(step, false); + } + + // prc => dst + bitonic_BLK_on_channel_store(step); + + add(reg_src, step * jcp_.blk_size * jcp_.data_size); + add(reg_dst, step * jcp_.blk_size * jcp_.data_size); + add(reg_dst_idx, step * jcp_.blk_size * sizeof(int)); + sub(reg_work_amount, step); + + jmp(topk_main_loop_label, T_NEAR); + } + L(topk_main_loop_end_label); + + // tail exists because working buffer has planar layout, though source buffer has blocked layout) + if (tail) { + Xbyak::Label topk_tail_loop_end_label; + cmp(reg_work_amount, tail); + jl(topk_tail_loop_end_label, T_NEAR); + + // src => prc + bitonic_BLK_on_channel_load(tail); + + bitonic_sort_vector(tail); + if (jcp_.sort_index) { + bitonic_sort_vector(tail, false); + } + + // prc => dst + bitonic_BLK_on_channel_store(tail); + + L(topk_tail_loop_end_label); + } + } + + inline void bitonic_sort_vector(int elt_num, bool cmp_val = true) { + if (cmp_val) { + mov(reg_i, jcp_.bitonic_idx_cnt); + mov(reg_aux, ptr[reg_params + GET_OFF(bitonic_idx_buf)]); + } else { + mov(reg_i, jcp_.bitonic_k_idx_cnt); + mov(reg_aux, ptr[reg_params + GET_OFF(bitonic_k_idx_buf)]); + } + + Xbyak::Label topk_main_loop_label; + Xbyak::Label topk_main_loop_end_label; + L(topk_main_loop_label); + { + cmp(reg_i, 0); + je(topk_main_loop_end_label, T_NEAR); + + bitonic_swap_vector(elt_num, cmp_val); + + add(reg_aux, 2 * sizeof(int)); + sub(reg_i, 2); + + jmp(topk_main_loop_label, T_NEAR); + } + L(topk_main_loop_end_label); + } + + inline void bitonic_BLK_on_channel_load(int elt_num) { + for (int i = 0; i < jcp_.axis_dim; i++) { + for (int j = 0; j < elt_num; j++) { + int offset = i / jcp_.blk_size * blk_stride + i % jcp_.blk_size + j * jcp_.blk_size; + + load_scalar(xmm_tmp, ptr[reg_src + offset * jcp_.data_size], data_type); + store_scalar(ptr[reg_prc + (i * jcp_.sort_stride + j) * jcp_.data_size], xmm_tmp, data_type); + + uni_vmovdqu(xmm_tmp, table_val(i)); + store_scalar(ptr[reg_prc_idx + (i * jcp_.sort_stride + j) * sizeof(int)], xmm_tmp, memory::data_type::s32, false); + } + } + } + + inline void bitonic_BLK_on_channel_store(int elt_num) { + for (int i = 0; i < jcp_.top_k; i++) { + for (int j = 0; j < elt_num; j++) { + int offset = i / jcp_.blk_size * blk_stride + i % jcp_.blk_size + j * jcp_.blk_size; + + load_scalar(xmm_tmp, ptr[reg_prc + (i * jcp_.sort_stride + j) * jcp_.data_size], data_type); + store_scalar(ptr[reg_dst + offset * jcp_.data_size], xmm_tmp, data_type); + + load_scalar(xmm_tmp, ptr[reg_prc_idx + (i * jcp_.sort_stride + j) * sizeof(int)], memory::data_type::s32); + store_scalar(ptr[reg_dst_idx + offset * sizeof(int)], xmm_tmp, memory::data_type::s32); + } + } + } + + inline void bitonic_get_addr(Xbyak::Reg64 reg_base, int data_size, int offset = 0) { + mov(reg_aux_idx.cvt32(), ptr[reg_aux + offset]); + mul_by_const(reg_aux_idx, reg_tmp_64, data_size); + add(reg_aux_idx, reg_base); + } + + inline void bitonic_swap_vector(int elt_num, bool cmp_val = true) { + bitonic_get_addr(reg_prc, jcp_.data_size, 0); + load_emitter->emit_code({static_cast(reg_aux_idx.getIdx())}, {static_cast(vmm_val_l.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + bitonic_get_addr(reg_prc, jcp_.data_size, sizeof(int)); + load_emitter->emit_code({static_cast(reg_aux_idx.getIdx())}, {static_cast(vmm_val_r.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + bitonic_get_addr(reg_prc_idx, sizeof(int), 0); + load_emitter->emit_code({static_cast(reg_aux_idx.getIdx())}, {static_cast(vmm_idx_l.getIdx())}, + std::make_shared(Precision::I32, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + bitonic_get_addr(reg_prc_idx, sizeof(int), sizeof(int)); + load_emitter->emit_code({static_cast(reg_aux_idx.getIdx())}, {static_cast(vmm_idx_r.getIdx())}, + std::make_shared(Precision::I32, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + + swap_vector(vmm_val_l, vmm_idx_l, vmm_val_r, vmm_idx_r, cmp_val); + + bitonic_get_addr(reg_prc, jcp_.data_size, 0); + store_emitter->emit_code({static_cast(vmm_val_l.getIdx())}, {static_cast(reg_aux_idx.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + bitonic_get_addr(reg_prc, jcp_.data_size, sizeof(int)); + store_emitter->emit_code({static_cast(vmm_val_r.getIdx())}, {static_cast(reg_aux_idx.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + bitonic_get_addr(reg_prc_idx, sizeof(int), 0); + store_emitter->emit_code({static_cast(vmm_idx_l.getIdx())}, {static_cast(reg_aux_idx.getIdx())}, + std::make_shared(Precision::FP32, Precision::I32, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + bitonic_get_addr(reg_prc_idx, sizeof(int), sizeof(int)); + store_emitter->emit_code({static_cast(vmm_idx_r.getIdx())}, {static_cast(reg_aux_idx.getIdx())}, + std::make_shared(Precision::FP32, Precision::I32, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + } + + inline void topk_heap_sorting() { + mov(reg_heap_seq_idx, ptr[reg_params + GET_OFF(idx_seq_buf)]); + mov(reg_heap_axis_dim, ptr[reg_params + GET_OFF(axis_dim)]); + mov(reg_heap_top_k, ptr[reg_params + GET_OFF(top_k)]); + + // init dst + mov(reg_i, 0); + sub(reg_heap_top_k, step); + topk_heap_load(reg_heap_k_sub_step, step); + add(reg_heap_top_k, step); + topk_heap_load(reg_heap_top_k, 1); + mov(reg_zero, 0); + + // heapify + Xbyak::Label topk_heapify_loop_label; + Xbyak::Label topk_heapify_loop_end_label; + reg_sub_shr(reg_end, reg_heap_top_k, 2, 1); + mov(reg_i, reg_end); + sub(reg_heap_top_k, 1); + L(topk_heapify_loop_label); + { + heapify_sub_tree(reg_i, reg_heap_k_sub_1); + + cmp(reg_i, 0); + je(topk_heapify_loop_end_label, T_NEAR); + + sub(reg_i, 1); + jmp(topk_heapify_loop_label, T_NEAR); + } + L(topk_heapify_loop_end_label); + add(reg_heap_top_k, 1); + + // update + Xbyak::Label topk_main_loop_label; + Xbyak::Label topk_main_loop_end_label; + mov(reg_i, reg_heap_top_k); + sub(reg_heap_top_k, 1); + L(topk_main_loop_label); + { + cmp(reg_i, reg_heap_axis_dim); + je(topk_main_loop_end_label, T_NEAR); + + Xbyak::Label topk_update_loop_end_label; + get_addr_by_reg_idx(reg_aux, reg_src, reg_i, jcp_.data_size); + load_scalar(xmm_val_p, ptr[reg_aux], data_type); + + get_addr_by_reg_idx(reg_aux, reg_heap_seq_idx, reg_i, sizeof(int)); + load_scalar(xmm_idx_p, ptr[reg_aux], memory::data_type::s32); + load_scalar(xmm_val_l, ptr[reg_dst], data_type); + load_scalar(xmm_idx_l, ptr[reg_dst_idx], memory::data_type::s32); + + heap_cmp_node(xmm_val_p, xmm_idx_p, xmm_val_l, xmm_idx_l); + JMP_TO_LABEL(topk_update_loop_end_label); + + store_scalar(ptr[reg_dst], xmm_val_p, data_type); + store_scalar(ptr[reg_dst_idx], xmm_idx_p, memory::data_type::s32); + heapify_sub_tree(reg_zero, reg_heap_k_sub_1); + + L(topk_update_loop_end_label); + + add(reg_i, 1); + jmp(topk_main_loop_label, T_NEAR); + } + L(topk_main_loop_end_label); + + // extract topk + if (jcp_.sort_index) { + // reheapify by index + Xbyak::Label topk_reheapify_loop_label; + Xbyak::Label topk_reheapify_loop_end_label; + mov(reg_i, reg_end); + L(topk_reheapify_loop_label); + { + heapify_sub_tree(reg_i, reg_heap_k_sub_1, false); + + cmp(reg_i, 0); + je(topk_reheapify_loop_end_label, T_NEAR); + + sub(reg_i, 1); + jmp(topk_reheapify_loop_label, T_NEAR); + } + L(topk_reheapify_loop_end_label); + + // extract by index + topk_heap_extract(false); + } else { + // extract by value + topk_heap_extract(); + } + } + + inline void topk_heap_load(Xbyak::Reg64 ®_end, int s) { + Xbyak::Label topk_init_loop_label; + Xbyak::Label topk_init_loop_end_label; + L(topk_init_loop_label); + { + if (s == step) { + cmp(reg_i, reg_end); + jg(topk_init_loop_end_label, T_NEAR); + } else { + cmp(reg_i, reg_end); + je(topk_init_loop_end_label, T_NEAR); + } + + get_addr_by_reg_idx(reg_heap_outer_aux, reg_src, reg_i, jcp_.data_size); + load_emitter->emit_code({static_cast(reg_heap_outer_aux.getIdx())}, {static_cast(vmm_tmp.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, s), + {}, {load_pool_gpr_idxs}); + get_addr_by_reg_idx(reg_heap_outer_aux, reg_dst, reg_i, jcp_.data_size); + store_emitter->emit_code({static_cast(vmm_tmp.getIdx())}, {static_cast(reg_heap_outer_aux.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, s), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + if (s == step) { + table_to_vmm(vmm_tmp, reg_heap_seq_idx, reg_i, 0, sizeof(int)); + } else { + get_addr_by_reg_idx(reg_heap_outer_aux, reg_heap_seq_idx, reg_i, sizeof(int)); + load_emitter->emit_code({static_cast(reg_heap_outer_aux.getIdx())}, {static_cast(vmm_tmp.getIdx())}, + std::make_shared(Precision::I32, Precision::I32, 1), + {}, {load_pool_gpr_idxs}); + } + get_addr_by_reg_idx(reg_heap_outer_aux, reg_dst_idx, reg_i, sizeof(int)); + store_emitter->emit_code({static_cast(vmm_tmp.getIdx())}, {static_cast(reg_heap_outer_aux.getIdx())}, + std::make_shared(Precision::I32, Precision::I32, s), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + + add(reg_i, s); + jmp(topk_init_loop_label, T_NEAR); + } + L(topk_init_loop_end_label); + } + + inline void topk_heap_extract(bool cmp_val = true) { + Xbyak::Label topk_extract_label; + Xbyak::Label topk_extract_end_label; + mov(reg_i, reg_heap_k_sub_1); + L(topk_extract_label); + { + cmp(reg_i, 0); + je(topk_extract_end_label, T_NEAR); + + heap_swap_root(reg_i); + sub(reg_i, 1); + heapify_sub_tree(reg_zero, reg_i_sub_1, cmp_val); + + jmp(topk_extract_label, T_NEAR); + } + L(topk_extract_end_label); + } + + inline void heapify_sub_tree(const Xbyak::Reg64 ®_idx, const Xbyak::Reg64 ®_valid, bool cmp_val = true) { + Xbyak::Label topk_heapify_loop_label; + Xbyak::Label topk_heapify_loop_end_label; + Xbyak::Label topk_lchild_loop_label; + Xbyak::Label topk_rchild_loop_label; + + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_i.cvt32()); + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_heap_top_k.cvt32()); + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_zero.cvt32()); + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_end.cvt32()); + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_heap_axis_dim.cvt32()); + sub(rsp, sizeof(int64_t)); + mov(ptr[rsp], reg_src); + sub(rsp, sizeof(int64_t)); + mov(ptr[rsp], reg_heap_seq_idx); + sub(rsp, sizeof(int64_t)); + mov(ptr[rsp], reg_params); + + // reg_idx and reg_valid should not be used + mov(reg_heapify_i, reg_idx); + mov(reg_heapify_valid, reg_valid); + + cmp(reg_heapify_valid, 0); + je(topk_heapify_loop_end_label, T_NEAR); + + reg_sub_shr(reg_heapify_end, reg_heapify_valid, 1, 1); + mov(reg_j, reg_heapify_i); + mov(reg_aux, reg_dst); + mov(reg_aux_idx, reg_dst_idx); + reg_mul_add(reg_aux, reg_heapify_tmp, reg_heapify_i, jcp_.data_size); + reg_mul_add(reg_aux_idx, reg_heapify_tmp, reg_heapify_i, sizeof(int)); + reg_mul_add(reg_offset, reg_heapify_i, 2, 1); + mul_by_const(reg_offset, reg_heapify_tmp, jcp_.data_size); + if (jcp_.data_size != sizeof(int)) { + reg_mul_add(reg_offset_idx, reg_heapify_i, 2, 1); + mul_by_const(reg_offset_idx, reg_heapify_tmp, sizeof(int)); + } + + L(topk_heapify_loop_label); + { + cmp(reg_j, reg_heapify_end); + jg(topk_heapify_loop_end_label, T_NEAR); + + load_scalar(xmm_val_p, ptr[reg_aux], data_type); + load_scalar(xmm_idx_p, ptr[reg_aux_idx], memory::data_type::s32); + + // compare lchild-rchild + mov(reg_prc, reg_dst); + add(reg_prc, reg_offset); + mov(reg_prc_idx, reg_dst_idx); + if (jcp_.data_size != sizeof(int)) + add(reg_prc_idx, reg_offset_idx); + else + add(reg_prc_idx, reg_offset); + load_scalar(xmm_val_l, ptr[reg_prc], data_type); + load_scalar(xmm_idx_l, ptr[reg_prc_idx], memory::data_type::s32); + add(reg_prc, jcp_.data_size); + add(reg_prc_idx, sizeof(int)); + + // if last valid parent has no rchild + mov(reg_heapify_tmp, reg_heapify_valid); + shr(reg_heapify_tmp, 1); + cmp(reg_j, reg_heapify_tmp); + je(topk_lchild_loop_label, T_NEAR); + + load_scalar(xmm_val_r, ptr[reg_prc], data_type); + load_scalar(xmm_idx_r, ptr[reg_prc_idx], memory::data_type::s32); + + heap_cmp_node(xmm_val_l, xmm_idx_l, xmm_val_r, xmm_idx_r, cmp_val); + JMP_TO_LABEL(topk_lchild_loop_label); + + // compare node-rchild + L(topk_rchild_loop_label); + { + heap_cmp_node(xmm_val_p, xmm_idx_p, xmm_val_r, xmm_idx_r, cmp_val); + JMP_TO_LABEL(topk_heapify_loop_end_label); + + heap_swap_node(xmm_val_p, xmm_idx_p, xmm_val_r, xmm_idx_r); + mov(reg_aux, reg_prc); + mov(reg_aux_idx, reg_prc_idx); + add(reg_offset, jcp_.data_size); + shl(reg_offset, 1); + add(reg_offset, jcp_.data_size); + if (jcp_.data_size != sizeof(int)) { + add(reg_offset_idx, sizeof(int)); + shl(reg_offset_idx, 1); + add(reg_offset_idx, sizeof(int)); + } + shl(reg_j, 1); + add(reg_j, 2); + jmp(topk_heapify_loop_label, T_NEAR); + } + + // compare node-lchild + L(topk_lchild_loop_label); + { + heap_cmp_node(xmm_val_p, xmm_idx_p, xmm_val_l, xmm_idx_l, cmp_val); + JMP_TO_LABEL(topk_heapify_loop_end_label); + + sub(reg_prc, jcp_.data_size); + sub(reg_prc_idx, sizeof(int)); + heap_swap_node(xmm_val_p, xmm_idx_p, xmm_val_l, xmm_idx_l); + mov(reg_aux, reg_prc); + mov(reg_aux_idx, reg_prc_idx); + shl(reg_offset, 1); + add(reg_offset, jcp_.data_size); + if (jcp_.data_size != sizeof(int)) { + shl(reg_offset_idx, 1); + add(reg_offset_idx, sizeof(int)); + } + shl(reg_j, 1); + add(reg_j, 1); + jmp(topk_heapify_loop_label, T_NEAR); + } + } + L(topk_heapify_loop_end_label); + + mov(reg_params, ptr[rsp]); + add(rsp, sizeof(int64_t)); + mov(reg_heap_seq_idx, ptr[rsp]); + add(rsp, sizeof(int64_t)); + mov(reg_src, ptr[rsp]); + add(rsp, sizeof(int64_t)); + mov(reg_heap_axis_dim.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + mov(reg_end.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + mov(reg_zero.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + mov(reg_heap_top_k.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + mov(reg_i.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + } + + inline void heap_cmp_node(Xmm xmm_val_a, Xmm xmm_idx_a, Xmm xmm_val_b, Xmm xmm_idx_b, bool cmp_val = true) { + if (isa == cpu::x64::avx512_common) { + if (cmp_val) + vcmpps(k_mask, xmm_val_a, xmm_val_b, heap_cmp_flg); + else + vcmpps(k_mask, xmm_idx_a, xmm_idx_b, _cmp_lt_os); + } else { + if (cmp_val) + uni_vcmpps(xmm_mask, xmm_val_a, xmm_val_b, heap_cmp_flg); + else + uni_vcmpps(xmm_mask, xmm_idx_a, xmm_idx_b, _cmp_lt_os); + } + } + + // n: node, c: child + inline void heap_swap_node(Xmm xmm_val_n, Xmm xmm_idx_n, Xmm xmm_val_c, Xmm xmm_idx_c) { + // swap store + store_scalar(ptr[reg_aux], xmm_val_c, data_type); + store_scalar(ptr[reg_aux_idx], xmm_idx_c, memory::data_type::s32); + store_scalar(ptr[reg_prc], xmm_val_n, data_type); + store_scalar(ptr[reg_prc_idx], xmm_idx_n, memory::data_type::s32); + } + + inline void heap_swap_root(const Xbyak::Reg64 ®_idx) { + get_addr_by_reg_idx(reg_aux, reg_dst, reg_idx, jcp_.data_size); + get_addr_by_reg_idx(reg_aux_idx, reg_dst_idx, reg_idx, sizeof(int)); + + load_scalar(xmm_val_p, ptr[reg_dst], data_type); + load_scalar(xmm_idx_p, ptr[reg_dst_idx], memory::data_type::s32); + load_scalar(xmm_val_l, ptr[reg_aux], data_type); + load_scalar(xmm_idx_l, ptr[reg_aux_idx], memory::data_type::s32); + + store_scalar(ptr[reg_aux], xmm_val_p, data_type); + store_scalar(ptr[reg_aux_idx], xmm_idx_p, memory::data_type::s32); + store_scalar(ptr[reg_dst], xmm_val_l, data_type); + store_scalar(ptr[reg_dst_idx], xmm_idx_l, memory::data_type::s32); + } + + inline void topk_bubble_vector() { + mov(reg_bubble_block_idx, ptr[reg_params + GET_OFF(idx_block_buf)]); + if (!jcp_.bubble_inplace) { + mov(reg_block_sort_stride, ptr[reg_params + GET_OFF(sort_stride)]); + mov(reg_bubble_axis_dim, ptr[reg_params + GET_OFF(axis_dim)]); + mov(reg_bubble_block_top_k, ptr[reg_params + GET_OFF(top_k)]); + } + + Xbyak::Label topk_main_loop_label; + Xbyak::Label topk_main_loop_end_label; + L(topk_main_loop_label); + { + cmp(reg_work_amount, step); + jl(topk_main_loop_end_label, T_NEAR); + + if (jcp_.bubble_inplace) { + topk_bubble_inplace(step); + } else { + topk_bubble(step); + } + + add(reg_src, step * jcp_.data_size); + add(reg_dst, step * jcp_.data_size); + add(reg_dst_idx, step * sizeof(int)); + sub(reg_work_amount, step); + + jmp(topk_main_loop_label, T_NEAR); + } + L(topk_main_loop_end_label); + + // tail + if (jcp_.bubble_inplace) { + if (tail) { + Xbyak::Label topk_tail_loop_end_label; + cmp(reg_work_amount, tail); + jl(topk_tail_loop_end_label, T_NEAR); + + topk_bubble_inplace(tail); + + L(topk_tail_loop_end_label); + } + } else { + Xbyak::Label topk_tail_loop_label; + Xbyak::Label topk_tail_loop_end_label; + L(topk_tail_loop_label); + { + cmp(reg_work_amount, 0); + je(topk_tail_loop_end_label, T_NEAR); + + topk_bubble(1); + + add(reg_src, jcp_.data_size); + add(reg_dst, jcp_.data_size); + add(reg_dst_idx, sizeof(int)); + sub(reg_work_amount, 1); + + jmp(topk_tail_loop_label, T_NEAR); + } + L(topk_tail_loop_end_label); + } + } + + inline void reg_add(const Xbyak::Reg64 ®_sum, const Xbyak::Reg64 ®_a, const Xbyak::Reg64 ®_b) { + mov(reg_sum, reg_a); + add(reg_sum, reg_b); + } + + inline void query_table_by_reg_idx(const Xbyak::Reg64 ®_table, const Xbyak::Reg64 ®_idx, int offset, size_t size) { + mov(reg_tmp, reg_idx); + add(reg_tmp, offset); + mul_by_const(reg_tmp, reg_tmp_64, size); + add(reg_tmp, reg_table); + } + + inline void table_to_vmm(Vmm vmm_src, const Xbyak::Reg64 ®_table, const Xbyak::Reg64 ®_idx, int offset, size_t size) { + query_table_by_reg_idx(reg_table, reg_idx, offset, size); + uni_vmovdqu(vmm_src, ptr[reg_tmp]); + } + + inline void table_to_xmm(Xmm xmm_src, const Xbyak::Reg64 ®_table, const Xbyak::Reg64 ®_idx, int offset, size_t size) { + query_table_by_reg_idx(reg_table, reg_idx, offset, size); + uni_vmovss(xmm_src, ptr[reg_tmp]); + } + + inline void get_addr_by_reg_idx(const Xbyak::Reg ®_out, const Xbyak::Reg ®_base, const Xbyak::Reg64 ®_in, int value) { + mov(reg_out, reg_in); + mul_by_const(reg_out, reg_tmp_64, value); + add(reg_out, reg_base); + } + + inline void get_addr_by_reg_idx(const Xbyak::Reg ®_out, const Xbyak::Reg ®_base, const Xbyak::Reg64 ®_in, int value, + const Xbyak::Reg64 ®_value) { + mov(reg_out, reg_in); + imul(reg_out, reg_value); + mul_by_const(reg_out, reg_tmp_64, value); + add(reg_out, reg_base); + } + + inline void get_addr_by_reg_idx(const Xbyak::Reg ®_out, const Xbyak::Reg ®_base, const Xbyak::Reg64 ®_in, + const Xbyak::Reg64 ®_value) { + mov(reg_out, reg_in); + imul(reg_out, reg_value); + add(reg_out, reg_base); + } + + inline void reg_mul_add(const Xbyak::Reg ®_out, const Xbyak::Reg64 ®_in, int mul_val, int add_val) { + mov(reg_out, reg_in); + mul_by_const(reg_out, reg_tmp_64, mul_val); + add(reg_out, add_val); + } + + inline void reg_mul_add(const Xbyak::Reg ®_out, const Xbyak::Reg ®_tmp, const Xbyak::Reg64 ®_in, int mul_val) { + mov(reg_tmp, reg_in); + mul_by_const(reg_tmp, reg_tmp_64, mul_val); + add(reg_out, reg_tmp); + } + + inline void reg_mul_add(const Xbyak::Reg ®_out, int mul_val, const Xbyak::Reg64 ®_base) { + mul_by_const(reg_out, reg_tmp_64, mul_val); + add(reg_out, reg_base); + } + + inline void reg_sub_shr(const Xbyak::Reg ®_out, const Xbyak::Reg64 ®_in, int sub_val, int shr_val) { + mov(reg_out, reg_in); + sub(reg_out, sub_val); + shr(reg_out, shr_val); + } + + inline void reg_sub_mul(const Xbyak::Reg ®_out, const Xbyak::Reg64 ®_in, int sub_val, int mul_val) { + mov(reg_out, reg_in); + sub(reg_out, sub_val); + mul_by_const(reg_out, reg_tmp_64, mul_val); + } + + inline void reg_shl(const Xbyak::Reg ®_out, int rate) { + switch (rate) { + case 1: + break; + case 2: + shl(reg_out, 1); + break; + case 4: + shl(reg_out, 2); + break; + default: + assert(!"incorrect data size rate"); + } + } + + inline void reg_shr(const Xbyak::Reg ®_out, int rate) { + switch (rate) { + case 1: + break; + case 2: + shr(reg_out, 1); + break; + case 4: + shr(reg_out, 2); + break; + default: + assert(!"incorrect data size rate"); + } + } + + inline void reg_div_blk_size(const Xbyak::Reg ®_out, const Xbyak::Reg64 ®_in, int blk_size) { + mov(reg_out, reg_in); + switch (blk_size) { + case 8: + shr(reg_out, 3); + break; + case 16: + shr(reg_out, 4); + break; + default: + assert(!"incorrect blk_size"); + } + } + + inline void reg_mod_blk_size(const Xbyak::Reg ®_out, const Xbyak::Reg64 ®_in, int blk_size) { + mov(reg_out, reg_in); + reg_div_blk_size(reg_tmp_64, reg_in, blk_size); + switch (blk_size) { + case 8: + shl(reg_tmp_64, 3); + break; + case 16: + shl(reg_tmp_64, 4); + break; + default: + assert(!"incorrect blk_size"); + } + sub(reg_out, reg_tmp_64); + } + + inline void reg_calc_offset_by_channel_idx(const Xbyak::Reg ®_out, const Xbyak::Reg64 ®_stride, + const Xbyak::Reg64 ®_channel_idx, int blk_size) { + reg_div_blk_size(reg_out, reg_channel_idx, blk_size); + imul(reg_out, reg_stride); + reg_mod_blk_size(reg_tmp, reg_channel_idx, blk_size); + add(reg_out, reg_tmp); + } + + inline void topk_bubble(int elt_num) { + reg_shl(reg_block_sort_stride, jcp_.data_size); + + // init dst + Xbyak::Label topk_init_loop_label; + Xbyak::Label topk_init_loop_end_label; + mov(reg_i, 0); + L(topk_init_loop_label); + { + cmp(reg_i, reg_bubble_block_top_k); + je(topk_init_loop_end_label, T_NEAR); + + get_addr_by_reg_idx(reg_tmp, reg_src, reg_block_sort_stride_byte, reg_i); + load_emitter->emit_code({static_cast(reg_tmp.getIdx())}, {static_cast(vmm_tmp.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + get_addr_by_reg_idx(reg_tmp, reg_dst, reg_block_sort_stride_byte, reg_i); + store_emitter->emit_code({static_cast(vmm_tmp.getIdx())}, {static_cast(reg_tmp.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + + table_to_vmm(vmm_tmp, reg_bubble_block_idx, reg_i, 0, vlen); + get_addr_by_reg_idx(reg_tmp, reg_dst_idx, reg_block_sort_stride_byte, sizeof(int) / jcp_.data_size, reg_i); + store_emitter->emit_code({static_cast(vmm_tmp.getIdx())}, {static_cast(reg_tmp.getIdx())}, + std::make_shared(Precision::I32, Precision::I32, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + + add(reg_i, 1); + jmp(topk_init_loop_label, T_NEAR); + } + L(topk_init_loop_end_label); + + // sort + topk_bubble_vector_sort(elt_num); + + // update + Xbyak::Label topk_update_loop_label; + Xbyak::Label topk_update_loop_end_label; + mov(reg_i, reg_bubble_block_top_k); + L(topk_update_loop_label); + { + cmp(reg_i, reg_bubble_axis_dim); + je(topk_update_loop_end_label, T_NEAR); + + get_addr_by_reg_idx(reg_tmp, reg_src, reg_block_sort_stride_byte, reg_i); + load_emitter->emit_code({static_cast(reg_tmp.getIdx())}, {static_cast(vmm_val_r.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + + table_to_vmm(vmm_idx_r, reg_bubble_block_idx, reg_i, 0, vlen); + uni_vcvtdq2ps(vmm_idx_r, vmm_idx_r); + + sub(rsp, sizeof(int64_t)); + mov(ptr[rsp], reg_bubble_block_idx); + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_bubble_axis_dim.cvt32()); + + Xbyak::Label topk_update_inner_loop_label; + Xbyak::Label topk_update_inner_loop_end_label; + mov(reg_block_r, reg_bubble_block_top_k); + sub(reg_bubble_block_top_k, 1); + L(topk_update_inner_loop_label); + { + cmp(reg_block_r, 0); + je(topk_update_inner_loop_end_label, T_NEAR); + + mov(reg_block_l, reg_block_r); + sub(reg_block_l, 1); + bubble_swap_vector(reg_block_l, reg_block_r, elt_num); + + sub(reg_block_r, 1); + jmp(topk_update_inner_loop_label, T_NEAR); + } + L(topk_update_inner_loop_end_label); + add(reg_bubble_block_top_k, 1); + + mov(reg_bubble_axis_dim.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + mov(reg_bubble_block_idx, ptr[rsp]); + add(rsp, sizeof(int64_t)); + + add(reg_i, 1); + jmp(topk_update_loop_label, T_NEAR); + } + L(topk_update_loop_end_label); + + if (jcp_.sort_index) { + topk_bubble_vector_sort(elt_num, false); + } + + reg_shr(reg_block_sort_stride, jcp_.data_size); + } + + inline void topk_bubble_vector_sort(int elt_num, bool cmp_val = true) { + sub(rsp, sizeof(int64_t)); + mov(ptr[rsp], reg_bubble_block_idx); + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_bubble_axis_dim.cvt32()); + + Xbyak::Label topk_sort_loop_label; + Xbyak::Label topk_sort_loop_end_label; + mov(reg_i, 0); + sub(reg_bubble_block_top_k, 1); + L(topk_sort_loop_label); + { + cmp(reg_i, reg_bubble_block_k_sub_1); + je(topk_sort_loop_end_label, T_NEAR); + + Xbyak::Label topk_sort_inner_loop_label; + Xbyak::Label topk_sort_inner_loop_end_label; + mov(reg_block_r, reg_bubble_block_k_sub_1); + L(topk_sort_inner_loop_label); + { + cmp(reg_block_r, reg_i); + je(topk_sort_inner_loop_end_label, T_NEAR); + + mov(reg_block_l, reg_block_r); + sub(reg_block_l, 1); + bubble_swap_vector(reg_block_l, reg_block_r, elt_num, cmp_val); + + sub(reg_block_r, 1); + jmp(topk_sort_inner_loop_label, T_NEAR); + } + L(topk_sort_inner_loop_end_label); + + add(reg_i, 1); + jmp(topk_sort_loop_label, T_NEAR); + } + L(topk_sort_loop_end_label); + add(reg_bubble_block_top_k, 1); + + mov(reg_bubble_axis_dim.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + mov(reg_bubble_block_idx, ptr[rsp]); + add(rsp, sizeof(int64_t)); + } + + inline void topk_bubble_inplace(int elt_num) { + // load + for (int i = 0; i < jcp_.top_k; i++) { + load_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val(i).getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num, i * jcp_.sort_stride * jcp_.data_size), + {}, {load_pool_gpr_idxs}); + uni_vmovdqu(vmm_idx(i), table_val(i)); + uni_vcvtdq2ps(vmm_idx(i), vmm_idx(i)); + } + // sort + for (int i = 0; i < jcp_.top_k - 1; i++) { + for (int j = jcp_.top_k - 1; j > i; j--) { + swap_vector(vmm_val(j - 1), vmm_idx(j - 1), vmm_val(j), vmm_idx(j)); + } + } + for (int i = jcp_.top_k; i < jcp_.axis_dim; i++) { + load_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val(jcp_.top_k).getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num, i * jcp_.sort_stride * jcp_.data_size), + {}, {load_pool_gpr_idxs}); + uni_vmovdqu(vmm_idx(jcp_.top_k), table_val(i)); + uni_vcvtdq2ps(vmm_idx(jcp_.top_k), vmm_idx(jcp_.top_k)); + for (int j = jcp_.top_k; j > 0; j--) { + swap_vector(vmm_val(j - 1), vmm_idx(j - 1), vmm_val(j), vmm_idx(j)); + } + } + if (jcp_.sort_index) { + for (int i = 0; i < jcp_.top_k - 1; i++) { + for (int j = jcp_.top_k - 1; j > i; j--) { + swap_vector(vmm_val(j - 1), vmm_idx(j - 1), vmm_val(j), vmm_idx(j), false); + } + } + } + // store + for (int i = 0; i < jcp_.top_k; i++) { + store_emitter->emit_code({static_cast(vmm_val(i).getIdx())}, {static_cast(reg_dst.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, elt_num, i * jcp_.sort_stride * jcp_.data_size), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + store_emitter->emit_code({static_cast(vmm_idx(i).getIdx())}, {static_cast(reg_dst_idx.getIdx())}, + std::make_shared(Precision::FP32, Precision::I32, elt_num, i * jcp_.sort_stride * sizeof(int)), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + } + } + + inline void topk_bubble_BLK_on_channel_horiz() { + mov(reg_bubble_axis_dim, ptr[reg_params + GET_OFF(axis_dim)]); + mov(reg_seq_sort_stride, ptr[reg_params + GET_OFF(sort_stride)]); + + // load and sort + mov(reg_i, 0); + Xbyak::Label topk_load_sort_label; + Xbyak::Label topk_load_sort_end_label; + cmp(reg_bubble_axis_dim, jcp_.blk_size); + jge(topk_load_sort_label, T_NEAR); + + load_scalar(xmm_val(0), ptr[reg_src], data_type); + uni_vmovss(xmm_idx(0), table_bubble_seq_idx(0)); + uni_vcvtdq2ps(xmm_idx(0), xmm_idx(0)); + jmp(topk_load_sort_end_label, T_NEAR); + + L(topk_load_sort_label); + { + mov(reg_bubble_seq_idx, ptr[reg_params + GET_OFF(idx_seq_buf)]); + load_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val(0).getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, step, 0), + {}, {load_pool_gpr_idxs}); + uni_vmovdqu(vmm_idx(0), table_bubble_seq_idx(0)); + uni_vcvtdq2ps(vmm_idx(0), vmm_idx(0)); + if (isa == cpu::x64::sse41) { + load_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val(1).getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, step, 4 * jcp_.data_size), + {}, {load_pool_gpr_idxs}); + uni_vmovdqu(vmm_idx(1), table_bubble_seq_idx(4)); + uni_vcvtdq2ps(vmm_idx(1), vmm_idx(1)); + swap_vector(vmm_val(0), vmm_idx(0), vmm_val(1), vmm_idx(1)); + } + + Xbyak::Label topk_iter_label; + Xbyak::Label topk_iter_end_label; + mov(reg_i, jcp_.blk_size); + sub(reg_bubble_axis_dim, jcp_.blk_size); + L(topk_iter_label); + { + cmp(reg_i, reg_bubble_axis_dim); + jg(topk_iter_end_label, T_NEAR); + + get_addr_by_reg_idx(reg_aux, reg_src, reg_i, jcp_.data_size, reg_seq_sort_stride); + load_emitter->emit_code({static_cast(reg_aux.getIdx())}, {static_cast(vmm_val(1).getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, step), + {}, {load_pool_gpr_idxs}); + table_to_vmm(vmm_idx(1), reg_bubble_seq_idx, reg_i, 0, sizeof(int)); + uni_vcvtdq2ps(vmm_idx(1), vmm_idx(1)); + swap_vector(vmm_val(0), vmm_idx(0), vmm_val(1), vmm_idx(1)); + if (isa == cpu::x64::sse41) { + add(reg_aux, 4 * jcp_.data_size); + load_emitter->emit_code({static_cast(reg_aux.getIdx())}, {static_cast(vmm_val(1).getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, step), + {}, {load_pool_gpr_idxs}); + table_to_vmm(vmm_idx(1), reg_bubble_seq_idx, reg_i, 4, sizeof(int)); + uni_vcvtdq2ps(vmm_idx(1), vmm_idx(1)); + swap_vector(vmm_val(0), vmm_idx(0), vmm_val(1), vmm_idx(1)); + } + + add(reg_i, jcp_.blk_size); + jmp(topk_iter_label, T_NEAR); + } + L(topk_iter_end_label); + add(reg_bubble_axis_dim, jcp_.blk_size); + + horiz_process(); + } + L(topk_load_sort_end_label); + + Xbyak::Label topk_tail_label; + Xbyak::Label topk_tail_end_label; + get_addr_by_reg_idx(reg_aux, reg_src, reg_i, jcp_.data_size, reg_seq_sort_stride); + L(topk_tail_label); + { + cmp(reg_i, reg_bubble_axis_dim); + je(topk_tail_end_label, T_NEAR); + + load_scalar(xmm_val(1), ptr[reg_aux], data_type); + table_to_xmm(xmm_idx(1), reg_bubble_seq_idx, reg_i, 0, sizeof(int)); + uni_vcvtdq2ps(xmm_idx(1), xmm_idx(1)); + bubble_swap_xmm(xmm_val(0), xmm_idx(0), xmm_val(1), xmm_idx(1)); + + add(reg_i, 1); + add(reg_aux, jcp_.data_size); + jmp(topk_tail_label, T_NEAR); + } + L(topk_tail_end_label); + + // store + store_scalar(ptr[reg_dst], xmm_val(0), data_type); + store_scalar(ptr[reg_dst_idx], xmm_idx(0), memory::data_type::s32); + } + + // dst: xmm_val(0) and xmm_idx(0) + // aux: xmm_val(2/3/4) and xmm_idx(2/3/4) + inline void horiz_process() { + if (isa == cpu::x64::sse41) { + horize_top1(); + } else if (isa == cpu::x64::avx2) { + Xbyak::Ymm ymm_val_dst = Xbyak::Ymm(vmm_val(0).getIdx()); + vextractf128(xmm_val(2), ymm_val_dst, 0); + vextractf128(xmm_val(3), ymm_val_dst, 1); + Xbyak::Ymm ymm_idx_dst = Xbyak::Ymm(vmm_idx(0).getIdx()); + vextractf128(xmm_idx(2), ymm_idx_dst, 0); + vextractf128(xmm_idx(3), ymm_idx_dst, 1); + bubble_swap_xmm(xmm_val(2), xmm_idx(2), xmm_val(3), xmm_idx(3)); + uni_vmovups(xmm_val(0), xmm_val(2)); + uni_vmovups(xmm_idx(0), xmm_idx(2)); + horize_top1(); + } else { + Xbyak::Zmm zmm_val_dst = Xbyak::Zmm(vmm_val(0).getIdx()); + vextractf32x4(xmm_val(2), zmm_val_dst, 0); + vextractf32x4(xmm_val(3), zmm_val_dst, 1); + Xbyak::Zmm zmm_idx_dst = Xbyak::Zmm(vmm_idx(0).getIdx()); + vextractf32x4(xmm_idx(2), zmm_idx_dst, 0); + vextractf32x4(xmm_idx(3), zmm_idx_dst, 1); + bubble_swap_xmm(xmm_val(2), xmm_idx(2), xmm_val(3), xmm_idx(3)); + vextractf32x4(xmm_val(3), zmm_val_dst, 2); + vextractf32x4(xmm_val(4), zmm_val_dst, 3); + vextractf32x4(xmm_idx(3), zmm_idx_dst, 2); + vextractf32x4(xmm_idx(4), zmm_idx_dst, 3); + bubble_swap_xmm(xmm_val(3), xmm_idx(3), xmm_val(4), xmm_idx(4)); + bubble_swap_xmm(xmm_val(2), xmm_idx(2), xmm_val(3), xmm_idx(3)); + uni_vmovups(xmm_val(0), xmm_val(2)); + uni_vmovups(xmm_idx(0), xmm_idx(2)); + horize_top1(); + } + } + + // dst: xmm_val(0) and xmm_idx(0) + // aux: xmm_val(3) and xmm_idx(3) + inline void horize_top1() { + uni_vmovshdup(xmm_val(3), xmm_val(0)); // dst:1,2,3,4; aux:2,2,4,4 + uni_vmovshdup(xmm_idx(3), xmm_idx(0)); + bubble_swap_xmm(xmm_val(0), xmm_idx(0), xmm_val(3), xmm_idx(3)); // dst:f(1,2),f(2,2),f(3,4),f(4,4) + uni_vmovhlps(xmm_val(3), xmm_val(3), xmm_val(0)); // aux:f(3,4),f(4,4),4,4 + uni_vmovhlps(xmm_idx(3), xmm_idx(3), xmm_idx(0)); + bubble_swap_xmm(xmm_val(0), xmm_idx(0), xmm_val(3), xmm_idx(3)); // dst:f(1,2,3,4),... + } + + inline void topk_bubble_BLK_on_channel_verti() { + if (jcp_.bubble_inplace) { + topk_bubble_BLK_on_channel_inplace(); + } else { + topk_bubble_BLK_on_channel(); + } + } + + inline void topk_bubble_BLK_on_channel() { + mov(reg_bubble_seq_idx, ptr[reg_params + GET_OFF(idx_seq_buf)]); + mov(reg_bubble_axis_dim, ptr[reg_params + GET_OFF(axis_dim)]); + mov(reg_seq_sort_stride, ptr[reg_params + GET_OFF(sort_stride)]); + mov(reg_bubble_seq_top_k, ptr[reg_params + GET_OFF(top_k)]); + + mul_by_const(reg_seq_sort_stride, reg_tmp_64, jcp_.blk_size); + + // init dst + Xbyak::Label topk_init_next_label; + Xbyak::Label topk_init_iter_label; + Xbyak::Label topk_init_iter_end_label; + mov(reg_i, 0); + mov(reg_sub_idx, 0); + mov(reg_aux, 0); + L(topk_init_iter_label); + { + cmp(reg_i, reg_bubble_seq_top_k); + je(topk_init_iter_end_label, T_NEAR); + + reg_add(reg_tmp, reg_src, reg_aux); + load_scalar(xmm_tmp, ptr[reg_tmp], data_type); + reg_add(reg_tmp, reg_dst, reg_aux); + store_scalar(ptr[reg_tmp], xmm_tmp, data_type); + + table_to_xmm(xmm_tmp, reg_bubble_seq_idx, reg_i, 0, sizeof(int)); + get_addr_by_reg_idx(reg_tmp, reg_dst_idx, reg_aux, sizeof(int) / jcp_.data_size); + store_scalar(ptr[reg_tmp], xmm_tmp, memory::data_type::s32, false); + + add(reg_sub_idx, 1); + cmp(reg_sub_idx, jcp_.blk_size); + jl(topk_init_next_label, T_NEAR); + + mov(reg_sub_idx, 0); + reg_sub_mul(reg_tmp, reg_blk_stride, jcp_.blk_size, jcp_.data_size); + add(reg_aux, reg_tmp); + + L(topk_init_next_label); + { + add(reg_i, 1); + add(reg_aux, jcp_.data_size); + jmp(topk_init_iter_label, T_NEAR); + } + } + L(topk_init_iter_end_label); + + // sort + topk_bubble_BLK_on_channel_sort(); + + // update + Xbyak::Label topk_next_label; + Xbyak::Label topk_iter_label; + Xbyak::Label topk_iter_end_label; + mov(reg_i, reg_bubble_seq_top_k); + reg_calc_offset_by_channel_idx(reg_seq_tmp, reg_blk_stride, reg_bubble_seq_top_k, jcp_.blk_size); + get_addr_by_reg_idx(reg_aux, reg_src, reg_seq_tmp, jcp_.data_size); + reg_mod_blk_size(reg_sub_idx, reg_bubble_seq_top_k, jcp_.blk_size); + L(topk_iter_label); + { + cmp(reg_i, reg_bubble_axis_dim); + je(topk_iter_end_label, T_NEAR); + + load_scalar(xmm_val_r, ptr[reg_aux], data_type); + table_to_xmm(xmm_idx_r, reg_bubble_seq_idx, reg_i, 0, sizeof(int)); + uni_vcvtdq2ps(xmm_idx_r, xmm_idx_r); + + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_prc.cvt32()); + + Xbyak::Label topk_update_inner_loop_label; + Xbyak::Label topk_update_inner_loop_end_label; + mov(reg_seq_r, reg_bubble_seq_top_k); + sub(reg_bubble_seq_top_k, 1); + L(topk_update_inner_loop_label); + { + cmp(reg_seq_r, 0); + je(topk_update_inner_loop_end_label, T_NEAR); + + mov(reg_seq_l, reg_seq_r); + sub(reg_seq_l, 1); + bubble_swap_by_index(reg_seq_l, reg_seq_r); + + sub(reg_seq_r, 1); + jmp(topk_update_inner_loop_label, T_NEAR); + } + L(topk_update_inner_loop_end_label); + add(reg_bubble_seq_top_k, 1); + + mov(reg_prc.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + + add(reg_sub_idx, 1); + cmp(reg_sub_idx, jcp_.blk_size); + jl(topk_next_label, T_NEAR); + + mov(reg_sub_idx, 0); + reg_sub_mul(reg_tmp, reg_blk_stride, jcp_.blk_size, jcp_.data_size); + add(reg_aux, reg_tmp); + + L(topk_next_label); + { + add(reg_i, 1); + add(reg_aux, jcp_.data_size); + jmp(topk_iter_label, T_NEAR); + } + } + L(topk_iter_end_label); + + if (jcp_.sort_index) { + topk_bubble_BLK_on_channel_sort(false); + } + } + + inline void topk_bubble_BLK_on_channel_sort(bool cmp_val = true) { + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_prc.cvt32()); + + Xbyak::Label topk_sort_loop_label; + Xbyak::Label topk_sort_loop_end_label; + mov(reg_i, 0); + sub(reg_bubble_seq_top_k, 1); + L(topk_sort_loop_label); + { + cmp(reg_i, reg_bubble_seq_k_sub_1); + je(topk_sort_loop_end_label, T_NEAR); + + Xbyak::Label topk_sort_inner_loop_label; + Xbyak::Label topk_sort_inner_loop_end_label; + mov(reg_seq_r, reg_bubble_seq_k_sub_1); + L(topk_sort_inner_loop_label); + { + cmp(reg_seq_r, reg_i); + je(topk_sort_inner_loop_end_label, T_NEAR); + + mov(reg_seq_l, reg_seq_r); + sub(reg_seq_l, 1); + bubble_swap_by_index(reg_seq_l, reg_seq_r, cmp_val); + + sub(reg_seq_r, 1); + jmp(topk_sort_inner_loop_label, T_NEAR); + } + L(topk_sort_inner_loop_end_label); + + add(reg_i, 1); + jmp(topk_sort_loop_label, T_NEAR); + } + L(topk_sort_loop_end_label); + add(reg_bubble_seq_top_k, 1); + + mov(reg_prc.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + } + + inline void topk_bubble_BLK_on_channel_inplace() { + // load + for (int i = 0; i < jcp_.top_k; i++) { + int offset = i / jcp_.blk_size * blk_stride + i % jcp_.blk_size; + load_scalar(xmm_val(i), ptr[reg_src + offset * jcp_.data_size], data_type); + uni_vmovdqu(xmm_idx(i), table_val(i)); + uni_vcvtdq2ps(xmm_idx(i), xmm_idx(i)); + } + // sort + for (int i = 0; i < jcp_.top_k - 1; i++) { + for (int j = jcp_.top_k - 1; j > i; j--) { + bubble_swap_xmm(xmm_val(j - 1), xmm_idx(j - 1), xmm_val(j), xmm_idx(j)); + } + } + for (int i = jcp_.top_k; i < jcp_.axis_dim; i++) { + int offset = i / jcp_.blk_size * blk_stride + i % jcp_.blk_size; + load_scalar(xmm_val(jcp_.top_k), ptr[reg_src + offset * jcp_.data_size], data_type); + uni_vmovdqu(xmm_idx(jcp_.top_k), table_val(i)); + uni_vcvtdq2ps(xmm_idx(jcp_.top_k), xmm_idx(jcp_.top_k)); + for (int j = jcp_.top_k; j > 0; j--) { + bubble_swap_xmm(xmm_val(j - 1), xmm_idx(j - 1), xmm_val(j), xmm_idx(j)); + } + } + if (jcp_.sort_index) { + for (int i = 0; i < jcp_.top_k - 1; i++) { + for (int j = jcp_.top_k - 1; j > i; j--) { + bubble_swap_xmm(xmm_val(j - 1), xmm_idx(j - 1), xmm_val(j), xmm_idx(j), false); + } + } + } + // store + for (int i = 0; i < jcp_.top_k; i++) { + int offset = i / jcp_.blk_size * blk_stride + i % jcp_.blk_size; + store_scalar(ptr[reg_dst + offset * jcp_.data_size], xmm_val(i), data_type); + store_scalar(ptr[reg_dst_idx + offset * sizeof(int)], xmm_idx(i), memory::data_type::s32); + } + } + + inline void bubble_swap_vector(const Xbyak::Reg64 ®_l, const Xbyak::Reg64 ®_r, int elt_num, bool cmp_val = true) { + mov(reg_tmp_64, reg_block_sort_stride_byte); + imul(reg_tmp_64, reg_l); + + // load l + mov(reg_tmp, reg_tmp_64); + add(reg_tmp, reg_dst); + load_emitter->emit_code({static_cast(reg_tmp.getIdx())}, {static_cast(vmm_val_l.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + reg_shl(reg_tmp_64, sizeof(int) / jcp_.data_size); + mov(reg_tmp, reg_tmp_64); + add(reg_tmp, reg_dst_idx); + reg_shr(reg_tmp_64, sizeof(int) / jcp_.data_size); + load_emitter->emit_code({static_cast(reg_tmp.getIdx())}, {static_cast(vmm_idx_l.getIdx())}, + std::make_shared(Precision::I32, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + + // load r + Xbyak::Label topk_load_jmp_label; + cmp(reg_r, reg_bubble_block_k_sub_1); + jg(topk_load_jmp_label, T_NEAR); + { + add(reg_tmp_64, reg_block_sort_stride_byte); + mov(reg_tmp, reg_tmp_64); + add(reg_tmp, reg_dst); + load_emitter->emit_code({static_cast(reg_tmp.getIdx())}, {static_cast(vmm_val_r.getIdx())}, + std::make_shared(jcp_.precision, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + reg_shl(reg_tmp_64, sizeof(int) / jcp_.data_size); + mov(reg_tmp, reg_tmp_64); + add(reg_tmp, reg_dst_idx); + reg_shr(reg_tmp_64, sizeof(int) / jcp_.data_size); + load_emitter->emit_code({static_cast(reg_tmp.getIdx())}, {static_cast(vmm_idx_r.getIdx())}, + std::make_shared(Precision::I32, Precision::FP32, elt_num), + {}, {load_pool_gpr_idxs}); + sub(reg_tmp_64, reg_block_sort_stride_byte); + } + L(topk_load_jmp_label); + + swap_vector(vmm_val_l, vmm_idx_l, vmm_val_r, vmm_idx_r, cmp_val); + + // store l + mov(reg_tmp, reg_tmp_64); + add(reg_tmp, reg_dst); + store_emitter->emit_code({static_cast(vmm_val_l.getIdx())}, {static_cast(reg_tmp.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + reg_shl(reg_tmp_64, sizeof(int) / jcp_.data_size); + mov(reg_tmp, reg_tmp_64); + add(reg_tmp, reg_dst_idx); + reg_shr(reg_tmp_64, sizeof(int) / jcp_.data_size); + store_emitter->emit_code({static_cast(vmm_idx_l.getIdx())}, {static_cast(reg_tmp.getIdx())}, + std::make_shared(Precision::FP32, Precision::I32, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + + // store r + Xbyak::Label topk_store_jmp_label; + cmp(reg_r, reg_bubble_block_k_sub_1); + jg(topk_store_jmp_label, T_NEAR); + { + add(reg_tmp_64, reg_block_sort_stride_byte); + mov(reg_tmp, reg_tmp_64); + add(reg_tmp, reg_dst); + store_emitter->emit_code({static_cast(vmm_val_r.getIdx())}, {static_cast(reg_tmp.getIdx())}, + std::make_shared(Precision::FP32, jcp_.precision, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + reg_shl(reg_tmp_64, sizeof(int) / jcp_.data_size); + mov(reg_tmp, reg_tmp_64); + add(reg_tmp, reg_dst_idx); + reg_shr(reg_tmp_64, sizeof(int) / jcp_.data_size); + store_emitter->emit_code({static_cast(vmm_idx_r.getIdx())}, {static_cast(reg_tmp.getIdx())}, + std::make_shared(Precision::FP32, Precision::I32, elt_num), + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + } + L(topk_store_jmp_label); + } + + inline void swap_vector(Vmm vmm_val_a, Vmm vmm_idx_a, Vmm vmm_val_b, Vmm vmm_idx_b, bool cmp_val = true) { + if (isa == cpu::x64::avx512_common) { + if (cmp_val) + vcmpps(k_mask, vmm_val_a, vmm_val_b, cmp_flg); + else + vcmpps(k_mask, vmm_idx_a, vmm_idx_b, _cmp_nle_us); + + uni_vmovups(vmm_tmp, vmm_val_a); + vblendmps(vmm_val_a | k_mask, vmm_val_a, vmm_val_b); + vblendmps(vmm_val_b | k_mask, vmm_val_b, vmm_tmp); + + uni_vmovups(vmm_tmp, vmm_idx_a); + vblendmps(vmm_idx_a | k_mask, vmm_idx_a, vmm_idx_b); + vblendmps(vmm_idx_b | k_mask, vmm_idx_b, vmm_tmp); + } else { + if (cmp_val) + uni_vcmpps(vmm_mask, vmm_val_a, vmm_val_b, cmp_flg); + else + uni_vcmpps(vmm_mask, vmm_idx_a, vmm_idx_b, _cmp_nle_us); + + uni_vmovups(vmm_tmp, vmm_val_a); + uni_vblendvps(vmm_val_a, vmm_val_a, vmm_val_b, vmm_mask); + uni_vblendvps(vmm_val_b, vmm_val_b, vmm_tmp, vmm_mask); + + uni_vmovups(vmm_tmp, vmm_idx_a); + uni_vblendvps(vmm_idx_a, vmm_idx_a, vmm_idx_b, vmm_mask); + uni_vblendvps(vmm_idx_b, vmm_idx_b, vmm_tmp, vmm_mask); + } + } + + inline void bubble_swap_by_index(const Xbyak::Reg64 ®_l, const Xbyak::Reg64 ®_r, bool cmp_val = true) { + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_i.cvt32()); + sub(rsp, sizeof(int)); + mov(ptr[rsp], reg_prc_idx.cvt32()); + sub(rsp, sizeof(int64_t)); + mov(ptr[rsp], reg_aux); + + mov(reg_aux, reg_blk_stride); + reg_calc_offset_by_channel_idx(reg_offset_l, reg_aux, reg_l, jcp_.blk_size); + reg_calc_offset_by_channel_idx(reg_offset_r, reg_aux, reg_r, jcp_.blk_size); + + get_addr_by_reg_idx(reg_aux, reg_dst, reg_offset_l, jcp_.data_size); + load_scalar(xmm_val_l, ptr[reg_aux], data_type); + get_addr_by_reg_idx(reg_aux, reg_dst_idx, reg_offset_l, sizeof(int)); + load_scalar(xmm_idx_l, ptr[reg_aux], memory::data_type::s32); + + Xbyak::Label topk_load_jmp_label; + cmp(reg_r, reg_bubble_seq_k_sub_1); + jg(topk_load_jmp_label, T_NEAR); + { + get_addr_by_reg_idx(reg_aux, reg_dst, reg_offset_r, jcp_.data_size); + load_scalar(xmm_val_r, ptr[reg_aux], data_type); + get_addr_by_reg_idx(reg_aux, reg_dst_idx, reg_offset_r, sizeof(int)); + load_scalar(xmm_idx_r, ptr[reg_aux], memory::data_type::s32); + } + L(topk_load_jmp_label); + + bubble_swap_xmm(xmm_val_l, xmm_idx_l, xmm_val_r, xmm_idx_r, cmp_val); + + get_addr_by_reg_idx(reg_aux, reg_dst, reg_offset_l, jcp_.data_size); + store_scalar(ptr[reg_aux], xmm_val_l, data_type); + get_addr_by_reg_idx(reg_aux, reg_dst_idx, reg_offset_l, sizeof(int)); + store_scalar(ptr[reg_aux], xmm_idx_l, memory::data_type::s32); + + Xbyak::Label topk_store_jmp_label; + cmp(reg_r, reg_bubble_seq_k_sub_1); + jg(topk_store_jmp_label, T_NEAR); + { + get_addr_by_reg_idx(reg_aux, reg_dst, reg_offset_r, jcp_.data_size); + store_scalar(ptr[reg_aux], xmm_val_r, data_type); + get_addr_by_reg_idx(reg_aux, reg_dst_idx, reg_offset_r, sizeof(int)); + store_scalar(ptr[reg_aux], xmm_idx_r, memory::data_type::s32); + } + L(topk_store_jmp_label); + + mov(reg_aux, ptr[rsp]); + add(rsp, sizeof(int64_t)); + mov(reg_prc_idx.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + mov(reg_i.cvt32(), ptr[rsp]); + add(rsp, sizeof(int)); + } + + inline void bubble_swap_xmm(Xmm xmm_val_a, Xmm xmm_idx_a, Xmm xmm_val_b, Xmm xmm_idx_b, bool cmp_val = true) { + if (isa == cpu::x64::avx512_common) { + if (cmp_val) + vcmpps(k_mask, xmm_val_a, xmm_val_b, cmp_flg); + else + vcmpps(k_mask, xmm_idx_a, xmm_idx_b, _cmp_nle_us); + + uni_vmovups(xmm_tmp, xmm_val_a); + vblendmps(xmm_val_a | k_mask, xmm_val_a, xmm_val_b); + vblendmps(xmm_val_b | k_mask, xmm_val_b, xmm_tmp); + + uni_vmovups(xmm_tmp, xmm_idx_a); + vblendmps(xmm_idx_a | k_mask, xmm_idx_a, xmm_idx_b); + vblendmps(xmm_idx_b | k_mask, xmm_idx_b, xmm_tmp); + } else { + if (cmp_val) + uni_vcmpps(xmm_mask, xmm_val_a, xmm_val_b, cmp_flg); + else + uni_vcmpps(xmm_mask, xmm_idx_a, xmm_idx_b, _cmp_nle_us); + + uni_vmovups(xmm_tmp, xmm_val_a); + uni_vblendvps(xmm_val_a, xmm_val_a, xmm_val_b, xmm_mask); + uni_vblendvps(xmm_val_b, xmm_val_b, xmm_tmp, xmm_mask); + + uni_vmovups(xmm_tmp, xmm_idx_a); + uni_vblendvps(xmm_idx_a, xmm_idx_a, xmm_idx_b, xmm_mask); + uni_vblendvps(xmm_idx_b, xmm_idx_b, xmm_tmp, xmm_mask); + } + } + + inline void load_scalar(Xmm xmm_src, const Xbyak::Address &op, memory::data_type src_dt, bool cvt_dt = true) { + switch (src_dt) { + case memory::data_type::f32: + case memory::data_type::s32: + uni_vmovss(xmm_src, op); + break; + case memory::data_type::bf16: + uni_vpinsrw(xmm_src, xmm_src, op, 0x0); + uni_vpslld(xmm_src, xmm_src, 16); + break; + case memory::data_type::s8: + movsx(reg_tmp_32, op); + uni_vmovq(xmm_src, reg_tmp_64); + break; + case memory::data_type::u8: + movzx(reg_tmp_32, op); + uni_vmovq(xmm_src, reg_tmp_64); + break; + default: + assert(!"unknown src_dt"); + } + + if (cvt_dt && !isFloatCompatible(src_dt)) { + uni_vcvtdq2ps(xmm_src, xmm_src); + } + } + + inline void store_scalar(const Xbyak::Address &op, Xmm xmm_dst, memory::data_type dst_dt, bool cvt_dt = true) { + if (cvt_dt && !isFloatCompatible(dst_dt)) { + uni_vcvtps2dq(xmm_dst, xmm_dst); + } + + switch (dst_dt) { + case memory::data_type::f32: + case memory::data_type::s32: + uni_vmovss(op, xmm_dst); + break; + case memory::data_type::bf16: + uni_vpsrld(xmm_dst, xmm_dst, 16); + uni_vpextrw(op, xmm_dst, 0x0); + break; + case memory::data_type::s8: + uni_vpackssdw(xmm_dst, xmm_dst, xmm_dst); + uni_vpacksswb(xmm_dst, xmm_dst, xmm_dst); + uni_vmovq(reg_tmp_64, xmm_dst); + mov(op, reg_tmp_8); + break; + case memory::data_type::u8: + uni_vpackusdw(xmm_dst, xmm_dst, xmm_dst); + uni_vpackuswb(xmm_dst, xmm_dst, xmm_dst); + uni_vmovq(reg_tmp_64, xmm_dst); + mov(op, reg_tmp_8); + break; + default: + assert(!"unknown dst_dt"); + } + } + + void prepare_idx_table() { + auto broadcast_int = [&](int val) { + for (size_t d = 0; d < vlen / sizeof(float); ++d) { + dd(val); + } + }; + + align(64); + L(l_table); + + // 00000000 11111111 22222222 ... + for (int i = 0; i < jcp_.axis_dim; i++) { + broadcast_int(i); + } + } +}; bool MKLDNNTopKNode::isSupportedOperation(const std::shared_ptr& op, std::string& errorMessage) noexcept { try { @@ -25,12 +1791,21 @@ bool MKLDNNTopKNode::isSupportedOperation(const std::shared_ptr(topKOp->get_input_node_shared_ptr(TOPK_K)); + if (!topKConst) { + errorMessage = "Second tensor is not constant in static shape mode"; + return false; + } + } + if (topKOp->get_mode() != ngraph::op::TopKMode::MAX && - topKOp->get_mode() != ngraph::op::TopKMode::MIN) { + topKOp->get_mode() != ngraph::op::TopKMode::MIN) { errorMessage = "Unsupported mode."; return false; } - if (!MKLDNNPlugin::one_of(topKOp->get_sort_type(), ngraph::op::TopKSortType::NONE, + if (!one_of(topKOp->get_sort_type(), ngraph::op::TopKSortType::NONE, ngraph::op::TopKSortType::SORT_VALUES, ngraph::op::TopKSortType::SORT_INDICES)) { errorMessage = "Unsupported sort type."; @@ -42,316 +1817,518 @@ bool MKLDNNTopKNode::isSupportedOperation(const std::shared_ptr& op, const mkldnn::engine& eng, - MKLDNNWeightsSharing::Ptr &cache) : MKLDNNNode(op, eng, cache) { +MKLDNNTopKNode::MKLDNNTopKNode(const std::shared_ptr& op, const mkldnn::engine& eng, MKLDNNWeightsSharing::Ptr &cache) + : MKLDNNNode(op, eng, cache) { std::string errorMessage; - if (!isSupportedOperation(op, errorMessage)) { + if (isSupportedOperation(op, errorMessage)) { + errorPrefix = "TopK layer with name '" + getName() + "'"; + + auto topKOp = ngraph::as_type_ptr(op); + + auto in_dims = topKOp->get_input_partial_shape(TOPK_DATA); + auto out_dims = topKOp->get_output_partial_shape(TOPK_DATA); + auto out_idx_dims = topKOp->get_output_partial_shape(TOPK_INDEX); + auto in_dims_size = in_dims.size(); + + if (!isDynamicNgraphNode(op)) { + auto topKConst = std::dynamic_pointer_cast(topKOp->get_input_node_shared_ptr(TOPK_K)); + if (!topKConst) { + IE_THROW() << errorPrefix << "gets non-constant second tensor in static shape mode!"; + } + } + + axis = topKOp->get_axis(); + mode_max = topKOp->get_mode() == ngraph::op::TopKMode::MAX; + sort_index = topKOp->get_sort_type() == ngraph::op::TopKSortType::SORT_INDICES; + + top_k = 0; + vec_idx_seq.clear(); + vec_idx_block.clear(); + + if (inputShapes.size() != 2 || outputShapes.size() < 2) + IE_THROW() << errorPrefix << " gets incorrect number of input/output edges!"; + + if (getInputShapeAtPort(TOPK_DATA).getRank() != getOutputShapeAtPort(TOPK_DATA).getRank()) + IE_THROW() << errorPrefix << " gets incorrect number of input/output dimensions!"; + + if (getInputShapeAtPort(TOPK_K).getRank() != 1) + IE_THROW() << errorPrefix << " gets incorrect index vector dimension! Index vector should be 1 dimension."; + + if (out_dims != out_idx_dims) + IE_THROW() << errorPrefix << " gets incorrect output tensor dimension sizes!"; + + if (axis < 0) + axis += in_dims_size; + if (axis < 0 || axis >= static_cast(in_dims_size)) + IE_THROW() << errorPrefix << " gets incorrect input parameters dimensions and axis number!"; + } else { IE_THROW(NotImplemented) << errorMessage; } - auto topK1Op = ngraph::as_type_ptr(op); - - axis = topK1Op->get_axis(); - mode_max = topK1Op->get_mode() == ngraph::op::TopKMode::MAX; - sort_value = topK1Op->get_sort_type() == ngraph::op::TopKSortType::SORT_VALUES; } +void MKLDNNTopKNode::getSupportedDescriptors() {} + void MKLDNNTopKNode::initSupportedPrimitiveDescriptors() { if (!supportedPrimitiveDescriptors.empty()) return; - std::vector outDataConf; - outDataConf.reserve(outputShapes.size()); - outDataConf.emplace_back(LayoutType::ncsp, Precision::FP32); - for (int i = 1; i < outputShapes.size(); ++i) - outDataConf.emplace_back(LayoutType::ncsp, Precision::I32); - - addSupportedPrimDesc({{LayoutType::ncsp, Precision::FP32}, - {LayoutType::ncsp, Precision::I32}}, - outDataConf, - impl_desc_type::ref_any); -} - -void MKLDNNTopKNode::execute(mkldnn::stream strm) { - const float *src = reinterpret_cast(getParentEdgeAt(TOPK_DATA)->getMemoryPtr()->GetPtr()); - src_k = reinterpret_cast(getParentEdgeAt(TOPK_K)->getMemoryPtr()->GetPtr())[0]; - float* dst_data = nullptr; - int* dst_idx = nullptr; - - if (outputShapes.size() == 1) { - if (getOriginalOutputPrecisionAtPort(0) == Precision::FP32) { - dst_data = reinterpret_cast(getChildEdgesAtPort(0)[0]->getMemoryPtr()->GetPtr()); - } else { - dst_idx = reinterpret_cast(getChildEdgesAtPort(0)[0]->getMemoryPtr()->GetPtr()); - } - const VectorDims& dstDims = getChildEdgesAtPort(0)[0]->getMemory().getStaticDims(); - - if (dstDims[axis] != static_cast(src_k)) { - std::string errorMsg = "Output tensor dimension mismatch"; - IE_THROW() << errorMsg; - } - } else if (outputShapes.size() == 2) { - dst_data = reinterpret_cast(getChildEdgesAtPort(TOPK_VALUE)[0]->getMemoryPtr()->GetPtr()); - const VectorDims& dst_data_dims = getChildEdgesAtPort(TOPK_VALUE)[0]->getMemory().getStaticDims(); - - dst_idx = reinterpret_cast(getChildEdgesAtPort(TOPK_INDEX)[0]->getMemoryPtr()->GetPtr()); - const VectorDims& dst_idx_dims = getChildEdgesAtPort(TOPK_INDEX)[0]->getMemory().getStaticDims(); - - if (dst_idx_dims[axis] != static_cast(src_k) || dst_data_dims[axis] != static_cast(src_k)) { - std::string errorMsg = "Output tensors dimension mismatch"; - IE_THROW() << errorMsg; - } + impl_desc_type impl_type; + if (mayiuse(cpu::x64::avx512_common)) { + impl_type = impl_desc_type::jit_avx512; + } else if (mayiuse(cpu::x64::avx2)) { + impl_type = impl_desc_type::jit_avx2; + } else if (mayiuse(cpu::x64::sse41)) { + impl_type = impl_desc_type::jit_sse42; } else { - std::string errorMsg = "Output tensors amount mismatch"; - IE_THROW() << errorMsg; + impl_type = impl_desc_type::ref; } - const VectorDims& in_dims = getParentEdgeAt(TOPK_DATA)->getMemory().getStaticDims(); + jit_mode = mayiuse(cpu::x64::sse41); + compile_kernel = jit_mode; - if (in_dims[axis] < static_cast(src_k)) - src_k = in_dims[axis]; + static const Precision supportedPrecision[] = { + Precision::FP32, + Precision::BF16, + Precision::I32, + Precision::I8, + Precision::U8 + }; - bool is_last_dim = false; - for (int j = in_dims.size() - 1; j >= 0; j--) { - if (in_dims[j] != 1) { - if (static_cast(j) == axis) - is_last_dim = true; - break; - } - } - - dim = static_cast(in_dims[axis]); - before_num = count(in_dims, 0, axis); - - if (src_k == 1) { - if (is_last_dim) { - if (mode_max) - top1(src, dst_data, dst_idx, in_dims); - else - top1(src, dst_data, dst_idx, in_dims); + Precision dataPrecision = getOriginalOutputPrecisionAtPort(TOPK_DATA); + if (dataPrecision == Precision::BF16 && !mayiuse(avx512_core)) + IE_THROW() << errorPrefix << " gets incorrect isa for BF16! AVX512 must be supported!"; + bool precisionSupported = std::find(std::begin(supportedPrecision), std::end(supportedPrecision), dataPrecision) + != std::end(supportedPrecision); + if (!precisionSupported) { + if (dataPrecision.is_float()) { + dataPrecision = Precision::FP32; } else { - if (mode_max) - top1_axis(src, dst_data, dst_idx, in_dims); - else - top1_axis(src, dst_data, dst_idx, in_dims); - } - } else { - if (is_last_dim) { - if (mode_max) - topk(src, dst_data, dst_idx, in_dims); - else - topk(src, dst_data, dst_idx, in_dims); - } else { - if (mode_max) - topk_axis(src, dst_data, dst_idx, in_dims); - else - topk_axis(src, dst_data, dst_idx, in_dims); + dataPrecision = Precision::I32; } } -} -bool MKLDNNTopKNode::created() const { - return getType() == TopK; -} + std::vector> dataFomats{ + {LayoutType::ncsp, LayoutType::ncsp}, + {LayoutType::nspc, LayoutType::nspc}, + {LayoutType::nCsp16c, LayoutType::nCsp16c}, + {LayoutType::nCsp8c, LayoutType::nCsp8c} + }; -bool MKLDNNTopKNode::needPrepareParams() const { - return false; -} - -void MKLDNNTopKNode::executeDynamicImpl(mkldnn::stream strm) { - execute(strm); + for (const auto &df : dataFomats) { + addSupportedPrimDesc({{df.first, dataPrecision}, {LayoutType::ncsp, Precision::I32}}, + {{df.second, dataPrecision}, {df.second, Precision::I32}}, + impl_type); + } } bool MKLDNNTopKNode::needShapeInfer() const { - const int kValue = reinterpret_cast(getParentEdgeAt(TOPK_K)->getMemoryPtr()->GetPtr())[0]; - return inputShapesModified() || kValue != src_k; + const int src_k = reinterpret_cast(getParentEdgeAt(TOPK_K)->getMemoryPtr()->GetPtr())[0]; + return inputShapesModified() || src_k != top_k; } std::vector MKLDNNTopKNode::shapeInfer() const { return MKLDNNNode::shapeInferGeneric(PortMask(1)); } -template class Compare2> -void MKLDNNTopKNode::top1_axis(const float* src_data, float* dst_data, int* dst_idx, VectorDims in_dims) { - int after_num = count(in_dims, axis + 1, in_dims.size()); - int first_index = 0; +bool MKLDNNTopKNode::needPrepareParams() const { + const int src_k = reinterpret_cast(getParentEdgeAt(TOPK_K)->getMemoryPtr()->GetPtr())[0]; + return inputShapesModified() || top_k != src_k; +} -#if defined(HAVE_SSE) || defined(HAVE_AVX2) || defined(HAVE_AVX512F) - parallel_for2d(before_num, after_num / block_size, [&](int i0, int ib1) { - int s_index = i0 * dim * after_num + ib1 * block_size; - vec_type_f vmax_val = _mm_uni_loadu_ps(src_data + s_index); - vec_type_i vindex_max_val = _mm_uni_setzero_si(); - for (int i2 = 1; i2 < dim; i2++) { - s_index += after_num; - vec_type_f vsrc = _mm_uni_loadu_ps(src_data + s_index); - vmask_type vmask = Compare1::cmp_ps(vsrc, vmax_val); - vmax_val = _mm_uni_blendv_ps(vmax_val, vsrc, vmask); +void MKLDNNTopKNode::prepareParams() { + auto &dstMemPtr = getChildEdgeAt(TOPK_DATA)->getMemoryPtr(); + auto &srcMemPtr = getParentEdgeAt(TOPK_DATA)->getMemoryPtr(); + if (!dstMemPtr || !dstMemPtr->GetPrimitivePtr()) + IE_THROW() << errorPrefix << " has not allocated destination memory."; + if (!srcMemPtr || !srcMemPtr->GetPrimitivePtr()) + IE_THROW() << errorPrefix << " has not allocate input memory."; + if (getSelectedPrimitiveDescriptor() == nullptr) + IE_THROW() << errorPrefix << " has nullable preferable primitive descriptor"; - vec_type_i vindex_cur_val = _mm_uni_set1_epi32(i2); -#if defined(HAVE_AVX512F) - vindex_max_val = _mm512_mask_blend_epi32(vmask, vindex_max_val, vindex_cur_val); -#else - vindex_max_val = _mm_uni_blendv_epi8(vindex_max_val, vindex_cur_val, _mm_uni_castps_si(vmask)); -#endif + src_dims = srcMemPtr->getDesc().getShape().getDims(); + dst_dims = dstMemPtr->getDesc().getShape().getDims(); + + if (isDynamicNode()) { + const int src_k = reinterpret_cast(getParentEdgeAt(TOPK_K)->getMemoryPtr()->GetPtr())[0]; + if (src_k > src_dims[axis]) + IE_THROW() << errorPrefix << " gets top_k out of range!"; + if (top_k != src_k) { + top_k = src_k; + } + } else { + top_k = reinterpret_cast(getParentEdgeAt(TOPK_K)->getMemoryPtr()->GetPtr())[0]; + } + + if (jit_mode) { + if (srcMemPtr->getDesc().hasLayoutType(LayoutType::ncsp)) { + layout = TopKLayoutType::topk_ncsp; + } else if (srcMemPtr->getDesc().hasLayoutType(LayoutType::nspc)) { + layout = TopKLayoutType::topk_nspc; + } else { + layout = TopKLayoutType::topk_blocked; + } + + auto selectedPD = getSelectedPrimitiveDescriptor(); + auto data_type = MKLDNNExtensionUtils::IEPrecisionToDataType(selectedPD->getConfig().inConfs[TOPK_DATA].desc->getPrecision()); + data_size = MKLDNNExtensionUtils::sizeOfDataType(data_type); + + topk_innermost = (layout == TopKLayoutType::topk_ncsp && axis == static_cast(dst_dims.size() - 1)) || + ((layout == TopKLayoutType::topk_nspc || layout == TopKLayoutType::topk_blocked) && axis == 1); + + if (mayiuse(cpu::x64::avx512_common)) { + blk_size = 16; + count_xmm = 16; // only 16 vector registers are valid in sse instructions + } else if (mayiuse(cpu::x64::sse41)) { + blk_size = 8; + count_xmm = 16; + } + + auto layout_dims = dstMemPtr->GetDescWithType()->getBlockDims(); + calc_dims_size(layout_dims); + + axis_dim = src_dims[axis]; + + // [case 1]: if 2 * (top_k + 1) + 2 <= count_xmm, thus top_k is small enough that the vector registers are sufficient + // to keep all necessary data for sorting, no need to load and store frequently, use inplace bubble sort; + // (horizotal sorting cases not included) + // [case 2]: only when topk is imposed on innermost dimsension of planar(ncsp/nspc) layout, should heap sort be used; + // [case 3]: by default, use bitonic sort when alg_cost_bitonic < alg_cost_bubble, otherwise use bubble sort. + // alg_cost_bitonic = (N / 4) * logN * (logN + 1) + // alg_cost_bubble = K * (K - 1) / 2 + (N - K) * K + // where, N = axis_dim, K = topk_k + // the above two alg_costs are not the exact implementation costs, yet it's proper to use them to decide + // which algorithm should be used for specific N and K. + if (!isDynamicNode()) { + if (top_k <= count_xmm / 2 - 2) { + algorithm = TopKAlgorithm::topk_bubble_sort; + bubble_inplace = layout == TopKLayoutType::topk_blocked && topk_innermost && top_k == 1 ? false : true; + } else if ((layout == TopKLayoutType::topk_ncsp || layout == TopKLayoutType::topk_nspc) && topk_innermost) { + algorithm = TopKAlgorithm::topk_heap_sort; + } else { + auto log_axis_dim = log2(axis_dim); + size_t alg_cost_bitonic = static_cast((axis_dim / 4.0f) * log_axis_dim * (log_axis_dim + 1)); + size_t alg_cost_bubble = top_k * (top_k - 1) / 2 + (axis_dim - top_k) * top_k; + if (alg_cost_bitonic < alg_cost_bubble) { + algorithm = TopKAlgorithm::topk_bitonic_sort; + } else { + algorithm = TopKAlgorithm::topk_bubble_sort; + bubble_inplace = false; + } } - if (dst_data) - _mm_uni_storeu_ps(dst_data + i0 * after_num + ib1 * block_size, vmax_val); - if (dst_idx) - _mm_uni_storeu_si(reinterpret_cast(dst_idx + i0 * after_num + ib1 * block_size), vindex_max_val); + } else { + if ((layout == TopKLayoutType::topk_ncsp || layout == TopKLayoutType::topk_nspc) && topk_innermost) { + algorithm = TopKAlgorithm::topk_heap_sort; + } else { + algorithm = TopKAlgorithm::topk_bubble_sort; + bubble_inplace = false; + } + } + + prepare_original_idx(); + + if (compile_kernel) { + auto jcp = jit_topk_config_params(); + jcp.precision = selectedPD->getConfig().inConfs[TOPK_DATA].desc->getPrecision(); + jcp.data_size = data_size; + jcp.blk_size = blk_size; + jcp.layout = layout; + jcp.top_k = top_k; + jcp.axis_dim = axis_dim; + jcp.mode_max = mode_max; + jcp.sort_index = sort_index; + jcp.topk_innermost = topk_innermost; + jcp.algorithm = algorithm; + jcp.bubble_inplace = bubble_inplace; + jcp.sort_stride = static_cast(I); + jcp.work_amount = static_cast(I); + + if (algorithm == TopKAlgorithm::topk_bitonic_sort) { + size_t src_count = srcMemPtr->GetDescWithType()->getPaddedElementsCount(); + vec_process_ptr.resize(src_count * data_size); + vec_process_idx_ptr.resize(src_count * sizeof(int32_t)); + + calc_bitonic_idx(axis_dim, jcp.bitonic_idx_cnt, true); + if (sort_index) { + calc_bitonic_idx(top_k, jcp.bitonic_k_idx_cnt, false); + } + } + + if (mayiuse(cpu::x64::avx512_common)) { + topk_kernel.reset(new jit_uni_topk_kernel_f32(jcp)); + } else if (mayiuse(cpu::x64::avx2)) { + topk_kernel.reset(new jit_uni_topk_kernel_f32(jcp)); + } else if (mayiuse(cpu::x64::sse41)) { + topk_kernel.reset(new jit_uni_topk_kernel_f32(jcp)); + } + + if (topk_kernel) + topk_kernel->create_ker(); + + compile_kernel = false; + } + } else { //reference mode + int j; + for (j = src_dims.size() - 1; j >= 0; j--) { + if (src_dims[j] != 1) + break; + } + dim = static_cast(src_dims[axis]); + before_num = count(src_dims, 0, axis); + } +} + +void MKLDNNTopKNode::executeDynamicImpl(mkldnn::stream strm) { + execute(strm); +} + +void MKLDNNTopKNode::execute(mkldnn::stream strm) { + auto &srcMemPtr = getParentEdgeAt(TOPK_DATA)->getMemoryPtr(); + auto &dstMemPtr = getChildEdgeAt(TOPK_DATA)->getMemoryPtr(); + auto &dstIndexesMemPtr = getChildEdgeAt(TOPK_INDEX)->getMemoryPtr(); + + const uint8_t *src_data = reinterpret_cast(srcMemPtr->GetPtr()); + uint8_t *dst_data = reinterpret_cast(dstMemPtr->GetPtr()); + uint8_t *dst_idx = reinterpret_cast(dstIndexesMemPtr->GetPtr()); + + if (jit_mode) { + topk_process(src_data, dst_data, dst_idx); + } else { + if (layout == TopKLayoutType::topk_ncsp) { + auto in_ptr = reinterpret_cast(src_data); + auto out_ptr = reinterpret_cast(dst_data); + auto out_idx_ptr = reinterpret_cast(dst_idx); + topk_ref(in_ptr, out_ptr, out_idx_ptr); + } else { + IE_THROW() << errorPrefix << "only support plain layout on machine w/o sse42."; + } + } +} + +void MKLDNNTopKNode::topk_process(const uint8_t *in_ptr, uint8_t *out_ptr, uint8_t *out_idx_ptr) { + uint8_t *process_ptr = vec_process_ptr.data(); + uint8_t *process_idx_ptr = vec_process_idx_ptr.data(); + + // [blocked layout with topk on C] + if (layout == TopKLayoutType::topk_blocked && topk_innermost) { + size_t IA = div_up(src_dims[1], blk_size); + size_t OA = div_up(dst_dims[1], blk_size); + if (algorithm == TopKAlgorithm::topk_bubble_sort) { + parallel_for2d(O, I, [&](size_t o, size_t i) { + const uint8_t *in_ptr_a = in_ptr + (o * IA * I + i) * blk_size * data_size; + uint8_t *out_ptr_a = out_ptr + (o * OA * I + i) * blk_size * data_size; + uint8_t *out_idx_ptr_a = out_idx_ptr + (o * OA * I + i) * blk_size * sizeof(int32_t); + size_t work_amount = 1; + topk_kernel_process(in_ptr_a, out_ptr_a, out_idx_ptr_a, NULL, NULL, work_amount); + }); + } else if (algorithm == TopKAlgorithm::topk_bitonic_sort) { + parallel_for(O, [&](size_t o) { + const uint8_t *in_ptr_a = in_ptr + o * IA * I * blk_size * data_size; + uint8_t *process_ptr_a = process_ptr + o * IA * I * blk_size * data_size; + uint8_t *process_idx_ptr_a = process_idx_ptr + o * IA * I * blk_size * sizeof(int32_t); + uint8_t *out_ptr_a = out_ptr + o * OA * I * blk_size * data_size; + uint8_t *out_idx_ptr_a = out_idx_ptr + o * OA * I * blk_size * sizeof(int32_t); + size_t work_amount = I; + topk_kernel_process(in_ptr_a, out_ptr_a, out_idx_ptr_a, process_ptr_a, process_idx_ptr_a, work_amount); + }); + } + } else { // [planar layout] [blocked layout with topk on non-C] + parallel_for2d(O, I / blk_size, [&](size_t o, size_t k) { + const uint8_t *in_ptr_a = in_ptr + (o * A * I + k * blk_size) * data_size; + uint8_t *process_ptr_a = process_ptr + (o * A * I + k * blk_size) * data_size; + uint8_t *process_idx_ptr_a = process_idx_ptr + (o * A * I + k * blk_size) * sizeof(int32_t); + uint8_t *out_ptr_a = out_ptr + (o * top_k * I + k * blk_size) * data_size; + uint8_t *out_idx_ptr_a = out_idx_ptr + (o * top_k * I + k * blk_size) * sizeof(int32_t); + size_t work_amount = blk_size; + topk_kernel_process(in_ptr_a, out_ptr_a, out_idx_ptr_a, process_ptr_a, process_idx_ptr_a, work_amount); }); - first_index = after_num / block_size * block_size; -#endif - int rest = after_num - first_index; - parallel_for2d(before_num, rest, [&](int i0, int i1) { - int index_max_val = 0; - int s_index = i0 * dim * after_num + first_index + i1; - float max_val = src_data[s_index]; - for (int i2 = 1; i2 < dim; i2++) { - s_index += after_num; - if (Compare2()(src_data[s_index], max_val)) { - max_val = src_data[s_index]; - index_max_val = i2; - } + + size_t tail_start = I / blk_size * blk_size; + size_t work_amount = I - tail_start; + if (work_amount) { + parallel_for(O, [&](size_t o) { + const uint8_t *in_ptr_a = in_ptr + (o * A * I + tail_start) * data_size; + uint8_t *process_ptr_a = process_ptr + (o * A * I + tail_start) * data_size; + uint8_t *process_idx_ptr_a = process_idx_ptr + (o * A * I + tail_start) * sizeof(int32_t); + uint8_t *out_ptr_a = out_ptr + (o * top_k * I + tail_start) * data_size; + uint8_t *out_idx_ptr_a = out_idx_ptr + (o * top_k * I + tail_start) * sizeof(int32_t); + topk_kernel_process(in_ptr_a, out_ptr_a, out_idx_ptr_a, process_ptr_a, process_idx_ptr_a, work_amount); + }); } - if (dst_data) - dst_data[i0 * after_num + first_index + i1] = max_val; - if (dst_idx) - dst_idx[i0 * after_num + first_index + i1] = index_max_val; - }); + } } -template