/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once
#ifdef PADDLE_WITH_CUDA
#include <cuda.h>
#endif
#ifdef PADDLE_WITH_HIP
#include <hip/hip_runtime.h>
#endif
#include <stdio.h>

#include "paddle/phi/common/bfloat16.h"
#include "paddle/phi/common/complex.h"
#include "paddle/phi/common/float16.h"

template <typename T>
using complex = phi::dtype::complex<T>;

namespace phi {

#define CUDA_ATOMIC_WRAPPER(op, T) \
  __device__ __forceinline__ T CudaAtomic##op(T *address, const T val)

#define USE_CUDA_ATOMIC(op, T) \
  CUDA_ATOMIC_WRAPPER(op, T) { return atomic##op(address, val); }

// Default thread count per block(or block size).
// TODO(typhoonzero): need to benchmark against setting this value
//                    to 1024.
constexpr int PADDLE_CUDA_NUM_THREADS = 512;

// For atomicAdd.
USE_CUDA_ATOMIC(Add, float);
USE_CUDA_ATOMIC(Add, int);
USE_CUDA_ATOMIC(Add, unsigned int);

CUDA_ATOMIC_WRAPPER(Add, bool) {
  size_t offset = reinterpret_cast<size_t>(address) & 3;
  uint32_t *address_as_ui =
      reinterpret_cast<uint32_t *>(reinterpret_cast<char *>(address) - offset);
  uint32_t old = *address_as_ui;
  uint32_t shift = offset * 8;
  uint32_t old_byte;
  uint32_t newval;
  uint32_t assumed;

  do {
    assumed = old;
    old_byte = (old >> shift) & 0xff;
    newval = static_cast<uint8_t>(val + static_cast<bool>(old_byte));
    newval = (old & ~(0x000000ff << shift)) | (newval << shift);
    old = atomicCAS(address_as_ui, assumed, newval);
  } while (assumed != old);

  return static_cast<bool>(old & 0xff);
}

CUDA_ATOMIC_WRAPPER(Add, uint8_t) {
  size_t offset = reinterpret_cast<size_t>(address) & 3;
  uint32_t *address_as_ui =
      reinterpret_cast<uint32_t *>(reinterpret_cast<char *>(address) - offset);
  uint32_t old = *address_as_ui;
  uint32_t shift = offset * 8;
  uint32_t old_byte;
  uint32_t newval;
  uint32_t assumed;

  do {
    assumed = old;
    old_byte = (old >> shift) & 0xff;
    newval = static_cast<uint8_t>(val + static_cast<uint8_t>(old_byte));
    newval = (old & ~(0x000000ff << shift)) | (newval << shift);
    old = atomicCAS(address_as_ui, assumed, newval);
  } while (assumed != old);

  return static_cast<uint8_t>(old & 0xff);
}

CUDA_ATOMIC_WRAPPER(Add, int8_t) {
  size_t offset = reinterpret_cast<size_t>(address) & 3;
  uint32_t *address_as_ui =
      reinterpret_cast<uint32_t *>(reinterpret_cast<char *>(address) - offset);
  uint32_t old = *address_as_ui;
  uint32_t shift = offset * 8;
  uint32_t old_byte;
  uint32_t newval;
  uint32_t assumed;

  do {
    assumed = old;
    old_byte = (old >> shift) & 0xff;
    newval = static_cast<int8_t>(val + static_cast<int8_t>(old_byte));
    newval = (old & ~(0x000000ff << shift)) | (newval << shift);
    old = atomicCAS(address_as_ui, assumed, newval);
  } while (assumed != old);

  return static_cast<int8_t>(old & 0xff);
}

CUDA_ATOMIC_WRAPPER(Add, int16_t) {
  size_t offset = reinterpret_cast<size_t>(address) & 2;
  uint32_t *address_as_ui =
      reinterpret_cast<uint32_t *>(reinterpret_cast<char *>(address) - offset);
  bool is_32_align = offset;
  uint32_t old = *address_as_ui;
  uint32_t old_bytes;
  uint32_t newval;
  uint32_t assumed;

  do {
    assumed = old;
    old_bytes = is_32_align ? old >> 16 : old & 0xffff;
    newval = static_cast<uint16_t>(val + static_cast<int16_t>(old_bytes));
    newval = is_32_align ? (old & 0xffff) | (newval << 16)
                         : (old & 0xffff0000) | newval;
    old = atomicCAS(address_as_ui, assumed, newval);
  } while (assumed != old);

  return static_cast<int16_t>(old & 0xffff);
}
// CUDA API uses unsigned long long int, we cannot use uint64_t here.
// It because unsigned long long int is not necessarily uint64_t
USE_CUDA_ATOMIC(Add, unsigned long long int);  // NOLINT

CUDA_ATOMIC_WRAPPER(Add, int64_t) {
  // Here, we check long long int must be int64_t.
  static_assert(sizeof(int64_t) == sizeof(long long int),  // NOLINT
                "long long should be int64");
  return CudaAtomicAdd(
      reinterpret_cast<unsigned long long int *>(address),  // NOLINT
      static_cast<unsigned long long int>(val));            // NOLINT
}

#if defined(__HIPCC__) || (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 600)
USE_CUDA_ATOMIC(Add, double);
#else
CUDA_ATOMIC_WRAPPER(Add, double) {
  unsigned long long int *address_as_ull =                  // NOLINT
      reinterpret_cast<unsigned long long int *>(address);  // NOLINT
  unsigned long long int old = *address_as_ull, assumed;    // NOLINT

  do {
    assumed = old;
    old = atomicCAS(address_as_ull,
                    assumed,
                    __double_as_longlong(val + __longlong_as_double(assumed)));

    // Note: uses integer comparison to avoid hang in case of NaN
  } while (assumed != old);

  return __longlong_as_double(old);
}
#endif

// NOTE(zhangbo): cuda do not have atomicCAS for __nv_bfloat16.
inline __device__ uint32_t bf16_add_to_low_half(uint32_t val, float x) {
  phi::dtype::bfloat16 low_half;
  // the bfloat16 in lower 16bits
  low_half.x = static_cast<uint16_t>(val & 0xFFFFu);
  low_half =
      static_cast<phi::dtype::bfloat16>(static_cast<float>(low_half) + x);
  return (val & 0xFFFF0000u) | low_half.x;
}

inline __device__ uint32_t bf16_add_to_high_half(uint32_t val, float x) {
  phi::dtype::bfloat16 high_half;
  // the bfloat16 in higher 16bits
  high_half.x = static_cast<uint16_t>(val >> 16);
  high_half =
      static_cast<phi::dtype::bfloat16>(static_cast<float>(high_half) + x);
  return (val & 0xFFFFu) | (static_cast<uint32_t>(high_half.x) << 16);
}

#if CUDA_VERSION >= 11000 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
static __device__ __forceinline__ phi::dtype::bfloat16 CUDABF16ToPDBF16(
    __nv_bfloat16 x) {
  return *reinterpret_cast<phi::dtype::bfloat16 *>(&x);
}

static __device__ __forceinline__ __nv_bfloat16
PDBF16ToCUDABF16(phi::dtype::bfloat16 x) {
  return *reinterpret_cast<__nv_bfloat16 *>(&x);
}

CUDA_ATOMIC_WRAPPER(Add, phi::dtype::bfloat16) {
  return CUDABF16ToPDBF16(atomicAdd(reinterpret_cast<__nv_bfloat16 *>(address),
                                    PDBF16ToCUDABF16(val)));
}
#else
CUDA_ATOMIC_WRAPPER(Add, phi::dtype::bfloat16) {
  // concrete packed bfloat16 value may exist in lower or higher 16bits
  // of the 32bits address.
  uint32_t *address_as_ui = reinterpret_cast<uint32_t *>(
      reinterpret_cast<char *>(address) -
      (reinterpret_cast<uintptr_t>(address) & 0x02));
  float val_f = static_cast<float>(val);
  uint32_t old = *address_as_ui;
  uint32_t sum;
  uint32_t newval;
  uint32_t assumed;
  if (((uintptr_t)address & 0x02) == 0) {
    // the bfloat16 value stay at lower 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(
          address_as_ui, assumed, bf16_add_to_low_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::bfloat16 ret;
    ret.x = old & 0xFFFFu;
    return ret;
  } else {
    // the bfloat16 value stay at higher 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(
          address_as_ui, assumed, bf16_add_to_high_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::bfloat16 ret;
    ret.x = old >> 16;
    return ret;
  }
}
#endif

CUDA_ATOMIC_WRAPPER(Add, complex<float>) {
  float *real = reinterpret_cast<float *>(address);
  float *imag = real + 1;
  return complex<float>(CudaAtomicAdd(real, val.real),
                        CudaAtomicAdd(imag, val.imag));
}

CUDA_ATOMIC_WRAPPER(Add, complex<double>) {
  double *real = reinterpret_cast<double *>(address);
  double *imag = real + 1;
  return complex<double>(CudaAtomicAdd(real, val.real),
                         CudaAtomicAdd(imag, val.imag));
}

// NOTE(dzhwinter): cuda do not have atomicCAS for half.
// Just use the half address as a unsigned value address and
// do the atomicCAS. According to the value store at high 16 bits
// or low 16 bits, then do a different sum and CAS.
// Given most warp-threads will failed on the atomicCAS, so this
// implemented should be avoided in high concurrency. It's will be
// slower than the way convert value into 32bits and do a full atomicCAS.

// convert the value into float and do the add arithmetic.
// then store the result into a uint32.
inline __device__ uint32_t add_to_low_half(uint32_t val, float x) {
  phi::dtype::float16 low_half;
  // the float16 in lower 16bits
  low_half.x = static_cast<uint16_t>(val & 0xFFFFu);
  low_half = static_cast<phi::dtype::float16>(static_cast<float>(low_half) + x);
  return (val & 0xFFFF0000u) | low_half.x;
}

inline __device__ uint32_t add_to_high_half(uint32_t val, float x) {
  phi::dtype::float16 high_half;
  // the float16 in higher 16bits
  high_half.x = static_cast<uint16_t>(val >> 16);
  high_half =
      static_cast<phi::dtype::float16>(static_cast<float>(high_half) + x);
  return (val & 0xFFFFu) | (static_cast<uint32_t>(high_half.x) << 16);
}

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
static __device__ __forceinline__ phi::dtype::float16 CUDAFP16ToPDFP16(
    __half x) {
  return *reinterpret_cast<phi::dtype::float16 *>(&x);
}

static __device__ __forceinline__ __half
PDFP16ToCUDAFP16(phi::dtype::float16 x) {
  return *reinterpret_cast<__half *>(&x);
}

CUDA_ATOMIC_WRAPPER(Add, phi::dtype::float16) {
  return CUDAFP16ToPDFP16(
      atomicAdd(reinterpret_cast<__half *>(address), PDFP16ToCUDAFP16(val)));
}
#else
CUDA_ATOMIC_WRAPPER(Add, phi::dtype::float16) {
  // concrete packed float16 value may exist in lower or higher 16bits
  // of the 32bits address.
  uint32_t *address_as_ui = reinterpret_cast<uint32_t *>(
      reinterpret_cast<char *>(address) -
      (reinterpret_cast<uintptr_t>(address) & 0x02));
  float val_f = static_cast<float>(val);
  uint32_t old = *address_as_ui;
  uint32_t sum;
  uint32_t newval;
  uint32_t assumed;
  if (((uintptr_t)address & 0x02) == 0) {
    // the float16 value stay at lower 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(address_as_ui, assumed, add_to_low_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::float16 ret;
    ret.x = old & 0xFFFFu;
    return ret;
  } else {
    // the float16 value stay at higher 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(address_as_ui, assumed, add_to_high_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::float16 ret;
    ret.x = old >> 16;
    return ret;
  }
}
#endif

template <typename T, bool IsAvailable, typename NVType, typename NVVec2Type>
struct VecAtomicAddHelperBase {
  static constexpr auto kIsAvailable = IsAvailable;
  using NVT = NVType;
  using NVVec2T = NVVec2Type;
};

template <typename T>
struct VecAtomicAddHelper : VecAtomicAddHelperBase<T, false, void, void> {};

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 700
template <>
struct VecAtomicAddHelper<phi::dtype::float16>
    : VecAtomicAddHelperBase<phi::dtype::float16, true, __half, __half2> {};
#endif

#if CUDA_VERSION >= 11000 && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
template <>
struct VecAtomicAddHelper<phi::dtype::bfloat16>
    : VecAtomicAddHelperBase<phi::dtype::bfloat16,
                             true,
                             __nv_bfloat16,
                             __nv_bfloat162> {};
#endif

// The performance of "atomicAdd(half* )" is bad, but for "atomicAdd(half2* )"
// is good. So for fp16 type, we can use "atomicAdd(half2* )" to speed up.
template <typename T,
          typename std::enable_if<VecAtomicAddHelper<T>::kIsAvailable>::type * =
              nullptr>
__device__ __forceinline__ void fastAtomicAdd(T *tensor,
                                              size_t index,
                                              const size_t numel,
                                              T value) {
  // whether the address is 32-byte aligned.
  using NVT = typename VecAtomicAddHelper<T>::NVT;
  using NVVec2T = typename VecAtomicAddHelper<T>::NVVec2T;
  NVT *target_addr = reinterpret_cast<NVT *>(tensor + index);
  bool aligned_half2 =
      (reinterpret_cast<std::uintptr_t>(target_addr) % sizeof(NVVec2T) == 0);

  if (aligned_half2 && index < (numel - 1)) {
    NVVec2T value2;
    value2.x = *reinterpret_cast<NVT *>(&value);
    value2.y = 0.0;
    atomicAdd(reinterpret_cast<NVVec2T *>(target_addr), value2);

  } else if (!aligned_half2 && index > 0) {
    NVVec2T value2;
    value2.x = 0.0;
    value2.y = *reinterpret_cast<NVT *>(&value);
    atomicAdd(reinterpret_cast<NVVec2T *>(target_addr - 1), value2);

  } else {
    atomicAdd(reinterpret_cast<NVT *>(tensor) + index,
              *reinterpret_cast<NVT *>(&value));
  }
}

template <typename T,
          typename std::enable_if<!VecAtomicAddHelper<T>::kIsAvailable>::type
              * = nullptr>
__device__ __forceinline__ void fastAtomicAdd(T *arr,
                                              size_t index,
                                              const size_t numel,
                                              T value) {
  CudaAtomicAdd(arr + index, value);
}

// For atomicMul.
CUDA_ATOMIC_WRAPPER(Mul, int) {
  int res = *address, old = res;  // NOLINT
  do {
    old = res;
    res = atomicCAS(address,     // NOLINT
                    old,         // NOLINT
                    val * old);  // NOLINT
  } while (old != res);
  return res;
}

CUDA_ATOMIC_WRAPPER(Mul, unsigned int) {
  unsigned int res = *address, old = res;  // NOLINT
  do {
    old = res;
    res = atomicCAS(address,     // NOLINT
                    old,         // NOLINT
                    val * old);  // NOLINT
  } while (old != res);
  return res;
}
// CUDA API uses unsigned long long int, we cannot use uint64_t here.
// It because unsigned long long int is not necessarily uint64_t
CUDA_ATOMIC_WRAPPER(Mul, unsigned long long int) {  // NOLINT
  unsigned long long int old = *address, assumed;   // NOLINT

  do {
    assumed = old;
    old = atomicCAS(address, assumed, val * assumed);
  } while (assumed != old);
  return old;
}

CUDA_ATOMIC_WRAPPER(Mul, int64_t) {
  // Here, we check long long int must be int64_t.
  static_assert(sizeof(int64_t) == sizeof(long long int),  // NOLINT
                "long long should be int64");
  long long int res = *address, old = res;  // NOLINT
  do {
    old = res;
    res = (long long int)atomicCAS(                                  // NOLINT
        (unsigned long long int *)address,                           // NOLINT
        (unsigned long long int)old,                                 // NOLINT
        (unsigned long long int)val * (unsigned long long int)old);  // NOLINT
  } while (old != res);
  return res;
}

CUDA_ATOMIC_WRAPPER(Mul, float) {
  int *const address_as_i = reinterpret_cast<int *>(address);
  int old = *address_as_i, assumed;

  do {
    assumed = old;
    old = atomicCAS(
        address_as_i, assumed, __float_as_int(val * __int_as_float(assumed)));
  } while (assumed != old);

  return __int_as_float(old);
}

__device__ __forceinline__ uint32_t __loadAligned(const uintptr_t base_addr,
                                                  uint32_t mask,
                                                  uint32_t shift) {
  // get 4B aligned address
  uint32_t aligned_value = *reinterpret_cast<const uint32_t *>(base_addr);
  return (aligned_value & mask) >> shift;
}

CUDA_ATOMIC_WRAPPER(Mul, uint8_t) {
  // get 4D aligned base address
  uintptr_t base_addr = reinterpret_cast<uintptr_t>(address) & (~3);
  uint32_t offset = reinterpret_cast<uintptr_t>(address) - base_addr;
  uint32_t shift = offset * 8;
  uint32_t mask = 0xFFU << shift;

  uint32_t old32 = __loadAligned(base_addr, mask, shift), assumed32 = 0;

  do {
    assumed32 = old32;
    uint8_t current = static_cast<uint8_t>((old32 & mask) >> shift);
    uint8_t new_val = current * val;
    uint32_t new32 =
        (old32 & ~mask) | (static_cast<uint32_t>(new_val) << shift);

    old32 =
        atomicCAS(reinterpret_cast<uint32_t *>(base_addr), assumed32, new32);
  } while (assumed32 != old32);

  return static_cast<uint8_t>((old32 & mask) >> shift);
}

CUDA_ATOMIC_WRAPPER(Mul, int16_t) {
  // get 4D aligned base address
  uintptr_t base_addr = reinterpret_cast<uintptr_t>(address) & (~3);
  uint32_t offset = (reinterpret_cast<uintptr_t>(address) - base_addr) / 2;
  uint32_t shift = offset * 16;
  uint32_t mask = 0xFFFFU << shift;

  uint32_t old32 = __loadAligned(base_addr, mask, shift), assumed32 = 0;

  do {
    assumed32 = old32;
    int16_t current = static_cast<int16_t>((old32 & mask) >> shift);
    int16_t new_val = current * val;
    uint32_t new32 =
        (old32 & ~mask) | (static_cast<uint32_t>(new_val) << shift);

    old32 =
        atomicCAS(reinterpret_cast<uint32_t *>(base_addr), assumed32, new32);
  } while (assumed32 != old32);

  return static_cast<int16_t>((old32 & mask) >> shift);
}

CUDA_ATOMIC_WRAPPER(Mul, double) {
  unsigned long long int *const address_as_ull =            // NOLINT
      reinterpret_cast<unsigned long long int *>(address);  // NOLINT
  unsigned long long int old = *address_as_ull, assumed;    // NOLINT

  do {
    assumed = old;

    old = atomicCAS(address_as_ull,
                    assumed,
                    __double_as_longlong(val * __longlong_as_double(assumed)));
  } while (assumed != old);

  return __longlong_as_double(old);
}

inline __device__ uint32_t mul_to_low_half(uint32_t val, float x) {
  phi::dtype::float16 low_half;
  // The float16 in lower 16bits
  low_half.x = static_cast<uint16_t>(val & 0xFFFFu);
  low_half = static_cast<phi::dtype::float16>(static_cast<float>(low_half) * x);
  return (val & 0xFFFF0000u) | low_half.x;
}

inline __device__ uint32_t mul_to_high_half(uint32_t val, float x) {
  phi::dtype::float16 high_half;
  // The float16 in higher 16bits
  high_half.x = static_cast<uint16_t>(val >> 16);
  high_half =
      static_cast<phi::dtype::float16>(static_cast<float>(high_half) * x);
  return (val & 0xFFFFu) | (static_cast<uint32_t>(high_half.x) << 16);
}

CUDA_ATOMIC_WRAPPER(Mul, phi::dtype::float16) {
  uint32_t *address_as_ui = reinterpret_cast<uint32_t *>(
      reinterpret_cast<char *>(address) -
      (reinterpret_cast<uintptr_t>(address) & 0x02));
  float val_f = static_cast<float>(val);
  uint32_t old = *address_as_ui;
  uint32_t assumed;
  if (((uintptr_t)address & 0x02) == 0) {
    // The float16 value stay at lower 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(address_as_ui, assumed, mul_to_low_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::float16 ret;
    ret.x = old & 0xFFFFu;
    return ret;
  } else {
    // The float16 value stay at higher 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(address_as_ui, assumed, mul_to_high_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::float16 ret;
    ret.x = old >> 16;
    return ret;
  }
}

inline __device__ uint32_t bf16_mul_to_low_half(uint32_t val, float x) {
  phi::dtype::bfloat16 low_half;
  // The bfloat16 in lower 16bits
  low_half.x = static_cast<uint16_t>(val & 0xFFFFu);
  low_half =
      static_cast<phi::dtype::bfloat16>(static_cast<float>(low_half) * x);
  return (val & 0xFFFF0000u) | low_half.x;
}

inline __device__ uint32_t bf16_mul_to_high_half(uint32_t val, float x) {
  phi::dtype::bfloat16 high_half;
  // The bfloat16 in higher 16bits
  high_half.x = static_cast<uint16_t>(val >> 16);
  high_half =
      static_cast<phi::dtype::bfloat16>(static_cast<float>(high_half) * x);
  return (val & 0xFFFFu) | (static_cast<uint32_t>(high_half.x) << 16);
}

CUDA_ATOMIC_WRAPPER(Mul, phi::dtype::bfloat16) {
  uint32_t *address_as_ui = reinterpret_cast<uint32_t *>(
      reinterpret_cast<char *>(address) -
      (reinterpret_cast<uintptr_t>(address) & 0x02));
  float val_f = static_cast<float>(val);
  uint32_t old = *address_as_ui;
  uint32_t assumed;
  if (((uintptr_t)address & 0x02) == 0) {
    // The bfloat16 value stay at lower 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(
          address_as_ui, assumed, bf16_mul_to_low_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::bfloat16 ret;
    ret.x = old & 0xFFFFu;
    return ret;
  } else {
    // The bfloat16 value stay at higher 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(
          address_as_ui, assumed, bf16_mul_to_high_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::bfloat16 ret;
    ret.x = old >> 16;
    return ret;
  }
}

// For atomicMax
USE_CUDA_ATOMIC(Max, int);
USE_CUDA_ATOMIC(Max, unsigned int);
// CUDA API uses unsigned long long int, we cannot use uint64_t here.
// It because unsigned long long int is not necessarily uint64_t
#if defined(__HIPCC__) || (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 350)
USE_CUDA_ATOMIC(Max, unsigned long long int);  // NOLINT
#else
CUDA_ATOMIC_WRAPPER(Max, unsigned long long int) {  // NOLINT
  if (*address >= val) {
    return *address;
  }

  unsigned long long int old = *address, assumed;  // NOLINT

  do {
    assumed = old;
    if (assumed >= val) {
      break;
    }

    old = atomicCAS(address, assumed, val);
  } while (assumed != old);
}
#endif

CUDA_ATOMIC_WRAPPER(Max, int64_t) {
  // Here, we check long long int must be int64_t.
  static_assert(sizeof(int64_t) == sizeof(long long int),  // NOLINT
                "long long should be int64");
  long long int res = *address;  // NOLINT
  while (val > res) {
    long long int old = res;                                           // NOLINT
    res = (long long int)atomicCAS((unsigned long long int *)address,  // NOLINT
                                   (unsigned long long int)old,        // NOLINT
                                   (unsigned long long int)val);       // NOLINT
    if (res == old) {
      break;
    }
  }
  return res;
}

CUDA_ATOMIC_WRAPPER(Max, float) {
  if (*address >= val) {
    return *address;
  }

  int *const address_as_i = reinterpret_cast<int *>(address);
  int old = *address_as_i, assumed;

  do {
    assumed = old;
    if (__int_as_float(assumed) >= val) {
      break;
    }

    old = atomicCAS(address_as_i, assumed, __float_as_int(val));
  } while (assumed != old);

  return __int_as_float(old);
}

CUDA_ATOMIC_WRAPPER(Max, double) {
  if (*address >= val) {
    return *address;
  }

  unsigned long long int *const address_as_ull =            // NOLINT
      reinterpret_cast<unsigned long long int *>(address);  // NOLINT
  unsigned long long int old = *address_as_ull, assumed;    // NOLINT

  do {
    assumed = old;
    if (__longlong_as_double(assumed) >= val) {
      break;
    }

    old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val));
  } while (assumed != old);

