// Copyright (c) Microsoft Corporation.
// SPDX-License-Identifier: Apache-2.0

// DeepSpeed Team

#include "dequantization_utils.h"
#include "memory_access_utils.h"
#include "quantization_utils.h"
#include "reduction_utils.h"

using rop = reduce::ROpType;

namespace swiz_quant {
constexpr int max_threads = 512;
constexpr int min_threads = 32;

constexpr int step_granularity = 2;
constexpr int h_per_step = step_granularity * quantize::h_per_load;
}  // namespace swiz_quant

template <int numBits, int totalChunks, int threads, quantize::Type quantType>
__global__ void swizzled_quant_kernel(int8_t* quantized_data,
                                      float* quantized_scales,
                                      const __half* uncompressed_data,
                                      int elems_per_group,
                                      int nodes,
                                      int devices_per_node)
{
    cg::thread_block tb = cg::this_thread_block();
    cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);

    // Indexing offsets, same as normal quantization for in-case
    const int block_rank = blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y;
    const int block_offset = block_rank * elems_per_group;
    const int elem_offset = tb.thread_index().x * quantize::h_per_load;
    const int base_offset = block_offset + elem_offset;
    const int stride = tb.size() * quantize::h_per_load;
    const __half* input_base = uncompressed_data + base_offset;

    // Local buffer
    __half2 local_buffer[totalChunks * quantize::h2_per_load];

    quantize::GroupStats<quantType> stats;
#pragma unroll
    for (int i = 0; i < totalChunks; i++) {
        __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load;

        mem_access::load_global<quantize::granularity>(
            iteration_buffer, input_base + i * stride, elem_offset + i * stride < elems_per_group);

#pragma unroll
        for (int j = 0; j < quantize::h2_per_load; j++) { stats.update(iteration_buffer[j]); }
    }

    auto params = stats.template get_params<numBits, threads>(tb, warp);

    const int partition_id = blockIdx.z;
    const int partition_offset = partition_id / devices_per_node;
    const int partition_base = (partition_id % devices_per_node) * nodes;
    const int pipelining_offset = blockIdx.y * (devices_per_node * nodes);
    const int output_partition = (pipelining_offset + partition_base + partition_offset);

    constexpr int out_scalar_effect = 8 / numBits;
    const int out_block_rank = output_partition * gridDim.x + blockIdx.x;
    const int out_block_offset = out_block_rank * elems_per_group / out_scalar_effect;
    const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect;
    int8_t* out_base = quantized_data + out_base_offset;

    const int out_stride = stride / out_scalar_effect;
    constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect;

    if (tb.thread_index().x == 0) { params.store(quantized_scales, out_block_rank); }

#pragma unroll
    for (int i = 0; i < totalChunks; i++) {
        if (i * stride + elem_offset < elems_per_group) {
            int8_t local_output[quantize::h_per_load / out_scalar_effect];
            quantize::_chunk<numBits, quantType>(
                local_output, local_buffer + i * quantize::h2_per_load, params);
            mem_access::store_global<num_int8_out>(out_base + i * out_stride, local_output);
        }
    }
}

#define LAUNCH_SWIZZLE_QUANT(total_chunks, threads)                                           \
    swizzled_quant_kernel<numBits, total_chunks, threads, qType><<<grid, block, 0, stream>>>( \
        q_data, q_scales, input_data, elems_per_group, nodes, devices_per_node);

/*
Swizzled quantization reorganizes the quantized groups in order to better facilitate
communication. As an example of the partitioning scheme we have the following example
of 2 node, 4 device swizzling:

 --- --- --- --- --- --- --- ---
| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
 --- --- --- --- --- --- --- ---
becomes
 --- --- --- --- --- --- --- ---
| 0 | 4 | 1 | 5 | 2 | 6 | 3 | 7 |
 --- --- --- --- --- --- --- ---

Multiple quantization groups may be mapped into a single partition. In order to better support
later pipelining, we may also perform an additional slicing. In two-way slicing, for instance,
the first halves of each partition are concatenated.
*/

