/*!
 *  Copyright (c) 2020 by Contributors
 * \file array/cuda/fp16.cuh
 * \brief float16 related functions.
 * \note this file is modified from TVM project:
 *       https://github.com/apache/tvm/blob/e561007f0c330e3d14c2bc8a3ef40fb741db9004/src/target/source/literal/cuda_half_t.h.
 */
#ifndef DGL_ARRAY_CUDA_FP16_CUH_
#define DGL_ARRAY_CUDA_FP16_CUH_


#ifdef USE_FP16
#include <cuda_fp16.h>
#include <algorithm>

static __device__ __forceinline__ half max(half a, half b) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
  return __hgt(__half(a), __half(b)) ? a : b;
#else
  return __half(max(float(a), float(b)));  // NOLINT
#endif
}

static __device__ __forceinline__ half min(half a, half b) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 530
  return __hlt(__half(a), __half(b)) ? a : b;
#else
  return __half(min(float(a), float(b)));  // NOLINT
#endif
}

#ifdef __CUDACC__
// Arithmetic FP16 operations for architecture >= 5.3 are already defined in cuda_fp16.h
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ < 530)
__device__ __forceinline__ __half operator+(const __half& lh, const __half& rh) {
  return __half(float(lh) + float(rh));  // NOLINT
}
__device__ __forceinline__ __half operator-(const __half& lh, const __half& rh) {
  return __half(float(lh) - float(rh));  // NOLINT
}
__device__ __forceinline__ __half operator*(const __half& lh, const __half& rh) {
  return __half(float(lh) * float(rh));  // NOLINT
}
__device__ __forceinline__ __half operator/(const __half& lh, const __half& rh) {
  return __half(float(lh) / float(rh));  // NOLINT
}

__device__ __forceinline__ __half& operator+=(__half& lh, const __half& rh) {  // NOLINT
  lh = __half(float(lh) + float(rh)); return lh;  // NOLINT
}
__device__ __forceinline__ __half& operator-=(__half& lh, const __half& rh) {  // NOLINT
  lh = __half(float(lh) - float(rh)); return lh;  // NOLINT
}
__device__ __forceinline__ __half& operator*=(__half& lh, const __half& rh) {  // NOLINT
  lh = __half(float(lh) * float(rh)); return lh;  // NOLINT
}
__device__ __forceinline__ __half& operator/=(__half& lh, const __half& rh) {  // NOLINT
  lh = __half(float(lh) / float(rh)); return lh;  // NOLINT
}

__device__ __forceinline__ __half& operator++(__half& h) {  // NOLINT
  h = __half(float(h) + 1.0f); return h;  // NOLINT
}
__device__ __forceinline__ __half& operator--(__half& h) {  // NOLINT
  h = __half(float(h) - 1.0f); return h;  // NOLINT
}
__device__ __forceinline__ __half  operator++(__half& h, int) {  // NOLINT
  __half ret = h; h = __half(float(h) + 1.0f); return ret;  // NOLINT
}
__device__ __forceinline__ __half  operator--(__half& h, int) {  // NOLINT
  __half ret = h; h = __half(float(h) - 1.0f); return ret;  // NOLINT
}

__device__ __forceinline__ __half operator+(const __half& h) { return h; }
__device__ __forceinline__ __half operator-(const __half& h) {
  return __half(-float(h));  // NOLINT
}

__device__ __forceinline__ bool operator==(const __half& lh, const __half& rh) {
  return float(lh) == float(rh);  // NOLINT
}
__device__ __forceinline__ bool operator!=(const __half& lh, const __half& rh) {
  return float(lh) != float(rh);  // NOLINT
}
__device__ __forceinline__ bool operator> (const __half& lh, const __half& rh) {
  return float(lh) >  float(rh);  // NOLINT
}
__device__ __forceinline__ bool operator< (const __half& lh, const __half& rh) {
  return float(lh) <  float(rh);  // NOLINT
}
__device__ __forceinline__ bool operator>=(const __half& lh, const __half& rh) {
  return float(lh) >= float(rh);  // NOLINT
}
__device__ __forceinline__ bool operator<=(const __half& lh, const __half& rh) {
  return float(lh) <= float(rh);  // NOLINT
}
#endif  // __CUDA_ARCH__ < 530
#endif  // __CUDACC__

#endif  // USE_FP16

#endif  // DGL_ARRAY_CUDA_FP16_CUH_