  return __longlong_as_double(old);
}

inline __device__ uint32_t max_to_low_half(uint32_t val, float x) {
  phi::dtype::float16 low_half;
  // The float16 in lower 16bits
  low_half.x = static_cast<uint16_t>(val & 0xFFFFu);
  low_half =
      static_cast<phi::dtype::float16>(max(static_cast<float>(low_half), x));
  return (val & 0xFFFF0000u) | low_half.x;
}

inline __device__ uint32_t max_to_high_half(uint32_t val, float x) {
  phi::dtype::float16 high_half;
  // The float16 in higher 16bits
  high_half.x = static_cast<uint16_t>(val >> 16);
  high_half =
      static_cast<phi::dtype::float16>(max(static_cast<float>(high_half), x));
  return (val & 0xFFFFu) | (static_cast<uint32_t>(high_half.x) << 16);
}

CUDA_ATOMIC_WRAPPER(Max, phi::dtype::float16) {
  if (*address >= val) {
    return *address;
  }
  uint32_t *address_as_ui = reinterpret_cast<uint32_t *>(
      reinterpret_cast<char *>(address) -
      (reinterpret_cast<uintptr_t>(address) & 0x02));
  float val_f = static_cast<float>(val);
  uint32_t old = *address_as_ui;
  uint32_t assumed;
  if (((uintptr_t)address & 0x02) == 0) {
    // The float16 value stay at lower 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(address_as_ui, assumed, max_to_low_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::float16 ret;
    ret.x = old & 0xFFFFu;
    return ret;
  } else {
    // The float16 value stay at higher 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(address_as_ui, assumed, max_to_high_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::float16 ret;
    ret.x = old >> 16;
    return ret;
  }
}

