// Copyright 2022 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 TILE_HEIGHT & (TILE_HEIGHT-1) == 0 and TILE_HEIGHT != 0
$assert TILE_WIDTH & (TILE_WIDTH-1) == 0 and TILE_WIDTH != 0
$assert TILE_HEIGHT in [1, 2, 4]
$assert TILE_WIDTH in [1, 2, 4]

#include <assert.h>

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

void xnn_x24_transposec_ukernel__${TILE_HEIGHT}x${TILE_WIDTH}_scalar(
    const void *input,
    void * output,
    size_t input_stride,
    size_t output_stride,
    size_t block_width,
    size_t block_height)
{
  assert(output_stride >= block_height * 3);
  assert(input_stride >= block_width * 3);

  $if TILE_HEIGHT == 1:
    const size_t input_reset = ${TILE_WIDTH * 3} - block_height * input_stride;
    const size_t output_reset = ${TILE_WIDTH} * output_stride - block_height * 3;
  $else:
    const size_t input_reset = ${TILE_WIDTH * 3} - round_down_po2(block_height, ${TILE_HEIGHT}) * input_stride;
    const size_t output_reset = ${TILE_WIDTH} * output_stride - block_height * 3;
  const size_t input_offset = ${TILE_HEIGHT} * input_stride;

  const uint8_t* i0 = (const uint8_t*) input;
  $for N in range(1, TILE_HEIGHT):
    const uint8_t* i${N} = (const uint8_t*) ((uintptr_t) i${N-1} + input_stride);

  uint8_t* o0 = (uint8_t*) output;
  $for N in range(1, TILE_WIDTH):
    uint8_t* o${N} = (uint8_t*) ((uintptr_t) o${N-1} + output_stride);

  do {
    $if TILE_WIDTH > 1:
      if XNN_UNPREDICTABLE(block_width < 2) {
        o1 = o0;
      }
    $for N in range(2, TILE_WIDTH, 2):
      if XNN_UNPREDICTABLE(block_width <= ${N}) {
        o${N} = o0;
      }
      if XNN_UNPREDICTABLE(block_width < ${N+2}) {
        o${N+1} = o0;
      }
    size_t bh = block_height;
    for (; bh >= ${TILE_HEIGHT}; bh -= ${TILE_HEIGHT}) {
      $for M in reversed(range(TILE_WIDTH)):
        $POS = 0
        $for N in range(TILE_HEIGHT):
          o${M}[${POS}] = i${N}[${M * 3}];
          o${M}[${POS + 1}] = i${N}[${M * 3 + 1}];
          o${M}[${POS + 2}] = i${N}[${M * 3 + 2}];
          $POS += 3
        o${M} += ${POS};
      $for N in range(TILE_HEIGHT):
        i${N} = (const uint8_t*) ((uintptr_t) i${N} + input_offset);
    }
    $if TILE_HEIGHT > 1:
      const uint8_t* i = i0;
      $if TILE_HEIGHT > 2:
        if (bh & 2) {
          $for M in reversed(range(TILE_WIDTH)):
            o${M}[0] = i0[${M * 3}];
            o${M}[1] = i0[${M * 3 + 1}];
            o${M}[2] = i0[${M * 3 + 2}];
            o${M}[3] = i1[${M * 3}];
            o${M}[4] = i1[${M * 3 + 1}];
            o${M}[5] = i1[${M * 3 + 2}];
            o${M} += 6;
          i = i2;
        }
      if (bh & 1) {
        $for M in reversed(range(TILE_WIDTH)):
          o${M}[0] = i[${M * 3}];
          o${M}[1] = i[${M * 3 + 1}];
          o${M}[2] = i[${M * 3 + 2}];
          o${M} += 3;
      }

    i0 = (const uint8_t*) ((uintptr_t) i0 + input_reset);
    $for N in range(1, TILE_HEIGHT):
      i${N} = (const uint8_t*) ((uintptr_t) i${N-1} + input_stride);
    $for N in range(TILE_WIDTH):
      o${N} = (uint8_t*) ((uintptr_t) o${N} + output_reset);
    block_width = doz(block_width, ${TILE_WIDTH});
  } while (block_width != 0);
}