#include <ATen/ATen.h>

#include "torch_npu/csrc/core/OverflowUtils.h"
#include "torch_npu/csrc/core/npu/sys_ctrl/npu_sys_ctrl.h"
#include "op_plugin/OpInterface.h"

namespace torch_npu {
namespace utils {

OverflowUtil::OverflowUtil() {}

OverflowUtil::~OverflowUtil() {}

void OverflowUtil::EnableOverflowNpu() {
  auto result = c10_npu::NpuSysCtrl::GetInstance().OverflowSwitchEnable();
  return;
}

bool OverflowUtil::CheckOverflowNpu() {
  auto options = at::TensorOptions(c10::DeviceType::PrivateUse1).dtype(at::kFloat);
  at::Tensor tmp = at::empty({8}, options);
  auto floatStatus = op_plugin::npu_alloc_float_status(tmp);
  auto result = op_plugin::npu_get_float_status(floatStatus);
  if (result.cpu()[0].item().toInt() != 0) {
    return true;
  }
  return false;
}

void OverflowUtil::ClearOverflowNpu() {
  auto options = at::TensorOptions(c10::DeviceType::PrivateUse1).dtype(at::kFloat);
  at::Tensor tmp = at::empty({8}, options);
  auto floatStatus = op_plugin::npu_alloc_float_status(tmp);
  auto result = op_plugin::npu_clear_float_status(floatStatus);
  return;
}

}
}