/// \file vec.cuh
/// \brief Aligned vector types for coalesced global memory access.
///
/// `AlignedVector<T, N>` wraps `N` elements of type `T` in a naturally
/// aligned struct so that the compiler emits wide (vectorized) load/store
/// instructions (e.g. `LDG.128`). The maximum supported vector width is
/// 256 bits (32 bytes), matching CUDA's widest vector load.

#pragma once
#include <sgl_kernel/utils.cuh>

#include <cstddef>
#include <cstdint>

namespace device {

namespace details {

/// \brief Maps byte-width to the corresponding unsigned integer type.
template <std::size_t N>
struct uint_trait {};

template <>
struct uint_trait<1> {
  using type = uint8_t;
};

template <>
struct uint_trait<2> {
  using type = uint16_t;
};

template <>
struct uint_trait<4> {
  using type = uint32_t;
};

template <>
struct uint_trait<8> {
  using type = uint64_t;
};

/// \brief Alias: maps `sizeof(T)` to matching unsigned int type.
template <typename T>
using sized_int = typename uint_trait<sizeof(T)>::type;

}  // namespace details

/// \brief Raw aligned storage for `N` elements of type `T`.
template <typename T, std::size_t N>
struct alignas(sizeof(T) * N) AlignedStorage {
  T data[N];
};

/**
 * \brief Aligned vector for vectorized memory access on GPU.
 *
 * Stores `N` elements of type `T` with natural alignment so that a single
 * `load`/`store` call compiles to a wide memory transaction.
 *
 * \tparam T Element type (e.g. `fp16_t`, `bf16_t`, `float`).
 * \tparam N Number of elements. Must be a power of two and
 *           `sizeof(T) * N <= 32` (256 bits).
 *
 * Example:
 * \code
 *   AlignedVector<fp16_t, 8> vec;  // 16 bytes, 128-bit aligned
 *   vec.load(input_ptr, tid);      // vectorized load
 *   vec[0] = vec[0] + 1;
 *   vec.store(output_ptr, tid);    // vectorized store
 * \endcode
 */
template <typename T, std::size_t N>
struct AlignedVector {
 private:
  static_assert(
      (N > 0 && (N & (N - 1)) == 0) && sizeof(T) * N <= kMaxVecBytes,
      "CUDA vector size exceeds arch limit: max 16 bytes on pre-Blackwell/AMD, "
      "32 bytes on Blackwell or greater");
  using element_t = typename details::sized_int<T>;
  using storage_t = AlignedStorage<element_t, N>;

 public:
  /// \brief Vectorized load from `ptr` at the given element `offset`.
  SGL_DEVICE void load(const void* ptr, int64_t offset = 0) {
    m_storage = reinterpret_cast<const storage_t*>(ptr)[offset];
  }
  /// \brief Vectorized store to `ptr` at the given element `offset`.
  SGL_DEVICE void store(void* ptr, int64_t offset = 0) const {
    reinterpret_cast<storage_t*>(ptr)[offset] = m_storage;
  }
  /// \brief Fill all N elements with the same `value`.
  SGL_DEVICE void fill(T value) {
    const auto store_value = *reinterpret_cast<element_t*>(&value);
#pragma unroll
    for (std::size_t i = 0; i < N; ++i) {
      m_storage.data[i] = store_value;
    }
  }

  SGL_DEVICE auto operator[](std::size_t idx) -> T& {
    return reinterpret_cast<T*>(&m_storage)[idx];
  }
  SGL_DEVICE auto operator[](std::size_t idx) const -> T {
    return reinterpret_cast<const T*>(&m_storage)[idx];
  }
  SGL_DEVICE auto data() -> T* {
    return reinterpret_cast<T*>(&m_storage);
  }
  SGL_DEVICE auto data() const -> const T* {
    return reinterpret_cast<const T*>(&m_storage);
  }

 private:
  storage_t m_storage;
};

}  // namespace device