template <int numBits, quantize::Type qType>
void launch_swizzled_quant_impl(int8_t* q_data,
                                float* q_scales,
                                const __half* input_data,
                                int groups,
                                int elems_per_group,
                                int pipelining,
                                int nodes,
                                int devices_per_node,
                                cudaStream_t stream)
{
    const int one_step_threads =
        next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step);
    const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads
                                                                         : swiz_quant::max_threads;
    const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads
                                                                : max_threads;

    dim3 block(threads);
    const int groups_per_partition = groups / (nodes * devices_per_node);
    assert(groups_per_partition % pipelining == 0);
    const int contiguous_groups = groups_per_partition / pipelining;
    const int partitions = nodes * devices_per_node;
    dim3 grid(contiguous_groups, pipelining, partitions);

    const int elems_per_step = threads * swiz_quant::h_per_step;
    const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step);
    const int total_unroll = external_unroll * swiz_quant::step_granularity;

    assert(total_unroll % 2 == 0);

    if (threads == 32) {
        LAUNCH_SWIZZLE_QUANT(2, 32);
    } else if (threads == 64) {
        LAUNCH_SWIZZLE_QUANT(2, 64);
    } else if (threads == 128) {
        LAUNCH_SWIZZLE_QUANT(2, 128);
    } else if (threads == 256) {
        LAUNCH_SWIZZLE_QUANT(2, 256);
    } else if (threads == 512) {
        if (total_unroll == 2) {
            LAUNCH_SWIZZLE_QUANT(2, 512);
        } else if (total_unroll == 4) {
            LAUNCH_SWIZZLE_QUANT(4, 512);
        } else if (total_unroll == 6) {
            LAUNCH_SWIZZLE_QUANT(6, 512);
        } else if (total_unroll == 8) {
            LAUNCH_SWIZZLE_QUANT(8, 512);
        } else if (total_unroll == 10) {
            LAUNCH_SWIZZLE_QUANT(10, 512);
        }
    }
}

#define DISPATCH_SWIZZLE_QUANT(num_bits, qtype)                   \
    launch_swizzled_quant_impl<num_bits, qtype>(q_data,           \
                                                q_scales,         \
                                                input_data,       \
                                                groups,           \
                                                elems_per_group,  \
                                                pipelining,       \
                                                nodes,            \
                                                devices_per_node, \
                                                stream);

void launch_swizzled_quant(int8_t* q_data,
                           float* q_scales,
                           const __half* input_data,
                           int num_bits,
                           quantize::Type q_type,
                           int groups,
                           int elems_per_group,
                           int pipelining,
                           int nodes,
                           int devices_per_node,
                           cudaStream_t stream)
{
    if (num_bits == 4) {
        if (q_type == quantize::Type::Asymmetric) {
            DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Asymmetric);
        } else if (q_type == quantize::Type::Symmetric) {
            DISPATCH_SWIZZLE_QUANT(4, quantize::Type::Symmetric);
        }
    } else if (num_bits == 8) {
        if (q_type == quantize::Type::Asymmetric) {
            DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Asymmetric);
        } else if (q_type == quantize::Type::Symmetric) {
            DISPATCH_SWIZZLE_QUANT(8, quantize::Type::Symmetric);
        }
    }
}

