From ade90b204362af44910ccc4986df1dcbb0428734 Mon Sep 17 00:00:00 2001 From: Chenhu Wang Date: Thu, 17 Nov 2022 19:31:08 +0800 Subject: [PATCH] [CPU] MVN nspc kernel optimization (#13268) --- src/plugins/intel_cpu/src/nodes/mvn.cpp | 865 ++++++++++++++---- src/plugins/intel_cpu/src/nodes/mvn.h | 15 +- .../functional/single_layer_tests/mvn.cpp | 8 +- 3 files changed, 697 insertions(+), 191 deletions(-) diff --git a/src/plugins/intel_cpu/src/nodes/mvn.cpp b/src/plugins/intel_cpu/src/nodes/mvn.cpp index b4dd367ee00..d4d046f240f 100644 --- a/src/plugins/intel_cpu/src/nodes/mvn.cpp +++ b/src/plugins/intel_cpu/src/nodes/mvn.cpp @@ -110,15 +110,17 @@ struct jit_uni_mvn_mean_variance_kernel_f32 : public jit_uni_mvn_mean_variance_k } void generate() override { - tail_step = jcp_.planar_layout ? (jcp_.D * jcp_.H * jcp_.W) - ((jcp_.D * jcp_.H * jcp_.W) / vector_step) * vector_step : + tail_step = jcp_.layout == MVNLayoutType::mvn_planar ? (jcp_.D * jcp_.H * jcp_.W) - ((jcp_.D * jcp_.H * jcp_.W) / vector_step) * vector_step : jcp_.C - (jcp_.C / vector_step) * vector_step; Precision dst_prc = isFloatCompatible(jcp_.src_prc) ? Precision::FP32 : Precision::I32; load_vector_emitter.reset(new jit_load_emitter(this, isa, jcp_.src_prc, dst_prc, vector_step)); load_tail_emitter.reset(new jit_load_emitter(this, isa, jcp_.src_prc, dst_prc, tail_step)); - load_tail_with_fill_emitter.reset(new jit_load_emitter(this, isa, jcp_.src_prc, dst_prc, tail_step, Precision::FP32, true)); + load_tail_with_fill_emitter.reset(new jit_load_emitter(this, isa, jcp_.src_prc, dst_prc, tail_step, Precision::FP32, true, "zero")); + load_scalar_with_fill_emitter.reset(new jit_load_emitter(this, isa, jcp_.src_prc, dst_prc, scalar_step, Precision::FP32, true, "zero")); this->preamble(); + mov(reg_table, l_table); mov(reg_src, ptr[reg_params + GET_OFF(src)]); if (jcp_.normalize_variance) { mov(reg_mean, ptr[reg_params + GET_OFF(mean)]); @@ -133,7 +135,7 @@ struct jit_uni_mvn_mean_variance_kernel_f32 : public jit_uni_mvn_mean_variance_k mov(reg_oc_off, ptr[reg_params + GET_OFF(oc_off)]); if (jcp_.normalize_variance) { - if (jcp_.planar_layout || jcp_.across_channels) { + if (jcp_.layout == MVNLayoutType::mvn_planar || jcp_.across_channels) { uni_vbroadcastss(vmm_mean, ptr[reg_mean]); } else { uni_vmovups(vmm_mean, ptr[reg_mean]); @@ -142,37 +144,24 @@ struct jit_uni_mvn_mean_variance_kernel_f32 : public jit_uni_mvn_mean_variance_k load_pool_gpr_idxs = {static_cast(reg_load_store_mask.getIdx()), static_cast(reg_load_table.getIdx())}; - if (jcp_.planar_layout) { + if (jcp_.layout == MVNLayoutType::mvn_planar) { worker_unroll(); if (tail_step != 0) { - worker_tail_planar(); + worker_partial(false, true); } // hsum+store if (!jcp_.normalize_variance && !isFloatCompatible(jcp_.src_prc)) uni_vcvtdq2ps(vmm_sum, vmm_sum); Vmm vmm_dst = jcp_.normalize_variance ? vmm_variance : vmm_sum; - if (isa == cpu::x64::sse41) { - hsum_store(vmm_dst); - } else if (isa == cpu::x64::avx2) { - Xbyak::Ymm ymm_sum = Xbyak::Ymm(vmm_dst.getIdx()); - vextractf128(xmm_aux1, ymm_sum, 0); - vextractf128(xmm_aux2, ymm_sum, 1); - uni_vaddps(xmm_aux1, xmm_aux1, xmm_aux2); - hsum_store(xmm_aux1); - } else { - Xbyak::Zmm zmm_sum = Xbyak::Zmm(vmm_dst.getIdx()); - vextractf32x4(xmm_aux1, zmm_sum, 0); - vextractf32x4(xmm_aux2, zmm_sum, 1); - uni_vaddps(xmm_aux1, xmm_aux1, xmm_aux2); - vextractf32x4(xmm_aux2, zmm_sum, 2); - vextractf32x4(xmm_aux3, zmm_sum, 3); - uni_vaddps(xmm_aux2, xmm_aux2, xmm_aux3); - uni_vaddps(xmm_aux1, xmm_aux1, xmm_aux2); - hsum_store(xmm_aux1); - } + reduce_sum_store_vmm(vmm_dst.getIdx()); + } else if (jcp_.layout == MVNLayoutType::mvn_by_channel) { + if (jcp_.across_channels) + nspc_ac_ker(); + else + nspc_pc_ker(); } else { - // blk+nspc + // blk int repeats = (isa == cpu::x64::sse41) ? 2 : 1; // block size is also 8 on cpu::x64::sse41 with two step process int sse42_step = 4; for (int i = 0; i < repeats; i++) { @@ -257,6 +246,9 @@ struct jit_uni_mvn_mean_variance_kernel_f32 : public jit_uni_mvn_mean_variance_k load_vector_emitter->emit_data(); load_tail_emitter->emit_data(); load_tail_with_fill_emitter->emit_data(); + load_scalar_with_fill_emitter->emit_data(); + + prepare_table(); } private: @@ -266,6 +258,7 @@ private: const int vlen = cpu_isa_traits::vlen; const int vector_step = vlen / sizeof(float); int tail_step = 0; + int scalar_step = 1; Xbyak::Reg64 reg_src = r8; Xbyak::Reg64 reg_mean = r9; @@ -279,24 +272,229 @@ private: Xbyak::Reg64 reg_aux = r15; Xbyak::Reg64 reg_oc_off = rax; + Xbyak::Reg64 reg_table = rdx; + Xbyak::Label l_table; - Vmm vmm_val = Vmm(0); - Vmm vmm_mean = Vmm(1); - Vmm vmm_variance = Vmm(2); + Vmm vmm_val = Vmm(1); + Vmm vmm_mean = Vmm(2); + Vmm vmm_variance = Vmm(3); Vmm vmm_sum = vmm_mean; - Xbyak::Xmm xmm_aux1 = Xbyak::Xmm(3); - Xbyak::Xmm xmm_aux2 = Xbyak::Xmm(4); - Xbyak::Xmm xmm_aux3 = Xbyak::Xmm(5); - Vmm vmm_zero = Vmm(6); + Xbyak::Xmm xmm_aux1 = Xbyak::Xmm(4); + Xbyak::Xmm xmm_aux2 = Xbyak::Xmm(5); + Xbyak::Xmm xmm_aux3 = Xbyak::Xmm(6); + Vmm vmm_zero = Vmm(0); + // 8-15 for unloop Xbyak::Opmask k_mask = Xbyak::Opmask(7); std::unique_ptr load_vector_emitter = nullptr; std::unique_ptr load_tail_emitter = nullptr; std::unique_ptr load_tail_with_fill_emitter = nullptr; + std::unique_ptr load_scalar_with_fill_emitter = nullptr; std::vector load_pool_gpr_idxs; + // nspc across channel + inline void nspc_ac_ker() { + Xbyak::Label loop_label; + Xbyak::Label loop_end_label; + Xbyak::Label scalar_loop_label; + Xbyak::Label scalar_loop_end_label; + L(loop_label); + { + cmp(reg_work_amount, vector_step); + jl(loop_end_label, T_NEAR); + + worker_full_size(); + add(reg_src, vector_step * jcp_.src_data_size); + + sub(reg_work_amount, vector_step); + jmp(loop_label, T_NEAR); + } + L(loop_end_label); + + L(scalar_loop_label); + { + cmp(reg_work_amount, 1); + jl(scalar_loop_end_label, T_NEAR); + + worker_partial(true, true); + add(reg_src, scalar_step * jcp_.src_data_size); + + sub(reg_work_amount, scalar_step); + jmp(scalar_loop_label, T_NEAR); + } + L(scalar_loop_end_label); + + if (!jcp_.normalize_variance && !isFloatCompatible(jcp_.src_prc)) + uni_vcvtdq2ps(vmm_sum, vmm_sum); + Vmm vmm_dst = jcp_.normalize_variance ? vmm_variance : vmm_sum; + reduce_sum_store_vmm(vmm_dst.getIdx()); + } + + // nspc per channel with unroll + inline void nspc_pc_ker() { + // 4 unroll vector + size_t unroll_size = 4; + size_t vec_num = div_up(jcp_.C, vector_step); + unroll_size = vec_num >= unroll_size ? unroll_size : vec_num; + size_t unroll_number = div_up(vec_num, unroll_size); + + int ur_base = 4; + Xbyak::Reg64 reg_src_aux = reg_stride; + Xbyak::Reg64 reg_work_amount_bk = rbx; + mov(reg_work_amount_bk, reg_work_amount); + for (int ur_num = 0; ur_num < unroll_number; ur_num++) { + // 4-15 for unroll. 4-7 for src, 8-11 for m/v sum, 12-15 for mean + int ur_offset_elt = ur_num * unroll_size * vector_step; + int ur_offset = ur_offset_elt * sizeof(float); + size_t unroll_size_rt = std::min(vec_num - ur_num * unroll_size, unroll_size); + size_t elt_num = std::min(jcp_.C - ur_num * unroll_size * vector_step, unroll_size * vector_step); + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vpxor(Vmm(ur_base + 4 + ur_size), Vmm(ur_base + 4 + ur_size), Vmm(ur_base + 4 + ur_size)); + } + if (jcp_.normalize_variance) { + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vmovups(Vmm(ur_base + 8 + ur_size), ptr[reg_mean + ur_offset + ur_size * vector_step * sizeof(float)]); + } + } + + mov(reg_src_aux, reg_src); + mov(reg_work_amount, reg_work_amount_bk); + + Xbyak::Label loop_label; + Xbyak::Label loop_end_label; + L(loop_label); + { + cmp(reg_work_amount, 0); + jle(loop_end_label, T_NEAR); + + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + bool is_tails = ur_offset_elt + ur_size * vector_step + vector_step > jcp_.C; + if (is_tails) { + load_tail_emitter->emit_code({static_cast(reg_src_aux.getIdx())}, + {static_cast(ur_base + ur_size)}, {}, {load_pool_gpr_idxs}); + add(reg_src_aux, tail_step * jcp_.src_data_size); + } else { + load_vector_emitter->emit_code({static_cast(reg_src_aux.getIdx())}, + {static_cast(ur_base + ur_size)}, {}, {load_pool_gpr_idxs}); + add(reg_src_aux, vector_step * jcp_.src_data_size); + } + } + add(reg_src_aux, (jcp_.C - elt_num) * jcp_.src_data_size); + prefetcht0(ptr[reg_src_aux]); + + if (jcp_.normalize_variance) { + if (!isFloatCompatible(jcp_.src_prc)) { + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vcvtdq2ps(Vmm(ur_base + ur_size), Vmm(ur_base + ur_size)); + } + } + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vsubps(Vmm(ur_base + ur_size), Vmm(ur_base + ur_size), Vmm(ur_base + 8 + ur_size)); + } + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vfmadd231ps(Vmm(ur_base + 4 + ur_size), Vmm(ur_base + ur_size), Vmm(ur_base + ur_size)); + } + } else { + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + if (!isFloatCompatible(jcp_.src_prc)) + uni_vpaddd(Vmm(ur_base + 4 + ur_size), Vmm(ur_base + 4 + ur_size), Vmm(ur_base + ur_size)); + else + uni_vaddps(Vmm(ur_base + 4 + ur_size), Vmm(ur_base + 4 + ur_size), Vmm(ur_base + ur_size)); + } + } + + sub(reg_work_amount, 1); + jmp(loop_label, T_NEAR); + } + L(loop_end_label); + + // store sum/variance + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + if (jcp_.normalize_variance) { + uni_vmovups(ptr[reg_variance + ur_offset + ur_size * vector_step * sizeof(float)], Vmm(ur_base + 4 + ur_size)); + } else { + if (!isFloatCompatible(jcp_.src_prc)) + uni_vcvtdq2ps(Vmm(ur_base + 4 + ur_size), Vmm(ur_base + 4 + ur_size)); + uni_vmovups(ptr[reg_sum + ur_offset + ur_size * vector_step * sizeof(float)], Vmm(ur_base + 4 + ur_size)); + } + } + + add(reg_src, unroll_size_rt * vector_step * jcp_.src_data_size); + } + } + + inline void worker_unroll(bool is_tail = false) { + // if mean(sum) for continous data, then fast pass for major part + if (!jcp_.normalize_variance && jcp_.layout == MVNLayoutType::mvn_planar) { + Vmm vmm_one = Vmm(15); + // i8/u8 fast path + if (mayiuse(avx512_core_vnni) && jcp_.src_data_size == 1) { + uni_vmovups(vmm_one, ptr[reg_table]); + Xbyak::Label loop_8bit_label; + Xbyak::Label loop_8bit_end_label; + L(loop_8bit_label); + { + cmp(reg_work_amount, 4); + jl(loop_8bit_end_label, T_NEAR); + + if (jcp_.src_prc == Precision::I8) { + vpdpbusd(vmm_sum, vmm_one, ptr[reg_src]); + } else { + uni_vmovdqu(vmm_val, ptr[reg_src]); + vpdpbusd(vmm_sum, vmm_val, vmm_one); + } + + add(reg_src, vlen); + sub(reg_work_amount, 4); + + jmp(loop_8bit_label, T_NEAR); + } + L(loop_8bit_end_label); + } + // bf16 fast path + if (mayiuse(avx512_core_bf16) && jcp_.src_prc == Precision::BF16) { + uni_vmovups(vmm_one, ptr[reg_table]); + Xbyak::Label loop_bf16_label; + Xbyak::Label loop_bf16_end_label; + L(loop_bf16_label); + { + cmp(reg_work_amount, 2); + jl(loop_bf16_end_label, T_NEAR); + + vdpbf16ps(vmm_sum, vmm_one, ptr[reg_src]); + + add(reg_src, vlen); + sub(reg_work_amount, 2); + + jmp(loop_bf16_label, T_NEAR); + } + L(loop_bf16_end_label); + } + } + + Xbyak::Label loop_label; + Xbyak::Label loop_end_label; + L(loop_label); + { + cmp(reg_work_amount, 0); + jle(loop_end_label, T_NEAR); + + if (jcp_.layout != MVNLayoutType::mvn_planar && is_tail) { + worker_partial(false, false); + } else { + worker_full_size(); + } + + add(reg_src, reg_stride); + sub(reg_work_amount, 1); + + jmp(loop_label, T_NEAR); + } + L(loop_end_label); + } + inline void worker_full_size() { load_vector_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val.getIdx())}, {}, {load_pool_gpr_idxs}); @@ -317,75 +515,42 @@ private: } } - inline void worker_tail_blk() { - load_tail_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val.getIdx())}, - {}, {load_pool_gpr_idxs}); - - if (jcp_.normalize_variance) { - // all with float - if (!isFloatCompatible(jcp_.src_prc)) - uni_vcvtdq2ps(vmm_val, vmm_val); - - uni_vsubps(vmm_val, vmm_val, vmm_mean); - uni_vfmadd231ps(vmm_variance, vmm_val, vmm_val); - } else { - // for sum, int execute prc for int-family data type - if (!isFloatCompatible(jcp_.src_prc)) - uni_vpaddd(vmm_sum, vmm_sum, vmm_val); - else - uni_vaddps(vmm_sum, vmm_sum, vmm_val); - } - } - - inline void worker_unroll(bool is_tail = false) { - Xbyak::Label loop_label; - Xbyak::Label loop_end_label; - L(loop_label); - { - cmp(reg_work_amount, 0); - jle(loop_end_label, T_NEAR); - - if (!jcp_.planar_layout && is_tail) { - worker_tail_blk(); - } else { - worker_full_size(); - } - - add(reg_src, reg_stride); - sub(reg_work_amount, 1); - - jmp(loop_label, T_NEAR); - } - L(loop_end_label); - } - - inline void worker_tail_planar() { - load_tail_with_fill_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val.getIdx())}, + // needed and supported case: 1. scalar with zero pad. 2. tails w/ or w/o zero pad + inline void worker_partial(bool is_scalar, bool is_zero_pad) { + if (is_scalar) { + load_scalar_with_fill_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val.getIdx())}, {}, {load_pool_gpr_idxs}); - + } else { + if (is_zero_pad) + load_tail_with_fill_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val.getIdx())}, + {}, {load_pool_gpr_idxs}); + else + load_tail_emitter->emit_code({static_cast(reg_src.getIdx())}, {static_cast(vmm_val.getIdx())}, + {}, {load_pool_gpr_idxs}); + } if (jcp_.normalize_variance) { if (!isFloatCompatible(jcp_.src_prc)) uni_vcvtdq2ps(vmm_val, vmm_val); - uni_vsubps(vmm_val, vmm_val, vmm_mean); - - uni_vpxor(vmm_zero, vmm_zero, vmm_zero); - if (isa == cpu::x64::sse41) { - uint8 imm = 1; - imm = ~((imm << tail_step) - imm); - blendps(vmm_val, vmm_zero, imm); - } else if (isa == cpu::x64::avx2) { - uint8 imm = 1; - imm = ~((imm << tail_step) - imm); - vblendps(vmm_val, vmm_val, vmm_zero, imm); - } else if (isa == cpu::x64::avx512_core) { - uint64_t tail_mask = 1; - tail_mask = ~((tail_mask << tail_step) - tail_mask); - mov(reg_aux, tail_mask); - kmovq(k_mask, reg_aux); - vblendmps(vmm_val | k_mask, vmm_val, vmm_zero); + if (is_zero_pad) { + int elt_num = is_scalar ? 1 : tail_step; + uni_vpxor(vmm_zero, vmm_zero, vmm_zero); + if (isa == cpu::x64::sse41) { + uint8 imm = 1; + imm = ~((imm << elt_num) - imm); + blendps(vmm_val, vmm_zero, imm); + } else if (isa == cpu::x64::avx2) { + uint8 imm = 1; + imm = ~((imm << elt_num) - imm); + vblendps(vmm_val, vmm_val, vmm_zero, imm); + } else if (isa == cpu::x64::avx512_core) { + uint64_t tail_mask = 1; + tail_mask = ~((tail_mask << elt_num) - tail_mask); + mov(reg_aux, tail_mask); + kmovq(k_mask, reg_aux); + vblendmps(vmm_val | k_mask, vmm_val, vmm_zero); + } } - uni_vfmadd231ps(vmm_variance, vmm_val, vmm_val); } else { if (!isFloatCompatible(jcp_.src_prc)) @@ -395,9 +560,9 @@ private: } } - inline void hsum_store(Xbyak::Xmm xmm_sum) { - uni_vmovshdup(xmm_aux3, xmm_sum); // sum:1,2,3,4; aux3:2,2,4,4 - uni_vaddps(xmm_sum, xmm_sum, xmm_aux3); // sum:1+2,2+2,3+4,4+4 + inline void reduce_sum_store_xmm(Xbyak::Xmm xmm_sum) { + uni_vmovshdup(xmm_aux3, xmm_sum); // sum:1,2,3,4; aux3:2,2,4,4 + uni_vaddps(xmm_sum, xmm_sum, xmm_aux3); // sum:1+2,2+2,3+4,4+4 uni_vmovhlps(xmm_aux3, xmm_aux3, xmm_sum); // aux3:3+4,4+4,4,4 uni_vaddps(xmm_sum, xmm_sum, xmm_aux3); // sum:1+2+3+4,... if (jcp_.normalize_variance) { @@ -406,6 +571,50 @@ private: uni_vmovss(ptr[reg_sum], xmm_sum); } } + + inline void reduce_sum_store_vmm(int vmm_idx) { + if (isa == cpu::x64::sse41) { + reduce_sum_store_xmm(Xmm(vmm_idx)); + } else if (isa == cpu::x64::avx2) { + Xbyak::Ymm ymm_sum = Xbyak::Ymm(vmm_idx); + vextractf128(xmm_aux1, ymm_sum, 0); + vextractf128(xmm_aux2, ymm_sum, 1); + uni_vaddps(xmm_aux1, xmm_aux1, xmm_aux2); + reduce_sum_store_xmm(xmm_aux1); + } else { + Xbyak::Zmm zmm_sum = Xbyak::Zmm(vmm_idx); + vextractf32x4(xmm_aux1, zmm_sum, 0); + vextractf32x4(xmm_aux2, zmm_sum, 1); + uni_vaddps(xmm_aux1, xmm_aux1, xmm_aux2); + vextractf32x4(xmm_aux2, zmm_sum, 2); + vextractf32x4(xmm_aux3, zmm_sum, 3); + uni_vaddps(xmm_aux2, xmm_aux2, xmm_aux3); + uni_vaddps(xmm_aux1, xmm_aux1, xmm_aux2); + reduce_sum_store_xmm(xmm_aux1); + } + } + + void prepare_table() { + const unsigned int cvals[] = { + // 4 * 8 = 32 bit + 0x01010101, // one byte + 0x3f803f80 // one bf16 + }; + + align(64); + L(l_table); + + if (mayiuse(avx512_core_vnni) && (jcp_.src_prc == Precision::U8 || jcp_.src_prc == Precision::I8)) { + for (size_t d = 0; d < vector_step; ++d) { + dd(cvals[0]); + } + } + if (mayiuse(avx512_core_bf16) && jcp_.src_prc == Precision::BF16) { + for (size_t d = 0; d < vector_step; ++d) { + dd(cvals[1]); + } + } + } }; // mean,variance->mvn @@ -413,7 +622,22 @@ template struct jit_uni_mvn_kernel_f32 : public jit_uni_mvn_kernel, public jit_generator { DECLARE_CPU_JIT_AUX_FUNCTIONS(jit_uni_mvn_kernel_f32) - explicit jit_uni_mvn_kernel_f32(jit_mvn_config_params jcp, const dnnl_primitive_attr &attr) : jit_uni_mvn_kernel(jcp, attr), jit_generator(jit_name()) {} + explicit jit_uni_mvn_kernel_f32(jit_mvn_config_params jcp, const dnnl_primitive_attr &attr) : jit_uni_mvn_kernel(jcp, attr), jit_generator(jit_name()) { + const auto &p = attr_.post_ops_; + bool opt_scaleshift_applicable = jcp_.layout == MVNLayoutType::mvn_by_channel && isa == cpu::x64::avx512_core && !jcp_.across_channels; + if (opt_scaleshift_applicable) { + for (int i = 0; i < p.len(); i++) { + auto &post_op = p.entry_[i]; + if (post_op.is_depthwise()) { + if (0 == i && post_op.depthwise.alg == alg_kind::depthwise_scale_shift) { + optimized_scaleshift_num = 1; + } else if (1 == i && optimized_scaleshift_num == 1 && post_op.depthwise.alg == alg_kind::depthwise_scale_shift) { + optimized_scaleshift_num = 2; + } + } + } + } + } void create_ker() override { jit_generator::create_kernel(); @@ -436,7 +660,7 @@ struct jit_uni_mvn_kernel_f32 : public jit_uni_mvn_kernel, public jit_generator } } - tail_step = jcp_.planar_layout ? (jcp_.D * jcp_.H * jcp_.W) - ((jcp_.D * jcp_.H * jcp_.W) / vector_step) * vector_step : + tail_step = jcp_.layout == MVNLayoutType::mvn_planar ? (jcp_.D * jcp_.H * jcp_.W) - ((jcp_.D * jcp_.H * jcp_.W) / vector_step) * vector_step : jcp_.C - (jcp_.C / vector_step) * vector_step; load_vector_emitter.reset(new jit_load_emitter(this, isa, jcp_.src_prc, Precision::FP32, vector_step)); @@ -457,7 +681,7 @@ struct jit_uni_mvn_kernel_f32 : public jit_uni_mvn_kernel, public jit_generator mov(reg_dst_stride, ptr[reg_params + GET_OFF(dst_stride)]); mov(reg_oc_off, ptr[reg_params + GET_OFF(oc_off)]); - if (jcp_.planar_layout || jcp_.across_channels) { + if (jcp_.layout == MVNLayoutType::mvn_planar || jcp_.across_channels) { uni_vbroadcastss(vmm_mean, ptr[reg_mean]); if (jcp_.normalize_variance) uni_vbroadcastss(vmm_variance_inv, ptr[reg_variance_inv]); @@ -473,13 +697,18 @@ struct jit_uni_mvn_kernel_f32 : public jit_uni_mvn_kernel, public jit_generator store_pool_gpr_idxs = {static_cast(reg_load_store_mask.getIdx())}; store_pool_vec_idxs = {static_cast(vmm_zero.getIdx()), static_cast(vmm_val.getIdx())}; - if (jcp_.planar_layout) { + if (jcp_.layout == MVNLayoutType::mvn_planar) { worker_mvn_unroll(); if (tail_step != 0) { worker_mvn(true); } + } else if (jcp_.layout == MVNLayoutType::mvn_by_channel) { + if (jcp_.across_channels) + norm_nspc_ac_ker(); + else + norm_nspc_pc_ker(); } else { - // blk+nspc + // blk int repeats = (isa == cpu::x64::sse41) ? 2 : 1; // block size is also 8 on cpu::x64::sse41 for (int i = 0; i < repeats; i++) { int offset_sse42 = i * 4; @@ -567,13 +796,13 @@ private: Xbyak::Reg64 reg_load_table = r15; Xbyak::Reg64 reg_load_store_mask = rbp; - Vmm vmm_val = Vmm(1); - Vmm vmm_mean = Vmm(0); - Vmm vmm_variance_inv = Vmm(2); - Vmm vmm_zero = Vmm(3); + Vmm vmm_val = Vmm(3); + Vmm vmm_mean = Vmm(4); + Vmm vmm_variance_inv = Vmm(5); + Vmm vmm_zero = Vmm(2); - Vmm vmm_d_weights = Vmm(5); - Vmm vmm_d_bias = Vmm(6); + Vmm vmm_d_weights = Vmm(0); + Vmm vmm_d_bias = Vmm(1); std::unique_ptr load_vector_emitter = nullptr; std::unique_ptr load_tail_emitter = nullptr; @@ -588,6 +817,174 @@ private: std::vector store_pool_vec_idxs; std::vector load_pool_gpr_idxs; + // nspc norm per channel with unroll + inline void norm_nspc_pc_ker() { + // 4 unroll vector + size_t unroll_size = 4; + size_t vec_num = div_up(jcp_.C, vector_step); + unroll_size = vec_num >= unroll_size ? unroll_size : vec_num; + size_t unroll_number = div_up(vec_num, unroll_size); + + int ur_base = 4; + Xbyak::Reg64 reg_src_aux = reg_src_stride; + Xbyak::Reg64 reg_dst_aux = reg_dst_stride; + // 2 abi + Xbyak::Reg64 reg_work_amount_bk = rcx; + Xbyak::Reg64 reg_oc_off_bk = rdi; + mov(reg_oc_off_bk, reg_oc_off); + mov(reg_work_amount_bk, reg_work_amount); + for (int ur_num = 0; ur_num < unroll_number; ur_num++) { + // 4-15 for unroll. 4-7 for src, 8-11 for m, 12-15 for v + int ur_offset_elt = ur_num * unroll_size * vector_step; + int ur_offset = ur_offset_elt * sizeof(float); + size_t unroll_size_rt = std::min(vec_num - ur_num * unroll_size, unroll_size); + size_t elt_num = std::min(jcp_.C - ur_num * unroll_size * vector_step, unroll_size * vector_step); + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vmovups(Vmm(ur_base + 4 + ur_size), ptr[reg_mean + ur_offset + ur_size * vector_step * sizeof(float)]); + } + if (jcp_.normalize_variance) { + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vmovups(Vmm(ur_base + 8 + ur_size), ptr[reg_variance_inv + ur_offset + ur_size * vector_step * sizeof(float)]); + } + } + // optimized scaleshift + size_t post_ops_data_offset = 0; + for (int i = 0; i < optimized_scaleshift_num; i++) { + mov(reg_d_weights, ptr[reg_post_ops_data + post_ops_data_offset]); + add(reg_d_weights, ur_offset); + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vmovups(Vmm(16 + i * 4 + ur_size), ptr[reg_d_weights]); + add(reg_d_weights, vector_step * sizeof(float)); + } + mov(reg_d_bias, ptr[reg_post_ops_data + post_ops_data_offset]); + add(reg_d_bias, ur_offset + jcp_.C * sizeof(float)); + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vmovups(Vmm(24 + i * 4 + ur_size), ptr[reg_d_bias]); + add(reg_d_bias, vector_step * sizeof(float)); + } + post_ops_data_offset += sizeof(float*); + } + + mov(reg_src_aux, reg_src); + mov(reg_dst_aux, reg_dst); + mov(reg_work_amount, reg_work_amount_bk); + mov(reg_oc_off, reg_oc_off_bk); + + Xbyak::Label loop_label; + Xbyak::Label loop_end_label; + L(loop_label); + { + cmp(reg_work_amount, 0); + jle(loop_end_label, T_NEAR); + + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + bool is_tails = ur_offset_elt + ur_size * vector_step + vector_step > jcp_.C; + if (is_tails) { + load_tail_emitter->emit_code({static_cast(reg_src_aux.getIdx())}, + {static_cast(ur_base + ur_size)}, {}, {load_pool_gpr_idxs}); + add(reg_src_aux, tail_step * jcp_.src_data_size); + } else { + load_vector_emitter->emit_code({static_cast(reg_src_aux.getIdx())}, + {static_cast(ur_base + ur_size)}, {}, {load_pool_gpr_idxs}); + add(reg_src_aux, vector_step * jcp_.src_data_size); + } + } + add(reg_src_aux, (jcp_.C - elt_num) * jcp_.src_data_size); + prefetcht0(ptr[reg_src_aux]); + + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vsubps(Vmm(ur_base + ur_size), Vmm(ur_base + ur_size), Vmm(ur_base + 4 + ur_size)); + } + if (jcp_.normalize_variance) { + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vmulps(Vmm(ur_base + ur_size), Vmm(ur_base + ur_size), Vmm(ur_base + 8 + ur_size)); + } + } + + for (int i = 0; i < optimized_scaleshift_num; i++) { + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + uni_vfmadd132ps(Vmm(ur_base + ur_size), Vmm(24 + i * 4 + ur_size), Vmm(16 + i * 4 + ur_size)); + } + } + + if (attr_.post_ops_.len() != 0) { + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + apply_post_ops(jcp_.dst_prc, ur_base + ur_size, false); + bool is_tails = ur_offset_elt + ur_size * vector_step + vector_step > jcp_.C; + if (is_tails) + add(reg_oc_off, tail_step * sizeof(float)); + else + add(reg_oc_off, vector_step * sizeof(float)); + } + } + + for (int ur_size = 0; ur_size < unroll_size_rt; ur_size++) { + bool is_tails = ur_offset_elt + ur_size * vector_step + vector_step > jcp_.C; + if (is_tails) { + store_tail_emitter->emit_code({static_cast(ur_base + ur_size)}, {static_cast(reg_dst_aux.getIdx())}, + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + add(reg_dst_aux, tail_step * jcp_.dst_data_size); + } else { + store_vector_emitter->emit_code({static_cast(ur_base + ur_size)}, {static_cast(reg_dst_aux.getIdx())}, + {store_pool_vec_idxs}, {store_pool_gpr_idxs}); + add(reg_dst_aux, vector_step * jcp_.dst_data_size); + } + } + + add(reg_dst_aux, (jcp_.C - elt_num) * jcp_.dst_data_size); + sub(reg_oc_off, elt_num * sizeof(float)); + sub(reg_work_amount, 1); + jmp(loop_label, T_NEAR); + } + L(loop_end_label); + + add(reg_src, unroll_size_rt * vector_step * jcp_.src_data_size); + add(reg_dst, unroll_size_rt * vector_step * jcp_.dst_data_size); + add(reg_oc_off_bk, unroll_size_rt * vector_step * sizeof(float)); + } + } + + inline void norm_nspc_ac_ker() { + Xbyak::Reg64 reg_oc_off_bk = reg_src_stride; + if (attr_.post_ops_.len() != 0) { + mov(reg_oc_off_bk, reg_oc_off); + } + + size_t vec_num = div_up(jcp_.C, vector_step); + + Xbyak::Label loop_label; + Xbyak::Label loop_end_label; + L(loop_label); + { + cmp(reg_work_amount, 0); + jle(loop_end_label, T_NEAR); + + if (attr_.post_ops_.len() != 0) { + mov(reg_oc_off, reg_oc_off_bk); + } + + for (int v_num = 0; v_num < vec_num; v_num++) { + bool is_tail = (v_num * vector_step + vector_step > jcp_.C) ? true : false; + worker_mvn(is_tail); + if (is_tail) { + add(reg_src, tail_step * jcp_.src_data_size); + add(reg_dst, tail_step * jcp_.dst_data_size); + if (attr_.post_ops_.len() != 0) + add(reg_oc_off, tail_step * sizeof(float)); + } else { + add(reg_src, vector_step * jcp_.src_data_size); + add(reg_dst, vector_step * jcp_.dst_data_size); + if (attr_.post_ops_.len() != 0) + add(reg_oc_off, vector_step * sizeof(float)); + } + } + + sub(reg_work_amount, 1); + jmp(loop_label, T_NEAR); + } + L(loop_end_label); + } + inline void worker_mvn(bool is_tail) { const auto& load_emitter = is_tail ? load_tail_emitter : load_vector_emitter; const auto& store_emitter = is_tail ? store_tail_emitter : store_vector_emitter; @@ -599,7 +996,7 @@ private: if (jcp_.normalize_variance) uni_vmulps(vmm_val, vmm_val, vmm_variance_inv); - apply_post_ops(jcp_.dst_prc, jcp_.planar_layout); + apply_post_ops(jcp_.dst_prc, vmm_val.getIdx(), jcp_.layout == MVNLayoutType::mvn_planar); store_emitter->emit_code({static_cast(vmm_val.getIdx())}, {static_cast(reg_dst.getIdx())}, {store_pool_vec_idxs}, {store_pool_gpr_idxs}); @@ -625,7 +1022,7 @@ private: L(mvn_loop_end_label); } - void apply_post_ops(InferenceEngine::Precision dst_prc, bool is_broadcast) { + void apply_post_ops(InferenceEngine::Precision dst_prc, size_t vmm_idx, bool is_broadcast) { const auto &p = attr_.post_ops_; int eltwise_inj_idx = 0; int depthwise_inj_idx = 0; @@ -634,30 +1031,34 @@ private: for (int i = 0; i < p.len(); i++) { auto& post_op = p.entry_[i]; if (post_op.is_eltwise()) { - eltwise_injectors[eltwise_inj_idx]->compute_vector_range(vmm_val.getIdx(), vmm_val.getIdx() + 1); + eltwise_injectors[eltwise_inj_idx]->compute_vector_range(vmm_idx, vmm_idx + 1); eltwise_inj_idx++; } else if (post_op.is_depthwise()) { + if (post_op.depthwise.alg == alg_kind::depthwise_scale_shift && i < optimized_scaleshift_num) { + post_ops_data_offset += depthwise_injectors[depthwise_inj_idx]->memoryStep(); + depthwise_inj_idx++; + continue; + } mov(reg_d_weights, ptr[reg_post_ops_data + post_ops_data_offset]); add(reg_d_weights, reg_oc_off); depthwise_injectors[depthwise_inj_idx]->compute_vector_range( - vmm_val.getIdx(), vmm_val.getIdx() + 1, reg_d_weights, reg_d_weights, is_broadcast); + vmm_idx, vmm_idx + 1, reg_d_weights, reg_d_weights, is_broadcast); post_ops_data_offset += depthwise_injectors[depthwise_inj_idx]->memoryStep(); depthwise_inj_idx++; } else if (post_op.is_quantization()) { bool do_dequantization = post_op.quantization.alg == alg_kind::quantization_quantize_dequantize; bool do_rounding = do_dequantization || isFloatCompatible(dst_prc) || i != p.len() - 1; - int s_idx = vmm_val.getIdx(); quantization_injectors[quantization_inj_idx]->init_crop_ptrs(reg_post_ops_data + post_ops_data_offset, reg_oc_off); - quantization_injectors[quantization_inj_idx]->compute_crop(s_idx, s_idx + 1, 0, 0, is_broadcast); + quantization_injectors[quantization_inj_idx]->compute_crop(vmm_idx, vmm_idx + 1, 0, 0, is_broadcast); quantization_injectors[quantization_inj_idx]->init_input_scale_shift_ptrs(reg_post_ops_data + post_ops_data_offset, reg_oc_off); - quantization_injectors[quantization_inj_idx]->compute_input_scale_shift(s_idx, s_idx + 1, 0, do_rounding, 0, is_broadcast); + quantization_injectors[quantization_inj_idx]->compute_input_scale_shift(vmm_idx, vmm_idx + 1, 0, do_rounding, 0, is_broadcast); quantization_injectors[quantization_inj_idx]->init_output_scale_shift_ptrs(reg_post_ops_data + post_ops_data_offset, reg_oc_off); - quantization_injectors[quantization_inj_idx]->compute_output_scale_shift(s_idx, s_idx + 1, 0, 0, is_broadcast); + quantization_injectors[quantization_inj_idx]->compute_output_scale_shift(vmm_idx, vmm_idx + 1, 0, 0, is_broadcast); post_ops_data_offset += quantization_injectors[quantization_inj_idx]->memoryStep(); quantization_inj_idx++; @@ -855,7 +1256,7 @@ MVN::MVNJitExecutor::MVNJitExecutor(const MVNAttrs& mvnAttrs, jcp.dst_prc = mvnAttrs.dst_prc; jcp.src_data_size = src_data_size; jcp.dst_data_size = dst_data_size; - jcp.planar_layout = mvnAttrs.layout == MVNLayoutType::planar; + jcp.layout = mvnAttrs.layout; jcp.normalize_variance = mvnAttrs.normalizeVariance_; jcp.across_channels = mvnAttrs.execAcrossChannels_; int N = 0; @@ -900,8 +1301,10 @@ void MVN::MVNJitExecutor::exec(const uint8_t *src_data, uint8_t *dst_data, const if (!mvn_mean_kernel || (mvnAttrs.normalizeVariance_ && !mvn_variance_kernel) || !mvn_kernel) { IE_THROW() << "MVN layer doesn't create kernel to execute on sse41 above platform."; } - if (mvnAttrs.layout == MVNLayoutType::planar) { + if (mvnAttrs.layout == MVNLayoutType::mvn_planar) { mvn_pln(src_data, dst_data, post_ops_data_); + } else if (mvnAttrs.layout == MVNLayoutType::mvn_by_channel) { + mvn_nspc(src_data, dst_data, post_ops_data_); } else { mvn_blk(src_data, dst_data, post_ops_data_); } @@ -930,11 +1333,11 @@ void MVN::prepareParams() { mvnAttrs.src_prc = selectedPD->getConfig().inConfs[0].getMemDesc()->getPrecision(); mvnAttrs.dst_prc = selectedPD->getConfig().outConfs[0].getMemDesc()->getPrecision(); if (getParentEdgeAt(0)->getMemory().getDesc().hasLayoutType(LayoutType::ncsp)) { - mvnAttrs.layout = MVNLayoutType::planar; + mvnAttrs.layout = MVNLayoutType::mvn_planar; } else if (getParentEdgeAt(0)->getMemory().getDesc().hasLayoutType(LayoutType::nspc)) { - mvnAttrs.layout = MVNLayoutType::by_channel; + mvnAttrs.layout = MVNLayoutType::mvn_by_channel; } else { - mvnAttrs.layout = MVNLayoutType::block; + mvnAttrs.layout = MVNLayoutType::mvn_block; } MVNKey key = {mvnAttrs, dnnl::primitive_attr()}; @@ -1043,9 +1446,9 @@ void MVN::MVNJitExecutor::mvn_pln(const uint8_t* src_data, uint8_t* dst_data, co size_t src_stride_size = static_cast(blk_size * src_data_size); size_t dst_stride_size = static_cast(blk_size * dst_data_size); - parallel_for(N, [&](int b) { - size_t cb = b * C3; - if (mvnAttrs.execAcrossChannels_) { + if (mvnAttrs.execAcrossChannels_) { + parallel_for(N, [&](int b) { + size_t cb = b * C3; // Calculate mean value for one instance in batch // Parallel sum for each channel float C3inv = 1.f / static_cast(C3); @@ -1088,6 +1491,7 @@ void MVN::MVNJitExecutor::mvn_pln(const uint8_t* src_data, uint8_t* dst_data, co variance /= sqrtf(variance_temp * C3inv + mvnAttrs.epsValue_); else if (mvnAttrs.epsMode_ == OUTSIDE_SQRT) variance /= sqrtf(variance_temp * C3inv) + mvnAttrs.epsValue_; + // mvn for one instance in batch parallel_for(C, [&](int c) { size_t cc = cb + c * C2; @@ -1119,48 +1523,50 @@ void MVN::MVNJitExecutor::mvn_pln(const uint8_t* src_data, uint8_t* dst_data, co (*mvn_kernel)(&arg); }); } - } else { // per channel + }); + } else { + parallel_for2d(N, C, [&](size_t b, size_t c) { + size_t cb = b * C3; + size_t cc = cb + c * C2; float C2inv = 1.f / static_cast(C2); - parallel_for(C, [&](size_t c) { - // mean for this channel - float mean = 0.f; - size_t cc = cb + c * C2; - // the same arg for three kernels - auto arg = jit_mvn_call_args(); - arg.src = src_data + cc * src_data_size; - arg.dst = dst_data + cc * dst_data_size; - arg.sum = static_cast(&mean); - arg.src_stride = src_stride_size; - arg.dst_stride = dst_stride_size; - arg.work_amount = static_cast(C2 / blk_size); - arg.oc_off = static_cast(c * sizeof(float)); - arg.post_op_data = post_ops_data_; - (*mvn_mean_kernel)(&arg); - mean *= C2inv; + // mean for this channel + float mean = 0.f; + // the same arg for three kernels + auto arg = jit_mvn_call_args(); + arg.src = src_data + cc * src_data_size; + arg.dst = dst_data + cc * dst_data_size; + arg.sum = static_cast(&mean); + arg.src_stride = src_stride_size; + arg.dst_stride = dst_stride_size; + arg.work_amount = static_cast(C2 / blk_size); + arg.oc_off = static_cast(c * sizeof(float)); + arg.post_op_data = post_ops_data_; + (*mvn_mean_kernel)(&arg); - if (mvnAttrs.normalizeVariance_) { - // variance for this channel - float variance = 0.f; - arg.mean = static_cast(&mean); - arg.variance = static_cast(&variance); - (*mvn_variance_kernel)(&arg); + mean *= C2inv; - if (mvnAttrs.epsMode_ == INSIDE_SQRT) - variance = 1.f / sqrtf(variance * C2inv + mvnAttrs.epsValue_); - else if (mvnAttrs.epsMode_ == OUTSIDE_SQRT) - variance = 1.f / (sqrtf(variance * C2inv) + mvnAttrs.epsValue_); + if (mvnAttrs.normalizeVariance_) { + // variance for this channel + float variance = 0.f; + arg.mean = static_cast(&mean); + arg.variance = static_cast(&variance); + (*mvn_variance_kernel)(&arg); - // mvn for this channel - (*mvn_kernel)(&arg); - } else { - // mvn for this channel - arg.mean = static_cast(&mean); - (*mvn_kernel)(&arg); - } - }); - } - }); + if (mvnAttrs.epsMode_ == INSIDE_SQRT) + variance = 1.f / sqrtf(variance * C2inv + mvnAttrs.epsValue_); + else if (mvnAttrs.epsMode_ == OUTSIDE_SQRT) + variance = 1.f / (sqrtf(variance * C2inv) + mvnAttrs.epsValue_); + + // mvn for this channel + (*mvn_kernel)(&arg); + } else { + // mvn for this channel + arg.mean = static_cast(&mean); + (*mvn_kernel)(&arg); + } + }); + } } void MVN::MVNRefExecutor::mvn_ref(const uint8_t* src_data, uint8_t* dst_data) { @@ -1261,6 +1667,107 @@ void MVN::MVNRefExecutor::mvn_ref(const uint8_t* src_data, uint8_t* dst_data) { }); } +void MVN::MVNJitExecutor::mvn_nspc(const uint8_t* src_data, uint8_t* dst_data, const void *post_ops_data_) { + size_t blk_size = 1; // channel blk for memory layout + if (mayiuse(cpu::x64::avx512_core)) { + blk_size = 16; + } else if (mayiuse(cpu::x64::avx2)) { + blk_size = 8; + } else { + blk_size = 4; + } + + size_t N = 1; size_t C = 1; size_t D = 1; size_t H = 1; size_t W = 1; + std::tie(N, C, D, H, W) = mvnAttrs.shape5D; + + size_t threads_num = parallel_get_num_threads(); + size_t aux_buffer_size = mvnAttrs.execAcrossChannels_ ? 1 : rnd_up(C, blk_size); + parallel_for(N, [&](size_t b) { + std::vector mean_buffer(aux_buffer_size * threads_num); + std::vector variance_buffer; + if (mvnAttrs.normalizeVariance_) { + variance_buffer.resize(aux_buffer_size * threads_num); + } + size_t b_offset = b * C * D * H * W; + + // kernel_type: 0 for mean, 1 for variance, 2 for normalization + auto worker = [&](const bool across_channel, const int kernel_type) { + parallel_nt(0, [&](const int ithr, const int nthr) { + size_t start = 0, end = 0; + splitter(D * H * W, nthr, ithr, start, end); + + auto arg = jit_mvn_call_args(); + arg.src = src_data + (b_offset + (start * C)) * src_data_size; + if (0 == kernel_type) { + arg.sum = &mean_buffer[aux_buffer_size * ithr]; + } else if (1 == kernel_type) { + arg.mean = &mean_buffer[0]; + arg.variance = &variance_buffer[aux_buffer_size * ithr]; + } else if (2 == kernel_type) { + arg.dst = dst_data + (b_offset + (start * C)) * dst_data_size; + arg.mean = &mean_buffer[0]; + if (mvnAttrs.normalizeVariance_) + arg.variance = &variance_buffer[0]; + arg.oc_off = 0; + arg.post_op_data = post_ops_data_; + } + arg.work_amount = (across_channel && kernel_type != 2) ? (end - start) * C : (end - start); + + if (0 == kernel_type) { + (*mvn_mean_kernel)(&arg); + } else if (1 == kernel_type) { + (*mvn_variance_kernel)(&arg); + } else if (2 == kernel_type) { + (*mvn_kernel)(&arg); + } + }); + }; + + if (mvnAttrs.execAcrossChannels_) { + float size_inv = 1.f / static_cast(C * D * H * W); + worker(true, 0); + for (size_t i = 1; i < threads_num; i++) { + mean_buffer[0] += mean_buffer[i]; + } + mean_buffer[0] *= size_inv; + if (mvnAttrs.normalizeVariance_) { + worker(true, 1); + for (size_t i = 1; i < threads_num; i++) { + variance_buffer[0] += variance_buffer[i]; + } + if (mvnAttrs.epsMode_ == INSIDE_SQRT) + variance_buffer[0] = 1.f / sqrtf(variance_buffer[0] * size_inv + mvnAttrs.epsValue_); + else if (mvnAttrs.epsMode_ == OUTSIDE_SQRT) + variance_buffer[0] = 1.f / (sqrtf(variance_buffer[0] * size_inv) + mvnAttrs.epsValue_); + } + worker(true, 2); + } else { // for per_channel + float size_inv = 1.f / static_cast(D * H * W); + worker(false, 0); + for (size_t i = 1; i < threads_num; i++) { + for (size_t c = 0; c < C; c++) + mean_buffer[c] += mean_buffer[c + aux_buffer_size * i]; + } + for (size_t c = 0; c < C; c++) + mean_buffer[c] *= size_inv; + if (mvnAttrs.normalizeVariance_) { + worker(false, 1); + for (size_t i = 1; i < threads_num; i++) { + for (size_t c = 0; c < C; c++) + variance_buffer[c] += variance_buffer[c + aux_buffer_size * i]; + } + for (size_t c = 0; c < C; c++) { + if (mvnAttrs.epsMode_ == INSIDE_SQRT) + variance_buffer[c] = 1.f / sqrtf(variance_buffer[c] * size_inv + mvnAttrs.epsValue_); + else if (mvnAttrs.epsMode_ == OUTSIDE_SQRT) + variance_buffer[c] = 1.f / (sqrtf(variance_buffer[c] * size_inv) + mvnAttrs.epsValue_); + } + } + worker(false, 2); + } + }); +} + void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, const void *post_ops_data_) { size_t blk_size = 1; // channel blk for memory layout if (mayiuse(cpu::x64::avx512_core)) { @@ -1274,8 +1781,7 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co size_t CB = div_up(C, blk_size); - bool is_nhwc = mvnAttrs.layout == MVNLayoutType::by_channel; - size_t C0 = is_nhwc ? W * C : W * blk_size; + size_t C0 = W * blk_size; size_t C1 = C0 * H; size_t C2 = C1 * D; size_t C3 = C2 * CB; @@ -1286,18 +1792,17 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co std::vector mean_buffer(aux_buffer_size * threads_num); std::vector variance_buffer(aux_buffer_size * threads_num); - size_t src_stride_size = is_nhwc ? static_cast(C * src_data_size) : static_cast(blk_size * src_data_size); - size_t dst_stride_size = is_nhwc ? static_cast(C * dst_data_size) : static_cast(blk_size * dst_data_size); + size_t src_stride_size = static_cast(blk_size * src_data_size); + size_t dst_stride_size = static_cast(blk_size * dst_data_size); for (size_t b = 0lu; b < N; b++) { - size_t b_offset = is_nhwc ? b * C5 : b * C3; + size_t b_offset = b * C3; if (mvnAttrs.execAcrossChannels_) { // mean for this instance in batch float C5inv = 1.f / static_cast(C5); float mean_temp = 0.0f; mean_temp = parallel_sum3d(CB, D, H, mean_temp, [&](size_t cb, size_t d, size_t h)->float { - size_t src_offset = is_nhwc ? b_offset + d * C1 + h * C0 + cb * blk_size - : b_offset + cb * C2 + d * C1 + h * C0; + size_t src_offset = b_offset + cb * C2 + d * C1 + h * C0; float mean_internal = 0.0f; ///////////////////////////////// @@ -1332,8 +1837,7 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co // variance: sum((x-mean)*(x-mean)) for one instance in batch float variance_temp = 0.0f; variance_temp = parallel_sum3d(CB, D, H, variance_temp, [&](size_t cb, size_t d, size_t h)->float { - size_t src_offset = is_nhwc ? b_offset + d * C1 + h * C0 + cb * blk_size - : b_offset + cb * C2 + d * C1 + h * C0; + size_t src_offset = b_offset + cb * C2 + d * C1 + h * C0; float variance_internal = 0.0f; auto variance_buffer_ptr = &variance_buffer[blk_size * parallel_get_thread_num()]; @@ -1361,10 +1865,10 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co variance /= sqrtf(variance_temp * C5inv + mvnAttrs.epsValue_); else if (mvnAttrs.epsMode_ == OUTSIDE_SQRT) variance /= sqrtf(variance_temp * C5inv) + mvnAttrs.epsValue_; + // mvn for one instance in batch parallel_for3d(CB, D, H, [&](size_t cb, size_t d, size_t h) { - size_t src_offset = is_nhwc ? b_offset + d * C1 + h * C0 + cb * blk_size - : b_offset + cb * C2 + d * C1 + h * C0; + size_t src_offset = b_offset + cb * C2 + d * C1 + h * C0; auto arg = jit_mvn_call_args(); arg.src = src_data + src_offset * src_data_size; arg.dst = dst_data + src_offset * dst_data_size; @@ -1380,8 +1884,7 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co } else { // mvn for one instance in batch parallel_for3d(CB, D, H, [&](size_t cb, size_t d, size_t h) { - size_t src_offset = is_nhwc ? b_offset + d * C1 + h * C0 + cb * blk_size - : b_offset + cb * C2 + d * C1 + h * C0; + size_t src_offset = b_offset + cb * C2 + d * C1 + h * C0; auto arg = jit_mvn_call_args(); arg.src = src_data + src_offset * src_data_size; arg.dst = dst_data + src_offset * dst_data_size; @@ -1403,8 +1906,7 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co // keep the compute order the same as planar parallel_for2d(D, H, [&](size_t thr_idx, size_t d, size_t h) { for (size_t cb = 0; cb < CB; cb++) { - size_t src_offset = is_nhwc ? b_offset + d * C1 + h * C0 + cb * blk_size - : b_offset + cb * C2 + d * C1 + h * C0; + size_t src_offset = b_offset + cb * C2 + d * C1 + h * C0; auto mean_buffer_ptr = &mean_buffer[blk_size * cb + aux_buffer_size * thr_idx]; auto arg = jit_mvn_call_args(); @@ -1431,8 +1933,7 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co parallel_for2d(D, H, [&](size_t thr_idx, size_t d, size_t h) { for (size_t cb = 0; cb < CB; cb++) { - size_t src_offset = is_nhwc ? b_offset + d * C1 + h * C0 + cb * blk_size - : b_offset + cb * C2 + d * C1 + h * C0; + size_t src_offset = b_offset + cb * C2 + d * C1 + h * C0; auto mean_buffer_ptr = &mean_buffer[blk_size * cb]; auto variance_buffer_ptr = &variance_buffer[blk_size * cb + aux_buffer_size * thr_idx]; @@ -1460,8 +1961,7 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co parallel_for2d(D, H, [&](size_t d, size_t h) { for (size_t cb = 0; cb < CB; cb++) { - size_t src_offset = is_nhwc ? b_offset + d * C1 + h * C0 + cb * blk_size - : b_offset + cb * C2 + d * C1 + h * C0; + size_t src_offset = b_offset + cb * C2 + d * C1 + h * C0; auto mean_buffer_ptr = &mean_buffer[blk_size * cb]; auto variance_buffer_ptr = &variance_buffer[blk_size * cb]; @@ -1482,8 +1982,7 @@ void MVN::MVNJitExecutor::mvn_blk(const uint8_t* src_data, uint8_t* dst_data, co // normalizeVariance_ == false parallel_for2d(D, H, [&](size_t d, size_t h) { for (size_t cb = 0; cb < CB; cb++) { - size_t src_offset = is_nhwc ? b_offset + d * C1 + h * C0 + cb * blk_size - : b_offset + cb * C2 + d * C1 + h * C0; + size_t src_offset = b_offset + cb * C2 + d * C1 + h * C0; auto mean_buffer_ptr = &mean_buffer[blk_size * cb]; auto arg = jit_mvn_call_args(); diff --git a/src/plugins/intel_cpu/src/nodes/mvn.h b/src/plugins/intel_cpu/src/nodes/mvn.h index ffab7316118..ebd4feef42e 100644 --- a/src/plugins/intel_cpu/src/nodes/mvn.h +++ b/src/plugins/intel_cpu/src/nodes/mvn.h @@ -14,8 +14,14 @@ namespace ov { namespace intel_cpu { namespace node { +enum MVNLayoutType { + mvn_planar, + mvn_block, + mvn_by_channel +}; + struct jit_mvn_config_params { - bool planar_layout; + MVNLayoutType layout; bool across_channels; bool normalize_variance; InferenceEngine::Precision src_prc; @@ -69,6 +75,7 @@ struct jit_uni_mvn_kernel { jit_mvn_config_params jcp_; const dnnl_primitive_attr &attr_; + int optimized_scaleshift_num = 0; }; class MVN : public Node { @@ -102,11 +109,6 @@ public: OUTSIDE_SQRT }; - enum MVNLayoutType { - planar, - block, - by_channel - }; struct MVNAttrs { MVNLayoutType layout; std::tuple shape5D; @@ -152,6 +154,7 @@ private: private: void mvn_pln(const uint8_t *in_ptr_, uint8_t *out_ptr_, const void *post_ops_data_); void mvn_blk(const uint8_t *in_ptr_, uint8_t *out_ptr_, const void *post_ops_data_); + void mvn_nspc(const uint8_t *in_ptr_, uint8_t *out_ptr_, const void *post_ops_data_); std::shared_ptr mvn_mean_kernel; std::shared_ptr mvn_variance_kernel; diff --git a/src/plugins/intel_cpu/tests/functional/single_layer_tests/mvn.cpp b/src/plugins/intel_cpu/tests/functional/single_layer_tests/mvn.cpp index 7f876c3f672..d9e039674bd 100644 --- a/src/plugins/intel_cpu/tests/functional/single_layer_tests/mvn.cpp +++ b/src/plugins/intel_cpu/tests/functional/single_layer_tests/mvn.cpp @@ -433,13 +433,17 @@ const std::vector inputShapesStatic_3D = { const std::vector inputShapesStatic_4D = { {1, 7, 3, 5}, {1, 15, 9, 5}, - {4, 41, 6, 9} + {4, 41, 6, 9}, + // cover channel case 4*16*2+16+3=147 + {1, 147, 2, 2} }; const std::vector inputShapesStatic_5D = { {1, 32, 8, 1, 6}, {1, 9, 1, 15, 9}, - {6, 64, 6, 1, 18} + {6, 64, 6, 1, 18}, + // cover channel case 4*16*2+16+9=153 + {6, 153, 2, 2, 2} }; const auto Mvn2DStatic = ::testing::Combine(