// 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 != 0
$assert KR == 4

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

#include <wasm_simd128.h>

#include "xnnpack/packw.h"


void xnn_x32_packw_gemm_goi_ukernel_x${NR}c${KR}__wasmsimd_u4(
  size_t g,
  size_t nc,
  size_t kc,
  size_t nr,
  size_t kr,
  size_t sr,
  const uint32_t* weights,
  const uint32_t* bias,
  const void* scale,
  uint32_t* packed_weights,
  size_t extra_bytes,
  const void* params)
{
  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);

  do {
    // NC main loop multiple of ${NR}
    const uint32_t* w0 = (const uint32_t*) weights;
    size_t n = nc;

    for (; n >= ${NR}; n -= ${NR}) {
      if XNN_LIKELY(bias != NULL) {
        $for N in range(NR):
          packed_weights[${N}] = bias[${N}];
        bias += ${NR};
      } else {
        $for N in range(NR):
          packed_weights[${N}] = 0;
      }
      packed_weights += ${NR};

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

      // KC main loop multiple of ${NR}x${KR}
      size_t k = kc;
      for (; k >= ${KR}; k -= ${KR}) {
        // Read blocks of 2x4
        // a b c d
        // e f g h
        $for N in range(NR):
          const v128_t v${N} = wasm_v128_load(w${N});
          w${N} += ${KR};

        wasm_v128_store(packed_weights, v0);
        $for N in range(1,NR):
          wasm_v128_store(packed_weights + ${N*KR}, v${N});
        packed_weights += ${NR*KR};
      }

      // KC remainder (1..${KR-1})
      if XNN_UNLIKELY(k != 0) {
        assert(k >= 1);
        assert(k <= ${KR-1});
        switch (k) {
          case 1:
          {
            // Read blocks of 2x1
            // a
            // e
            $for N in range(NR):
              const v128_t v${N} = wasm_v128_load32_zero(w${N});
              ++w${N};
            wasm_v128_store(packed_weights, v0);
            $for N in range(1,NR):
              wasm_v128_store(packed_weights + ${N*KR}, v${N});
            packed_weights += ${NR*KR};
            break;
          }
          case 2:
          {
            // Read blocks of 2x2
            // a b
            // e f
            $for N in range(NR):
              const v128_t v${N} = wasm_v128_load64_zero(w${N});
              w${N} += 2;
            wasm_v128_store(packed_weights, v0);
            $for N in range(1,NR):
              wasm_v128_store(packed_weights + ${N*KR}, v${N});
            packed_weights += ${NR*KR};
            break;
          }
          case 3:
          {
            // Read blocks of 2x3
            // a b c
            // e f g
            $for N in range(NR):
              v128_t v${N} = wasm_v128_load64_zero(w${N});
              v${N} = wasm_v128_load32_lane(w${N} + 2, v${N}, 2);
              w${N} += 3;
            wasm_v128_store(packed_weights, v0);
            $for N in range(1,NR):
              wasm_v128_store(packed_weights + ${N*KR}, v${N});
            packed_weights += ${NR*KR};
            break;
          }
          default:
            XNN_UNREACHABLE;
        }
      }
      packed_weights = (uint32_t*) ((uintptr_t) packed_weights + extra_bytes);
      w0 = w${NR-1};
    }

    // NC remainder (1..${NR-1})
    if XNN_UNLIKELY(n != 0) {
      assert(n >= 1);
      assert(n <= ${NR-1});
      if XNN_LIKELY(bias != NULL) {
        size_t nb = n;
        do {
          *packed_weights++  = *bias++;
        } while (--nb != 0);
        packed_weights += (${NR} - n);
      } else {
        $for N in range(NR):
          packed_weights[${N}] = 0;
        packed_weights += ${NR};
      }

      // NR remainder has less than ${NR} rows so last row is not loaded
      $for N in range(1, NR-1):
        const uint32_t* 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${KR}
      size_t k = kc;
      for (; k >= ${KR}; k -= ${KR}) {
        // Read blocks of 2x4
        // a b c d
        // e f g h
        $for N in range(NR-1):
          const v128_t v${N} = wasm_v128_load(w${N});
          w${N} += ${KR};

        wasm_v128_store(packed_weights, v0);
        $for N in range(1,NR):
          wasm_v128_store(packed_weights + ${N*KR}, v${min(N,NR-2)});
        packed_weights += ${NR*KR};
      }

      // KC remainder (1..${KR-1})
      if XNN_UNLIKELY(k != 0) {
        assert(k >= 1);
        assert(k <= ${KR-1});
        switch (k) {
          case 1:
          {
            // Read blocks of 1x1
            // a
            $for N in range(NR-1):
              const v128_t v${N} = wasm_v128_load32_zero(w${N});
              ++w${N};
            wasm_v128_store(packed_weights, v0);
            $for N in range(1,NR):
              wasm_v128_store(packed_weights + ${N*KR}, v${min(N,NR-2)});
            packed_weights += ${NR*KR};
            break;
          }
          case 2:
          {
            // Read blocks of 1x2
            // a b
            $for N in range(NR-1):
              const v128_t v${N} = wasm_v128_load64_zero(w${N});
              w${N} += 2;
            wasm_v128_store(packed_weights, v0);
            $for N in range(1,NR):
              wasm_v128_store(packed_weights + ${N*KR}, v${min(N,NR-2)});
            packed_weights += ${NR*KR};
            break;
          }
          case 3:
          {
            // Read blocks of 1x3
            // a b c
            $for N in range(NR-1):
              v128_t v${N} = wasm_v128_load64_zero(w${N});
              v${N} = wasm_v128_load32_lane(w${N} + 2, v${N}, 2);
              w${N} += 3;
            wasm_v128_store(packed_weights, v0);
            $for N in range(1,NR):
              wasm_v128_store(packed_weights + ${N*KR}, v${min(N,NR-2)});
            packed_weights += ${NR*KR};
            break;
          }
          default:
            XNN_UNREACHABLE;
        }
      }
      packed_weights = (uint32_t*) ((uintptr_t) packed_weights + extra_bytes);
    }
    weights += nc * kc;
  } while (--g != 0);
}
