#ifdef MNN_SUPPORT_FP16
#pragma OPENCL EXTENSION cl_khr_fp16 : enable
#endif

#define GLOBAL_SIZE_3_DIMS \
    __private const int global_size_dim0, __private const int global_size_dim1, __private const int global_size_dim2,

#define DEAL_NON_UNIFORM_DIM3(input1, input2, input3)                                             \
    if (input1 >= global_size_dim0 || input2 >= global_size_dim1 || input3 >= global_size_dim2) { \
        return;                                                                                   \
    }

#define DEAL_HEAD_DIM_NOT_ALIGN \
    if(hd * 4 + 3 >= head_dim) {\
        temp_0.w = (FLOAT)0;\
        temp_1.w = (FLOAT)0;\
        temp_2.w = (FLOAT)0;\
        temp_3.w = (FLOAT)0;\
    }\
    if(hd * 4 + 2 >= head_dim) {\
        temp_0.z = (FLOAT)0;\
        temp_1.z = (FLOAT)0;\
        temp_2.z = (FLOAT)0;\
        temp_3.z = (FLOAT)0;\
    }\
    if(hd * 4 + 1 >= head_dim) {\
        temp_0.y = (FLOAT)0;\
        temp_1.y = (FLOAT)0;\
        temp_2.y = (FLOAT)0;\
        temp_3.y = (FLOAT)0;\
    }

#define DEAL_SEQ_LEN_NOT_ALIGN \
    if(4 * sl + 3 >= seq_len) {\
        temp_3 = (FLOAT4)0;\
    }\
    if(4 * sl + 2 >= seq_len) {\
        temp_2 = (FLOAT4)0;\
    }\
    if(4 * sl + 1 >= seq_len) {\
        temp_1 = (FLOAT4)0;\
    }

