#include <ATen/ATen.h>
#include <torch/library.h>

#if defined(BACKEND_CUDA)
#include <c10/cuda/CUDAGuard.h>
#elif defined(BACKEND_HIP)
#include <c10/hip/HIPGuard.h>
#else
#error "BACKEND_CUDA or BACKEND_HIP must be defined"
#endif

at::Tensor guard_loop(at::Tensor tensor, int64_t loops) {
  const auto device = static_cast<c10::DeviceIndex>(tensor.device().index());
  for (int64_t i = 0; i < loops; ++i) {
#if defined(BACKEND_CUDA)
    c10::cuda::CUDAGuard guard(device);
#else
    c10::hip::HIPGuard guard(device);
#endif
  }
  return tensor;
}

TORCH_LIBRARY(fastpt_c_overhead_mre, m) {
  m.def("guard_loop(Tensor tensor, int loops) -> Tensor", guard_loop);
}