inline __device__ uint32_t bf16_max_to_low_half(uint32_t val, float x) {
  phi::dtype::bfloat16 low_half;
  // The bfloat16 in lower 16bits
  low_half.x = static_cast<uint16_t>(val & 0xFFFFu);
  low_half =
      static_cast<phi::dtype::bfloat16>(max(static_cast<float>(low_half), x));
  return (val & 0xFFFF0000u) | low_half.x;
}

inline __device__ uint32_t bf16_max_to_high_half(uint32_t val, float x) {
  phi::dtype::bfloat16 high_half;
  // The bfloat16 in higher 16bits
  high_half.x = static_cast<uint16_t>(val >> 16);
  high_half =
      static_cast<phi::dtype::bfloat16>(max(static_cast<float>(high_half), x));
  return (val & 0xFFFFu) | (static_cast<uint32_t>(high_half.x) << 16);
}

CUDA_ATOMIC_WRAPPER(Max, phi::dtype::bfloat16) {
  if (*address >= val) {
    return *address;
  }
  uint32_t *address_as_ui = reinterpret_cast<uint32_t *>(
      reinterpret_cast<char *>(address) -
      (reinterpret_cast<uintptr_t>(address) & 0x02));
  float val_f = static_cast<float>(val);
  uint32_t old = *address_as_ui;
  uint32_t assumed;
  if (((uintptr_t)address & 0x02) == 0) {
    // The bfloat16 value stay at lower 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(
          address_as_ui, assumed, bf16_max_to_low_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::bfloat16 ret;
    ret.x = old & 0xFFFFu;
    return ret;
  } else {
    // The bfloat16 value stay at higher 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(
          address_as_ui, assumed, bf16_max_to_high_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::bfloat16 ret;
    ret.x = old >> 16;
    return ret;
  }
}

