/*
 * Copyright (c) 2023-2024 Arm Limited.
 *
 * SPDX-License-Identifier: MIT
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to
 * deal in the Software without restriction, including without limitation the
 * rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
 * sell copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
#if defined(__aarch64__)

#include "src/core/NEON/kernels/NEReorderKernel.h"

#include "arm_compute/core/Helpers.h"
#include "arm_compute/core/Validate.h"
#include "arm_compute/runtime/Scheduler.h"

#include "src/common/utils/Log.h"
#include "src/core/NEON/kernels/arm_gemm/transform.hpp"

namespace arm_compute
{

void NEReorderKernel::run(const Window &window, const ThreadInfo &info)
{
    ARM_COMPUTE_UNUSED(info);
    ARM_COMPUTE_ERROR_ON_UNCONFIGURED_KERNEL(this);
    ARM_COMPUTE_ERROR_ON_INVALID_SUBWINDOW(INEKernel::window(), window);
    switch (_input->info()->data_type())
    {
        case DataType::F32:
        {
            const int ksize_rows_elements = _xmax * _ksize;
            const int jump_rows           = ksize_rows_elements * window.x().start();
            const int k_start             = window.x().start() * _ksize;
            const int k_end               = std::min(window.x().end() * _ksize, _kmax);
            const int stride              = _kmax;
            if (k_start < k_end)
            {
                switch (_output_wf)
                {
                    case WeightFormat::OHWIo4:
                    {
                        switch (_output->info()->data_type())
                        {
                            case DataType::F32:
                                arm_gemm::Transform<4, 1, true, arm_gemm::VLType::None>(
                                    reinterpret_cast<float *>(_output->buffer()) + jump_rows,
                                    reinterpret_cast<float *>(_input->buffer()), stride, k_start, k_end, 0, _xmax);
                                break;
                            case DataType::BFLOAT16:
                                arm_gemm::Transform<4, 4, true, arm_gemm::VLType::None>(
                                    reinterpret_cast<bfloat16 *>(_output->buffer()) + jump_rows,
                                    reinterpret_cast<float *>(_input->buffer()), stride, k_start, k_end, 0, _xmax);
                                break;
                            default:
                                ARM_COMPUTE_ERROR("Unsupported data type!");
                        }
                        break;
                    }
#if defined(ARM_COMPUTE_ENABLE_SVE)
                    case WeightFormat::OHWIo8:
                    {
                        switch (_output->info()->data_type())
                        {
                            case DataType::F32:
                                arm_gemm::Transform<1, 1, true, arm_gemm::VLType::SVE>(
                                    reinterpret_cast<float *>(_output->buffer()) + jump_rows,
                                    reinterpret_cast<float *>(_input->buffer()), stride, k_start, k_end, 0, _xmax);
                                break;
                            case DataType::BFLOAT16:
                                arm_gemm::Transform<2, 4, true, arm_gemm::VLType::SVE>(
                                    reinterpret_cast<bfloat16 *>(_output->buffer()) + jump_rows,
                                    reinterpret_cast<float *>(_input->buffer()), stride, k_start, k_end, 0, _xmax);
                                break;
                            default:
                                ARM_COMPUTE_ERROR("Unsupported data type!");
                        }
                        break;
                    }
#endif /* ARM_COMPUTE_ENABLE_SVE */
                    default:
                    {
                        ARM_COMPUTE_ERROR("Unsupported data type!");
                        break;
                    }
                }
            }
            break;
        }
        default:
            ARM_COMPUTE_ERROR("Unsupported data type!");
    }
}

NEReorderKernel::NEReorderKernel()
    : _input(nullptr),
      _output(nullptr),
      _ksize(0),
      _kmax(0),
      _xmax(0),
      _input_wf(WeightFormat::ANY),
      _output_wf(WeightFormat::ANY)
{
}