template <int numBits, int totalChunks, int threads, quantize::Type quantType>
__global__ void loco_swizzled_quant_kernel(int8_t* quantized_data,
                                           float* quantized_scales,
                                           const __half* uncompressed_data,
                                           __half* error_feedback,
                                           const float err_beta,
                                           int groups,
                                           int elems_per_group,
                                           int pipelining,
                                           int nodes,
                                           int devices_per_node)
{
    cg::thread_block tb = cg::this_thread_block();
    cg::thread_block_tile<hw_warp_size> warp = cg::tiled_partition<hw_warp_size>(tb);

    // Indexing offsets, same as normal quantization for in-case
    const int block_rank_data =
        blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y;
    const int block_offset_data = block_rank_data * elems_per_group;
    const int elem_offset = tb.thread_index().x * quantize::h_per_load;
    const int base_offset_data = block_offset_data + elem_offset;
    const int stride = tb.size() * quantize::h_per_load;
    const __half* uncompressed_data_base = uncompressed_data + base_offset_data;

    const int partition_id = blockIdx.z;
    const int partition_offset = partition_id / devices_per_node;
    const int partition_base = (partition_id % devices_per_node) * nodes;
    const int pipelining_offset = blockIdx.y * (devices_per_node * nodes);
    const int output_partition = (pipelining_offset + partition_base + partition_offset);
    const int block_rank_err = output_partition * gridDim.x + blockIdx.x;

    const int block_offset_err = block_rank_err * elems_per_group;
    const int base_offset_err = block_offset_err + elem_offset;
    __half* error_feedback_base = error_feedback + base_offset_err;

    __half2 local_buffer[totalChunks * quantize::h2_per_load];
    __half2 err_buffer[totalChunks * quantize::h2_per_load];

    quantize::GroupStats<quantType> stats;

#pragma unroll
    for (int i = 0; i < totalChunks; i++) {
        __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load;
        __half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load;
        const int i_stride = i * stride;
        bool do_loads = (elem_offset + i_stride) < elems_per_group;

        mem_access::load_global<quantize::granularity>(
            iteration_buffer, uncompressed_data_base + i_stride, do_loads);

        mem_access::load_global<quantize::granularity>(
            iter_err_buffer, error_feedback_base + i_stride, do_loads);

#pragma unroll
        for (int j = 0; j < quantize::h2_per_load; j++) {
            iteration_buffer[j] = __hadd2(iteration_buffer[j], iter_err_buffer[j]);
            stats.update(iteration_buffer[j]);
        }
    }

    auto params = stats.template get_params<numBits, threads>(tb, warp);

    // Initialize dequantization parameters based on params
    auto de_params = params;
    de_params.scale = 1.0f / params.scale;
    if constexpr (quantType == quantize::Type::Asymmetric) { de_params.offset = params.offset; }

    if (threadIdx.x == 0) { params.store(quantized_scales, block_rank_err); }

    constexpr int out_scalar_effect = 8 / numBits;
    const int out_block_offset = block_rank_err * elems_per_group / out_scalar_effect;
    const int out_base_offset = out_block_offset + elem_offset / out_scalar_effect;
    int8_t* out_base = quantized_data + out_base_offset;

    const int out_stride = stride / out_scalar_effect;
    constexpr int num_int8_out = quantize::h_per_load / out_scalar_effect;

#pragma unroll
    for (int i = 0; i < totalChunks; i++) {
        const int i_stride = i * stride;
        __half2* iteration_buffer = local_buffer + i * quantize::h2_per_load;
        __half2* iter_err_buffer = err_buffer + i * quantize::h2_per_load;

        if (i_stride + elem_offset < elems_per_group) {
            int8_t local_output[quantize::h_per_load / out_scalar_effect];
            quantize::_chunk<numBits, quantType>(local_output, iteration_buffer, params);
            mem_access::store_global<num_int8_out>(out_base + i * out_stride, local_output);

            // Dequantize the quantized output to compute the dequantized value
            __half2 dequant_buffer[quantize::h2_per_load];
            dequantize::chunk<numBits, quantType>(dequant_buffer, local_output, de_params);

// Compute new error: sum - dequant_buffer
#pragma unroll
            for (int k = 0; k < quantize::h2_per_load; k++) {
                // __half2 to float2
                float2 iter_buf_f = __half22float2(iteration_buffer[k]);
                float2 dequant_buf_f = __half22float2(dequant_buffer[k]);

                // Update within float precision
                float2 new_error_f;
                new_error_f.x = iter_buf_f.x - dequant_buf_f.x;
                new_error_f.y = iter_buf_f.y - dequant_buf_f.y;

                float2 iter_err_buf_f = __half22float2(iter_err_buffer[k]);

                iter_err_buf_f.x = err_beta * iter_err_buf_f.x + (1.0f - err_beta) * new_error_f.x;
                iter_err_buf_f.y = err_beta * iter_err_buf_f.y + (1.0f - err_beta) * new_error_f.y;

                // float2 back to __half2
                iter_err_buffer[k] = __float22half2_rn(iter_err_buf_f);
            }
            __half2* error_feedback_base_h2 = reinterpret_cast<__half2*>(error_feedback_base);
            mem_access::store_global<quantize::granularity>(error_feedback_base_h2 + i_stride / 2,
                                                            iter_err_buffer);
        }
    }
}

