text: /*******************************************************************************
  * Copyright 2018 Intel Corporation
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  *******************************************************************************/
  
  #include "math_utils.hpp"
  #include "mkldnn_thread.hpp"
  #include "simple_q10n.hpp"
  #include "gemm_x8s8s32x_inner_product.hpp"
  
  namespace mkldnn {
  namespace impl {
  namespace cpu {
  
  using namespace math;
  using namespace memory_format;
  using namespace memory_tracking::names;
  
  template<data_type_t src_type, data_type_t dst_type>
  gemm_x8s8s32x_inner_product_fwd_t<src_type, dst_type>::pp_kernel_t::pp_kernel_t(
          const pd_t *pd, bool dst_is_acc)
      : ker_(nullptr), OC_(pd->OC())
      , bias_data_type_(data_type::undef), bias_data_type_size_(0)
      , scale_idx_mult_(0), rmode_(round_mode::nearest)
      , do_bias_(false), do_relu_(false)
  {
      using namespace types;
  
      scale_idx_mult_ = (pd->attr()->output_scales_.mask_ == (1 << 1));
      rmode_ = pd->attr()->round_mode_;
  
      auto &post_ops = pd->attr()->post_ops_;
      do_relu_ = post_ops.len_ == 1;
      do_bias_ = pd->with_bias();
      bias_data_type_ = pd->desc()->bias_desc.data_type;
      if (do_bias_) {
          assert(bias_data_type_ != data_type::undef);
          bias_data_type_size_ = data_type_size(bias_data_type_);
      }
  
      if (!mayiuse(avx512_core))
          // use fallback code for older CPUs since they do not have optimized
          // x8s8s32 GEMM anyways. The configuration variables above are used by
          // the fallback code.
          return;
      else
          generate();
  }
  
  template<data_type_t src_type, data_type_t dst_type>
  void gemm_x8s8s32x_inner_product_fwd_t<src_type, dst_type>::pp_kernel_t::generate()
  {
      using namespace Xbyak;
      using namespace utils;
      using namespace round_mode;
  
      // TODO: clean-up
      Reg64 reg_param = abi_param1;
      Reg64 reg_dst = rdx;
      Reg64 reg_acc = rax;
      Reg64 reg_bias = rbx;
      Reg64 reg_scales = rsi;
  
      Reg64 reg_len = r8;
      Reg64 reg_tmp = rcx; // intentional for shifting purposes
      Reg64 reg_oc_offset = r9;
      Reg64 reg_rem_mask = r10;
      Opmask kreg_rem_mask = k1;
      Opmask kreg_relu_cmp = k2;
  
      const size_t vlen = cpu_isa_traits<avx512_common>::vlen / sizeof(float);
  
      Zmm vreg_zero = Zmm(0);
      Zmm vreg_scale = Zmm(1);
      Zmm vreg_nslope = Zmm(2);
  
      auto vreg_dst = [&](int idx) { return Zmm(3 + idx * 2 + 0); };
      auto vreg_bias = [&](int idx) { return Zmm(3 + idx * 2 + 1); };
  
      preamble();
  
  #define PARAM_OFF(x) offsetof(ker_args, x)
      mov(reg_dst, ptr[reg_param + PARAM_OFF(dst)]);
      mov(reg_acc, ptr[reg_param + PARAM_OFF(acc)]);
      mov(reg_bias, ptr[reg_param + PARAM_OFF(bias)]);
      mov(reg_scales, ptr[reg_param + PARAM_OFF(scales)]);
      mov(reg_len, ptr[reg_param + PARAM_OFF(len)]);
      mov(reg_oc_offset, ptr[reg_param + PARAM_OFF(oc_offset)]);
      vbroadcastss(vreg_nslope, ptr[reg_param + PARAM_OFF(nslope)]);
      if (scale_idx_mult_ == 0)
          vbroadcastss(vreg_scale, dword[reg_scales]);
  #undef PARAM_OFF
  
      if (do_relu_ || dst_type == data_type::u8)
          vxorps(vreg_zero, vreg_zero, vreg_zero);
  
      // Load accumulated value, convert to float, apply bias (if any), scaling,
      // and relu (if any); then convert to destination type and store
      auto compute = [&](size_t offset, int idx, bool apply_mask) {
          auto acc_addr = ptr[reg_acc + offset * sizeof(acc_data_t)];
  
          if (scale_idx_mult_ > 0) {
              assert(scale_idx_mult_ == 1);
              auto scale_addr = ptr[reg_scales + offset * sizeof(float)];
              auto vreg_scale_ = vreg_scale;
              if (apply_mask)
                  vreg_scale_ = vreg_scale_ | kreg_rem_mask;
              vmovups(vreg_scale, scale_addr);
          }
  
          auto vreg_dst_ = vreg_dst(idx);
          if (apply_mask)
              vreg_dst_ = vreg_dst_ | kreg_rem_mask;
          vcvtdq2ps(vreg_dst_, acc_addr);
  
          if (do_bias_) {
              auto bias_addr = ptr[reg_bias + offset * bias_data_type_size_];
              auto vreg_bias_ = vreg_bias(idx);
              if (apply_mask)
                  vreg_bias_ = vreg_bias_ | kreg_rem_mask;
  
              switch (bias_data_type_) {
              case data_type::s8:
                  vpmovsxbd(vreg_bias_, bias_addr);
                  break;
              case data_type::u8:
                  vpmovzxbd(vreg_bias_, bias_addr);
                  break;
              case data_type::s32:
              case data_type::f32:
                  vmovups(vreg_bias_, bias_addr);
                  break;
              default: assert(!"unimplemented");
              }
              if (bias_data_type_ != data_type::f32)
                  vcvtdq2ps(vreg_bias(idx), vreg_bias(idx));
              vaddps(vreg_dst(idx), vreg_dst(idx), vreg_bias(idx));
          }
  
          vmulps(vreg_dst(idx), vreg_dst(idx), vreg_scale);
          if (do_relu_) {
              vcmpps(kreg_relu_cmp, vreg_dst(idx), vreg_zero, _cmp_lt_os);
              vmulps(vreg_dst(idx) | kreg_relu_cmp, vreg_dst(idx), vreg_nslope);
          }
  
          if (dst_type == data_type::u8)
              vmaxps(vreg_dst(idx), vreg_dst(idx), vreg_zero);
  
          if (dst_type != data_type::f32) {
              auto rmode_control = (rmode_ == nearest ? T_rn_sae : T_rd_sae);
              vcvtps2dq(vreg_dst(idx) | rmode_control, vreg_dst(idx));
          }
  
          auto dst_addr = ptr[reg_dst + offset * sizeof(dst_data_t)];
          switch (dst_type) {
          case data_type::s8:
              vpmovsdb(dst_addr, vreg_dst_);
              break;
          case data_type::u8:
              vpmovusdb(dst_addr, vreg_dst_);
              break;
          case data_type::f32:
          case data_type::s32:
              vmovups(dst_addr, vreg_dst_);
              break;
          default: assert(!"unimplemented");
          }
      };
  
      // Advance all pointers by an immediate
      auto advance_ptrs_imm = [&](size_t offset) {
          add(reg_dst, offset * sizeof(dst_data_t));
          add(reg_acc, offset * sizeof(acc_data_t));
          if (scale_idx_mult_) {
              assert(scale_idx_mult_ == 1);
              add(reg_scales, offset * sizeof(float));
          }
          if (do_bias_)
              add(reg_bias, offset * bias_data_type_size_);
      };
  
      // Advance all pointers by a value stored in a register
      auto advance_ptrs_reg = [&](Reg64 offset) {
          lea(reg_dst, ptr[reg_dst + offset * sizeof(dst_data_t)]);
          lea(reg_acc, ptr[reg_acc + offset * sizeof(acc_data_t)]);
          if (scale_idx_mult_) {
              assert(scale_idx_mult_ == 1);
              lea(reg_scales, ptr[reg_scales + offset * sizeof(float)]);
          }
          if (do_bias_)
              lea(reg_bias, ptr[reg_bias + offset * bias_data_type_size_]);
      };
  
      // Rewind pointers that point to data that is indixed by output channel
      // (bias or per-oc scaling factors)
      auto rewind_ptrs = [&]() {
          if (do_bias_)
              sub(reg_bias, OC_ * bias_data_type_size_);
          if (scale_idx_mult_) {
              assert(scale_idx_mult_ == 1);
              sub(reg_scales, OC_ * sizeof(float));
          }
      };
  
      //      <-------------------- OC ------------------------------->
      //
      // ^    +....................+----------------------------------+
      // |    :   not accessed     |          Prologue loop           |
      // |    +--------------------+----------------------------------+
      //      |                                                       |
      // M    |                 Main loop (unrolled)                  |
      // B    |                                                       |
      //      +--------------------------------+----------------------+
      // |    |       Epilogue loop            |      not accessed    :
      // v    +--------------------------------+......................+
  
      Label prologue_end;
      cmp(reg_oc_offset, 0);
      je(prologue_end, T_NEAR);
  
      // Prologue loop
      {
          mov(reg_tmp, OC_);
          sub(reg_tmp, reg_oc_offset);
          cmp(reg_tmp, reg_len);
          cmovg(reg_tmp, reg_len);
          sub(reg_len, reg_tmp);
  
          Label prologue_loop, prologue_loop_tail, prologue_loop_end;
          cmp(reg_tmp, vlen);
          jle(prologue_loop_tail, T_NEAR); // Skips for reg_tmp == 16 too (?)
          L(prologue_loop); {
              compute(0, 0, false);
              advance_ptrs_imm(vlen);
              sub(reg_tmp, vlen);
              cmp(reg_tmp, vlen);
              jge(prologue_loop, T_NEAR);
          }
  
          L(prologue_loop_tail);
          mov(reg_rem_mask, 1);
          shl(reg_rem_mask, cl); // cl == reg_tmp because reg_tmp <= vlen here
          sub(reg_rem_mask, 1);
          jz(prologue_loop_end, T_NEAR);
  
          kmovq(kreg_rem_mask, reg_rem_mask);
          compute(0, 0, true);
          advance_ptrs_reg(reg_tmp);
  
          L(prologue_loop_end);
          rewind_ptrs();
      }
      L(prologue_end);
  
      // Main loop
      Label main_loop_end;
      {
          cmp(reg_len, OC_);
          jle(main_loop_end, T_NEAR);
  
          Label main_loop;
          L(main_loop); {
              size_t def_unroll = 4;
              size_t max_unroll = 13;
  
              size_t OC_loop, OC_tail;
              if (OC_ < max_unroll * vlen) {
                  // Fully unroll small loops
                  OC_loop = 0;
                  OC_tail = OC_;
              } else {
                  OC_loop = vlen * def_unroll;
                  OC_tail = OC_ % OC_loop;
              }
  
              assert(!!OC_loop || !!OC_tail);
  
              if (OC_tail % vlen) {
                  int vlen_tail = OC_tail % vlen;
                  unsigned tail_mask = (1 << vlen_tail) - 1;
                  mov(reg_tmp, tail_mask);
                  kmovq(kreg_rem_mask, reg_tmp);
              }
  
              if (OC_loop) {
                  mov(reg_tmp, rnd_dn(OC_, OC_loop));
                  Label oc_loop;
                  L(oc_loop); {
                      for (size_t offset = 0; offset < OC_loop; offset += vlen)
                          compute(offset, offset / vlen, false);
                      advance_ptrs_imm(OC_loop);
                      sub(reg_tmp, OC_loop);
                      jnz(oc_loop);
                  }
              }
  
              if (OC_tail) {
                  for (size_t offset = 0; offset < OC_tail; offset += vlen) {
                      bool use_mask = (offset + vlen) > OC_tail;
                      compute(offset, offset / vlen, use_mask);
                  }
                  advance_ptrs_imm(OC_tail);
              }
  
              rewind_ptrs();
              sub(reg_len, OC_);
              cmp(reg_len, OC_);
              jge(main_loop, T_NEAR);
          }
      }
      L(main_loop_end);
  
      // Epilogue loop
      Label epilogue_end;
      {
          cmp(reg_len, 0);
          je(epilogue_end, T_NEAR);
  
          Label epilogue_loop, epilogue_loop_tail;
          cmp(reg_len, vlen);
          jle(epilogue_loop_tail, T_NEAR); // Skips for reg_len == 16 (?)
          L(epilogue_loop); {
              compute(0, 0, false);
              sub(reg_len, vlen);
              advance_ptrs_imm(vlen);
              cmp(reg_len, vlen);
              jge(epilogue_loop, T_NEAR);
          }
  
          L(epilogue_loop_tail);
          mov(reg_tmp, reg_len); // reg_tmp is rcx, and we need cl for the shift
          mov(reg_rem_mask, 1);
          shl(reg_rem_mask, cl); // reg_tmp == rcx and reg_tail < vlen == 16
          sub(reg_rem_mask, 1);
          jz(epilogue_end, T_NEAR);
          kmovq(kreg_rem_mask, reg_rem_mask);
          compute(0, 0, true);
      }
  
      L(epilogue_end);
  
      postamble();
  
      ker_ = getCode<decltype(ker_)>();
  }
  
  template<data_type_t src_type, data_type_t dst_type>
  void gemm_x8s8s32x_inner_product_fwd_t<src_type, dst_type>::pp_kernel_t::operator ()(
          dst_data_t *dst, const acc_data_t *acc,
          const char *bias, const float *scales, float nslope,
          size_t start, size_t end)
  {
      using math::get_bias;
  
      if (end <= start)
          return;
  
      if (ker_) {
          // JIT
          ker_args args;
          size_t oc_offset = start % OC_;
          args.dst = dst + start;
          args.acc = acc + start;
          args.bias = bias + oc_offset * bias_data_type_size_;
          args.scales = scales + scale_idx_mult_ * oc_offset;
          args.nslope = nslope;
          args.len = end - start;
          args.oc_offset = oc_offset;
          ker_(&args);
      } else {
          // Fallback
          size_t oc = start % OC_;
          for (size_t i = start; i < end; i++) {
              float d = (float)acc[i];
              float b = get_bias(bias, oc, bias_data_type_);
              d = d + b;
              d *= scales[oc * scale_idx_mult_];
              if (do_relu_ && d < 0)
                  d *= nslope;
              dst[i] = qz_a1b0<float, dst_data_t>()(d, rmode_);
              oc = (oc == OC_ - 1) ? 0 : oc + 1;
          }
      }
  };
  
  template <data_type_t src_type, data_type_t dst_type>
  void gemm_x8s8s32x_inner_product_fwd_t<src_type, dst_type
          >::execute_forward() const {
      auto src = reinterpret_cast<const src_data_t *>(this->input_memory(0));
      auto weights = reinterpret_cast<const wei_data_t *>(this->input_memory(1));
      auto bias = reinterpret_cast<const char *>(this->input_memory(2));
      auto dst = reinterpret_cast<dst_data_t *>(this->memory());
  
      const int MB = pd()->MB();
      const int OC = pd()->OC();
  
      bool wei_tr = utils::one_of(pd()->weights_pd()->desc()->format,
               oihw, oidhw, oi);
  
      const int M = OC;
      const int N = MB;
      const int K = pd()->IC_total_padded();
      const int8_t off_a = 0, off_b = 0;
      const int32_t off_c = 0;
  
      const float *scales = pd()->attr()->output_scales_.scales_;
  
      const auto &post_ops = pd()->attr()->post_ops_;
      const bool do_relu = post_ops.len_ == 1;
      const float nslope = do_relu ? post_ops.entry_[0].eltwise.alpha : 0.f;
  
      acc_data_t *acc = pd()->dst_is_acc_
          ? (acc_data_t *)dst
          : scratchpad().template get<acc_data_t>(key_iprod_int_dat_in_acc_dt);
  
      const float onef = 1.0, zerof = 0.0;
  
      if (src_type == data_type::u8) {
          mkldnn_gemm_s8u8s32(wei_tr ? "T" : "N", "N", "F", &M, &N, &K, &onef,
                  weights, wei_tr ? &K : &M, &off_a, (uint8_t *)src, &K, &off_b, &zerof,
                  acc, &M, &off_c);
      } else if (src_type == data_type::s8) {
          mkldnn_gemm_s8s8s32(wei_tr ? "T" : "N", "N", "F", &M, &N, &K, &onef,
                  weights, wei_tr ? &K : &M, &off_a, (int8_t *)src, &K, &off_b, &zerof,
                  acc, &M, &off_c);
      } else {
          assert(!"incorrect src type");
      }
  
      const bool force_sequential = MB * OC < 2000;
      parallel(force_sequential ? 1 : 0, [&](int ithr, int nthr) {
              size_t start, end;
              balance211((size_t)OC * MB, nthr, ithr, start, end);
              (*pp_kernel_)(dst, acc, bias, scales, nslope, start, end);
              });
  }
  
  using namespace data_type;
  
  template struct gemm_x8s8s32x_inner_product_fwd_t<u8, f32>;
  template struct gemm_x8s8s32x_inner_product_fwd_t<u8, s32>;
  template struct gemm_x8s8s32x_inner_product_fwd_t<u8, s8>;
  template struct gemm_x8s8s32x_inner_product_fwd_t<u8, u8>;
  template struct gemm_x8s8s32x_inner_product_fwd_t<s8, f32>;
  template struct gemm_x8s8s32x_inner_product_fwd_t<s8, s32>;
  template struct gemm_x8s8s32x_inner_product_fwd_t<s8, s8>;
  template struct gemm_x8s8s32x_inner_product_fwd_t<s8, u8>;
  }
  }
  }
  