/*
 * Copyright (c) 2018-2021, NVIDIA CORPORATION.
 *
 * 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

#include <cuda_fp16.h>
#include "cuda_utils.cuh"

namespace raft {

template <typename math_, int VecLen>
struct IOType {};
template <>
struct IOType<bool, 1> {
  static_assert(sizeof(bool) == sizeof(int8_t),
                "IOType bool size assumption failed");
  typedef int8_t Type;
};
template <>
struct IOType<bool, 2> {
  typedef int16_t Type;
};
template <>
struct IOType<bool, 4> {
  typedef int32_t Type;
};
template <>
struct IOType<bool, 8> {
  typedef int2 Type;
};
template <>
struct IOType<bool, 16> {
  typedef int4 Type;
};
template <>
struct IOType<int8_t, 1> {
  typedef int8_t Type;
};
template <>
struct IOType<int8_t, 2> {
  typedef int16_t Type;
};
template <>
struct IOType<int8_t, 4> {
  typedef int32_t Type;
};
template <>
struct IOType<int8_t, 8> {
  typedef int2 Type;
};
template <>
struct IOType<int8_t, 16> {
  typedef int4 Type;
};
template <>
struct IOType<uint8_t, 1> {
  typedef uint8_t Type;
};
template <>
struct IOType<uint8_t, 2> {
  typedef uint16_t Type;
};
template <>
struct IOType<uint8_t, 4> {
  typedef uint32_t Type;
};
template <>
struct IOType<uint8_t, 8> {
  typedef uint2 Type;
};
template <>
struct IOType<uint8_t, 16> {
  typedef uint4 Type;
};
template <>
struct IOType<int16_t, 1> {
  typedef int16_t Type;
};
template <>
struct IOType<int16_t, 2> {
  typedef int32_t Type;
};
template <>
struct IOType<int16_t, 4> {
  typedef int2 Type;
};
template <>
struct IOType<int16_t, 8> {
  typedef int4 Type;
};
template <>
struct IOType<uint16_t, 1> {
  typedef uint16_t Type;
};
template <>
struct IOType<uint16_t, 2> {
  typedef uint32_t Type;
};
template <>
struct IOType<uint16_t, 4> {
  typedef uint2 Type;
};
template <>
struct IOType<uint16_t, 8> {
  typedef uint4 Type;
};
template <>
struct IOType<__half, 1> {
  typedef __half Type;
};
template <>
struct IOType<__half, 2> {
  typedef __half2 Type;
};
template <>
struct IOType<__half, 4> {
  typedef uint2 Type;
};
template <>
struct IOType<__half, 8> {
  typedef uint4 Type;
};
template <>
struct IOType<__half2, 1> {
  typedef __half2 Type;
};
template <>
struct IOType<__half2, 2> {
  typedef uint2 Type;
};
template <>
struct IOType<__half2, 4> {
  typedef uint4 Type;
};
template <>
struct IOType<int32_t, 1> {
  typedef int32_t Type;
};
template <>
struct IOType<int32_t, 2> {
  typedef uint2 Type;
};
template <>
struct IOType<int32_t, 4> {
  typedef uint4 Type;
};
template <>
struct IOType<uint32_t, 1> {
  typedef uint32_t Type;
};
template <>
struct IOType<uint32_t, 2> {
  typedef uint2 Type;
};
template <>
struct IOType<uint32_t, 4> {
  typedef uint4 Type;
};
template <>
struct IOType<float, 1> {
  typedef float Type;
};
template <>
struct IOType<float, 2> {
  typedef float2 Type;
};
template <>
struct IOType<float, 4> {
  typedef float4 Type;
};
template <>
struct IOType<int64_t, 1> {
  typedef int64_t Type;
};
template <>
struct IOType<int64_t, 2> {
  typedef uint4 Type;
};
template <>
struct IOType<uint64_t, 1> {
  typedef uint64_t Type;
};
template <>
struct IOType<uint64_t, 2> {
  typedef uint4 Type;
};
template <>
struct IOType<unsigned long long, 1> {
  typedef unsigned long long Type;
};
template <>
struct IOType<unsigned long long, 2> {
  typedef uint4 Type;
};
template <>
struct IOType<double, 1> {
  typedef double Type;
};
template <>
struct IOType<double, 2> {
  typedef double2 Type;
};

/**
     * @struct TxN_t
     *
     * @brief Internal data structure that is used to define a facade for vectorized
     * loads/stores across the most common POD types. The goal of his file is to
     * provide with CUDA programmers, an easy way to have compiler issue vectorized
     * load or store instructions to memory (either global or shared). Vectorized
     * accesses to memory are important as they'll utilize its resources
     * efficiently,
     * when compared to their non-vectorized counterparts. Obviously, for whatever
     * reasons if one is unable to issue such vectorized operations, one can always
     * fallback to using POD types.
     *
     * Concept of vectorized accesses : Threads process multiple elements
     * to speed up processing. These are loaded in a single read thanks
     * to type promotion. It is then reinterpreted as a vector elements
     * to perform the kernel's work.
     *
     * Caution : vectorized accesses requires input adresses to be memory aligned
     * according not to the input type but to the promoted type used for reading.
     *
     * Example demonstrating the use of load operations, performing math on such
     * loaded data and finally storing it back.
     * @code{.cu}
     * TxN_t<uint8_t,8> mydata1, mydata2;
     * int idx = (threadIdx.x + (blockIdx.x * blockDim.x)) * mydata1.Ratio;
     * mydata1.load(ptr1, idx);
     * mydata2.load(ptr2, idx);
     * #pragma unroll
     * for(int i=0;i<mydata1.Ratio;++i) {
     *     mydata1.val.data[i] += mydata2.val.data[i];
     * }
     * mydata1.store(ptr1, idx);
     * @endcode
     *
     * By doing as above, the interesting thing is that the code effectively remains
     * almost the same, in case one wants to upgrade to TxN_t<uint16_t,16> type.
     * Only change required is to replace variable declaration appropriately.
     *
     * Obviously, it's caller's responsibility to take care of pointer alignment!
     *
     * @tparam math_ the data-type in which the compute/math needs to happen
     * @tparam veclen_ the number of 'math_' types to be loaded/stored per
     * instruction
     */
