#include <ATen/native/BinaryOps.h>

#include <limits>

#include <ATen/AccumulateType.h>
#include <ATen/Dispatch.h>
#include <ATen/native/DispatchStub.h>
#include <ATen/native/TensorIterator.h>
#include <ATen/native/cuda/Loops.cuh>

// NOTE: CUDA on Windows requires that the enclosing function
// of a __device__ lambda not have internal linkage.

namespace at {
namespace native {

void sigmoid_backward_kernel_cuda(TensorIterator& iter) {
  AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "sigmoid_backward_cuda", [&]() {
    AT_SKIP_BFLOAT16_IF_NOT_ROCM(scalar_t, "sigmoid_backward_cuda", [&] {
      gpu_kernel(iter, []GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
        return a * (scalar_t(1.) - b) * b;
      });
    });
  });
}

void logit_backward_kernel_cuda(TensorIterator& iter, Scalar eps_scalar) {
  AT_DISPATCH_FLOATING_TYPES_AND2(
      at::ScalarType::Half,
      at::ScalarType::BFloat16,
      iter.dtype(),
      "logit_cuda",
      [&]() {
        AT_SKIP_BFLOAT16_IF_NOT_ROCM(scalar_t, "logit_cuda", [&] {
          using T_ACC = acc_type<scalar_t, true>;
          const T_ACC eps = eps_scalar.to<T_ACC>();
          if (eps < T_ACC(0)) {
            gpu_kernel(
                iter, [] GPU_LAMBDA(scalar_t dy, scalar_t x) -> scalar_t {
                  const T_ACC dy_acc = static_cast<T_ACC>(dy);
                  const T_ACC x_acc = static_cast<T_ACC>(x);
                  return (x_acc < T_ACC(0) || x_acc > T_ACC(1))
                      ? std::numeric_limits<T_ACC>::quiet_NaN()
                      : dy_acc / (x_acc * (T_ACC(1) - x_acc));
                });
          } else {
            const T_ACC lo = eps;
            const T_ACC hi = T_ACC(1) - eps;
            gpu_kernel(
                iter, [lo, hi] GPU_LAMBDA(scalar_t dy, scalar_t x) -> scalar_t {
                  const T_ACC dy_acc = static_cast<T_ACC>(dy);
                  const T_ACC x_acc = static_cast<T_ACC>(x);
                  return (x_acc < lo || x_acc > hi)
                      ? T_ACC(0)
                      : dy_acc / (x_acc * (T_ACC(1) - x_acc));
                });
          }
        });
      });
}

void tanh_backward_kernel_cuda(TensorIterator& iter) {
  AT_DISPATCH_FLOATING_AND_COMPLEX_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, iter.dtype(), "tanh_backward_cuda", [&]() {
    AT_SKIP_BFLOAT16_IF_NOT_ROCM(scalar_t, "tanh_backward_cuda", [&] {
      gpu_kernel(iter, [] GPU_LAMBDA(scalar_t a, scalar_t b) -> scalar_t {
        return a * (scalar_t{1.} - b * b);
      });
    });
  });
}

REGISTER_DISPATCH(sigmoid_backward_stub, &sigmoid_backward_kernel_cuda);
REGISTER_DISPATCH(logit_backward_stub, &logit_backward_kernel_cuda);
REGISTER_DISPATCH(tanh_backward_stub, &tanh_backward_kernel_cuda);

} // namespace native
} // namespace at
