// Copyright 2020-2026 XMOS LIMITED.
// This Software is subject to the terms of the XMOS Public Licence: Version 1.

#if defined(__VX4B__)


#include "../asm_helper.h"

/*  

headroom_t vect_s32_sqrt(
    int32_t a[],
    const int32_t b[],
    const unsigned length,
    const right_shift_t b_shr,
    const unsigned depth);

*/


#define NSTACKVECTS     (3)
#define NSTACKWORDS     (8+8*(NSTACKVECTS))

#define FUNCTION_NAME   vect_s32_sqrt

// Temporary vector needed because there's no instruction to do vR[] * vR[]
#define STACK_VEC_TMP       (NSTACKWORDS-8-1)
// Holds the shifted values of b[] while we're solving it.
#define STACK_VEC_TARGET    (NSTACKWORDS-16-1)
// Holds the power of 2 that is currently being worked on inside hte inner loop.
// @todo If we had an instruction that set each vR[k] to the value of a register, this wouldn't be needed.
#define STACK_VEC_POW       (NSTACKWORDS-24-1)

#define STACK_DEPTH     0

#define a           a0
#define b           a1
#define length      a2
#define b_shr       a3

#define depth       s2
#define mask_vec    s3
#define _32         s4
#define _1          s5
#define tmp         s8

.text
.p2align 2


FUNCTION_NAME:
    xm.entsp (NSTACKWORDS)*4
    xm.stdsp  s3,s2,8
    xm.stdsp  s5,s4,16
    xm.stdsp  s7,s6,0
    sw a4, (STACK_DEPTH)*4                (sp)
// Set VPU mode to 32-bit
// (length << 2) is the length of the vector in bytes.
{ li t3, 0                              ; sw s8, 24                          (sp) }
{ slli length, length, 2                ; xm.vsetc t3                           }

// Maximum supported depth is 31
{ li tmp, 31                            ; lw t3, (STACK_DEPTH)*4                (sp) }
{ xm.assert t3                          ; nop                                   }
{ xm.sltu t3, tmp, t3                   ; li _1, 1                              }
{ li _32, 32                            ; xm.brff t3, .L_vect_loop_top          }
    { nop                               ; sw tmp, (STACK_DEPTH)*4                (sp) }


.L_vect_loop_top:

    // mask_vec is a byte mask for the elements of a[] that we're currently working on.
    // using VSTRPV with mask_vec prevents us from corrupting the headroom register.
    // depth is the number of MSBs that we're solving for
    { xm.mkmsk mask_vec, length         ; lw depth, (STACK_DEPTH)*4              (sp) }

    // First initialize the target vector using b[]
    // (Doing this first allows this function to operate in-place on b[] if desired)
    // @todo If we wanted to, we could do a VSIGN + VLMUL here to take an absolute value of each b[k],
    //       since this function will not work for any negative b[k].
        xm.vlashr b, b_shr
    { addi t3,sp, (STACK_VEC_TARGET)*4  ; add b, b, _32                         }
        xm.vstrpv t3, mask_vec

    // Initialize the result (a[]) with 0's
    { nop                               ; xm.vclrdr                             }
        xm.vstrpv a, mask_vec

    // VEC_POW[] is the bit we're currently solving for. Initialize to the first non-sign bit.
    // (The VSTD is to zero out the VEC_POW[] elements that are going to be masked out, because
    //  we're going to use VEC_POW[] later to update the headroom register)
    la t3, vpu_vec_0x40000000
    { addi t3,sp, (STACK_VEC_POW)*4     ; xm.vldr t3                            }
    { nop                               ; xm.vstd t3                            }
        xm.vstrpv t3, mask_vec 

    // This saves us a few cycles on the first iteration (because of loop alignment, we'd need a 
    // 'bu .L_sqrt_loop_top' here even if we didn't want to skip ahead). It's necessary because 
    // we don't want to right-shift VEC_POW[] on the first iteration (it's already 2^30), and we 
    // can't fix that by initializing VEC_POW[] to 0x80000000 above because that's negative and 
    // VLASHR is an arithmetic shift.
    { addi t3,sp, (STACK_VEC_TARGET)*4  ; xm.bu .L_first_iter                   }

    // Inner loop. Iteratively solving for the square root bit-by-bit
    // 12 instructions + 1 FNOP
    .p2align 4
    .L_sqrt_loop_top:

        // Load the next power of 2 and store it back to VEC_POW[]
            xm.vlashr t3, _1
            xm.vstrpv t3, mask_vec

        // Add the current power of 2 to each a[] to get the next value to be tested.
        // test[k] <-- a[k] + VEC_POW
        { addi t3,sp, (STACK_VEC_TMP)*4 ; xm.vladd a                            }


        // vR[] contains the values we're testing. Store it and square it
        // vR[k] <-- ( test[k] * test[k] ) >> 30
            xm.vstrpv t3, mask_vec
        { addi t3,sp, (STACK_VEC_TARGET)*4 ; xm.vlmul0 t3                          }

        .L_first_iter:

        // Subtract the squared test values from the target vector   
        // vR[k] <-- target[k] - (( test[k] * test[k] ) >> 30)
        { nop                           ; xm.vlsub t3                           }

        // If vR[k] is negative, the test value was too large, so we don't want to update those a[k]
        // for which vR[k] is negative.

        //  vR[k] = a[k] + MAX( signum( vR[k] ), 0 ) * VEC_POW[k]

        { addi depth, depth, -1         ; xm.vsign                              }
        { addi t3,sp, (STACK_VEC_POW)*4 ; xm.vpos                               }
        { nop                           ; xm.vlmul0 t3                          }
        { nop                           ; xm.vladd a                            }

        // Store the updated results in a[]
            xm.vstrpv a, mask_vec
        { nop                           ; xm.bt depth, .L_sqrt_loop_top         }
    .L_sqrt_loop_bot:

    // a[] now contains the results, but we haven't updated the headroom register because we've only
    // been using VSTRPV. So, update the headroom register
    // @todo Do we need to update the headroom register? Aren't we more or less guaranteed there's no
    // headroom, because we got rid of the headroom of b[]? Should work out the math on this later.
    
    // We used mask_vec when initializing VEC_POW[], so we can use that here to avoid corrupting
    // the headroom register with data that comes after a[]. t3 is already pointing at VEC_POW[].
        xm.vstrpv t3, mask_vec
    { sub length, length, _32           ; xm.vldr t3                            }

    // If (length - 32) < 1 we're done.
    { xm.slt tmp, length, _1            ; xm.vstr t3                            }
    { add a, a, _32                     ; nop                                   }
    beqz tmp, .L_vect_loop_top 
.L_vect_loop_bot:

.L_finish:

    xm.lddsp  s3,s2,8
    xm.lddsp  s5,s4,16
    xm.lddsp  s7,s6,0
{ li a0, 31                             ; xm.vgetc t3                           }
{ xm.zexti t3, 5                        ; lw s8, 24                      (sp)   }
{ sub a0, a0, t3                        ; xm.retsp (NSTACKWORDS)*4              } 


.L_func_end:


.global FUNCTION_NAME
.type FUNCTION_NAME,@function
.resource_const FUNCTION_NAME, "stack_frame_bytes", (NSTACKWORDS)*4
.resource_list_empty FUNCTION_NAME, "tail_callees"
.resource_list_empty FUNCTION_NAME, "callees"
.resource_list_empty FUNCTION_NAME, "parallel_callees"
.size FUNCTION_NAME, .L_func_end - FUNCTION_NAME


#endif //defined(__VX4B__)