// For atomicMin
USE_CUDA_ATOMIC(Min, int);
USE_CUDA_ATOMIC(Min, unsigned int);
// CUDA API uses unsigned long long int, we cannot use uint64_t here.
// It because unsigned long long int is not necessarily uint64_t
#if defined(__HIPCC__) || (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 350)
USE_CUDA_ATOMIC(Min, unsigned long long int);  // NOLINT
#else
CUDA_ATOMIC_WRAPPER(Min, unsigned long long int) {  // NOLINT
  if (*address <= val) {
    return *address;
  }

  unsigned long long int old = *address, assumed;  // NOLINT

  do {
    assumed = old;
    if (assumed <= val) {
      break;
    }

    old = atomicCAS(address, assumed, val);
  } while (assumed != old);
}
#endif

CUDA_ATOMIC_WRAPPER(Min, int64_t) {
  // Here, we check long long int must be int64_t.
  static_assert(sizeof(int64_t) == sizeof(long long int),  // NOLINT
                "long long should be int64");
  long long int res = *address;  // NOLINT
  while (val < res) {
    long long int old = res;                                           // NOLINT
    res = (long long int)atomicCAS((unsigned long long int *)address,  // NOLINT
                                   (unsigned long long int)old,        // NOLINT
                                   (unsigned long long int)val);       // NOLINT
    if (res == old) {
      break;
    }
  }
  return res;
}

