// 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 KBLOCK >= 1
$assert TYPE in ["int8_t"]

#include <assert.h>
#include <stddef.h>
#include <stdint.h>

#include "xnnpack/math.h"
#include "xnnpack/packw.h"
#include "xnnpack/unaligned.h"

$BITS = {"int8_t": 8}[TYPE]
$BTYPE = {"int8_t": "uint32_t"}[TYPE]
$WTYPE = {"int8_t": "int8_t"}[TYPE]
void xnn_x${BITS}_packw_gemm_goi_ukernel_x${NR}__scalar_u${KBLOCK}(
  size_t g,
  size_t nc,
  size_t kc,
  size_t nr,
  size_t kr,
  size_t sr,
  const ${WTYPE}* weights,
  const ${BTYPE}* bias,
  const void* scale,
  ${WTYPE}* packed_weights,
  size_t extra_bytes,
  const void* params)
{
  assert(g != 0);
  assert(nc != 0);
  assert(kc != 0);
  assert(nr == ${NR});   // This kernel is for NR=${NR}
  assert(kr == 1);
  assert(sr == 1);
  assert(weights != NULL);
  assert(packed_weights != NULL);

  ${TYPE}* out = (${TYPE}*) packed_weights;
  const ${BTYPE}* b = (const ${BTYPE}*) bias;

  do {
    // NC main loop multiple of ${NR}
    const ${TYPE}* w0 = (const ${TYPE}*) weights;
    size_t n = nc;
    for (;n >= ${NR}; n -= ${NR}) {
      if XNN_LIKELY(b != NULL) {
        $for N in range(NR):
          $if BTYPE == TYPE:
            out[${N}] = b[${N}];
          $else:
            unaligned_store_s32(out + ${N} * sizeof(int32_t), b[${N}]);
        b += ${NR};
      } else {
        $for N in range(NR):
          $if BTYPE == TYPE:
            out[${N}] = 0;
          $else:
            unaligned_store_s32(out + ${N} * sizeof(int32_t), 0);
      }
      $if BTYPE == TYPE:
        out += ${NR};
      $else:
        out += ${NR} * sizeof(${BTYPE});

      $for N in range(1, NR):
        const ${TYPE}* w${N} = w${N-1} + kc;

      // KC main loop multiple of ${NR}x${KBLOCK}
      size_t k = kc;
      for (; k >= ${KBLOCK}; k -= ${KBLOCK}) {
        $for N in range(NR):
          $for K in range(KBLOCK):
            const ${TYPE} v${N}${K} = w${N}[${K}];
          w${N} += ${KBLOCK};
        $for K in range(KBLOCK):
          $for N in range(NR):
            out[${N+K*NR}] = v${N}${K};
        out += ${NR*KBLOCK};
      }

      // KC remainder
      for (; k != 0; --k) {
        $for N in range(NR):
          const ${TYPE} v${N} = *w${N}++;
          out[${N}] = v${N};
        out += ${NR};
      }
      out = (${TYPE}*) ((uintptr_t) out + extra_bytes);
      w0 = w${NR-1};
    }

    // NC remainder (1..${NR-1})
    if XNN_UNLIKELY(n != 0) {
      if XNN_LIKELY(b != NULL) {
        size_t nb = n;
        do {
          $if BTYPE == TYPE:
            *out++ = *b++;
          $else:
            unaligned_store_s32(out, *b++);
            out += sizeof(${BTYPE});
        } while (--nb != 0);
      } else {
        size_t nb = n;
        do {
          $if BTYPE == TYPE:
            *out++ = 0;
          $else:
            unaligned_store_s32(out, 0);
            out += sizeof(${BTYPE});
        } while (--nb != 0);
      }
      $if BTYPE == TYPE:
        out += (${NR} - n);
      $else:
        out += (${NR} - n) * sizeof(${BTYPE});

      $if NR > 2:
        // NR remainder has less than ${NR} rows so last row is not loaded
      $for N in range(1, NR-1):
        const ${TYPE}* w${N} = w${N-1} + kc;
        $if N % 2 == 0:
          if XNN_UNPREDICTABLE(n <= ${N}) {
            w${N} = w${N-1};
          }
        $else:
          if XNN_UNPREDICTABLE(n < ${N+1}) {
            w${N} = w${N-1};
          }

      // KC main loop multiple of ${NR}x${KBLOCK}
      size_t k = kc;
      for (; k >= ${KBLOCK}; k -= ${KBLOCK}) {
        $for N in range(NR-1):
          $for K in range(KBLOCK):
            const ${TYPE} v${N}${K} = w${N}[${K}];
          w${N} += ${KBLOCK};
        $for K in range(KBLOCK):
          $for N in range(NR-1):
            out[${N+K*NR}] = v${N}${K};
        out += ${NR*KBLOCK};
      }

      // KC remainder of 1..${KBLOCK-1}
      for (; k != 0; --k) {
        $for N in range(NR-1):
          const ${TYPE} v${N} = *w${N}++;
          out[${N}] = v${N};
        out += ${NR};
      }
      out = (${TYPE}*) ((uintptr_t) out + extra_bytes);
    }
    weights += nc * kc;
  } while (--g != 0);
}