void NEReorderKernel::configure(const ITensor            *input,
                                ITensor                  *output,
                                arm_compute::WeightFormat input_wf,
                                arm_compute::WeightFormat output_wf)
{
    ARM_COMPUTE_LOG_PARAMS(input, output, input_wf, output_wf);
    ARM_COMPUTE_ERROR_ON_NULLPTR(input, output);
    ARM_COMPUTE_ERROR_THROW_ON(validate(input->info(), output->info(), input_wf, output_wf));

    // Set variables
    _input     = input;
    _output    = output;
    _input_wf  = input_wf;
    _output_wf = output_wf;

    // Setting parameters for transform
    auto dims = input->info()->num_dimensions();
    switch (dims)
    {
        case 2:
        {
            _xmax = input->info()->dimension(0); // Number of columns in input matrix
            _kmax = input->info()->dimension(1); // Number of rows in input matrix
            break;
        }
        case 4:
        {
            _xmax = input->info()->dimension(2); // Number of columns in input matrix
            _kmax = input->info()->dimension(3); // Number of rows in input matrix
            break;
        }
        default:
        {
            ARM_COMPUTE_ERROR("Only 2 or 4 dimensions supported.");
        }
    }

    // Configure kernel window
    // Window size is set by rows / _ksize
    Window win;
    int    window_size = 0;
    switch (_output_wf)
    {
#if defined(ARM_COMPUTE_ENABLE_SVE)
        case WeightFormat::OHWIo8:
        {
            _ksize      = 8;
            window_size = _kmax / _ksize;
            break;
        }
#endif /* ARM_COMPUTE_ENABLE_SVE */
        case WeightFormat::OHWIo4:
        {
            _ksize      = 4;
            window_size = _kmax / _ksize;
            break;
        }
        default:
        {
            ARM_COMPUTE_ERROR("Unsupported weight format.");
            break;
        }
    }
    if (_kmax % _ksize != 0)
    {
        window_size += 1;
    }

    win.set(Window::DimX, Window::Dimension(0, window_size, 1));

    INEKernel::configure(win);
}

Status NEReorderKernel::validate(const ITensorInfo        *input,
                                 const ITensorInfo        *output,
                                 arm_compute::WeightFormat input_wf,
                                 arm_compute::WeightFormat output_wf)
{
    ARM_COMPUTE_RETURN_ERROR_ON_NULLPTR(input, output);
    ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() == DataType::UNKNOWN);
    if (output->tensor_shape().total_size() != 0)
    {
        ARM_COMPUTE_RETURN_ERROR_ON(input->data_type() != DataType::F32);
        ARM_COMPUTE_RETURN_ERROR_ON(output->data_type() != DataType::F32 && output->data_type() != DataType::BFLOAT16);
        ARM_COMPUTE_RETURN_ERROR_ON_MISMATCHING_QUANTIZATION_INFO(input, output);
        // Only input WeightFormat OHWI supported
        ARM_COMPUTE_RETURN_ERROR_ON(input_wf != arm_compute::WeightFormat::OHWI);
        int  input_x_dim;
        int  input_k_dim;
        int  output_x_dim;
        int  output_k_dim;
        auto dims = output->num_dimensions();
        switch (dims)
        {
            case 2:
            {
                input_x_dim  = input->dimension(0);  // Number of columns in input matrix
                input_k_dim  = input->dimension(1);  // Number of rows in input matrix
                output_x_dim = output->dimension(0); // Number of columns in output matrix
                output_k_dim = output->dimension(1); // Number of rows in output matrix
                break;
            }
            case 4:
            {
                input_x_dim  = input->dimension(2);  // Number of columns in input matrix
                input_k_dim  = input->dimension(3);  // Number of rows in input matrix
                output_x_dim = output->dimension(2); // Number of columns in output matrix
                output_k_dim = output->dimension(3); // Number of rows in output matrix
                break;
            }
            default:
            {
                ARM_COMPUTE_RETURN_ERROR_MSG("Only 2 or 4 dimensions supported.");
            }
        }

        int ksize = 0;
        switch (output_wf)
        {
#if defined(ARM_COMPUTE_ENABLE_SVE)
            case WeightFormat::OHWIo8:
            {
                if (Scheduler::get().cpu_info().has_sve() && arm_gemm::utils::get_vector_length<float>() == 8)
                {
                    ksize = 8;
                }
                else
                {
                    ARM_COMPUTE_RETURN_ERROR_MSG("Unsupported weight format.");
                }
                break;
            }
#endif /* ARM_COMPUTE_ENABLE_SVE */
            case WeightFormat::OHWIo4:
            {
                ksize = 4;
                break;
            }
            default:
            {
                ARM_COMPUTE_RETURN_ERROR_MSG("Unsupported weight format.");
                break;
            }
        }

        // output k_dim needs to be same as input but multiple of ksize
        int32_t rnd_up_input_kdim = arm_compute::ceil_to_multiple<int32_t, int32_t>(input_k_dim, ksize);
        ARM_COMPUTE_RETURN_ERROR_ON(rnd_up_input_kdim != output_k_dim);
        // output x_dim needs to be same as input
        ARM_COMPUTE_RETURN_ERROR_ON(input_x_dim != output_x_dim);
    }
    return Status{};
}

} // namespace arm_compute

#endif // defined(__aarch64__)