CUDA_ATOMIC_WRAPPER(Min, float) {
  if (*address <= val) {
    return *address;
  }

  int *const address_as_i = reinterpret_cast<int *>(address);
  int old = *address_as_i, assumed;

  do {
    assumed = old;
    if (__int_as_float(assumed) <= val) {
      break;
    }

    old = atomicCAS(address_as_i, assumed, __float_as_int(val));
  } while (assumed != old);

  return __int_as_float(old);
}

CUDA_ATOMIC_WRAPPER(Min, double) {
  if (*address <= val) {
    return *address;
  }

  unsigned long long int *const address_as_ull =            // NOLINT
      reinterpret_cast<unsigned long long int *>(address);  // NOLINT
  unsigned long long int old = *address_as_ull, assumed;    // NOLINT

  do {
    assumed = old;
    if (__longlong_as_double(assumed) <= val) {
      break;
    }

    old = atomicCAS(address_as_ull, assumed, __double_as_longlong(val));
  } while (assumed != old);

  return __longlong_as_double(old);
}

inline __device__ uint32_t min_to_low_half(uint32_t val, float x) {
  phi::dtype::float16 low_half;
  // The float16 in lower 16bits
  low_half.x = static_cast<uint16_t>(val & 0xFFFFu);
  low_half =
      static_cast<phi::dtype::float16>(min(static_cast<float>(low_half), x));
  return (val & 0xFFFF0000u) | low_half.x;
}