__kernel void split_transpose_qkv(GLOBAL_SIZE_3_DIMS
                              __global const FLOAT *input, // [Batch, seqLen/4, mNumHead * 3 * mHeadDim, 4]
                              __global FLOAT *output_q, // [Batch * mNumHead, head_dim_pack_k, seq_len_pack_mn / qSeqSplitNum]
                              __global FLOAT *output_k, // [Batch * mNumHead, head_dim_pack_k, seq_len_pack_mn]
                              __global FLOAT *output_v, // [Batch * mNumHead, ROUND_UP(seqLen, tile), head_dim_pack_mn]
                              __private const int seq_len_pack_mn,
                              __private const int seq_len_piece,
                              __private const int head_dim_pack_mn,
                              __private const int head_dim_pack_k,
                              __private const int seq_len,
                              __private const int head_num,
                              __private const int head_dim,
                              __private const int batch,
                              __private const int seq_index
) {
    const int sl = get_global_id(0); // seqLen_4
    const int hd = get_global_id(1); // mHeadDim_4
    const int z = get_global_id(2); // Batch * mNumHead
    DEAL_NON_UNIFORM_DIM3(sl, hd, z);
    
    const int b = z / head_num;
    const int hn = z % head_num;
    
    const int offset_q = ((b * head_num + hn) * head_dim_pack_k + 4 * hd) * seq_len_piece + 4 * sl;

    if(seq_index > 0) {
        // fill output_q only
        if(sl * 4 >= seq_len || hd * 4 >= head_dim) {
            if(hd * 4 < head_dim_pack_k) {
                if(sl * 4 < seq_len_piece) {
                    vstore4((FLOAT4)0, 0, output_q + offset_q);
                    vstore4((FLOAT4)0, 0, output_q + offset_q + seq_len_piece);
                    vstore4((FLOAT4)0, 0, output_q + offset_q + seq_len_piece + seq_len_piece);
                    vstore4((FLOAT4)0, 0, output_q + offset_q + seq_len_piece + seq_len_piece + seq_len_piece);
                }
            }
            return;
        }
        
        const int offset_inp = ((((seq_index * seq_len_piece / 4 + sl) * batch + b) * head_num + hn) * 3 * head_dim + 4 * hd) * 4;

        if(sl * 4 < seq_len_piece) {
            FLOAT4 temp_0 = vload4(0, input + offset_inp);
            FLOAT4 temp_1 = vload4(0, input + offset_inp + 4);
            FLOAT4 temp_2 = vload4(0, input + offset_inp + 8);
            FLOAT4 temp_3 = vload4(0, input + offset_inp + 12);
            #ifdef HEADDIM_LEAVE
            DEAL_HEAD_DIM_NOT_ALIGN
            #endif
            #ifdef SEQLEN_LEAVE
            DEAL_SEQ_LEN_NOT_ALIGN
            #endif
            vstore4(temp_0, 0, output_q + offset_q);
            vstore4(temp_1, 0, output_q + offset_q + seq_len_piece);
            vstore4(temp_2, 0, output_q + offset_q + seq_len_piece + seq_len_piece);
            vstore4(temp_3, 0, output_q + offset_q + seq_len_piece + seq_len_piece + seq_len_piece);
        }
        return;
    }
    const int offset_k = ((b * head_num + hn) * head_dim_pack_k + 4 * hd) * seq_len_pack_mn + 4 * sl;

    const int offset_v = ((b * head_num + hn) * seq_len_pack_mn + 4 * sl) * head_dim_pack_mn + 4 * hd;
    if(sl * 4 >= seq_len || hd * 4 >= head_dim) {
        if(hd * 4 < head_dim_pack_k) {
            if(sl * 4 < seq_len_piece) {
                vstore4((FLOAT4)0, 0, output_q + offset_q);
                vstore4((FLOAT4)0, 0, output_q + offset_q + seq_len_piece);
                vstore4((FLOAT4)0, 0, output_q + offset_q + seq_len_piece + seq_len_piece);
                vstore4((FLOAT4)0, 0, output_q + offset_q + seq_len_piece + seq_len_piece + seq_len_piece);
            }
            vstore4((FLOAT4)0, 0, output_k + offset_k);
            vstore4((FLOAT4)0, 0, output_k + offset_k + seq_len_pack_mn);
            vstore4((FLOAT4)0, 0, output_k + offset_k + seq_len_pack_mn + seq_len_pack_mn);
            vstore4((FLOAT4)0, 0, output_k + offset_k + seq_len_pack_mn + seq_len_pack_mn + seq_len_pack_mn);
        }
        vstore4((FLOAT4)0, 0, output_v + offset_v);
        vstore4((FLOAT4)0, 0, output_v + offset_v + head_dim_pack_mn);
        vstore4((FLOAT4)0, 0, output_v + offset_v + head_dim_pack_mn + head_dim_pack_mn);
        vstore4((FLOAT4)0, 0, output_v + offset_v + head_dim_pack_mn + head_dim_pack_mn + head_dim_pack_mn);
        
        return;
    }
    

    const int offset_inp = (((sl * batch + b) * head_num + hn) * 3 * head_dim + 4 * hd) * 4;

    
    if(sl * 4 < seq_len_piece) {
        FLOAT4 temp_0 = vload4(0, input + offset_inp);
        FLOAT4 temp_1 = vload4(0, input + offset_inp + 4);
        FLOAT4 temp_2 = vload4(0, input + offset_inp + 8);
        FLOAT4 temp_3 = vload4(0, input + offset_inp + 12);
        #ifdef HEADDIM_LEAVE
        DEAL_HEAD_DIM_NOT_ALIGN
        #endif
        #ifdef SEQLEN_LEAVE
        DEAL_SEQ_LEN_NOT_ALIGN
        #endif
        vstore4(temp_0, 0, output_q + offset_q);
        vstore4(temp_1, 0, output_q + offset_q + seq_len_piece);
        vstore4(temp_2, 0, output_q + offset_q + seq_len_piece + seq_len_piece);
        vstore4(temp_3, 0, output_q + offset_q + seq_len_piece + seq_len_piece + seq_len_piece);
    }
    
    {
        // K
        FLOAT4 temp_0 = vload4(0, input + offset_inp + 4*head_dim);
        FLOAT4 temp_1 = vload4(0, input + offset_inp + 4*head_dim + 4);
        FLOAT4 temp_2 = vload4(0, input + offset_inp + 4*head_dim + 8);
        FLOAT4 temp_3 = vload4(0, input + offset_inp + 4*head_dim + 12);
        #ifdef HEADDIM_LEAVE
        DEAL_HEAD_DIM_NOT_ALIGN
        #endif
        #ifdef SEQLEN_LEAVE
        DEAL_SEQ_LEN_NOT_ALIGN
        #endif
        
        vstore4(temp_0, 0, output_k + offset_k);
        vstore4(temp_1, 0, output_k + offset_k + seq_len_pack_mn);
        vstore4(temp_2, 0, output_k + offset_k + seq_len_pack_mn + seq_len_pack_mn);
        vstore4(temp_3, 0, output_k + offset_k + seq_len_pack_mn + seq_len_pack_mn + seq_len_pack_mn);
        
        // V
        temp_0 = vload4(0, input + offset_inp + 8 * head_dim);
        temp_1 = vload4(0, input + offset_inp + 8 * head_dim + 4);
        temp_2 = vload4(0, input + offset_inp + 8 * head_dim + 8);
        temp_3 = vload4(0, input + offset_inp + 8 * head_dim + 12);
        #ifdef HEADDIM_LEAVE
        DEAL_HEAD_DIM_NOT_ALIGN
        #endif
        #ifdef SEQLEN_LEAVE
        DEAL_SEQ_LEN_NOT_ALIGN
        #endif
        
        vstore4((FLOAT4){temp_0.x, temp_1.x, temp_2.x, temp_3.x}, 0, output_v + offset_v);
        vstore4((FLOAT4){temp_0.y, temp_1.y, temp_2.y, temp_3.y}, 0, output_v + offset_v + head_dim_pack_mn);
        vstore4((FLOAT4){temp_0.z, temp_1.z, temp_2.z, temp_3.z}, 0, output_v + offset_v + head_dim_pack_mn + head_dim_pack_mn);
        vstore4((FLOAT4){temp_0.w, temp_1.w, temp_2.w, temp_3.w}, 0, output_v + offset_v + head_dim_pack_mn + head_dim_pack_mn + head_dim_pack_mn);
    }
}


