/*
 * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * 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 "common/kernels/kernel.h"
#include <algorithm>
#include <array>
#include <assert.h>
#include <cfloat>
#include <cstdio>
#include <math.h>
#include <stdio.h>
namespace nvinfer1
{
namespace plugin
{
// This macro is to control shared memory usage. If set to 1, kernel loads the whole feature map
// into shared memory for reuse; If set to 0, kernel loads data from global memory directly.
// Roi pooling performance is data dependent. You can test which value is better to your data.
// If all bboxes are very small, 0 is recommended, otherwise, shared memory will load many unused
// data; If bboxes have many overlaps, 1 is recommended to avoid duplicate loads.
// 1 requires larger shared memory size. It may fail if it is larger than CUDA allowed per-block
// shared memory size upper bound. Then you have to use 0.
#define ROIPOOLING_FEATURE_MAP_USE_SHMEM 1

template <typename T>
__device__ T getMax();

template <>
__device__ __forceinline__ int8_t getMax<int8_t>()
{
    return INT8_MAX;
}

template <>
__device__ __forceinline__ float getMax<float>()
{
    return FLT_MAX;
}

// ROI POOLING FORWARD KERNEL
template <typename DATA_T, typename ROI_T, bool INFER_ONLY, bool FM_IN_SMEM>
__global__ void ROIPoolingForwardKernelAligned(int32_t ROICount, const ROI_T* rois,
    int32_t N, // feature map size
    int32_t C, // feature map size
    int32_t H, // feature map size
    int32_t W, // feature map size
    const DATA_T* featureMap, const int32_t poolingH, const int32_t poolingW, const float spatialScale, DATA_T* top,
    int32_t* maxIds, int32_t fmapStep)
{
    extern __shared__ float smem[];
    DATA_T* feature_shr = (DATA_T*) &smem[0];
    int* rois_shr = nullptr;
    if (FM_IN_SMEM)
    {
        rois_shr = (int*) &feature_shr[H * W];
    }
    else
    {
        rois_shr = (int*) &feature_shr[0];
        feature_shr = nullptr;
    }

    const int batch = blockIdx.x / C;
    const int channel = blockIdx.x % C;

    // load ROIs to shared memory
    for (int j = threadIdx.x; j < ROICount; j += blockDim.x)
    {
        int offset = j << 2;
        float4 roi = reinterpret_cast<float4*>(const_cast<float*>(rois))[batch * ROICount + j];
        // spatialScale = 1.0 / featureStride
        // Convert the coordinates to feature map scale
        rois_shr[offset] = round(roi.x * spatialScale);     //roi_start_w
        rois_shr[offset + 1] = round(roi.y * spatialScale); //roi_start_h
        rois_shr[offset + 2] = round(roi.z * spatialScale) - round(roi.x * spatialScale); //roi_length_w
        rois_shr[offset + 3] = round(roi.w * spatialScale) - round(roi.y * spatialScale); // roi_length_h
    }

    // NC/xHW
    int fmapOffset = blockIdx.x / fmapStep * H * W * fmapStep + blockIdx.x % fmapStep;

    // Assumes #CTAs is just enough to cover all channels of all blocks
    const DATA_T* bottom_data_offset = featureMap + fmapOffset;
    if (FM_IN_SMEM)
    {
        // load the current channel to the shared memory
        for (int j = threadIdx.x; j < H * W; j += blockDim.x)
        {
            feature_shr[j] = bottom_data_offset[j * fmapStep];
        }
    }
    __syncthreads();

    for (int j = threadIdx.x; j < ROICount; j += blockDim.x)
    {
        const int offset = j << 2;
        // Force malformed ROIs to be 1x1
        int roi_start_w = rois_shr[offset];
        int roi_start_h = rois_shr[offset + 1];
        int roi_width = max(rois_shr[offset + 2] + 1, 1);
        int roi_height = max(rois_shr[offset + 3] + 1, 1);
        float bin_size_h = static_cast<float>(roi_height) / static_cast<float>(poolingH);
        float bin_size_w = static_cast<float>(roi_width) / static_cast<float>(poolingW);

        for (int ph = 0; ph < poolingH; ++ph)
        {
            for (int pw = 0; pw < poolingW; ++pw)
            {
                int hstart = static_cast<int>(floor(static_cast<float>(ph) * bin_size_h));
                int wstart = static_cast<int>(floor(static_cast<float>(pw) * bin_size_w));
                int hend = static_cast<int>(ceil(static_cast<float>(ph + 1) * bin_size_h));
                int wend = static_cast<int>(ceil(static_cast<float>(pw + 1) * bin_size_w));

                // Add roi offsets and clip to input boundaries
                // In fact, clipping should be done in the RPN, but just in case...
                hstart = min(max(hstart + roi_start_h, 0), H);
                hend = min(max(hend + roi_start_h, 0), H);
                wstart = min(max(wstart + roi_start_w, 0), W);
                wend = min(max(wend + roi_start_w, 0), W);
                bool is_empty = (hend <= hstart) || (wend <= wstart);

                // Define an empty pooling region to be zero
                DATA_T maxval = is_empty ? 0 : -getMax<DATA_T>();
                int maxId = -1;
                DATA_T data = 0;
                for (int h = hstart; h < hend; ++h)
                {
                    for (int w = wstart; w < wend; ++w)
                    {
                        int index = h * W + w;
                        if (FM_IN_SMEM)
                        {
                            data = feature_shr[index];
                        }
                        else
                        {
                            data = bottom_data_offset[index * fmapStep];
                        }
                        if (data > maxval)
                        {
                            maxval = data;
                            maxId = index;
                        }
                    }
                }
                top[(((batch * ROICount + j) * C + channel) * poolingH + ph) * poolingW + pw] = maxval;
                if (!INFER_ONLY)
                {
                    maxIds[(((batch * ROICount + j) * C + channel) * poolingH + ph) * poolingW + pw] = maxId;
                }
            } //for:pw
        }     //for:ph
    }         // for:j
}

template <typename DATA_T, DLayout_t DATA_L, typename ROI_T, bool INFER_ONLY>
pluginStatus_t ROIPoolingForwardKernelAlignedLauncher(cudaStream_t stream,
    const int R,        // TOTAL number of rois -> ~nmsMaxOut * N
    const int N,        // Batch size
    const int C,        // Channels
    const int H,        // Input feature map H
    const int W,        // Input feature map W
    const int poolingH, // Output feature map H
    const int poolingW, // Output feature map W
    const float spatialScale, const void* rois, const void* featureMap, void* top, int* maxIds, size_t deviceSmemSize)
{
    size_t roiShmemSize = (R / N) * 4 * sizeof(ROI_T);

#if ROIPOOLING_FEATURE_MAP_USE_SHMEM
    size_t shmemSize = H * W * sizeof(DATA_T) + roiShmemSize;
    const bool fmap_in_shmem = true;
#else
    size_t shmemSize = roiShmemSize;
    const bool fmap_in_shmem = false;
#endif

    if (shmemSize > deviceSmemSize)
    {
        return STATUS_BAD_PARAM;
    }

    // in the aligned version of ROI Pooling R should always be a multiple of N
    PLUGIN_ASSERT(R % N == 0);

    // NC/xHW
    int32_t fmapStep = 1;
    switch(DATA_L)
    {
    case NCHW: fmapStep = 1; break;
    case NC4HW:
        fmapStep = 4;
        PLUGIN_ASSERT((N * C) % 4 == 0);
        break;
    case NC32HW:
        fmapStep = 32;
        PLUGIN_ASSERT((N * C) % 32 == 0);
        break;
    default: PLUGIN_ASSERT(false);
    }

    if (shmemSize > 48 * 1024)
    {
        PLUGIN_CHECK(cudaFuncSetAttribute(&ROIPoolingForwardKernelAligned<DATA_T, ROI_T, INFER_ONLY, true>,
            cudaFuncAttributeMaxDynamicSharedMemorySize, shmemSize));
    }
    ROIPoolingForwardKernelAligned<DATA_T, ROI_T, INFER_ONLY, fmap_in_shmem><<<N * C, 256, shmemSize, stream>>>(R / N,
        (const ROI_T*) rois,
        N, // feature map size
        C, // feature map size
        H, // feature map size
        W, // feature map size
        (const DATA_T*) featureMap, poolingH, poolingW, spatialScale, (DATA_T*) top, maxIds, fmapStep);

    CSC(cudaGetLastError(), STATUS_FAILURE);
    return STATUS_SUCCESS;
}

// ROI POOLING LAUNCH CONFIG

typedef pluginStatus_t (*roiFwd)(cudaStream_t,
                                const int,   //R, // TOTAL number of rois -> ~nmsMaxOut * N
                                const int,   //N, // Batch size
                                const int,   //C, // Channels
                                const int,   //H, // Input feature map H
                                const int,   //W, // Input feature map W
                                const int,   //poolingH, // Output feature map H
                                const int,   //poolingW, // Output feature map W
                                const float, //spatialScale,
                                const void*, //rois,
                                const void*, //featureMap,
                                void*,       //top
                                int*,        //maxIds
                                size_t);     //device shmem size

// struct
struct roiFwdLaunchConfig
{
    DataType t_rois;
    DataType t_featureMap;
    DLayout_t l_featureMap;
    DataType t_top;
    DLayout_t l_top;
    bool inferOnly;
    roiFwd function;

    roiFwdLaunchConfig(
        DataType t_rois, DataType t_featureMap, DLayout_t l_featureMap, DataType t_top, DLayout_t l_top, bool inferOnly)
        : t_rois(t_rois)
        , t_featureMap(t_featureMap)
        , l_featureMap(l_featureMap)
        , t_top(t_top)
        , l_top(l_top)
        , inferOnly(inferOnly)
        , function(nullptr)
    {
    }

    roiFwdLaunchConfig(DataType t_rois,
                       DataType t_featureMap,
                       DLayout_t l_featureMap,
                       DataType t_top,
                       DLayout_t l_top,
                       bool inferOnly,
                       roiFwd function)
        : t_rois(t_rois)
        , t_featureMap(t_featureMap)
        , l_featureMap(l_featureMap)
        , t_top(t_top)
        , l_top(l_top)
        , inferOnly(inferOnly)
        , function(function)
    {
    }

    bool operator==(roiFwdLaunchConfig const& other) const
    {
        return (t_rois == other.t_rois)
            && (t_featureMap == other.t_featureMap)
            && (l_featureMap == other.l_featureMap)
            && (t_top == other.t_top)
            && (l_top == other.l_top)
            && (inferOnly == other.inferOnly);
    }
};

#define FLOAT32 nvinfer1::DataType::kFLOAT
#define INT8 nvinfer1::DataType::kINT8
static std::array<roiFwdLaunchConfig, 6> roiFwdLCOptions = {
    roiFwdLaunchConfig(FLOAT32, FLOAT32, NCHW, FLOAT32, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<float, NCHW, float, true>),
    roiFwdLaunchConfig(FLOAT32, FLOAT32, NCHW, FLOAT32, NCHW, false, ROIPoolingForwardKernelAlignedLauncher<float, NCHW, float, false>),
    roiFwdLaunchConfig(FLOAT32, INT8, NCHW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NCHW, float, true>),
    roiFwdLaunchConfig(FLOAT32, INT8, NC4HW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NC4HW, float, true>),
    roiFwdLaunchConfig(FLOAT32, INT8, NC32HW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NC32HW, float, true>),
    roiFwdLaunchConfig(FLOAT32, FLOAT32, NC4HW, FLOAT32, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<float, NC4HW, float, true>)};

// ROI INFERENCE
pluginStatus_t roiInference(cudaStream_t stream,
                           const int R,        // TOTAL number of rois -> ~nmsMaxOut * N
                           const int N,        // Batch size
                           const int C,        // Channels
                           const int H,        // Input feature map H
                           const int W,        // Input feature map W
                           const int poolingH, // Output feature map H
                           const int poolingW, // Output feature map W
                           const float spatialScale,
                           const nvinfer1::DataType t_rois,
                           const void* rois,
                           const nvinfer1::DataType t_featureMap,
                           const DLayout_t l_featureMap,
                           const void* featureMap,
                           const nvinfer1::DataType t_top,
                           const DLayout_t l_top,
                           void* top,
                           size_t deviceSmemSize)
{
    if (featureMap == NULL || rois == NULL || top == NULL)
    {
        return STATUS_BAD_PARAM;
    }

    DEBUG_PRINTF("&&&& ROIS %u\n", hash(rois, R * 4 * sizeof(float)));
    DEBUG_PRINTF("&&&& FMAP %u\n", hash(featureMap, N * C * H * W * sizeof(float)));

    roiFwdLaunchConfig rflc = roiFwdLaunchConfig(t_rois, t_featureMap, l_featureMap, t_top, l_top, true);
    ASSERT_PARAM(R > 0);

    for (unsigned i = 0; i < roiFwdLCOptions.size(); i++)
    {
        if (rflc == roiFwdLCOptions[i])
        {
            DEBUG_PRINTF("$$$$ ROI KERNEL %d\n", i);
            return roiFwdLCOptions[i].function(stream,
                                                R, N, C, H, W, poolingH, poolingW,
                                                spatialScale, rois, featureMap, top, NULL, deviceSmemSize);
        }
    }
    return STATUS_BAD_PARAM;
}
} // namespace plugin
} // namespace nvinfer1