template <typename math_, int veclen_>
struct TxN_t {
  /** underlying math data type */
  typedef math_ math_t;
  /** internal storage data type */
  typedef typename IOType<math_t, veclen_>::Type io_t;

  /** defines the number of 'math_t' types stored by this struct */
  static const int Ratio = veclen_;

  union {
    /** the vectorized data that is used for subsequent operations */
    math_t data[Ratio];
    /** internal data used to ensure vectorized loads/stores */
    io_t internal;
  } val;

  ///@todo: add default constructor

  /**
   * @brief Fill the contents of this structure with a constant value
   * @param _val the constant to be filled
   */
  DI void fill(math_t _val) {
#pragma unroll
    for (int i = 0; i < Ratio; ++i) {
      val.data[i] = _val;
    }
  }

  ///@todo: how to handle out-of-bounds!!?

  /**
   * @defgroup LoadsStores Global/Shared vectored loads or stores
   *
   * @brief Perform vectored loads/stores on this structure
   * @tparam idx_t index data type
   * @param ptr base pointer from where to load (or store) the data. It must
   *  be aligned to 'sizeof(io_t)'!
   * @param idx the offset from the base pointer which will be loaded
   *  (or stored) by the current thread. This must be aligned to 'Ratio'!
   *
   * @note: In case of loads, after a successful execution, the val.data will
   *  be populated with the desired data loaded from the pointer location. In
   * case of stores, the data in the val.data will be stored to that location.
   * @{
   */
  template <typename idx_t = int>
  DI void load(const math_t *ptr, idx_t idx) {
    const io_t *bptr = reinterpret_cast<const io_t *>(&ptr[idx]);
    val.internal = __ldg(bptr);
  }

  template <typename idx_t = int>
  DI void load(math_t *ptr, idx_t idx) {
    io_t *bptr = reinterpret_cast<io_t *>(&ptr[idx]);
    val.internal = *bptr;
  }

  template <typename idx_t = int>
  DI void store(math_t *ptr, idx_t idx) {
    io_t *bptr = reinterpret_cast<io_t *>(&ptr[idx]);
    *bptr = val.internal;
  }
  /** @} */
};

/** this is just to keep the compiler happy! */
template <typename math_>
struct TxN_t<math_, 0> {
  typedef math_ math_t;
  static const int Ratio = 1;

  union {
    math_t data[1];
  } val;

  DI void fill(math_t _val) {}
  template <typename idx_t = int>
  DI void load(const math_t *ptr, idx_t idx) {}
  template <typename idx_t = int>
  DI void load(math_t *ptr, idx_t idx) {}
  template <typename idx_t = int>
  DI void store(math_t *ptr, idx_t idx) {}
};

}  // namespace raft