#ifndef SOFTMAX_LOCAL_SIZE
    #define SOFTMAX_LOCAL_SIZE 512
#endif

// [outside, axis, inside] -> reduce: inside
__kernel void softmax_inside(GLOBAL_SIZE_3_DIMS
                            __global const FLOAT *input, // [batch * mNumHead, ROUND_UP(seqLen, tile), ROUND_UP(seqLen, tile)]
                            __global FLOAT *output,
                            __private const int inside_len,
                            __private const int4 shape // [batch * mNumHead, ROUND_UP(seqLen, tile), ROUND_UP(seqLen, tile)]
                            ) {
    const int inside = get_global_id(0);
    const int axis = get_global_id(1);
    const int outside = get_global_id(2);
    DEAL_NON_UNIFORM_DIM3(inside, axis, outside);

    const int offset = (outside * shape.y + axis) * shape.z + 0;

#if SOFTMAX_LOCAL_SIZE >= 4
    int lid = get_local_id(0);
    float local sum_mnn[SOFTMAX_LOCAL_SIZE];
    float local max_mnn[SOFTMAX_LOCAL_SIZE];

    /*Compute Max */
    float maxValue = (float)(-FLT_MAX);
    // clip to seq_len
    for (int i=lid; i<inside_len; i+=SOFTMAX_LOCAL_SIZE) {
        maxValue = fmax(maxValue, (float)input[offset+ i]);
    }
    max_mnn[lid] = maxValue;
    barrier(CLK_LOCAL_MEM_FENCE);
    #pragma unroll
    for(int i = SOFTMAX_LOCAL_SIZE/2; i > 0; i >>= 1){
        if (lid < i)
            max_mnn[lid] = fmax(max_mnn[lid], max_mnn[lid + i]);
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    maxValue = max_mnn[0];

    /*Compute Exp Sum*/
    float sumValue = 0;
    for (int i=lid; i<inside_len; i+=SOFTMAX_LOCAL_SIZE) {
        sumValue += exp((float)input[offset+ i] - maxValue);
    }
    sum_mnn[lid] = sumValue;
    barrier(CLK_LOCAL_MEM_FENCE);
    #pragma unroll
    for(int i = SOFTMAX_LOCAL_SIZE/2; i > 0; i >>= 1){
        if (lid < i)
            sum_mnn[lid] = sum_mnn[lid] + sum_mnn[lid + i];
        barrier(CLK_LOCAL_MEM_FENCE);
    }
    sumValue = sum_mnn[0];

    #ifdef OUTPUT_TRANSPOSE
    const int out_offset = (outside * shape.z + 0) * shape.y + axis;
    #endif
    /*Compute Result */
    for (int i=lid; i<inside_len; i+=SOFTMAX_LOCAL_SIZE) {
        float value = exp((float)input[offset+ i] - maxValue) / sumValue;
        #ifdef OUTPUT_TRANSPOSE
        output[out_offset+ i*shape.y] = value;
        #else
        output[offset+ i] = value;
        #endif
    }
    if(shape.z > inside_len){
        for(int i = lid + inside_len; i < shape.z; i+=SOFTMAX_LOCAL_SIZE){
            #ifdef OUTPUT_TRANSPOSE
            output[out_offset+ i*shape.y] = (FLOAT)0;
            #else
            output[offset+ i] = (FLOAT)0;
            #endif
        }
    }
#else
    /*Compute Max */
    float maxValue = (float)(-FLT_MAX);
    // clip to seq_len
    for (int i=0; i<inside_len; i++) {
        maxValue = fmax(maxValue, (float)input[offset+ i]);
    }

    /*Compute Exp Sum*/
    float sumValue = 0;
    for (int i=0; i<inside_len; i++) {
        sumValue += exp((float)input[offset+ i] - maxValue);
    }
    #ifdef OUTPUT_TRANSPOSE
    const int out_offset = (outside * shape.z + 0) * shape.y + axis;
    #endif
    /*Compute Result */
    for (int i=0; i<inside_len; i++) {
        float value = exp((float)input[offset+ i] - maxValue) / sumValue;
        #ifdef OUTPUT_TRANSPOSE
        output[out_offset+ i*shape.y] = value;
        #else
        output[offset+ i] = value;
        #endif
    }
    if(shape.z > inside_len){
        for(int i = inside_len; i < shape.z; i++){
            #ifdef OUTPUT_TRANSPOSE
            output[out_offset+ i*shape.y] = (FLOAT)0;
            #else
            output[offset+ i] = (FLOAT)0;
            #endif
        }
    }
#endif
}

// [N X Y4 4] -> [N Y X]
__kernel void trans_3d_buf(GLOBAL_SIZE_3_DIMS
                        __global const FLOAT* input,
                        __global FLOAT* output,
                        __private const int batch,
                        __private const int width,
                        __private const int height
) {
    int b = get_global_id(2);
    int w = get_global_id(0);
    int h = get_global_id(1);
    DEAL_NON_UNIFORM_DIM3(w, h, b);

    w = w << 3;
    h = h << 3;
    
    const int inp_offset = (b * width + w) * height + h;
    const int out_offset = (b * height + h) * width + w;

    FLOAT8 value_0 = vload8(0, input+inp_offset);
    FLOAT8 value_1 = vload8(0, input+inp_offset + height);
    FLOAT8 value_2 = vload8(0, input+inp_offset + height + height);
    FLOAT8 value_3 = vload8(0, input+inp_offset + height + height + height);
    FLOAT8 value_4 = vload8(0, input+inp_offset + (height << 2));
    FLOAT8 value_5 = vload8(0, input+inp_offset + height * 5);
    FLOAT8 value_6 = vload8(0, input+inp_offset + height * 6);
    FLOAT8 value_7 = vload8(0, input+inp_offset + height * 7);
    
    vstore8((FLOAT8){value_0.s0, value_1.s0, value_2.s0, value_3.s0, value_4.s0, value_5.s0, value_6.s0, value_7.s0}, 0, output + out_offset);
    vstore8((FLOAT8){value_0.s1, value_1.s1, value_2.s1, value_3.s1, value_4.s1, value_5.s1, value_6.s1, value_7.s1}, 0, output + out_offset + width);
    vstore8((FLOAT8){value_0.s2, value_1.s2, value_2.s2, value_3.s2, value_4.s2, value_5.s2, value_6.s2, value_7.s2}, 0, output + out_offset + width + width);
    vstore8((FLOAT8){value_0.s3, value_1.s3, value_2.s3, value_3.s3, value_4.s3, value_5.s3, value_6.s3, value_7.s3}, 0, output + out_offset + width + width + width);
    vstore8((FLOAT8){value_0.s4, value_1.s4, value_2.s4, value_3.s4, value_4.s4, value_5.s4, value_6.s4, value_7.s4}, 0, output + out_offset + (width << 2));
    vstore8((FLOAT8){value_0.s5, value_1.s5, value_2.s5, value_3.s5, value_4.s5, value_5.s5, value_6.s5, value_7.s5}, 0, output + out_offset + width * 5);
    vstore8((FLOAT8){value_0.s6, value_1.s6, value_2.s6, value_3.s6, value_4.s6, value_5.s6, value_6.s6, value_7.s6}, 0, output + out_offset + width * 6);
    vstore8((FLOAT8){value_0.s7, value_1.s7, value_2.s7, value_3.s7, value_4.s7, value_5.s7, value_6.s7, value_7.s7}, 0, output + out_offset + width * 7);
}

__kernel void clip_transpose_qkv(GLOBAL_SIZE_3_DIMS
                              __global const FLOAT *input, // [Batch * mNumHead, ROUND_UP(mHeadDim, tile), ROUND_UP(seqLen, tile)]
                              __global FLOAT *output, // [Batch, seqLen/4, mNumHead * mHeadDim, 4]
                              __private const int tile,
                              __private const int seq_len,
                              __private const int seq_len_piece,
                              __private const int head_num,
                              __private const int head_dim,
                              __private const int batch,
                              __private const int seq_index
) {
    
    const int sl = get_global_id(0); // seqLen_Piece_4
    const int hd = get_global_id(1); // mHeadDim_4
    const int z = get_global_id(2); // Batch * mNumHead
    DEAL_NON_UNIFORM_DIM3(sl, hd, z);
    
    const int b = z / head_num;
    const int hn = z % head_num;
    
    const int seq_len_4 = (seq_len + 3) / 4;
    
    if(seq_index * seq_len_piece / 4 + sl >= seq_len_4) {
        return;
    }
    const int seq_len_pack = seq_len_piece;//((seq_len + tile - 1) / tile) * tile;
    const int head_dim_pack = ((head_dim + tile - 1) / tile) * tile;
    
    const int offset_inp = ((b * head_num + hn) * head_dim_pack + 4 * hd) * seq_len_pack + 4 * sl;
    
    const int offset_out = ((((seq_index * seq_len_piece / 4 + sl) * batch + b) * head_num + hn) * head_dim + 4 * hd) * 4;

    // Q
    FLOAT4 temp_0 = vload4(0, input + offset_inp);
    FLOAT4 temp_1 = vload4(0, input + offset_inp + seq_len_pack);
    FLOAT4 temp_2 = vload4(0, input + offset_inp + 2 * seq_len_pack);
    FLOAT4 temp_3 = vload4(0, input + offset_inp + 3 * seq_len_pack);
    
    vstore4(temp_0, 0, output + offset_out);
    if(4 * hd + 1 > head_dim) return;
    vstore4(temp_1, 0, output + offset_out + 4);
    if(4 * hd + 2 > head_dim) return;
    vstore4(temp_2, 0, output + offset_out + 8);
    if(4 * hd + 3 > head_dim) return;
    vstore4(temp_3, 0, output + offset_out + 12);

}