// !!! This is a file automatically generated by hipify!!!
#include <ATen/dtk_macros.h>
      
#include <torch/extension.h>
#include <ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h>

void deformable_aggregation(
  float* output,
  const float* mc_ms_feat,
  const int* spatial_shape,
  const int* scale_start_index,
  const float* sample_location,
  const float* weights,
  int batch_size,
  int num_cams,
  int num_feat,
  int num_embeds,
  int num_scale,
  int num_anchors,
  int num_pts,
  int num_groups
);

void deformable_aggregation_ref(
  float* output,
  const float* mc_ms_feat,
  const int* spatial_shape,
  const int* scale_start_index,
  const float* sample_location,
  const float* weights,
  int batch_size,
  int num_cams,
  int num_feat,
  int num_embeds,
  int num_scale,
  int num_anchors,
  int num_pts,
  int num_groups
);  

at::Tensor deformable_aggregation_forward(
  const at::Tensor &_mc_ms_feat,
  const at::Tensor &_spatial_shape,
  const at::Tensor &_scale_start_index,
  const at::Tensor &_sampling_location,
  const at::Tensor &_weights
) {
  at::DeviceGuard guard(_mc_ms_feat.device());
  const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(_mc_ms_feat));
  int batch_size = _mc_ms_feat.size(0);
  int num_feat = _mc_ms_feat.size(1);
  int num_embeds = _mc_ms_feat.size(2);
  int num_cams = _spatial_shape.size(0);
  int num_scale = _spatial_shape.size(1);
  int num_anchors = _sampling_location.size(1);
  int num_pts = _sampling_location.size(2);
  int num_groups = _weights.size(5);

  const float* mc_ms_feat = _mc_ms_feat.data_ptr<float>();
  const int* spatial_shape = _spatial_shape.data_ptr<int>();
  const int* scale_start_index = _scale_start_index.data_ptr<int>();
  const float* sampling_location = _sampling_location.data_ptr<float>();
  const float* weights = _weights.data_ptr<float>();

  auto output = at::zeros({batch_size, num_anchors, num_embeds}, _mc_ms_feat.options());
  int warm_up = 10;
  int prof_cnt = 1;
  hipEvent_t start, stop;
  float milliseconds = 0;
  for (int _i = 0; _i < warm_up; _i++) {
    deformable_aggregation(
      output.data_ptr<float>(),
      mc_ms_feat, spatial_shape, scale_start_index, sampling_location, weights,
      batch_size, num_cams, num_feat, num_embeds, num_scale, num_anchors, num_pts, num_groups
    );
  }
  hipEventCreate(&start);
  hipEventCreate(&stop); 

  hipEventRecord(start);
  for (int _i = 0; _i < prof_cnt; _i++) {
    deformable_aggregation(
      output.data_ptr<float>(),
      mc_ms_feat, spatial_shape, scale_start_index, sampling_location, weights,
      batch_size, num_cams, num_feat, num_embeds, num_scale, num_anchors, num_pts, num_groups
    );
  }
  hipEventRecord(stop);

  hipEventSynchronize(stop);

  hipEventElapsedTime(&milliseconds, start, stop);

  printf("优化后Kernel 执行时间: %.3f 毫秒\n", milliseconds / prof_cnt);

  hipEventDestroy(start);
  hipEventDestroy(stop);

  return output;
}

at::Tensor deformable_aggregation_forward_ref(
  const at::Tensor &_mc_ms_feat,
  const at::Tensor &_spatial_shape,
  const at::Tensor &_scale_start_index,
  const at::Tensor &_sampling_location,
  const at::Tensor &_weights
) {
  at::DeviceGuard guard(_mc_ms_feat.device());
  const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(_mc_ms_feat));
  int batch_size = _mc_ms_feat.size(0);
  int num_feat = _mc_ms_feat.size(1);
  int num_embeds = _mc_ms_feat.size(2);
  int num_cams = _spatial_shape.size(0);
  int num_scale = _spatial_shape.size(1);
  int num_anchors = _sampling_location.size(1);
  int num_pts = _sampling_location.size(2);
  int num_groups = _weights.size(5);

  const float* mc_ms_feat = _mc_ms_feat.data_ptr<float>();
  const int* spatial_shape = _spatial_shape.data_ptr<int>();
  const int* scale_start_index = _scale_start_index.data_ptr<int>();
  const float* sampling_location = _sampling_location.data_ptr<float>();
  const float* weights = _weights.data_ptr<float>();

  auto output = at::zeros({batch_size, num_anchors, num_embeds}, _mc_ms_feat.options());
  int warm_up = 10;
  int prof_cnt = 1;
  hipEvent_t start, stop;
  float milliseconds = 0;

  for (int _i = 0; _i < warm_up; _i++) {
  deformable_aggregation_ref(
    output.data_ptr<float>(),
    mc_ms_feat, spatial_shape, scale_start_index, sampling_location, weights,
    batch_size, num_cams, num_feat, num_embeds, num_scale, num_anchors, num_pts, num_groups
  );
  }


  // 创建事件对象
  hipEventCreate(&start);
  hipEventCreate(&stop); 

  hipEventRecord(start);
  for (int _i = 0; _i < prof_cnt; _i++) {
    deformable_aggregation_ref(
      output.data_ptr<float>(),
      mc_ms_feat, spatial_shape, scale_start_index, sampling_location, weights,
      batch_size, num_cams, num_feat, num_embeds, num_scale, num_anchors, num_pts, num_groups
    );
  }

  hipEventRecord(stop);

  hipEventSynchronize(stop);

  hipEventElapsedTime(&milliseconds, start, stop);

  printf("优化前Kernel 执行时间: %.3f 毫秒\n", milliseconds / prof_cnt);

  hipEventDestroy(start);
  hipEventDestroy(stop);

  return output;
}


PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def(
    "deformable_aggregation_forward",
    &deformable_aggregation_forward,
    "deformable_aggregation_forward"
  );
  m.def(
    "deformable_aggregation_forward_ref",
    &deformable_aggregation_forward_ref,
    "deformable_aggregation_forward_ref"
  );

}

    