// !!! This is a file automatically generated by hipify!!!
#include "hip/hip_runtime.h"
#include <hip/hip_fp16.h> // 【修改 1】引入半精度支持

/*!
**************************************************************************
* Deformable DETR (Inference Only - Forward Pass)
**************************************************************************
*/

#include <cstdio>
#include <algorithm>
#include <cstring>

#define CUDA_KERNEL_LOOP(i, n)                          \
  for (int i = blockIdx.x * blockDim.x + threadIdx.x;   \
      i < (n);                                          \
      i += blockDim.x * gridDim.x)

const int CUDA_NUM_THREADS = 1024;
inline int GET_BLOCKS(const int N, const int num_threads)
{
  return (N + num_threads - 1) / num_threads;
}

// =================================================================================
// 核心数学计算：双线性插值 (Device 侧)
// 【修改 2】返回值强制改为 float，确保累加精度；坐标参数 h 和 w 也改为 float
// =================================================================================
template <typename scalar_t>
__device__ float ms_deform_attn_im2col_bilinear(const scalar_t* &bottom_data, 
                                                   const int &height, const int &width, const int &nheads, const int &channels,
                                                   const float &h, const float &w, const int &m, const int &c)
{
  const int h_low = floor(h);
  const int w_low = floor(w);
  const int h_high = h_low + 1;
  const int w_high = w_low + 1;

  // 这里的系数计算全程在 FP32 下进行
  const float lh = h - h_low;
  const float lw = w - w_low;
  const float hh = 1.0f - lh, hw = 1.0f - lw;

  const int w_stride = nheads * channels;
  const int h_stride = width * w_stride;
  const int h_low_ptr_offset = h_low * h_stride;
  const int h_high_ptr_offset = h_low_ptr_offset + h_stride;
  const int w_low_ptr_offset = w_low * w_stride;
  const int w_high_ptr_offset = w_low_ptr_offset + w_stride;
  const int base_ptr = m * channels + c;

  // 【修改 3】从显存读取 FP16 (scalar_t) 数据后，立即转换为 float
  float v1 = 0.0f;
  if (h_low >= 0 && w_low >= 0)
  {
    const int ptr1 = h_low_ptr_offset + w_low_ptr_offset + base_ptr;
    v1 = static_cast<float>(bottom_data[ptr1]);
  }
  
  float v2 = 0.0f;
  if (h_low >= 0 && w_high <= width - 1)
  {
    const int ptr2 = h_low_ptr_offset + w_high_ptr_offset + base_ptr;
    v2 = static_cast<float>(bottom_data[ptr2]);
  }
  
  float v3 = 0.0f;
  if (h_high <= height - 1 && w_low >= 0)
  {
    const int ptr3 = h_high_ptr_offset + w_low_ptr_offset + base_ptr;
    v3 = static_cast<float>(bottom_data[ptr3]);
  }
  
  float v4 = 0.0f;
  if (h_high <= height - 1 && w_high <= width - 1)
  {
    const int ptr4 = h_high_ptr_offset + w_high_ptr_offset + base_ptr;
    v4 = static_cast<float>(bottom_data[ptr4]);
  }

  const float w1 = hh * hw, w2 = hh * lw, w3 = lh * hw, w4 = lh * lw;

  // FP32 下的加权求和，防止溢出或下溢
  const float val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
  return val;
}

