// 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.

#include <xnnpack/assembly.h>

.syntax unified

// void xnn_cs16_bfly4_samples1_ukernel__asm_aarch32_neon_x2(
//     size_t batch,                         r0
//     size_t samples,                       (unused)
//     int16_t* data,                        r2
//     const int16_t* twiddle,               (unused)
//     size_t stride)                        (unused)

// d8-d15, r12-r11,r14(lr) need to be preserved if used. r13(sp),r15(pc) are reserved.

// Register usage
// vout0 r2 d0
// vout1    d1
// vout2    d2
// vout3    d3
// vtmp3    d4
// vtmp4    d5
// vtmp5    d6
// vtmp0    d7

// vdiv4    d16
// vnegr    d17

BEGIN_FUNCTION xnn_cs16_bfly4_samples1_ukernel__asm_aarch32_neon_x2
        .arm
#ifndef __APPLE__
        .arch       armv7-a
        .fpu        neon
#endif
        SUBS        r0, r0, 1               // batch - 1
        VMVN.U16    d16, 57344              // 8191
        VMOV.I32    d17, 0x0001ffff         // vnegr
        BLS         1f

        // Batch of 2 main loop
0:
        VLD4.32     {d0,d1,d2,d3}, [r2]     // input 2 batches
        SUBS        r0, r0, 2               // batch

        VQRDMULH.S16    d1, d1, d16         // vout1 /= 4
        VQRDMULH.S16    d3, d3, d16         // vout3 /= 4
        VQRDMULH.S16    d0, d0, d16         // vout0 /= 4
        VQRDMULH.S16    d2, d2, d16         // vout2 /= 4

        VSUB.I16    d5, d1, d3              // vtmp4 = vout1 - vout3
        VADD.I16    d4, d1, d3              // vtmp3 = vout1 + vout3

        VMUL.S16    d5, d5, d17             // vrev4 = vtmp4 -r, i
        VADD.I16    d7, d0, d2              // vtmp0 = vout0 + vout2
        VSUB.I16    d6, d0, d2              // vtmp5 = vout0 - vout2

        VADD.I16    d0, d7, d4              // vout0 = vtmp0 + vtmp3
        VSUB.I16    d2, d7, d4              // vout2 = vtmp0 - vtmp3

        VREV32.16   d5, d5                  // vrev4 = vtmp4 i, -r
        VADD.I16    d1, d6, d5              // vout1 = vtmp5 + vrev4
        VSUB.I16    d3, d6, d5              // vout3 = vtmp5 - vrev4

        VST4.32     {d0,d1,d2,d3}, [r2]!    // output 2 batches
        BHI         0b

        BXLO        lr                      // no remainder?  early return

        // Remainder batch of 1
1:
        VLD4.32     {d0[0],d1[0],d2[0],d3[0]}, [r2]     // input 1 batch

        VQRDMULH.S16    d1, d1, d16         // vout1 /= 4
        VQRDMULH.S16    d3, d3, d16         // vout3 /= 4
        VQRDMULH.S16    d0, d0, d16         // vout0 /= 4
        VQRDMULH.S16    d2, d2, d16         // vout2 /= 4

        VSUB.I16    d5, d1, d3              // vtmp4 = vout1 - vout3
        VADD.I16    d4, d1, d3              // vtmp3 = vout1 + vout3

        VMUL.S16    d5, d5, d17             // vrev4 = vtmp4 -r, i
        VADD.I16    d7, d0, d2              // vtmp0 = vout0 + vout2
        VSUB.I16    d6, d0, d2              // vtmp5 = vout0 - vout2

        VADD.I16    d0, d7, d4              // vout0 = vtmp0 + vtmp3
        VSUB.I16    d2, d7, d4              // vout2 = vtmp0 - vtmp3

        VREV32.16   d5, d5                  // vrev4 = vtmp4 i, -r
        VADD.I16    d1, d6, d5              // vout1 = vtmp5 + vrev4
        VSUB.I16    d3, d6, d5              // vout3 = vtmp5 - vrev4

        VST4.32     {d0[0],d1[0],d2[0],d3[0]}, [r2] // output 1 batch
        BX          lr

END_FUNCTION xnn_cs16_bfly4_samples1_ukernel__asm_aarch32_neon_x2

#ifdef __ELF__
.section ".note.GNU-stack","",%progbits
#endif