#define LAUNCH_LOCO_SWIZZLE_QUANT(total_chunks, threads)              \
    loco_swizzled_quant_kernel<numBits, total_chunks, threads, qType> \
        <<<grid, block, 0, stream>>>(output_data,                     \
                                     params,                          \
                                     input_data,                      \
                                     error_feedback,                  \
                                     err_beta,                        \
                                     groups,                          \
                                     elems_per_group,                 \
                                     pipelining,                      \
                                     nodes,                           \
                                     devices_per_node);

template <int numBits, quantize::Type qType>
void launch_loco_swizzled_quant_impl(int8_t* output_data,
                                     float* params,
                                     const __half* input_data,
                                     __half* error_feedback,
                                     const float err_beta,
                                     int groups,
                                     int elems_per_group,
                                     int pipelining,
                                     int nodes,
                                     int devices_per_node,
                                     cudaStream_t stream)
{
    const int one_step_threads =
        next_pow2((elems_per_group + swiz_quant::h_per_step - 1) / swiz_quant::h_per_step);
    const int max_threads = (one_step_threads < swiz_quant::max_threads) ? one_step_threads
                                                                         : swiz_quant::max_threads;
    const int threads = (max_threads < swiz_quant::min_threads) ? swiz_quant::min_threads
                                                                : max_threads;

    dim3 block(threads);
    const int groups_per_partition = groups / (nodes * devices_per_node);
    assert(groups_per_partition % pipelining == 0);
    const int contiguous_groups = groups_per_partition / pipelining;
    const int partitions = nodes * devices_per_node;
    dim3 grid(contiguous_groups, pipelining, partitions);

    const int elems_per_step = threads * swiz_quant::h_per_step;
    const int external_unroll = ((elems_per_group + elems_per_step - 1) / elems_per_step);
    const int total_unroll = external_unroll * swiz_quant::step_granularity;

    assert(total_unroll % 2 == 0);

    if (threads == 32) {
        LAUNCH_LOCO_SWIZZLE_QUANT(2, 32);
    } else if (threads == 64) {
        LAUNCH_LOCO_SWIZZLE_QUANT(2, 64);
    } else if (threads == 128) {
        LAUNCH_LOCO_SWIZZLE_QUANT(2, 128);
    } else if (threads == 256) {
        LAUNCH_LOCO_SWIZZLE_QUANT(2, 256);
    } else if (threads == 512) {
        if (total_unroll == 2) {
            LAUNCH_LOCO_SWIZZLE_QUANT(2, 512);
        } else if (total_unroll == 4) {
            LAUNCH_LOCO_SWIZZLE_QUANT(4, 512);
        } else if (total_unroll == 6) {
            LAUNCH_LOCO_SWIZZLE_QUANT(6, 512);
        } else if (total_unroll == 8) {
            LAUNCH_LOCO_SWIZZLE_QUANT(8, 512);
        } else if (total_unroll == 10) {
            LAUNCH_LOCO_SWIZZLE_QUANT(10, 512);
        }
    }
}

#define DISPATCH_LOCO_SWIZZLE_QUANT(num_bits, qtype)                   \
    launch_loco_swizzled_quant_impl<num_bits, qtype>(output_data,      \
                                                     params,           \
                                                     input_data,       \
                                                     error_feedback,   \
                                                     err_beta,         \
                                                     groups,           \
                                                     elems_per_group,  \
                                                     pipelining,       \
                                                     nodes,            \
                                                     devices_per_node, \
                                                     stream);

void launch_loco_swizzled_quant(int8_t* output_data,
                                float* params,
                                const __half* input_data,
                                __half* error_feedback,
                                const float err_beta,
                                int num_bits,
                                quantize::Type q_type,
                                int groups,
                                int elems_per_group,
                                int pipelining,
                                int nodes,
                                int devices_per_node,
                                cudaStream_t stream)
{
    if (num_bits == 4) {
        if (q_type == quantize::Type::Asymmetric) {
            DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Asymmetric);
        } else if (q_type == quantize::Type::Symmetric) {
            DISPATCH_LOCO_SWIZZLE_QUANT(4, quantize::Type::Symmetric);
        }
    } else if (num_bits == 8) {
        if (q_type == quantize::Type::Asymmetric) {
            DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Asymmetric);
        } else if (q_type == quantize::Type::Symmetric) {
            DISPATCH_LOCO_SWIZZLE_QUANT(8, quantize::Type::Symmetric);
        }
    }
}