/*
 * Copyright (c) 2019-2023, NVIDIA CORPORATION.  All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "src/fastertransformer/kernels/transpose_int8_kernels.h"

namespace fastertransformer {

// src is the result of batch MM, whose size is batch_size*head_num*(seq_len, size_per_head), CUBLASLT_ORDER_COL32
// dst is of m = batch_size*seq_len, k(n) = head_num*size_per_head, CUBLASLT_ORDER_COL32
// grid(seq_len, batch_size)
// block(size_per_head/4, head_num)
// assume size_per_head is multiples of 32
__global__ void transpose_COL32_kernel(char4*       dst,
                                       const int4*  src,
                                       const int    batch_size,
                                       const int    seq_len,
                                       const int    head_num,
                                       const int    size_per_head,
                                       const float* v_buf_addBias_deQFactor,
                                       const float* qk_afterSM_deQFactor,
                                       const float* out_scale_ptr,
                                       const int    batch_size_x_seq_len,
                                       const int    seq_len_x_size_per_head)
{
    const float scale      = __ldg(v_buf_addBias_deQFactor) * __ldg(qk_afterSM_deQFactor) * __ldg(out_scale_ptr);
    int         threadIdx4 = threadIdx.x << 2;
    int         batch_id   = blockIdx.y;
    int         seq_id     = blockIdx.x;
    int         head_id    = threadIdx.y;
    // get the (row, col) output layout of m*k
    // m = batch_size*seq_len
    // k = head_num*size_per_head
    int mk_row = batch_id * seq_len + seq_id;
    int mk_col = head_id * size_per_head + threadIdx4;
    // get the (row, col) layout of COL32; leading dimension = 32*m = 32*batch_size*seq_len
    int COL32_row = (mk_row << 5) + (mk_col & 31);
    // int COL32_col = mk_col >> 5;
    int outIdx = ((mk_col & 0xffffffe0) * batch_size_x_seq_len + COL32_row) >> 2;

    // get the (row, col) input layout of m'*k'
    // m' = seq_len
    // k' = size_per_head
    mk_row = seq_id;
    mk_col = threadIdx4;
    // get the (row, col) layout of COL32; leading dimension = 32*m' = 32*seq_len
    COL32_row = (mk_row << 5) + (mk_col & 31);
    // COL32_col = mk_col >> 5;

    int inIdx =
        ((batch_id * head_num + head_id) * seq_len_x_size_per_head + (mk_col & 0xffffffe0) * seq_len + COL32_row) >> 2;
    char4 tmp;

    int4 srcTmp4 = __ldg(src + inIdx);
    tmp.x        = float_to_int8_rn(srcTmp4.x * scale);
    tmp.y        = float_to_int8_rn(srcTmp4.y * scale);
    tmp.z        = float_to_int8_rn(srcTmp4.z * scale);
    tmp.w        = float_to_int8_rn(srcTmp4.w * scale);
    dst[outIdx]  = tmp;
}

void invokeTransposeCOL32(int8_t*      dst,
                          const int*   src,
                          const int    batch_size,
                          const int    seq_len,
                          const int    head_num,
                          const int    size_per_head,
                          const float* v_buf_addBias_deQFactor,
                          const float* qk_afterSM_deQFactor,
                          const float* out_scale_ptr,
                          cudaStream_t stream)
{
    assert(size_per_head % 32 == 0);
    transpose_COL32_kernel<<<dim3(seq_len, batch_size), dim3(size_per_head / 4, head_num), 0, stream>>>(
        (char4*)dst,
        (const int4*)src,
        batch_size,
        seq_len,
        head_num,
        size_per_head,
        v_buf_addBias_deQFactor,
        qk_afterSM_deQFactor,
        out_scale_ptr,
        batch_size * seq_len,
        seq_len * size_per_head);
}

// src is the result of batch MM, whose size is batch_size*head_num*(seq_len, size_per_head), CUBLASLT_ORDER_COL32
// dst is of m = batch_size*seq_len, k(n) = head_num*size_per_head, CUBLASLT_ORDER_COL32
// grid(seq_len, batch_size)
// block(size_per_head/4, head_num)
// assume size_per_head is multiples of 32
__global__ void transpose_COL32_kernel(int8_t*       dst,
                                       const int8_t* src,
                                       const int     batch_size,
                                       const int     seq_len,
                                       const int     head_num,
                                       const int     size_per_head,
                                       const float*  bmm2_deQFactor,
                                       const float*  out_scale_ptr,
                                       const int     batch_size_x_seq_len,
                                       const int     seq_len_x_size_per_head)
{
    int threadIdx4 = threadIdx.x << 2;
    int batch_id   = blockIdx.y;
    int seq_id     = blockIdx.x;
    int head_id    = threadIdx.y;
    // get the (row, col) output layout of m*k
    // m = batch_size*seq_len
    // k = head_num*size_per_head
    int mk_row = batch_id * seq_len + seq_id;
    int mk_col = head_id * size_per_head + threadIdx4;
    // get the (row, col) layout of COL32; leading dimension = 32*m = 32*batch_size*seq_len
    int COL32_row = (mk_row << 5) + (mk_col & 31);
    int COL32_col = mk_col >> 5;
    int outIdx    = ((COL32_col << 5) * batch_size_x_seq_len + COL32_row) >> 2;

    // get the (row, col) input layout of m'*k'
    // m' = seq_len
    // k' = size_per_head
    mk_row = seq_id;
    mk_col = threadIdx4;
    // get the (row, col) layout of COL32; leading dimension = 32*m' = 32*seq_len
    COL32_row = (mk_row << 5) + (mk_col & 31);
    COL32_col = mk_col >> 5;

    int inIdx =
        ((batch_id * head_num + head_id) * seq_len_x_size_per_head + (COL32_col << 5) * seq_len + COL32_row) >> 2;
    const char4* src_ptr4 = (const char4*)src;
    char4*       dst_ptr4 = (char4*)dst;
    dst_ptr4[outIdx]      = __ldg(src_ptr4 + inIdx);
}

void invokeTransposeCOL32(int8_t*       dst,
                          const int8_t* src,
                          const int     batch_size,
                          const int     seq_len,
                          const int     head_num,
                          const int     size_per_head,
                          const float*  bmm2_deQFactor,
                          const float*  out_scale_ptr,
                          cudaStream_t  stream)
{
    assert(size_per_head % 32 == 0);
    transpose_COL32_kernel<<<dim3(seq_len, batch_size), dim3(size_per_head / 4, head_num), 0, stream>>>(
        dst,
        src,
        batch_size,
        seq_len,
        head_num,
        size_per_head,
        bmm2_deQFactor,
        out_scale_ptr,
        batch_size * seq_len,
        seq_len * size_per_head);
}

// src is the result of batch MM, whose size is batch_size*head_num*(seq_len, size_per_head), CUBLASLT_ORDER_COL32
// dst is of m = valid_word_num, k(n) = head_num*size_per_head, CUBLASLT_ORDER_COL32
// grid(seq_len, batch_size)
// block(size_per_head/4, head_num)
// assume size_per_head is multiples of 32
__global__ void transpose_COL32_rebuild_padding_kernel(int8_t*        dst,
                                                       const int32_t* src,
                                                       const int*     sequence_id_map,
                                                       const int      valid_word_num,
                                                       const int      batch_size,
                                                       const int      seq_len,
                                                       const int      head_num,
                                                       const int      size_per_head,
                                                       const float*   v_buf_addBias_deQFactor,
                                                       const float*   qk_afterSM_deQFactor,
                                                       const float*   out_scale_ptr,
                                                       const int      seq_len_x_size_per_head)
{
    const float scale      = __ldg(v_buf_addBias_deQFactor) * __ldg(qk_afterSM_deQFactor) * __ldg(out_scale_ptr);
    int         threadIdx4 = threadIdx.x << 2;
    int         batch_id   = blockIdx.y;
    int         seq_id     = blockIdx.x;
    int         head_id    = threadIdx.y;
    // get the (row, col) output layout of m*k
    // m = valid_word_num
    // k = head_num*size_per_head
    int mk_row = __ldg(sequence_id_map + batch_id * seq_len + seq_id);
    if (mk_row >= 0) {
        int mk_col = head_id * size_per_head + threadIdx4;
        // get the (row, col) layout of COL32; leading dimension = 32*m = 32*valid_word_num
        int COL32_row = (mk_row << 5) + (mk_col & 31);
        int COL32_col = mk_col >> 5;
        int outIdx    = ((COL32_col << 5) * valid_word_num + COL32_row) >> 2;

        // get the (row, col) input layout of m'*k'
        // m' = seq_len
        // k' = size_per_head
        mk_row = seq_id;
        mk_col = threadIdx4;
        // get the (row, col) layout of COL32; leading dimension = 32*m' = 32*seq_len
        COL32_row = (mk_row << 5) + (mk_col & 31);
        COL32_col = mk_col >> 5;

        int inIdx = (batch_id * head_num + head_id) * seq_len_x_size_per_head + (COL32_col << 5) * seq_len + COL32_row;
        char4 tmp;
        tmp.x            = float_to_int8_rn(__ldg(src + inIdx) * scale);
        tmp.y            = float_to_int8_rn(__ldg(src + inIdx + 1) * scale);
        tmp.z            = float_to_int8_rn(__ldg(src + inIdx + 2) * scale);
        tmp.w            = float_to_int8_rn(__ldg(src + inIdx + 3) * scale);
        char4* dst_ptr4  = (char4*)dst;
        dst_ptr4[outIdx] = tmp;
    }
}

void invokeTransposeCOL32RebuildPadding(int8_t*      dst,
                                        const int*   src,
                                        const int*   sequence_id_map,
                                        const int    valid_word_num,
                                        const int    batch_size,
                                        const int    seq_len,
                                        const int    head_num,
                                        const int    size_per_head,
                                        const float* v_buf_addBias_deQFactor,
                                        const float* qk_afterSM_deQFactor,
                                        const float* out_scale_ptr,
                                        cudaStream_t stream)
{
    assert(size_per_head % 32 == 0);
    transpose_COL32_rebuild_padding_kernel<<<dim3(seq_len, batch_size), dim3(size_per_head / 4, head_num), 0, stream>>>(
        dst,
        src,
        sequence_id_map,
        valid_word_num,
        batch_size,
        seq_len,
        head_num,
        size_per_head,
        v_buf_addBias_deQFactor,
        qk_afterSM_deQFactor,
        out_scale_ptr,
        seq_len * size_per_head);
}

// src is the result of batch MM, whose size is batch_size*head_num*(seq_len, size_per_head), CUBLASLT_ORDER_COL32
// dst is of m = valid_word_num, k(n) = head_num*size_per_head, CUBLASLT_ORDER_COL32
// grid(seq_len, batch_size)
// block(size_per_head/4, head_num)
// assume size_per_head is multiples of 32
__global__ void transpose_COL32_rebuild_padding_kernel(int8_t*       dst,
                                                       const int8_t* src,
                                                       const int*    sequence_id_map,
                                                       const int     valid_word_num,
                                                       const int     batch_size,
                                                       const int     seq_len,
                                                       const int     head_num,
                                                       const int     size_per_head,
                                                       const float*  bmm2_deQFactor,
                                                       const float*  out_scale_ptr,
                                                       const int     seq_len_x_size_per_head)
{
    int threadIdx4 = threadIdx.x << 2;
    int batch_id   = blockIdx.y;
    int seq_id     = blockIdx.x;
    int head_id    = threadIdx.y;
    // get the (row, col) output layout of m*k
    // m = valid_word_num
    // k = head_num*size_per_head
    int mk_row = __ldg(sequence_id_map + batch_id * seq_len + seq_id);
    if (mk_row >= 0) {
        int mk_col = head_id * size_per_head + threadIdx4;
        // get the (row, col) layout of COL32; leading dimension = 32*m = 32*valid_word_num
        int COL32_row = (mk_row << 5) + (mk_col & 31);
        int COL32_col = mk_col >> 5;
        int outIdx    = ((COL32_col << 5) * valid_word_num + COL32_row) >> 2;

        // get the (row, col) input layout of m'*k'
        // m' = seq_len
        // k' = size_per_head
        mk_row = seq_id;
        mk_col = threadIdx4;
        // get the (row, col) layout of COL32; leading dimension = 32*m' = 32*seq_len
        COL32_row = (mk_row << 5) + (mk_col & 31);
        COL32_col = mk_col >> 5;

        int inIdx =
            ((batch_id * head_num + head_id) * seq_len_x_size_per_head + (COL32_col << 5) * seq_len + COL32_row) >> 2;

        const char4* src_ptr4 = (const char4*)src;

        char4* dst_ptr4  = (char4*)dst;
        dst_ptr4[outIdx] = __ldg(src_ptr4 + inIdx);
    }
}

void invokeTransposeCOL32RebuildPadding(int8_t*       dst,
                                        const int8_t* src,
                                        const int*    sequence_id_map,
                                        const int     valid_word_num,
                                        const int     batch_size,
                                        const int     seq_len,
                                        const int     head_num,
                                        const int     size_per_head,
                                        const float*  bmm2_deQFactor,
                                        const float*  out_scale_ptr,
                                        cudaStream_t  stream)
{
    assert(size_per_head % 32 == 0);
    transpose_COL32_rebuild_padding_kernel<<<dim3(seq_len, batch_size), dim3(size_per_head / 4, head_num), 0, stream>>>(
        dst,
        src,
        sequence_id_map,
        valid_word_num,
        batch_size,
        seq_len,
        head_num,
        size_per_head,
        bmm2_deQFactor,
        out_scale_ptr,
        seq_len * size_per_head);
}

// src is the result of batch MM, whose size is batch_size*head_num*(seq_len, size_per_head), CUBLASLT_ORDER_COL32
// dst is of m = batch_size*seq_len, k(n) = head_num*size_per_head, row major
// grid(seq_len, batch_size)
// block(size_per_head/4, head_num)
// assume size_per_head is multiples of 32
__global__ void transpose_COL32_ROW_kernel(int8_t*       dst,
                                           const int8_t* src,
                                           const int     batch_size,
                                           const int     seq_len,
                                           const int     head_num,
                                           const int     size_per_head,
                                           const float*  bmm2_deQFactor,
                                           const float*  out_scale_ptr,
                                           const int     head_num_x_size_per_head,
                                           const int     seq_len_x_size_per_head)
{
    int threadIdx4 = threadIdx.x << 2;
    int batch_id   = blockIdx.y;
    int seq_id     = blockIdx.x;
    int head_id    = threadIdx.y;
    // get the (row, col) output layout of m*k
    // m = batch_size*seq_len
    // k = head_num*size_per_head
    int mk_row = batch_id * seq_len + seq_id;
    int mk_col = head_id * size_per_head + threadIdx4;
    int outIdx = (mk_row * head_num_x_size_per_head + mk_col) >> 2;

    // get the (row, col) input layout of m'*k'
    // m' = seq_len
    // k' = size_per_head
    mk_row = seq_id;
    mk_col = threadIdx4;
    // get the (row, col) layout of COL32; leading dimension = 32*m' = 32*seq_len
    int COL32_row = (mk_row << 5) + (mk_col & 31);
    int COL32_col = mk_col >> 5;

    int inIdx =
        ((batch_id * head_num + head_id) * seq_len_x_size_per_head + (COL32_col << 5) * seq_len + COL32_row) >> 2;
    const char4* src_ptr4 = (const char4*)src;
    char4*       dst_ptr4 = (char4*)dst;
    dst_ptr4[outIdx]      = __ldg(src_ptr4 + inIdx);
}

void invokeTransposeCOL32ToRow(int8_t*       dst,
                               const int8_t* src,
                               const int     batch_size,
                               const int     seq_len,
                               const int     head_num,
                               const int     size_per_head,
                               const float*  bmm2_deQFactor,
                               const float*  out_scale_ptr,
                               cudaStream_t  stream)
{
    assert(size_per_head % 32 == 0);
    transpose_COL32_ROW_kernel<<<dim3(seq_len, batch_size), dim3(size_per_head / 4, head_num), 0, stream>>>(
        dst,
        src,
        batch_size,
        seq_len,
        head_num,
        size_per_head,
        bmm2_deQFactor,
        out_scale_ptr,
        head_num * size_per_head,
        seq_len * size_per_head);
}

// src is the result of batch MM, whose size is batch_size*head_num*(seq_len, size_per_head), CUBLASLT_ORDER_COL32
// dst is of m = valid_word_num, k(n) = head_num*size_per_head, row major
// grid(seq_len, batch_size)
// block(size_per_head/4, head_num)
// assume size_per_head is multiples of 32
__global__ void transpose_COL32_ROW_rebuild_padding_kernel(int8_t*       dst,
                                                           const int8_t* src,
                                                           const int*    sequence_id_map,
                                                           const int     valid_word_num,
                                                           const int     batch_size,
                                                           const int     seq_len,
                                                           const int     head_num,
                                                           const int     size_per_head,
                                                           const float*  bmm2_deQFactor,
                                                           const float*  out_scale_ptr,
                                                           const int     seq_len_x_size_per_head,
                                                           const int     head_num_x_size_per_head)
{
    int threadIdx4 = threadIdx.x << 2;
    int batch_id   = blockIdx.y;
    int seq_id     = blockIdx.x;
    int head_id    = threadIdx.y;
    // get the (row, col) output layout of m*k
    // m = valid_word_num
    // k = head_num*size_per_head
    int mk_row = __ldg(sequence_id_map + batch_id * seq_len + seq_id);
    if (mk_row >= 0) {
        int mk_col = head_id * size_per_head + threadIdx4;
        int outIdx = (mk_row * head_num_x_size_per_head + mk_col) >> 2;

        // get the (row, col) input layout of m'*k'
        // m' = seq_len
        // k' = size_per_head
        mk_row = seq_id;
        mk_col = threadIdx4;
        // get the (row, col) layout of COL32; leading dimension = 32*m' = 32*seq_len
        int COL32_row = (mk_row << 5) + (mk_col & 31);
        int COL32_col = mk_col >> 5;

        int inIdx =
            ((batch_id * head_num + head_id) * seq_len_x_size_per_head + (COL32_col << 5) * seq_len + COL32_row) >> 2;

        const char4* src_ptr4 = (const char4*)src;

        char4* dst_ptr4  = (char4*)dst;
        dst_ptr4[outIdx] = __ldg(src_ptr4 + inIdx);
    }
}

void invokeTransposeCOL32ToRowRebuildPadding(int8_t*       dst,
                                             const int8_t* src,
                                             const int*    sequence_id_map,
                                             const int     valid_word_num,
                                             const int     batch_size,
                                             const int     seq_len,
                                             const int     head_num,
                                             const int     size_per_head,
                                             const float*  bmm2_deQFactor,
                                             const float*  out_scale_ptr,
                                             cudaStream_t  stream)
{
    assert(size_per_head % 32 == 0);
    transpose_COL32_ROW_rebuild_padding_kernel<<<dim3(seq_len, batch_size),
                                                 dim3(size_per_head / 4, head_num),
                                                 0,
                                                 stream>>>(dst,
                                                           src,
                                                           sequence_id_map,
                                                           valid_word_num,
                                                           batch_size,
                                                           seq_len,
                                                           head_num,
                                                           size_per_head,
                                                           bmm2_deQFactor,
                                                           out_scale_ptr,
                                                           seq_len * size_per_head,
                                                           head_num * size_per_head);
}

__global__ void transpose_int8_tensor(int8_t* arr_t, const int8_t* arr, size_t dim_0, size_t dim_1)
{
    extern __shared__ int8_t shmem[];

    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    int local_id = threadIdx.y * blockDim.x + threadIdx.x;

    if (x >= dim_1 || y >= dim_0) {
        return;
    }

    shmem[local_id] = arr[y * dim_1 + x];
    __syncthreads();

    int new_x = local_id % blockDim.y;
    int new_y = local_id / blockDim.y;

    int x_t                  = blockIdx.y * blockDim.y + new_x;
    int y_t                  = blockIdx.x * blockDim.x + new_y;
    arr_t[y_t * dim_0 + x_t] = shmem[new_x * blockDim.x + new_y];
}

void invokeTransposeInt8Tensor(const Tensor& x_t, const Tensor& x, cudaStream_t stream)
{
    dim3 block_dim(32, 16);
    dim3 grid_dim((x.shape[1] + block_dim.x - 1) / block_dim.x, (x.shape[0] + block_dim.y - 1) / block_dim.y);

    const size_t shmem_size = block_dim.x * block_dim.y * sizeof(int8_t);

    transpose_int8_tensor<<<grid_dim, block_dim, shmem_size, stream>>>(
        x_t.getPtr<int8_t>(), x.getPtr<int8_t>(), x.shape[0], x.shape[1]);
}

}  // namespace fastertransformer
