src/qb4-packw/c8-aarch64-neondot.c.in

// Copyright 2023 Google LLC
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.


$assert NR > 1
$assert KR > 1
#include <arm_neon.h>
#include <assert.h>
#include <stddef.h>
#include <stdint.h>
#include <string.h>

#include "src/xnnpack/common.h"
#include "src/xnnpack/microparams.h"
#include "src/xnnpack/packw.h"

// convert a vector from packed nibbles to planar, and accumulate sum
static XNN_INTRINSIC
int8x16_t xnn_packed2planar(
  int32x4_t *vacc,
  const int8x16_t v,
  const int8x16_t vmask,
  const int8x16_t vones)
{
  const int8x16_t vl = vandq_s8(v, vmask);  // isolate lower int 4
  const int8x16_t vh = vshlq_n_s8(v, 4);// isolate upper int 4
  *vacc = vdotq_s32(*vacc, vh, vones);
  *vacc = vdotq_s32(*vacc, vl, vones);
  const int8x16_t v01 = vzip1q_s8(vh, vl);
  const int8x16_t v23 = vzip2q_s8(vh, vl);
  const uint8x16_t v02 = vreinterpretq_u8_s64(vuzp1q_s64(vreinterpretq_s64_s8(v01), vreinterpretq_s64_s8(v23)));
  const uint8x16_t v13 = vreinterpretq_u8_s64(vuzp2q_s64(vreinterpretq_s64_s8(v01), vreinterpretq_s64_s8(v23)));
  const uint8x16_t vl02 = vshrq_n_u8(v02, 4);
  const uint8x16_t v0123 = vorrq_u8(vl02, v13);
  return vreinterpretq_s8_u8(v0123);
}