inline __device__ uint32_t min_to_high_half(uint32_t val, float x) {
  phi::dtype::float16 high_half;
  // The float16 in higher 16bits
  high_half.x = static_cast<uint16_t>(val >> 16);
  high_half =
      static_cast<phi::dtype::float16>(min(static_cast<float>(high_half), x));
  return (val & 0xFFFFu) | (static_cast<uint32_t>(high_half.x) << 16);
}

CUDA_ATOMIC_WRAPPER(Min, phi::dtype::float16) {
  if (*address <= val) {
    return *address;
  }
  uint32_t *address_as_ui = reinterpret_cast<uint32_t *>(
      reinterpret_cast<char *>(address) -
      (reinterpret_cast<uintptr_t>(address) & 0x02));
  float val_f = static_cast<float>(val);
  uint32_t old = *address_as_ui;
  uint32_t assumed;
  if (((uintptr_t)address & 0x02) == 0) {
    // The float16 value stay at lower 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(address_as_ui, assumed, min_to_low_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::float16 ret;
    ret.x = old & 0xFFFFu;
    return ret;
  } else {
    // The float16 value stay at higher 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(address_as_ui, assumed, min_to_high_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::float16 ret;
    ret.x = old >> 16;
    return ret;
  }
}

inline __device__ uint32_t bf16_min_to_low_half(uint32_t val, float x) {
  phi::dtype::bfloat16 low_half;
  // The bfloat16 in lower 16bits
  low_half.x = static_cast<uint16_t>(val & 0xFFFFu);
  low_half =
      static_cast<phi::dtype::bfloat16>(min(static_cast<float>(low_half), x));
  return (val & 0xFFFF0000u) | low_half.x;
}

