// Copyright 2021 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.
$from itertools import chain
$import math
$assert IN_PTRS in ["MULTI", "REUSE"]
$assert OUT_PTRS in ["MULTI", "SWITCH", "MOV", "DEC"]
$assert SIZE in [8, 16, 32]
$TILE_SIZE = int(128/SIZE)
$NUM_ITERS = int(math.log2(TILE_SIZE))
$LO_PERM=str(list(chain.from_iterable((i, i+TILE_SIZE) for i in range((TILE_SIZE>>1)))))[1:-1]
$HI_PERM=str(list(chain.from_iterable(((TILE_SIZE>>1)+i, (TILE_SIZE>>1)+i+TILE_SIZE) for i in range((TILE_SIZE>>1)))))[1:-1]

#include <wasm_simd128.h>

#include <assert.h>

#include <xnnpack/common.h>
#include <xnnpack/math.h>
#include <xnnpack/transpose.h>

void xnn_x${SIZE}_transposec_ukernel__${TILE_SIZE}x${TILE_SIZE}_${IN_PTRS.lower()}_${OUT_PTRS.lower()}_wasmsimd(
    const uint${SIZE}_t* input,
    uint${SIZE}_t* output,
    size_t input_stride,
    size_t output_stride,
    size_t block_width,
    size_t block_height) XNN_OOB_READS
{
  assert(output_stride >= block_height * sizeof(uint${SIZE}_t));
  assert(input_stride >= block_width * sizeof(uint${SIZE}_t));

  const size_t tile_height = ${TILE_SIZE};
  const size_t tile_width = ${TILE_SIZE};
  const size_t tile_hbytes = tile_height * sizeof(uint${SIZE}_t);
  const size_t tile_wbytes = tile_width * sizeof(uint${SIZE}_t);
  const size_t input_reset = tile_wbytes - round_down_po2(block_height, tile_height) * input_stride;
  $if IN_PTRS == "MULTI":
    const size_t input_offset = tile_height * input_stride;
  $if OUT_PTRS in ["MOV", "DEC"]:
    const size_t output_reset = tile_width * output_stride - round_down_po2(block_height, 2) * sizeof(uint${SIZE}_t) - tile_hbytes;
  $else:
    const size_t output_reset = tile_width * output_stride - round_down_po2(block_height, 2) * sizeof(uint${SIZE}_t);

  $if IN_PTRS == "MULTI":
    const uint${SIZE}_t* i0 = input;
    $for N in range(1, TILE_SIZE):
      const uint${SIZE}_t* i${N} = (const uint${SIZE}_t*) ((uintptr_t) i${N-1} + input_stride);
  $else:
    const uint${SIZE}_t* i0 = input;
  $if OUT_PTRS == "MULTI":
    uint${SIZE}_t* o0 = (uint${SIZE}_t*) output;
    $for N in range(1, TILE_SIZE):
      uint${SIZE}_t* o${N} = (uint${SIZE}_t*) ((uintptr_t) o${N-1} + output_stride);
  $elif OUT_PTRS == "SWITCH":
    uint${SIZE}_t* o = (uint${SIZE}_t*) output;
  $else:
    uint${SIZE}_t* o = (uint${SIZE}_t*) ((uintptr_t) output - tile_hbytes);
  $if OUT_PTRS != "MULTI":
    const size_t minus_output_stride = -output_stride;

  do {
    $if OUT_PTRS == "MULTI":
      if XNN_UNPREDICTABLE(block_width < 2) {
        o1 = o0;
      }
      $for N in range(2, TILE_SIZE, 2):
        if XNN_UNPREDICTABLE(block_width <= ${N}) {
          o${N} = o0;
        }
        if XNN_UNPREDICTABLE(block_width < ${N+2}) {
          o${N+1} = o0;
        }
    $elif OUT_PTRS in ["MOV", "DEC"]:
      const size_t rem = min(block_width - 1, ${TILE_SIZE-1});
      const size_t oN_stride = rem * output_stride;
      const size_t oN_offset = oN_stride + tile_hbytes;
    $else:
      const size_t rem = min(block_width - 1, ${TILE_SIZE-1});
      const size_t oN_stride = rem * output_stride;
    size_t bh = block_height;
    for (; bh >= ${TILE_SIZE}; bh -= ${TILE_SIZE}) {
      $for N in range(TILE_SIZE):
        $if IN_PTRS == "REUSE":
          const v128_t v${NUM_ITERS}_${N} = wasm_v128_load(i0);
          i0 = (uint${SIZE}_t*) ((uintptr_t) i0 + input_stride);
        $else:
          const v128_t v${NUM_ITERS}_${N} = wasm_v128_load(i${N});
          i${N} = (uint${SIZE}_t*) ((uintptr_t) i${N} + input_offset);

      $for M in range(NUM_ITERS):
        $for N in range(TILE_SIZE >> 1):
          const v128_t v${NUM_ITERS-M-1}_${2*N} = wasm_v${SIZE}x${TILE_SIZE}_shuffle(v${NUM_ITERS-M}_${N}, v${NUM_ITERS-M}_${N+int(TILE_SIZE/2)}, ${LO_PERM});
          const v128_t v${NUM_ITERS-M-1}_${2*N+1} = wasm_v${SIZE}x${TILE_SIZE}_shuffle(v${NUM_ITERS-M}_${N}, v${NUM_ITERS-M}_${N+int(TILE_SIZE/2)}, ${HI_PERM});

      $if OUT_PTRS == "SWITCH":
        uint${SIZE}_t *oN = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
        switch (rem) {
          $for N in reversed(range(2, TILE_SIZE)):
            case ${N}:
              wasm_v128_store(oN, v0_${N});
              oN = (uint${SIZE}_t*) ((uintptr_t) oN + minus_output_stride);
          case 1:
            wasm_v128_store(oN, v0_1);
          case 0:
            wasm_v128_store(o, v0_0);
            o = (uint${SIZE}_t*) ((uintptr_t) o + tile_hbytes);
            break;
          default:
            XNN_UNREACHABLE;
        }
      $elif OUT_PTRS in ["MOV", "DEC"]:
        o = (uint${SIZE}_t*) ((uintptr_t) o + oN_offset);
        wasm_v128_store(o, v0_${TILE_SIZE-1});
        $if OUT_PTRS == "MOV":
          uint${SIZE}_t *oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
        $for N in reversed(range(2, TILE_SIZE-1, 2)):
          if XNN_UNPREDICTABLE(block_width > ${N+1}) {
            $if OUT_PTRS == "MOV":
              o = oN;
            $else:
              o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
          }
          wasm_v128_store(o, v0_${N});
          $if OUT_PTRS == "MOV":
            oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
          if XNN_UNPREDICTABLE(block_width >= ${N+1}) {
            $if OUT_PTRS == "MOV":
              o = oN;
            $else:
              o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
          }
          wasm_v128_store(o, v0_${N-1});
          $if OUT_PTRS == "MOV":
            oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
        if XNN_UNPREDICTABLE(block_width > 1) {
          $if OUT_PTRS == "MOV":
            o = oN;
          $else:
            o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
        }
        wasm_v128_store(o, v0_0);
      $else:
        $for N in reversed(range(TILE_SIZE)):
          wasm_v128_store(o${N}, v0_${N});
          o${N} = (uint${SIZE}_t*) ((uintptr_t) o${N} + tile_hbytes);
    }
    $if OUT_PTRS in ["MOV", "DEC"]:
      o = (uint${SIZE}_t*) ((uintptr_t) o + tile_hbytes);

    if (bh != 0) {
      $if IN_PTRS == "REUSE":
        const v128_t v${NUM_ITERS}_0 = wasm_v128_load(i0);
        $for N in range(1, TILE_SIZE - 1, 2):
          const uint${SIZE}_t *i${N} = (const uint${SIZE}_t*) ((uintptr_t) i${N-1} + input_stride);
          if XNN_UNPREDICTABLE(bh < ${N+1}) {
            i${N} = i${N-1};
          }
          const v128_t v${NUM_ITERS}_${N} = wasm_v128_load(i${N});
          const uint${SIZE}_t *i${N+1} = (const uint${SIZE}_t*) ((uintptr_t) i${N} + input_stride);
          if XNN_UNPREDICTABLE(bh <= ${N+1}) {
            i${N+1} = i${N};
          }
          const v128_t v${NUM_ITERS}_${N+1} = wasm_v128_load(i${N+1});
      $else:
        const v128_t v${NUM_ITERS}_0 = wasm_v128_load(i0);
        $for N in range(1, TILE_SIZE - 1, 2):
          if XNN_UNPREDICTABLE(bh < ${N+1}) {
            i${N} = i0;
          }
          const v128_t v${NUM_ITERS}_${N} = wasm_v128_load(i${N});
          if XNN_UNPREDICTABLE(bh <= ${N+1}) {
            i${N+1} = i0;
          }
          const v128_t v${NUM_ITERS}_${N+1} = wasm_v128_load(i${N+1});
      const v128_t v${NUM_ITERS}_${TILE_SIZE-1} = wasm_v128_xor(v${NUM_ITERS}_0, v${NUM_ITERS}_0);

      $for M in range(NUM_ITERS-1):
        $for N in range(TILE_SIZE >> 1):
          const v128_t v${NUM_ITERS-M-1}_${2*N} = wasm_v${SIZE}x${TILE_SIZE}_shuffle(v${NUM_ITERS-M}_${N}, v${NUM_ITERS-M}_${N+int(TILE_SIZE/2)}, ${LO_PERM});
          const v128_t v${NUM_ITERS-M-1}_${2*N+1} = wasm_v${SIZE}x${TILE_SIZE}_shuffle(v${NUM_ITERS-M}_${N}, v${NUM_ITERS-M}_${N+int(TILE_SIZE/2)}, ${HI_PERM});

      $for N in range(TILE_SIZE >> 1):
        v128_t v0_${2*N} = wasm_v${SIZE}x${TILE_SIZE}_shuffle(v1_${N}, v1_${N+int(TILE_SIZE/2)}, ${LO_PERM});
        v128_t v0_${2*N+1} = wasm_v${SIZE}x${TILE_SIZE}_shuffle(v1_${N}, v1_${N+int(TILE_SIZE/2)}, ${HI_PERM});

      if (bh & ${TILE_SIZE>>1}) {
        $if OUT_PTRS == "SWITCH":
          uint${SIZE}_t* oN = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
          switch (rem) {
            $for N in reversed(range(2, TILE_SIZE)):
              case ${N}:
                *((double*) oN) = wasm_f64x2_extract_lane(v0_${N}, 0);
                oN = (uint${SIZE}_t*) ((uintptr_t) oN + minus_output_stride);
            case 1:
              *((double*) oN) = wasm_f64x2_extract_lane(v0_1, 0);
            case 0:
              $if NUM_ITERS > 1:
                *((double*) o) = wasm_f64x2_extract_lane(v0_0, 0);
                o += ${TILE_SIZE>>1};
              $else:
                *((double*) o) = wasm_f64x2_extract_lane(v0_0, 0);
              break;
            default:
              XNN_UNREACHABLE;
          }
        $elif OUT_PTRS in ["MOV", "DEC"]:
          o = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
          *((double*) o) = wasm_f64x2_extract_lane(v0_${TILE_SIZE-1}, 0);
          $if OUT_PTRS == "MOV":
            uint${SIZE}_t *oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
          $for N in reversed(range(2, TILE_SIZE, 2)):
            if XNN_UNPREDICTABLE(block_width > ${N+1}) {
              $if OUT_PTRS == "MOV":
                o = oN;
              $else:
                o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            }
            *((double*) o) = wasm_f64x2_extract_lane(v0_${N}, 0);
            $if OUT_PTRS == "MOV":
              oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            if XNN_UNPREDICTABLE(block_width >= ${N+1}) {
              $if OUT_PTRS == "MOV":
                o = oN;
              $else:
                o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            }
            *((double*) o) = wasm_f64x2_extract_lane(v0_${N-1}, 0);
            $if OUT_PTRS == "MOV":
              oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
          if XNN_UNPREDICTABLE(block_width > 1) {
            $if OUT_PTRS == "MOV":
              o = oN;
            $else:
              o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
          }
          $if NUM_ITERS > 1:
            *((double*) o) = wasm_f64x2_extract_lane(v0_0, 0);
            o += ${TILE_SIZE>>1};
          $else:
            *((double*) o) = wasm_f64x2_extract_lane(v0_0, 0);
        $else:
          $for N in reversed(range(TILE_SIZE)):
            $if NUM_ITERS>1:
              *((double*) o${N}) = wasm_f64x2_extract_lane(v0_${N}, 0);
              o${N} += ${TILE_SIZE>>1};
            $else:
              *((double*) o${N}) = wasm_f64x2_extract_lane(v0_${N}, 0);
        $if NUM_ITERS > 1:
          $for N in range(TILE_SIZE):
            v0_${N} = wasm_v64x2_shuffle(v0_${N}, v0_${N}, 1, 1);
      }

      $if NUM_ITERS>1:
        if (bh & ${TILE_SIZE>>2}) {
          $if OUT_PTRS == "SWITCH":
            uint${SIZE}_t* oN = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
            switch (rem) {
              $for N in reversed(range(2, TILE_SIZE)):
                case ${N}:
                  *((float*) oN) = wasm_f32x4_extract_lane(v0_${N}, 0);
                  oN = (uint${SIZE}_t*) ((uintptr_t) oN + minus_output_stride);
              case 1:
                *((float*) oN) = wasm_f32x4_extract_lane(v0_1, 0);
              case 0:
                *((float*) o) = wasm_f32x4_extract_lane(v0_0, 0);
                $if SIZE < 32:
                  o += ${TILE_SIZE>>2};
                break;
              default:
                XNN_UNREACHABLE;
            }
          $elif OUT_PTRS in ["MOV", "DEC"]:
            o = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
            *((float*) o) = wasm_f32x4_extract_lane(v0_${TILE_SIZE-1}, 0);
            $if OUT_PTRS == "MOV":
              uint${SIZE}_t *oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            $for N in reversed(range(2, TILE_SIZE, 2)):
              if XNN_UNPREDICTABLE(block_width > ${N+1}) {
                $if OUT_PTRS == "MOV":
                  o = oN;
                $else:
                  o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              }
              *((float*) o) = wasm_f32x4_extract_lane(v0_${N}, 0);
              $if OUT_PTRS == "MOV":
                oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              if XNN_UNPREDICTABLE(block_width >= ${N+1}) {
                $if OUT_PTRS == "MOV":
                  o = oN;
                $else:
                  o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              }
              *((float*) o) = wasm_f32x4_extract_lane(v0_${N-1}, 0);
              $if OUT_PTRS == "MOV":
                oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            if XNN_UNPREDICTABLE(block_width > 1) {
              $if OUT_PTRS == "MOV":
                o = oN;
              $else:
                o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            }
            *((float*) o) = wasm_f32x4_extract_lane(v0_0, 0);
            $if SIZE < 32:
              o += ${TILE_SIZE>>2};
          $else:
            $for N in reversed(range(TILE_SIZE)):
              *((float*) o${N}) = wasm_f32x4_extract_lane(v0_${N}, 0);
              $if SIZE < 32:
                o${N} += ${TILE_SIZE>>2};
          $if NUM_ITERS > 2:
            $for N in range(TILE_SIZE):
              v0_${N} = wasm_u64x2_shr(v0_${N}, 32);
        }
      $if NUM_ITERS>2:
        if (bh & ${TILE_SIZE>>3}) {
          $if OUT_PTRS == "SWITCH":
            uint${SIZE}_t* oN = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
            switch (rem) {
              $for N in reversed(range(2, TILE_SIZE)):
                case ${N}:
                  $if SIZE == 16:
                    *oN = wasm_i16x8_extract_lane(v0_${N}, 0);
                    oN = (uint${SIZE}_t*) ((uintptr_t) oN + minus_output_stride);
                  $else:
                    *((uint16_t*) oN) = wasm_i16x8_extract_lane(v0_${N}, 0);
                    oN = (uint${SIZE}_t*) ((uintptr_t) oN + minus_output_stride);
              case 1:
                $if SIZE == 16:
                  *oN = wasm_i16x8_extract_lane(v0_1, 0);
                $else:
                  *((uint16_t*) oN) = wasm_i16x8_extract_lane(v0_1, 0);
              case 0:
                $if SIZE == 16:
                  *o = wasm_i16x8_extract_lane(v0_0, 0);
                $else:
                  *((uint16_t*) o) = wasm_i16x8_extract_lane(v0_0, 0);
                  o += ${TILE_SIZE>>3};
                break;
              default:
                XNN_UNREACHABLE;
            }
          $elif OUT_PTRS in ["MOV", "DEC"]:
            o = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
            $if SIZE == 16:
              *o = wasm_i16x8_extract_lane(v0_${TILE_SIZE-1}, 0);
            $else:
              *((uint16_t*) o) = wasm_i16x8_extract_lane(v0_${TILE_SIZE-1}, 0);
            $if OUT_PTRS == "MOV":
              uint${SIZE}_t *oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            $for N in reversed(range(2, TILE_SIZE, 2)):
              if XNN_UNPREDICTABLE(block_width > ${N+1}) {
                $if OUT_PTRS == "MOV":
                  o = oN;
                $else:
                  o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              }
              $if SIZE == 16:
                *o = wasm_i16x8_extract_lane(v0_${N}, 0);
              $else:
                *((uint16_t*) o) = wasm_i16x8_extract_lane(v0_${N}, 0);
              $if OUT_PTRS == "MOV":
                oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              if XNN_UNPREDICTABLE(block_width >= ${N+1}) {
                $if OUT_PTRS == "MOV":
                  o = oN;
                $else:
                  o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              }
              $if SIZE == 16:
                *o = wasm_i16x8_extract_lane(v0_${N-1}, 0);
              $else:
                *((uint16_t*) o) = wasm_i16x8_extract_lane(v0_${N-1}, 0);
              $if OUT_PTRS == "MOV":
                oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            if XNN_UNPREDICTABLE(block_width > 1) {
              $if OUT_PTRS == "MOV":
                o = oN;
              $else:
                o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            }
            $if SIZE == 16:
              *o = wasm_i16x8_extract_lane(v0_0, 0);
            $else:
              *((uint16_t*) o) = wasm_i16x8_extract_lane(v0_0, 0);
              o += ${TILE_SIZE>>3};
          $else:
            $for N in reversed(range(TILE_SIZE)):
              $if SIZE == 16:
                *o${N} = wasm_i16x8_extract_lane(v0_${N}, 0);
              $else:
                *((uint16_t*) o${N}) = wasm_i16x8_extract_lane(v0_${N}, 0);
                o${N} += ${TILE_SIZE>>3};
          $if NUM_ITERS>3:
            $for N in range(TILE_SIZE):
              v0_${N} = wasm_u32x4_shr(v0_${N}, 16);
        }
      $if SIZE == 8:
        if (bh & 1) {
          $if OUT_PTRS == "SWITCH":
            uint${SIZE}_t* oN = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
            switch (rem) {
              $for N in reversed(range(2, TILE_SIZE)):
                case ${N}:
                  *oN = wasm_i8x16_extract_lane(v0_${N}, 0);
                  oN = (uint${SIZE}_t*) ((uintptr_t) oN + minus_output_stride);
              case 1:
                *oN = wasm_i8x16_extract_lane(v0_1, 0);
              case 0:
                *o = wasm_i8x16_extract_lane(v0_0, 0);
                break;
              default:
                XNN_UNREACHABLE;
            }
          $elif OUT_PTRS in ["MOV", "DEC"]:
            o = (uint${SIZE}_t*) ((uintptr_t) o + oN_stride);
            *o = wasm_i8x16_extract_lane(v0_${TILE_SIZE-1}, 0);
            $if OUT_PTRS == "MOV":
              uint${SIZE}_t *oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            $for N in reversed(range(2, TILE_SIZE, 2)):
              if XNN_UNPREDICTABLE(block_width > ${N+1}) {
                $if OUT_PTRS == "MOV":
                  o = oN;
                $else:
                  o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              }
              *o = wasm_i8x16_extract_lane(v0_${N}, 0);
              $if OUT_PTRS == "MOV":
                oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              if XNN_UNPREDICTABLE(block_width >= ${N+1}) {
                $if OUT_PTRS == "MOV":
                  o = oN;
                $else:
                  o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
              }
              *o = wasm_i8x16_extract_lane(v0_${N-1}, 0);
              $if OUT_PTRS == "MOV":
                oN = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            if XNN_UNPREDICTABLE(block_width > 1) {
              $if OUT_PTRS == "MOV":
                o = oN;
              $else:
                o = (uint${SIZE}_t*) ((uintptr_t) o + minus_output_stride);
            }
            *o = wasm_i8x16_extract_lane(v0_0, 0);
          $else:
            $for N in reversed(range(TILE_SIZE)):
              *o${N} = wasm_i8x16_extract_lane(v0_${N}, 0);
        }
    }

    $if IN_PTRS == "MULTI":
      i0 = (const uint${SIZE}_t*) ((uintptr_t) i0 + input_reset);
      $for N in range(1, TILE_SIZE):
        i${N} = (const uint${SIZE}_t*) ((uintptr_t) i${N-1} + input_stride);
    $else:
      i0 = (const uint${SIZE}_t*) ((uintptr_t) i0 + input_reset);
    $if OUT_PTRS == "MULTI":
      o0 = (uint${SIZE}_t*) ((uintptr_t) o0 + output_reset);
      $for N in range(1, TILE_SIZE):
        o${N} = (uint${SIZE}_t*) ((uintptr_t) o${N} + output_reset);
    $else:
      o = (uint${SIZE}_t*) ((uintptr_t) o + output_reset);
    block_width = doz(block_width, tile_width);
  } while (block_width != 0);
}