// =================================================================================
// 核心调度内核：循环多尺度特征与采样点 (Global 侧)
// =================================================================================
template <typename scalar_t>
__global__ void ms_deformable_im2col_gpu_kernel(const int n,
                                                const scalar_t *data_value, 
                                                const int64_t *data_spatial_shapes,
                                                const int64_t *data_level_start_index, 
                                                const scalar_t *data_sampling_loc,
                                                const scalar_t *data_attn_weight,
                                                const int batch_size, 
                                                const int spatial_size, 
                                                const int num_heads,
                                                const int channels, 
                                                const int num_levels,
                                                const int num_query,
                                                const int num_point,
                                                scalar_t *data_col)
{
  CUDA_KERNEL_LOOP(index, n)
  {
    int _temp = index;
    const int c_col = _temp % channels;
    _temp /= channels;
    const int sampling_index = _temp; 
    const int m_col = _temp % num_heads;
    _temp /= num_heads;
    const int q_col = _temp % num_query;
    _temp /= num_query;
    const int b_col = _temp;

    scalar_t *data_col_ptr = data_col + index;
    int data_weight_ptr = sampling_index * num_levels * num_point;
    int data_loc_w_ptr = data_weight_ptr << 1;
    const int qid_stride = num_heads * channels;
    const int data_value_ptr_init_offset = b_col * spatial_size * qid_stride;
    
    // 【修改 4】特征累加器必须是 float，防止多次累加造成 16 位浮点数溢出
    float col = 0.0f; 
    
    for (int l_col=0; l_col < num_levels; ++l_col)
    {
      const int level_start_id = data_level_start_index[l_col];
      const int spatial_h_ptr = l_col << 1;
      const int spatial_h = data_spatial_shapes[spatial_h_ptr];
      const int spatial_w = data_spatial_shapes[spatial_h_ptr + 1];
      const scalar_t *data_value_ptr = data_value + (data_value_ptr_init_offset + level_start_id * qid_stride);
      
      for (int p_col=0; p_col < num_point; ++p_col)
      {
        // 【修改 5】极其关键：读取采样坐标和注意力权重时，立刻强转为 float
        const float loc_w = static_cast<float>(data_sampling_loc[data_loc_w_ptr]);
        const float loc_h = static_cast<float>(data_sampling_loc[data_loc_w_ptr + 1]);
        const float weight = static_cast<float>(data_attn_weight[data_weight_ptr]);

        // 在 FP32 下计算实际坐标，保证精度
        const float h_im = loc_h * spatial_h - 0.5f;
        const float w_im = loc_w * spatial_w - 0.5f;

        if (h_im > -1.0f && w_im > -1.0f && h_im < spatial_h && w_im < spatial_w)
        {
          // 这里的 ms_deform_attn_im2col_bilinear 已经改写为返回 float
          col += ms_deform_attn_im2col_bilinear(data_value_ptr, spatial_h, spatial_w, num_heads, channels, h_im, w_im, m_col, c_col) * weight;
        }

        data_weight_ptr += 1;
        data_loc_w_ptr += 2;
      }
    }
    // 【修改 6】算完之后，把 float 强转回 FP16 (scalar_t) 写回显存
    *data_col_ptr = static_cast<scalar_t>(col);
  }
}

// =================================================================================
// C++ 主机侧调用入口 (Launcher)
// =================================================================================
template <typename scalar_t>
void ms_deformable_im2col_cuda(hipStream_t stream,
                              const scalar_t* data_value,
                              const int64_t* data_spatial_shapes, 
                              const int64_t* data_level_start_index, 
                              const scalar_t* data_sampling_loc,
                              const scalar_t* data_attn_weight,
                              const int batch_size,
                              const int spatial_size, 
                              const int num_heads, 
                              const int channels, 
                              const int num_levels, 
                              const int num_query,
                              const int num_point,
                              scalar_t* data_col)
{
  const int num_kernels = batch_size * num_query * num_heads * channels;
  const int num_actual_kernels = batch_size * num_query * num_heads * channels;
  const int num_threads = CUDA_NUM_THREADS;
  
  hipLaunchKernelGGL(( ms_deformable_im2col_gpu_kernel<scalar_t>)
      , dim3(GET_BLOCKS(num_actual_kernels, num_threads)), dim3(num_threads),
          0, stream, 
      num_kernels, data_value, data_spatial_shapes, data_level_start_index, data_sampling_loc, data_attn_weight, 
      batch_size, spatial_size, num_heads, channels, num_levels, num_query, num_point, data_col);
  
  hipError_t err = hipGetLastError();
  if (err != hipSuccess)
  {
    printf("error in ms_deformable_im2col_cuda: %s\n", hipGetErrorString(err));
  }
}