inline __device__ uint32_t bf16_min_to_high_half(uint32_t val, float x) {
  phi::dtype::bfloat16 high_half;
  // The bfloat16 in higher 16bits
  high_half.x = static_cast<uint16_t>(val >> 16);
  high_half =
      static_cast<phi::dtype::bfloat16>(min(static_cast<float>(high_half), x));
  return (val & 0xFFFFu) | (static_cast<uint32_t>(high_half.x) << 16);
}

CUDA_ATOMIC_WRAPPER(Min, phi::dtype::bfloat16) {
  if (*address <= val) {
    return *address;
  }
  uint32_t *address_as_ui = reinterpret_cast<uint32_t *>(
      reinterpret_cast<char *>(address) -
      (reinterpret_cast<uintptr_t>(address) & 0x02));
  float val_f = static_cast<float>(val);
  uint32_t old = *address_as_ui;
  uint32_t assumed;
  if (((uintptr_t)address & 0x02) == 0) {
    // The bfloat16 value stay at lower 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(
          address_as_ui, assumed, bf16_min_to_low_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::bfloat16 ret;
    ret.x = old & 0xFFFFu;
    return ret;
  } else {
    // The bfloat16 value stay at higher 16 bits of the address.
    do {
      assumed = old;
      old = atomicCAS(
          address_as_ui, assumed, bf16_min_to_high_half(assumed, val_f));
    } while (old != assumed);
    phi::dtype::bfloat16 ret;
    ret.x = old >> 16;
    return ret;
  }
}

#define DEFINE_ATOMIC_MINMAX_U8(OpType, operator)                             \
  __device__ __forceinline__ uint8_t CudaAtomic##OpType(uint8_t *address,     \
                                                        const uint8_t val) {  \
    uintptr_t base_addr = reinterpret_cast<uintptr_t>(address) & (~3);        \
    uint32_t offset_bytes = reinterpret_cast<uintptr_t>(address) - base_addr; \
    uint32_t shift = 0, mask = 0;                                             \
    if constexpr (sizeof(uint8_t) == 1) {                                     \
      shift = offset_bytes * 8;                                               \
      mask = 0xFFU << shift;                                                  \
    } else {                                                                  \
      shift = (offset_bytes / 2) * 16;                                        \
      mask = 0xFFFFU << shift;                                                \
    }                                                                         \
    uint8_t current = 0;                                                      \
    uint8_t new_val = 0;                                                      \
    uint32_t assumed32 = 0, old32 = __loadAligned(base_addr, mask, shift);    \
    do {                                                                      \
      assumed32 = old32;                                                      \
      current = static_cast<uint8_t>((old32 & mask) >> shift);                \
      new_val = operator(current, val);                                       \
      uint32_t new32 =                                                        \
          (old32 & ~mask) | (static_cast<uint32_t>(new_val) << shift);        \
      old32 = atomicCAS(                                                      \
          reinterpret_cast<uint32_t *>(base_addr), assumed32, new32);         \
    } while (assumed32 != old32);                                             \
    return current;                                                           \
  }