void xnn_qb4_packw_gemm_goi_ukernel_x${NR}c${KR}__aarch64_neondot(
  size_t g,
  size_t nc,
  size_t kc,
  size_t nr,
  size_t kr,
  size_t sr,
  size_t bl,
  size_t n_stride,
  const uint8_t* weights,
  const int32_t* bias,
  const void* scale,
  int8_t* packed_weights,
  size_t extra_bytes_bl,
  size_t extra_bytes_n,
  const void* params) XNN_OOB_READS
{
  assert(g != 0);
  assert(nc != 0);
  assert(kc != 0);
  assert(nr == ${NR});
  assert(kr == ${KR});
  assert(sr == 1);
  assert(weights != NULL);
  assert(packed_weights != NULL);
  assert(extra_bytes_bl == nr * sizeof(uint16_t));
  assert(extra_bytes_n == nr * sizeof(float));
  assert(params != NULL);
  assert(kc % bl == 0);
  size_t num_blocks = kc / bl;
  size_t weight_stride = (n_stride >> 1);
  const int8x16_t vmask = vmovq_n_s8(INT8_C(0xF0));
  const int8x16_t vones = vmovq_n_s8(UINT8_C(0x01));

  uint8_t* out = (uint8_t*) packed_weights;
  const int32_t* b = (const int32_t*) bias;
  const float32x4_t vzeropoint = vmovq_n_f32((float) (((const struct xnn_qs8_qc4w_packing_params*) params)->input_zero_point + 0));
  const float32x4_t vrecip_sixteen = vmovq_n_f32(1.0f/ 16.0f);
  const int8_t kzp = (int8_t) (((const struct xnn_qs8_qc4w_packing_params*) params)->kernel_zero_point);
  assert(kzp == 0 || kzp == 8);
  const int8x16_t vkernel_zero_point = vmovq_n_s8(kzp * 0x11);

  do {
    // NC main loop multiple of ${NR}
    const int8_t* w0 = (const int8_t*) weights;
    const uint16_t* s0 = (const uint16_t*) scale;
    int n = nc;
    for (;n > 0; n -= ${NR}) {
      float* packed_k_scaled_sum = (float*) out;
      $for N in range(0, NR, 4):
        float32x4_t packed_k_scaled_sums${N//4} = vdupq_n_f32(0.0f);
      out += ${NR} * sizeof(float);

      // KC/2 bytes is KC Nibbles
      $for N in range(1, NR):
        const int8_t* w${N} = w${N-1} + weight_stride;

      // scales
      $for N in range(1, NR):
        const uint16_t* s${N} = s${N-1} + num_blocks;

      if XNN_UNLIKELY(n < ${NR}){
        $for N in range(1, NR):
          $if $N % 2 == 0:
            if XNN_UNPREDICTABLE(n <= ${N}) {
              w${N} = w${N-1};
              s${N} = s${N-1};
            }
          $else:
            if XNN_UNPREDICTABLE(n < ${N+1}) {
              w${N} = w${N-1};
              s${N} = s${N-1};
            }
      }

      size_t kb = kc;
      // Process k by blocks (bl)
      for (; kb >= bl; kb-=bl) {
        // Initialize KSum as subtracting bl zero points (8)
        $for N in range(0, NR, 2):
          int32x4_t ksum${N//2} = vdupq_n_s32(0);
        size_t k = bl;

        // KC Main loop multiple of 16x32
        for(; k >= 32; k-=32) {
          $for N in range(0, NR):
            const int64x2_t w${N}x01 = vreinterpretq_s64_s8(veorq_s8(vld1q_s8(w${N}), vkernel_zero_point)); w${N} += ${2 * KR};

          $for N in range(0, NR, 2):
            int8x16_t v${N}${N+1}_0 = vreinterpretq_s8_s64(vzip1q_s64(w${N}x01, w${N+1}x01));
            int8x16_t v${N}${N+1}_1 = vreinterpretq_s8_s64(vzip2q_s64(w${N}x01, w${N+1}x01));

          $for I in range(0, 2):
            $for N in range(0, NR, 4):
              v${N}${N+1}_${I} = xnn_packed2planar(&ksum${N//2}, v${N}${N+1}_${I}, vmask, vones);
              v${N+2}${N+3}_${I} = xnn_packed2planar(&ksum${(N+2)//2}, v${N+2}${N+3}_${I}, vmask, vones);

          $for I in range(0, 2):
            $for N in range(0, NR, 2):
              vst1q_s8((int8_t*) &out[${(I*NR + N)*KR}], v${N}${N+1}_${I});

          out += ${NR*16};
        }

        $for N in range(0, NR, 8):
          uint16x8_t bf_scales${N//8} = vdupq_n_u16(0);
          $for N_I in range(0, 8):
            bf_scales${N//8} = vsetq_lane_u16(s${N + N_I}[0], bf_scales${N//8}, ${N_I});

        $for N in range(0, NR, 8):
          float32x4_t f_scales${N//4} = vreinterpretq_f32_u32(vshll_n_u16(vget_low_u16(bf_scales${N//8}), 16));
          float32x4_t f_scales${(N+4)//4} = vreinterpretq_f32_u32(vshll_n_u16(vget_high_u16(bf_scales${N//8}), 16));

        $for N in range(0, NR):
          s${N} += 1;

        $for N in range(0, NR, 4):
          f_scales${N//4} = vmulq_f32(f_scales${N//4}, vrecip_sixteen);

        $for N in range(0, NR, 4):
          float32x4_t f_ksum${N//4} = vcvtq_f32_s32(vpaddq_s32(ksum${N//2}, ksum${(N+2)//2}));
          f_ksum${N//4} = vmulq_f32(f_ksum${N//4}, vzeropoint);
          packed_k_scaled_sums${N//4} = vfmsq_f32(packed_k_scaled_sums${N//4}, f_ksum${N//4}, f_scales${N//4});

        $for N in range(0, NR, 4):
          vst1q_f32(&packed_k_scaled_sum[${N}], packed_k_scaled_sums${N//4});

        $for N in range(0, NR, 4):
          vst1_u16((uint16_t*)out+${N}, vreinterpret_u16_s16(vshrn_n_s32(vreinterpretq_s32_f32(f_scales${N//4}), 16)));

        out += ${NR} * sizeof(uint16_t);
      }


      if XNN_LIKELY(b != NULL){
        if (n >= ${NR}) {
          $for N in range(0, NR, 4):
            const int32x4_t b${N//4} = vld1q_s32(&b[${N}]);
            vst1q_s32((int32_t*)out + ${N}, b${N//4});
          b += ${NR};
        } else {
          memcpy(out, b, n * sizeof(int32_t));
          memset((int32_t*) out + n, 0, (${NR} - n) * sizeof(int32_t));
          b += n;
        }
      } else {
        $for N in range(0, NR, 4):
          vst1q_s32((int32_t*)out + ${N}, vdupq_n_s32(0));
      }
      out += ${NR} * sizeof(uint32_t);
      w0 = w15 + weight_stride - (kc >> 1);
      s0 = s15;
    }
    weights += nc * weight_stride;
  } while (--g != 0);
}