DEFINE_ATOMIC_MINMAX_U8(Min, min)
DEFINE_ATOMIC_MINMAX_U8(Max, max)

#undef DEFINE_ATOMIC_MINMAX_U8

#define DEFINE_LOW_HALF_OP_I16(op)                                         \
  inline __device__ int op##_to_low_half(int val, int16_t x) {             \
    int16_t low_half = op(static_cast<int16_t>(val & 0x0000FFFF), x);      \
    return (val & 0xFFFF0000) | (static_cast<int>(low_half) & 0x0000FFFF); \
  }

#define DEFINE_HIGH_HALF_OP_I16(op)                                  \
  inline __device__ int op##_to_high_half(int val, int16_t x) {      \
    int16_t high_half = op(static_cast<int16_t>(val >> 16), x);      \
    return (val & 0x0000FFFF) | (static_cast<int>(high_half) << 16); \
  }

DEFINE_LOW_HALF_OP_I16(min)
DEFINE_LOW_HALF_OP_I16(max)
DEFINE_HIGH_HALF_OP_I16(min)
DEFINE_HIGH_HALF_OP_I16(max)

#define DEFINE_ATOMIC_MINMAX_I16(OpType, op, bypass_op)                        \
  __device__ __forceinline__ int16_t CudaAtomic##OpType(int16_t *address,      \
                                                        const int16_t val) {   \
    if (*address bypass_op val) {                                              \
      return *address;                                                         \
    }                                                                          \
    int *address_as_ui = reinterpret_cast<int *>(                              \
        reinterpret_cast<char *>(address) -                                    \
        (reinterpret_cast<uintptr_t>(address) & 0x02));                        \
    int old = 0, assumed = 0;                                                  \
    if ((uintptr_t)address & 0x02) {                                           \
      old = *address_as_ui;                                                    \
      do {                                                                     \
        assumed = old;                                                         \
        old = atomicCAS(                                                       \
            address_as_ui, assumed, op##_to_high_half(assumed, val));          \
      } while (old != assumed);                                                \
      return static_cast<int16_t>(old >> 16);                                  \
    } else {                                                                   \
      old = *address_as_ui;                                                    \
      do {                                                                     \
        assumed = old;                                                         \
        old =                                                                  \
            atomicCAS(address_as_ui, assumed, op##_to_low_half(assumed, val)); \
      } while (old != assumed);                                                \
      return static_cast<int16_t>(old & 0x0000FFFF);                           \
    }                                                                          \
  }

DEFINE_ATOMIC_MINMAX_I16(Min, min, <=)
DEFINE_ATOMIC_MINMAX_I16(Max, max, >=)

#undef DEFINE_ATOMIC_MINMAX_I16
#undef DEFINE_LOW_HALF_OP_I16
#undef DEFINE_HIGH_HALF_OP_I16

#ifdef PADDLE_WITH_CUDA
/*
 * One thead block deals with elementwise atomicAdd for vector of len.
 * @in: [x1, x2, x3, ...]
 * @out:[y1+x1, y2+x2, y3+x3, ...]
 * */

template <typename T,
          typename std::enable_if<!VecAtomicAddHelper<T>::kIsAvailable>::type
              * = nullptr>
__device__ __forceinline__ void VectorizedAtomicAddPerBlock(
    const int64_t len, int tid, int threads_per_block, const T *in, T *out) {
  for (int i = tid; i < len; i += threads_per_block) {
    CudaAtomicAdd(&out[i], in[i]);
  }
}

// Note: assume that len is even. If len is odd, call fastAtomicAdd directly.
template <typename T,
          typename std::enable_if<VecAtomicAddHelper<T>::kIsAvailable>::type * =
              nullptr>
__device__ __forceinline__ void VectorizedAtomicAddPerBlock(
    const int64_t len, int tid, int threads_per_block, const T *in, T *out) {
  int i = 0;
  int loops = len / 2 * 2;

  using NVT = typename VecAtomicAddHelper<T>::NVT;
  using NVVec2T = typename VecAtomicAddHelper<T>::NVVec2T;
  bool aligned_half2 =
      (reinterpret_cast<std::uintptr_t>(out) % sizeof(NVVec2T) == 0);

  if (aligned_half2) {
    for (i = tid * 2; i < loops; i += threads_per_block * 2) {
      NVVec2T value2;
      T value_1 = in[i];
      T value_2 = in[i + 1];
      value2.x = *reinterpret_cast<NVT *>(&value_1);
      value2.y = *reinterpret_cast<NVT *>(&value_2);
      atomicAdd(reinterpret_cast<NVVec2T *>(&out[i]), value2);
    }
    for (; i < len; i += threads_per_block) {
      fastAtomicAdd(out, i, len, in[i]);
    }
  } else {
    for (int i = tid; i < len; i += threads_per_block) {
      fastAtomicAdd(out, i, len, in[i]);
    }
  }
}

#endif
}  // namespace phi